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 futures::StreamExt;
 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/error.rs

@@ -26,6 +26,9 @@ pub enum SessionError {
     #[error("role violation")]
     RoleViolation,
 
+    #[error("protocol violation: {0}")]
+    ProtocolViolation(&'static str),
+
     /// Some VarInt was too large and we were too lazy to handle it
     #[error("varint bounds exceeded")]
     BoundsExceeded(#[from] coding::BoundsExceeded),
@@ -52,6 +55,7 @@ impl SessionError {
         match self {
             // PROTOCOL_VIOLATION (0x3) - The role negotiated in the handshake was violated
             Self::RoleViolation => 0x3,
+            Self::ProtocolViolation(_) => 0x3,
             // INTERNAL_ERROR (0x1) - Generic internal errors
             Self::Session(_) => 0x1,
             Self::Read(_) => 0x1,