fix(transport): close async subscriptions on ingress stop

Signed-off-by: Roman Volosatovs <[email protected]>

Author: rvolosatovs

Cargo.lock

@@ -52,6 +52,7 @@ async fn loopback() -> anyhow::Result<()> {
                             .shutdown()
                             .await
                             .context("failed to shutdown stream")?;
+                        drop(outgoing);
                         info!("wrote `bar`");
                         anyhow::Ok(())
                     },
@@ -65,6 +66,7 @@ async fn loopback() -> anyhow::Result<()> {
                             .shutdown()
                             .await
                             .context("failed to shutdown stream")?;
+                        drop(nested_tx);
                         info!("wrote `client->server`");
                         anyhow::Ok(())
                     },
@@ -125,6 +127,7 @@ async fn loopback() -> anyhow::Result<()> {
                             .shutdown()
                             .await
                             .context("failed to shutdown stream")?;
+                        drop(outgoing);
                         info!("wrote `foo`");
                         anyhow::Ok(())
                     },
@@ -138,6 +141,7 @@ async fn loopback() -> anyhow::Result<()> {
                             .shutdown()
                             .await
                             .context("failed to shutdown stream")?;
+                        drop(nested_tx);
                         info!("wrote `server->client`");
                         anyhow::Ok(())
                     },

crates/transport-quic/tests/loopback.rs

@@ -1,6 +1,6 @@
 [package]
 name = "wrpc-transport"
-version = "0.28.2"
+version = "0.28.3"
 description = "wRPC core transport functionality"
 
 authors.workspace = true

crates/transport/Cargo.toml

@@ -52,11 +52,17 @@ where
     results_io.spawn({
         let index = Arc::clone(&index);
         async move {
-            if let Err(err) = ingress(rx, index, results_tx).await {
+            if let Err(err) = ingress(rx, &index, results_tx).await {
                 error!(?err, "result ingress failed");
             } else {
                 debug!("result ingress successfully complete");
             }
+            let Ok(mut index) = index.lock() else {
+                error!("failed to lock index trie");
+                return;
+            };
+            trace!("shutting down index trie");
+            index.close_tx();
         }
         .in_current_span()
     });

crates/transport/src/frame/conn/client.rs

@@ -29,12 +29,12 @@ pub use client::*;
 pub use server::*;
 
 /// Index trie containing async stream subscriptions
-#[derive(Default)]
+#[derive(Debug, Default)]
 enum IndexTrie {
     #[default]
     Empty,
     Leaf {
-        tx: mpsc::Sender<std::io::Result<Bytes>>,
+        tx: Option<mpsc::Sender<std::io::Result<Bytes>>>,
         rx: Option<mpsc::Receiver<std::io::Result<Bytes>>>,
     },
     IndexNode {
@@ -65,7 +65,7 @@ impl<'a>
         ),
     ) -> Self {
         match path {
-            [] => Self::Leaf { tx, rx },
+            [] => Self::Leaf { tx: Some(tx), rx },
             [None, path @ ..] => Self::WildcardNode {
                 tx: None,
                 rx: None,
@@ -165,7 +165,7 @@ impl IndexTrie {
         let Some((i, path)) = path.split_first() else {
             return match self {
                 Self::Empty => None,
-                Self::Leaf { tx, .. } => Some(tx.clone()),
+                Self::Leaf { tx, .. } => tx.clone(),
                 Self::IndexNode { tx, .. } | Self::WildcardNode { tx, .. } => tx.clone(),
             };
         };
@@ -182,6 +182,33 @@ impl IndexTrie {
         }
     }
 
+    /// Closes all senders in the trie
+    #[instrument(level = "trace", skip(self), ret(level = "trace"))]
+    fn close_tx(&mut self) {
+        match self {
+            Self::Empty => {}
+            Self::Leaf { tx, .. } => {
+                mem::take(tx);
+            }
+            Self::IndexNode {
+                tx, ref mut nested, ..
+            } => {
+                mem::take(tx);
+                for nested in nested.iter_mut().flatten() {
+                    nested.close_tx();
+                }
+            }
+            Self::WildcardNode {
+                tx, ref mut nested, ..
+            } => {
+                mem::take(tx);
+                if let Some(nested) = nested {
+                    nested.close_tx();
+                }
+            }
+        }
+    }
+
     /// Inserts `sender` and `receiver` under a `path` - returns `false` if it failed and `true` if it succeeded.
     /// Tree state after `false` is returned is undefined
     #[instrument(level = "trace", skip(self, sender, receiver), ret(level = "trace"))]
@@ -208,14 +235,10 @@ impl IndexTrie {
                     let mut nested = Vec::with_capacity(n);
                     nested.resize_with(n, Option::default);
                     nested[*i] = Some(Self::from((path, sender, receiver)));
-                    *self = Self::IndexNode {
-                        tx: Some(tx),
-                        rx,
-                        nested,
-                    };
+                    *self = Self::IndexNode { tx, rx, nested };
                 } else {
                     *self = Self::WildcardNode {
-                        tx: Some(tx),
+                        tx,
                         rx,
                         nested: Some(Box::new(Self::from((path, sender, receiver)))),
                     };
@@ -405,7 +428,7 @@ impl AsyncWrite for Outgoing {
 #[instrument(level = "trace", skip_all, ret(level = "trace"))]
 async fn ingress(
     mut rx: impl AsyncRead + Unpin,
-    index: Arc<std::sync::Mutex<IndexTrie>>,
+    index: &std::sync::Mutex<IndexTrie>,
     param_tx: mpsc::Sender<std::io::Result<Bytes>>,
 ) -> std::io::Result<()> {
     loop {

crates/transport/src/frame/conn/mod.rs

@@ -149,11 +149,17 @@ where
             params_io.spawn({
                 let index = Arc::clone(&index);
                 async move {
-                    if let Err(err) = ingress(rx, index, params_tx).await {
+                    if let Err(err) = ingress(rx, &index, params_tx).await {
                         error!(?err, "parameter ingress failed");
                     } else {
                         debug!("parameter ingress successfully complete");
                     }
+                    let Ok(mut index) = index.lock() else {
+                        error!("failed to lock index trie");
+                        return;
+                    };
+                    trace!("shutting down index trie");
+                    index.close_tx();
                 }
                 .in_current_span()
             });