GO_AWAY support with graceful shutdown

moq-relay-ietf/src/main.rs

@@ -68,6 +68,12 @@ pub struct Cli {
     /// Requires --dev to enable the web server. Only serves files by exact CID - no index.
     #[arg(long)]
     pub mlog_serve: bool,
+
+    /// The public URL we advertise to other origins.
+    /// The provided certificate must be valid for this address.
+    #[arg(long)]
+    #[arg(default_value = "https://localhost:4443")]
+    pub public_url: Option<Url>,
 }
 
 #[tokio::main]
@@ -105,6 +111,7 @@ async fn main() -> anyhow::Result<()> {
 
     // Create a QUIC server for media.
     let relay = Relay::new(RelayConfig {
+        public_url: cli.public_url,
         tls: tls.clone(),
         bind: cli.bind,
         qlog_dir: qlog_dir_for_relay,

moq-relay-ietf/src/relay.rs

@@ -4,6 +4,8 @@ use anyhow::Context;
 
 use futures::{stream::FuturesUnordered, FutureExt, StreamExt};
 use moq_native_ietf::quic;
+use moq_transport::session::SessionMigration;
+use tokio::sync::broadcast;
 use url::Url;
 
 use crate::{Api, Consumer, Locals, Producer, Remotes, RemotesConsumer, RemotesProducer, Session};
@@ -31,10 +33,14 @@ pub struct RelayConfig {
     /// Our hostname which we advertise to other origins.
     /// We use QUIC, so the certificate must be valid for this address.
     pub node: Option<Url>,
+
+    /// The public URL we advertise to other origins.
+    pub public_url: Option<Url>,
 }
 
 /// MoQ Relay server.
 pub struct Relay {
+    public_url: Option<Url>,
     quic: quic::Endpoint,
     announce_url: Option<Url>,
     mlog_dir: Option<PathBuf>,
@@ -83,6 +89,7 @@ impl Relay {
         });
 
         Ok(Self {
+            public_url: config.public_url,
             quic,
             announce_url: config.announce,
             mlog_dir: config.mlog_dir,
@@ -96,6 +103,59 @@ impl Relay {
     pub async fn run(self) -> anyhow::Result<()> {
         let mut tasks = FuturesUnordered::new();
 
+        // Setup SIGTERM handler and broadcast channel
+        #[cfg(unix)]
+        let mut signal_term =
+            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::terminate())?;
+        let mut signal_int =
+            tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())?;
+
+        let (signal_tx, mut signal_rx) = broadcast::channel::<SessionMigration>(16);
+
+        // Get server address early for the shutdown signal
+        let server_addr = self
+            .quic
+            .server
+            .as_ref()
+            .context("missing TLS certificate")?
+            .local_addr()?;
+        // FIXME(itzmanish): this gives [::]:4433, which is not a valid URL
+        let shutdown_uri = if let Some(public_url) = &self.public_url {
+            public_url.clone().into()
+        } else {
+            format!("https://{}", server_addr)
+        };
+
+        // Spawn task to listen for SIGTERM and broadcast shutdown
+        let signal_tx_clone = signal_tx.clone();
+        tasks.push(
+            async move {
+                log::info!("Listening for SIGTERM");
+                #[cfg(unix)]
+                {
+                    tokio::select! {
+                        _ = signal_term.recv() => {
+                            log::info!("Received SIGTERM");
+                        }
+                        _ = signal_int.recv() => {
+                            log::info!("Received SIGINT");
+                        }
+                    }
+                    log::info!("broadcasting shutdown to all sessions");
+
+                    if let Err(e) = signal_tx.send(SessionMigration { uri: shutdown_uri }) {
+                        log::error!("failed to broadcast shutdown: {}", e);
+                    }
+                }
+                #[cfg(not(unix))]
+                {
+                    std::future::pending::<()>().await;
+                }
+                Ok(())
+            }
+            .boxed(),
+        );
+
         // Start the remotes producer task, if any
         let remotes = self.remotes.map(|(producer, consumer)| {
             tasks.push(producer.run().boxed());
@@ -115,10 +175,13 @@ impl Relay {
                 .context("failed to establish forward connection")?;
 
             // Create the MoQ session over the connection
-            let (session, publisher, subscriber) =
-                moq_transport::session::Session::connect(session, None)
-                    .await
-                    .context("failed to establish forward session")?;
+            let (session, publisher, subscriber) = moq_transport::session::Session::connect(
+                session,
+                None,
+                Some(signal_tx_clone.subscribe()),
+            )
+            .await
+            .context("failed to establish forward session")?;
 
             // Create a normal looking session, except we never forward or register announces.
             let session = Session {
@@ -158,12 +221,12 @@ impl Relay {
                     let remotes = remotes.clone();
                     let forward = forward_producer.clone();
                     let api = self.api.clone();
+                    let session_signal_rx = signal_tx_clone.subscribe();
 
                     // Spawn a new task to handle the connection
                     tasks.push(async move {
-
                         // Create the MoQ session over the connection (setup handshake etc)
-                        let (session, publisher, subscriber) = match moq_transport::session::Session::accept(conn, mlog_path).await {
+                        let (session, publisher, subscriber) = match moq_transport::session::Session::accept(conn, mlog_path, Some(session_signal_rx)).await {
                             Ok(session) => session,
                             Err(err) => {
                                 log::warn!("failed to accept MoQ session: {}", err);
@@ -186,6 +249,30 @@ impl Relay {
                     }.boxed());
                 },
                 res = tasks.next(), if !tasks.is_empty() => res.unwrap()?,
+                _ = signal_rx.recv() => {
+                    log::info!("received shutdown signal, waiting for {} active tasks to complete", tasks.len());
+
+                    // Give sessions a moment to send GOAWAY messages
+                    tokio::time::sleep(tokio::time::Duration::from_millis(200)).await;
+
+                    // Stop accepting new connections and wait for existing tasks to complete
+                    log::info!("draining {} remaining tasks...", tasks.len());
+                    let shutdown_timeout = tokio::time::Duration::from_secs(20);
+                    let result = tokio::time::timeout(shutdown_timeout, async {
+                        // Actually poll tasks to completion
+                        while let Some(res) = tasks.next().await {
+                            if let Err(e) = res {
+                                log::warn!("task failed during shutdown: {:?}", e);
+                            }
+                        }
+                    }).await;
+
+                    match result {
+                        Ok(_) => log::info!("all tasks completed successfully"),
+                        Err(_) => log::warn!("timed out waiting for tasks after {}s", shutdown_timeout.as_secs()),
+                    }
+                    break Ok(());
+                }
             }
         }
     }

moq-relay-ietf/src/remote.rs

@@ -12,7 +12,9 @@
 use moq_native_ietf::quic;
 use moq_transport::coding::TrackNamespace;
 use moq_transport::serve::{Track, TrackReader, TrackWriter};
+use moq_transport::session::SessionMigration;
 use moq_transport::watch::State;
+use tokio::sync::broadcast;
 use url::Url;
 
 use crate::Api;

moq-relay-ietf/src/session.rs

@@ -1,7 +1,7 @@
-use futures::{stream::FuturesUnordered, FutureExt, StreamExt};
-use moq_transport::session::SessionError;
-
 use crate::{Consumer, Producer};
+use futures::{stream::FuturesUnordered, FutureExt, StreamExt};
+use moq_transport::session::{SessionError, SessionMigration};
+use tokio::sync::broadcast;
 
 pub struct Session {
     pub session: moq_transport::session::Session,

moq-transport/src/message/mod.rs

@@ -222,3 +222,9 @@ message_types! {
     PublishOk = 0x1e,
     PublishError = 0x1f,
 }
+
+pub enum MiscMessage {
+    GoAway,
+    MaxRequestId,
+    RequestsBlocked,
+}

moq-transport/src/session/mod.rs

@@ -12,10 +12,12 @@ mod writer;
 pub use announce::*;
 pub use announced::*;
 pub use error::*;
+use log::info;
 pub use publisher::*;
 pub use subscribe::*;
 pub use subscribed::*;
 pub use subscriber::*;
+use tokio::sync::broadcast;
 pub use track_status_requested::*;
 
 use reader::*;
@@ -24,8 +26,8 @@ use writer::*;
 use futures::{stream::FuturesUnordered, StreamExt};
 use std::sync::{atomic, Arc, Mutex};
 
-use crate::coding::KeyValuePairs;
-use crate::message::Message;
+use crate::coding::{KeyValuePairs, SessionUri};
+use crate::message::{GoAway, Message};
 use crate::mlog;
 use crate::watch::Queue;
 use crate::{message, setup};
@@ -49,6 +51,14 @@ pub struct Session {
     /// Optional mlog writer for MoQ Transport events
     /// Wrapped in Arc<Mutex<>> to share across send/recv tasks when enabled
     mlog: Option<Arc<Mutex<mlog::MlogWriter>>>,
+
+    /// Optional signal receiver for migration events
+    signal_rx: Option<broadcast::Receiver<SessionMigration>>,
+}
+
+#[derive(Clone, Debug)]
+pub struct SessionMigration {
+    pub uri: String,
 }
 
 impl Session {
@@ -66,6 +76,7 @@ impl Session {
         recver: Reader,
         first_requestid: u64,
         mlog: Option<mlog::MlogWriter>,
+        signal_rx: Option<broadcast::Receiver<SessionMigration>>,
     ) -> (Self, Option<Publisher>, Option<Subscriber>) {
         let next_requestid = Arc::new(atomic::AtomicU64::new(first_requestid));
         let outgoing = Queue::default().split();
@@ -93,6 +104,7 @@ impl Session {
             subscriber: subscriber.clone(),
             outgoing: outgoing.1,
             mlog: mlog_shared,
+            signal_rx,
         };
 
         (session, publisher, subscriber)
@@ -103,6 +115,7 @@ impl Session {
     pub async fn connect(
         mut session: web_transport::Session,
         mlog_path: Option<PathBuf>,
+        signal_rx: Option<broadcast::Receiver<SessionMigration>>,
     ) -> Result<(Session, Publisher, Subscriber), SessionError> {
         let mlog = mlog_path.and_then(|path| {
             mlog::MlogWriter::new(path)
@@ -135,7 +148,7 @@ impl Session {
         // TODO: emit server_setup_parsed event
 
         // We are the client, so the first request id is 0
-        let session = Session::new(session, sender, recver, 0, mlog);
+        let session = Session::new(session, sender, recver, 0, mlog, signal_rx);
         Ok((session.0, session.1.unwrap(), session.2.unwrap()))
     }
 
@@ -144,6 +157,7 @@ impl Session {
     pub async fn accept(
         mut session: web_transport::Session,
         mlog_path: Option<PathBuf>,
+        signal_rx: Option<broadcast::Receiver<SessionMigration>>,
     ) -> Result<(Session, Option<Publisher>, Option<Subscriber>), SessionError> {
         let mut mlog = mlog_path.and_then(|path| {
             mlog::MlogWriter::new(path)
@@ -188,7 +202,7 @@ impl Session {
             sender.encode(&server).await?;
 
             // We are the server, so the first request id is 1
-            Ok(Session::new(session, sender, recver, 1, mlog))
+            Ok(Session::new(session, sender, recver, 1, mlog, signal_rx))
         } else {
             Err(SessionError::Version(client.versions, server_versions))
         }
@@ -198,9 +212,32 @@ impl Session {
     /// inbound control messages, receiving and processing new inbound uni-directional QUIC streams,
     /// and receiving and processing QUIC datagrams received
     pub async fn run(self) -> Result<(), SessionError> {
+        let mut cloned_outgoing = self.outgoing.clone();
+
+        // Spawn a task that waits for shutdown signal and pushes GOAWAY
+        // This runs independently and doesn't affect the main session tasks
+        if let Some(mut signal_rx) = self.signal_rx {
+            tokio::spawn(async move {
+                if let Ok(info) = signal_rx.recv().await {
+                    log::info!(
+                        "received terminate/interrupt signal, sending GOAWAY: {:#?}",
+                        info
+                    );
+                    let msg = GoAway {
+                        uri: SessionUri(info.uri),
+                    };
+                    if let Err(e) = cloned_outgoing.push(Message::GoAway(msg)) {
+                        log::error!("failed to push GOAWAY: {:#?}", e);
+                    } else {
+                        log::info!("GOAWAY message queued successfully");
+                    }
+                }
+            });
+        }
+
         tokio::select! {
             res = Self::run_recv(self.recver, self.publisher, self.subscriber.clone(), self.mlog.clone()) => res,
-            res = Self::run_send(self.sender, self.outgoing, self.mlog.clone()) => res,
+            res = Self::run_send(self.sender, self.outgoing, self.subscriber.clone(), self.mlog.clone()) => res,
             res = Self::run_streams(self.webtransport.clone(), self.subscriber.clone()) => res,
             res = Self::run_datagrams(self.webtransport, self.subscriber) => res,
         }
@@ -210,6 +247,7 @@ impl Session {
     async fn run_send(
         mut sender: Writer,
         mut outgoing: Queue<message::Message>,
+        mut subscriber: Option<Subscriber>,
         mlog: Option<Arc<Mutex<mlog::MlogWriter>>>,
     ) -> Result<(), SessionError> {
         while let Some(msg) = outgoing.pop().await {
@@ -256,6 +294,12 @@ impl Session {
                 }
             }
 
+            if let Message::GoAway(_m) = &msg {
+                subscriber
+                    .iter_mut()
+                    .for_each(|s| s.handle_go_away().unwrap_or(()));
+            }
+
             sender.encode(&msg).await?;
         }
 
@@ -340,6 +384,18 @@ impl Session {
                 Err(msg) => msg,
             };
 
+            let msg = match msg {
+                Message::GoAway(goaway) => {
+                    info!("Received GOAWAY: {:?}", goaway);
+                    subscriber
+                        .as_mut()
+                        .ok_or(SessionError::RoleViolation)?
+                        .handle_go_away()?;
+                    continue;
+                }
+                _ => msg,
+            };
+
             // TODO GOAWAY, MAX_REQUEST_ID, REQUESTS_BLOCKED
             log::warn!("Unimplemented message type received: {:?}", msg);
             return Err(SessionError::unimplemented(&format!(

moq-transport/src/session/publisher.rs

@@ -75,14 +75,14 @@ impl Publisher {
     pub async fn accept(
         session: web_transport::Session,
     ) -> Result<(Session, Publisher), SessionError> {
-        let (session, publisher, _) = Session::accept(session, None).await?;
+        let (session, publisher, _) = Session::accept(session, None, None).await?;
         Ok((session, publisher.unwrap()))
     }
 
     pub async fn connect(
         session: web_transport::Session,
     ) -> Result<(Session, Publisher), SessionError> {
-        let (session, publisher, _) = Session::connect(session, None).await?;
+        let (session, publisher, _) = Session::connect(session, None, None).await?;
         Ok((session, publisher))
     }