refactor: remove signal_rx parameter from session.run() and handle migration signals internally

Author: itzmanish

moq-clock-ietf/src/main.rs

@@ -64,7 +64,7 @@ async fn main() -> anyhow::Result<()> {
             let clock_publisher = clock::Publisher::new_datagram(track_writer.datagrams()?);
 
             tokio::select! {
-                res = session.run(None) => res.context("session error")?,
+                res = session.run() => res.context("session error")?,
                 res = clock_publisher.run() => res.context("clock error")?,
                 res = publisher.announce(tracks_reader) => res.context("failed to serve tracks")?,
             }
@@ -80,7 +80,7 @@ async fn main() -> anyhow::Result<()> {
             let clock_publisher = clock::Publisher::new(track_writer.subgroups()?);
 
             tokio::select! {
-                res = session.run(None) => res.context("session error")?,
+                res = session.run() => res.context("session error")?,
                 res = clock_publisher.run() => res.context("clock error")?,
                 res = publisher.announce(tracks_reader) => res.context("failed to serve tracks")?,
             }
@@ -104,7 +104,7 @@ async fn main() -> anyhow::Result<()> {
         let clock_subscriber = clock::Subscriber::new(track_reader);
 
         tokio::select! {
-            res = session.run(None) => res.context("session error")?,
+            res = session.run() => res.context("session error")?,
             res = clock_subscriber.run() => res.context("clock error")?,
             res = subscriber.subscribe(track_writer) => res.context("failed to subscribe to track")?,
         }

moq-pub/src/main.rs

@@ -76,7 +76,7 @@ async fn main() -> anyhow::Result<()> {
         .context("failed to create MoQ Transport publisher")?;
 
     tokio::select! {
-        res = session.run(None) => res.context("session error")?,
+        res = session.run() => res.context("session error")?,
         res = run_media(media) => {
             res.context("media error")?
         },

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

@@ -33,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>,
@@ -85,6 +89,7 @@ impl Relay {
         });
 
         Ok(Self {
+            public_url: config.public_url,
             quic,
             announce_url: config.announce,
             mlog_dir: config.mlog_dir,
@@ -105,7 +110,7 @@ impl Relay {
         let mut signal_int =
             tokio::signal::unix::signal(tokio::signal::unix::SignalKind::interrupt())?;
 
-        let (signal_tx, signal_rx) = broadcast::channel::<SessionMigration>(16);
+        let (signal_tx, mut signal_rx) = broadcast::channel::<SessionMigration>(16);
 
         // Get server address early for the shutdown signal
         let server_addr = self
@@ -114,7 +119,12 @@ impl Relay {
             .as_ref()
             .context("missing TLS certificate")?
             .local_addr()?;
-        let shutdown_uri = format!("https://{}", server_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();
@@ -148,8 +158,7 @@ impl Relay {
 
         // Start the remotes producer task, if any
         let remotes = self.remotes.map(|(producer, consumer)| {
-            let signal_rx = signal_rx.resubscribe();
-            tasks.push(producer.run(signal_rx).boxed());
+            tasks.push(producer.run().boxed());
             consumer
         });
 
@@ -166,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 {
@@ -184,15 +196,7 @@ impl Relay {
 
             let forward_producer = session.producer.clone();
 
-            tasks.push(
-                async move {
-                    session
-                        .run(signal_tx_clone.subscribe())
-                        .await
-                        .context("forwarding failed")
-                }
-                .boxed(),
-            );
+            tasks.push(async move { session.run().await.context("forwarding failed") }.boxed());
 
             forward_producer
         } else {
@@ -202,7 +206,6 @@ impl Relay {
         // Start the QUIC server loop
         let mut server = self.quic.server.context("missing TLS certificate")?;
         log::info!("listening on {}", server.local_addr()?);
-        let mut cloned_signal_rx = signal_rx.resubscribe();
 
         loop {
             tokio::select! {
@@ -218,12 +221,12 @@ impl Relay {
                     let remotes = remotes.clone();
                     let forward = forward_producer.clone();
                     let api = self.api.clone();
-                    let session_signal_rx = signal_rx.resubscribe();
+                    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);
@@ -238,26 +241,35 @@ impl Relay {
                             consumer: subscriber.map(|subscriber| Consumer::new(subscriber, locals, api, forward)),
                         };
 
-                        if let Err(err) = session.run(session_signal_rx).await {
+                        if let Err(err) = session.run().await {
                             log::warn!("failed to run MoQ session: {}", err);
                         }
 
                         Ok(())
                     }.boxed());
                 },
                 res = tasks.next(), if !tasks.is_empty() => res.unwrap()?,
-                _ = cloned_signal_rx.recv() => {
-                    log::info!("received shutdown signal, shutting down. Active tasks: {}", tasks.len());
-                    // set a timeout for waiting for tasks to be empty
-                    // FIXME(itzmanish): make this configurable and revisit
-                    let timeout = tokio::time::timeout(tokio::time::Duration::from_secs(20), async move {
-                        while !tasks.is_empty() {
-                            // sleep 500ms before checking again
-                            tokio::time::sleep(tokio::time::Duration::from_millis(500)).await;
+                _ = 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);
+                            }
                         }
-                    });
-                    if let Err(e) = timeout.await {
-                        log::warn!("timed out waiting for tasks to be empty: {}", 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

@@ -74,18 +74,13 @@ impl RemotesProducer {
     }
 
     /// Run the remotes producer to serve remote requests.
-    pub async fn run(
-        mut self,
-        signal_rx: broadcast::Receiver<SessionMigration>,
-    ) -> anyhow::Result<()> {
+    pub async fn run(mut self) -> anyhow::Result<()> {
         let mut tasks = FuturesUnordered::new();
 
         loop {
             tokio::select! {
                 Some(mut remote) = self.next() => {
                     let url = remote.url.clone();
-                    // Each remote task needs its own receiver
-                    let remote_signal_rx = signal_rx.resubscribe();
 
                     // Spawn a task to serve the remote
                     tasks.push(async move {
@@ -94,7 +89,7 @@ impl RemotesProducer {
                         log::warn!("serving remote: {:?}", info);
 
                         // Run the remote producer
-                        if let Err(err) = remote.run(remote_signal_rx).await {
+                        if let Err(err) = remote.run().await {
                             log::warn!("failed serving remote: {:?}, error: {}", info, err);
                         }
 
@@ -232,16 +227,13 @@ impl RemoteProducer {
         Self { info, state }
     }
 
-    pub async fn run(
-        &mut self,
-        signal_rx: broadcast::Receiver<SessionMigration>,
-    ) -> anyhow::Result<()> {
+    pub async fn run(&mut self) -> anyhow::Result<()> {
         // TODO reuse QUIC and MoQ sessions
         let (session, _quic_client_initial_cid) = self.quic.connect(&self.url).await?;
         let (session, subscriber) = moq_transport::session::Subscriber::connect(session).await?;
 
         // Run the session
-        let mut session = session.run(Some(signal_rx)).boxed();
+        let mut session = session.run().boxed();
         let mut tasks = FuturesUnordered::new();
 
         let mut done = None;

moq-relay-ietf/src/session.rs

@@ -11,12 +11,9 @@ pub struct Session {
 
 impl Session {
     /// Run the session, producer, and consumer as necessary.
-    pub async fn run(
-        self,
-        signal_rx: broadcast::Receiver<SessionMigration>,
-    ) -> Result<(), SessionError> {
+    pub async fn run(self) -> Result<(), SessionError> {
         let mut tasks = FuturesUnordered::new();
-        tasks.push(self.session.run(Some(signal_rx)).boxed());
+        tasks.push(self.session.run().boxed());
 
         if let Some(producer) = self.producer {
             tasks.push(producer.run().boxed());

moq-sub/src/main.rs

@@ -45,7 +45,7 @@ async fn main() -> anyhow::Result<()> {
     let mut media = Media::new(subscriber, tracks, out, config.catalog).await?;
 
     tokio::select! {
-        res = session.run(None) => res.context("session error")?,
+        res = session.run() => res.context("session error")?,
         res = media.run() => res.context("media error")?,
     }