Merge pull request #196 from englishm-cloudflare/me/release-idle-upstream-subscriptions

Author: englishm

moq-relay-ietf/Cargo.toml

@@ -70,6 +70,11 @@ thiserror = "2.0.17"
 metrics = "0.24"
 metrics-exporter-prometheus = { version = "0.16", optional = true }
 
+[dev-dependencies]
+# test-util provides the paused-clock runtime used to test idle timeouts without
+# sleeping for real.
+tokio = { version = "1", features = ["full", "test-util"] }
+
 [features]
 default = []
 metrics-prometheus = ["dep:metrics-exporter-prometheus"]

moq-relay-ietf/src/bin/moq-relay-ietf/main.rs

@@ -13,6 +13,7 @@ use url::Url;
 use api_coordinator::{ApiCoordinator, ApiCoordinatorConfig};
 use file_coordinator::FileCoordinator;
 use moq_relay_ietf::{Coordinator, Relay, RelayConfig, SessionConfig, Web, WebConfig};
+use std::time::Duration;
 
 #[derive(Parser, Clone)]
 pub struct Cli {
@@ -36,6 +37,12 @@ pub struct Cli {
     #[arg(long, default_value_t = 100)]
     pub max_request_id: u64,
 
+    /// Seconds to keep a cached track with no subscribers before releasing its
+    /// upstream subscription. 0 disables eviction, holding upstream
+    /// subscriptions for the lifetime of the upstream session.
+    #[arg(long, default_value_t = 30)]
+    pub cache_idle_timeout: u64,
+
     /// Forward all PUBLISH_NAMESPACE messages to the provided server for auth/routing.
     /// If not provided, the relay accepts every unique namespace publish.
     #[arg(long)]
@@ -188,23 +195,26 @@ async fn main() -> anyhow::Result<()> {
     };
 
     // Create a QUIC server for media.
-    let relay = Relay::new(RelayConfig {
-        tls: tls.clone(),
-        bind: Some(cli.bind),
-        endpoints: vec![],
-        qlog_dir: qlog_dir_for_relay,
-        mlog_dir: mlog_dir_for_relay,
-        node: cli.node,
-        announce: cli.announce,
-        coordinator,
-        session: SessionConfig {
-            max_request_id: cli.max_request_id,
+    let relay = Relay::new_with_cache_idle_timeout(
+        RelayConfig {
+            tls: tls.clone(),
+            bind: Some(cli.bind),
+            endpoints: vec![],
+            qlog_dir: qlog_dir_for_relay,
+            mlog_dir: mlog_dir_for_relay,
+            node: cli.node,
+            announce: cli.announce,
+            coordinator,
+            session: SessionConfig {
+                max_request_id: cli.max_request_id,
+            },
+            // No connection tagger: the default binary treats every inbound
+            // connection as a public client. Embedders that run relay-to-relay
+            // meshes supply a tagger to mark internal peers.
+            connection_tagger: None,
         },
-        // No connection tagger: the default binary treats every inbound
-        // connection as a public client. Embedders that run relay-to-relay
-        // meshes supply a tagger to mark internal peers.
-        connection_tagger: None,
-    })?;
+        Duration::from_secs(cli.cache_idle_timeout),
+    )?;
 
     if cli.dev {
         // Create a web server too.

moq-relay-ietf/src/consumer.rs

@@ -12,7 +12,9 @@ use moq_transport::{
 };
 use tokio::sync::Semaphore;
 
-use crate::{metrics::GaugeGuard, Coordinator, Locals, Producer, RemoteManager, SessionContext};
+use crate::{
+    metrics::GaugeGuard, Coordinator, Locals, Producer, RemoteManager, SessionContext, TrackRequest,
+};
 
 const MAX_INBOUND_PUBLISH_TRACKS_PER_SESSION: usize = 1024;
 
@@ -254,11 +256,11 @@ impl Consumer {
                     tracing::info!(namespace = %ns, "PUBLISH_NAMESPACE closed");
                     return Ok(());
                 },
-                Some(track) = requests.recv() => {
+                Some(TrackRequest { writer, lease }) = requests.recv() => {
                     let mut subscriber = self.subscriber.clone();
 
                     tasks.push(async move {
-                        let info = track.clone();
+                        let info = writer.info.clone();
                         let namespace = info.namespace.to_utf8_path();
                         let track_name = info.name.clone();
                         tracing::info!(
@@ -267,15 +269,47 @@ impl Consumer {
                             "forwarding subscribe: {:?}", info
                         );
 
-                        if let Err(err) = subscriber.subscribe(track).await {
-                            tracing::warn!(
-                                namespace = %namespace,
-                                track = %track_name,
-                                error = %err,
-                                "failed forwarding subscribe: {:?}", info
-                            )
+                        // Hold the subscription explicitly rather than using
+                        // `subscribe()`, so it can be dropped — sending
+                        // UNSUBSCRIBE — once downstream interest goes away.
+                        let subscribe = match subscriber.subscribe_open(writer).await {
+                            Ok(subscribe) => subscribe,
+                            Err(err) => {
+                                tracing::warn!(
+                                    namespace = %namespace,
+                                    track = %track_name,
+                                    error = %err,
+                                    "failed forwarding subscribe: {:?}", info
+                                );
+                                return Ok(());
+                            }
+                        };
+
+                        tokio::select! {
+                            res = subscribe.closed() => {
+                                if let Err(err) = res {
+                                    tracing::warn!(
+                                        namespace = %namespace,
+                                        track = %track_name,
+                                        error = %err,
+                                        "failed forwarding subscribe: {:?}", info
+                                    )
+                                }
+                            }
+                            // The cached track went unwatched long enough to be
+                            // evicted, so stop pulling it. Dropping `subscribe`
+                            // below sends UNSUBSCRIBE upstream.
+                            _ = lease.released() => {
+                                tracing::info!(
+                                    namespace = %namespace,
+                                    track = %track_name,
+                                    "releasing upstream subscription for idle cached track"
+                                );
+                            }
                         }
 
+                        drop(subscribe);
+
                         Ok(())
                     }.boxed());
                 },
@@ -352,13 +386,13 @@ impl Consumer {
         }
 
         tracing::debug!(
-            namespace = %namespace.to_utf8_path(),
+            namespace = %namespace,
             track = %track_name,
             "PUBLISH registered as exact local track"
         );
 
         publish.closed().await?;
-        tracing::info!(namespace = %namespace.to_utf8_path(), track = %track_name, "PUBLISH closed");
+        tracing::info!(namespace = %namespace, track = %track_name, "PUBLISH closed");
 
         Ok(())
     }