fix(rs-nats): close streams on drop

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

Author: rvolosatovs

Cargo.lock

@@ -6,6 +6,7 @@ version = "0.11.0"
 authors.workspace = true
 categories.workspace = true
 edition.workspace = true
+homepage.workspace = true
 license.workspace = true
 repository.workspace = true
 
@@ -140,6 +141,6 @@ wrpc-cli = { version = "0.3", path = "./crates/cli", default-features = false }
 wrpc-introspect = { version = "0.3", default-features = false, path = "./crates/introspect" }
 wrpc-runtime-wasmtime = { version = "0.22", path = "./crates/runtime-wasmtime", default-features = false }
 wrpc-transport = { version = "0.26.8", path = "./crates/transport", default-features = false }
-wrpc-transport-nats = { version = "0.23.1", path = "./crates/transport-nats", default-features = false }
+wrpc-transport-nats = { version = "0.23.2", path = "./crates/transport-nats", default-features = false }
 wrpc-transport-quic = { version = "0.1.2", path = "./crates/transport-quic", default-features = false }
 wrpc-wasmtime-nats-cli = { version = "0.8", path = "./crates/wasmtime-nats-cli", default-features = false }

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "wrpc-transport-nats"
-version = "0.23.1"
+version = "0.23.2"
 description = "wRPC NATS transport"
 
 authors.workspace = true

crates/transport-nats/Cargo.toml

@@ -357,6 +357,10 @@ impl AsyncRead for Reader {
         match self.incoming.poll_next_unpin(cx) {
             Poll::Ready(Some(Message { mut payload, .. })) => {
                 trace!(?payload, "received message");
+                if payload.is_empty() {
+                    trace!("received stream shutdown message");
+                    return Poll::Ready(Ok(()));
+                }
                 if payload.len() > cap {
                     trace!(len = payload.len(), cap, "partially reading the message");
                     buf.put_slice(&payload.split_to(cap));
@@ -380,11 +384,16 @@ impl AsyncRead for Reader {
 pub struct SubjectWriter {
     nats: async_nats::Client,
     tx: Subject,
+    shutdown: bool,
 }
 
 impl SubjectWriter {
     fn new(nats: async_nats::Client, tx: Subject) -> Self {
-        Self { nats, tx }
+        Self {
+            nats,
+            tx,
+            shutdown: false,
+        }
     }
 }
 
@@ -395,6 +404,7 @@ impl wrpc_transport::Index<Self> for SubjectWriter {
         Ok(Self {
             nats: self.nats.clone(),
             tx,
+            shutdown: false,
         })
     }
 }
@@ -450,18 +460,34 @@ impl AsyncWrite for SubjectWriter {
 
     #[instrument(level = "trace", skip_all, ret, fields(subject = self.tx.as_str()))]
     fn poll_shutdown(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<std::io::Result<()>> {
-        trace!("writing empty buffer to shut down stream");
+        trace!("writing stream shutdown message");
         ready!(self.as_mut().poll_write(cx, &[]))?;
+        self.shutdown = true;
         Poll::Ready(Ok(()))
     }
 }
 
+impl Drop for SubjectWriter {
+    fn drop(&mut self) {
+        if !self.shutdown {
+            let nats = self.nats.clone();
+            let subject = mem::replace(&mut self.tx, Subject::from_static(""));
+            tokio::spawn(async move {
+                trace!("writing stream shutdown message");
+                if let Err(err) = nats.publish(subject, Bytes::default()).await {
+                    warn!(?err, "failed to publish stream shutdown message")
+                }
+            });
+        }
+    }
+}
+
 #[derive(Debug, Default)]
 pub enum RootParamWriter {
     #[default]
     Corrupted,
     Handshaking {
-        tx: SubjectWriter,
+        nats: async_nats::Client,
         sub: Subscriber,
         indexed: std::sync::Mutex<Vec<(Vec<usize>, oneshot::Sender<SubjectWriter>)>>,
         buffer: Bytes,
@@ -474,9 +500,9 @@ pub enum RootParamWriter {
 }
 
 impl RootParamWriter {
-    fn new(tx: SubjectWriter, sub: Subscriber, buffer: Bytes) -> Self {
+    fn new(nats: async_nats::Client, sub: Subscriber, buffer: Bytes) -> Self {
         Self::Handshaking {
-            tx,
+            nats,
             sub,
             indexed: std::sync::Mutex::default(),
             buffer,
@@ -520,7 +546,7 @@ impl RootParamWriter {
                         reply: Some(tx), ..
                     })) => {
                         let Self::Handshaking {
-                            tx: SubjectWriter { nats, .. },
+                            nats,
                             indexed,
                             buffer,
                             ..
@@ -899,7 +925,7 @@ impl wrpc_transport::Invoke for Client {
         .context("failed to send handshake")?;
         Ok((
             ParamWriter::Root(RootParamWriter::new(
-                SubjectWriter::new((*self.nats).clone(), param_tx),
+                (*self.nats).clone(),
                 handshake_rx,
                 params,
             )),

crates/transport-nats/src/lib.rs

@@ -598,13 +598,61 @@ where
         )
         .await
         .context("failed to serve `test.async`")?;
+    let reset_inv = srv
+        .serve_values::<(String,), (String,)>("test", "reset", [Box::default(); 0])
+        .await
+        .context("failed to serve `test.reset`")?;
     let sync_inv = srv
         .serve_values("test", "sync", [Box::default(); 0])
         .await
         .context("failed to serve `test.sync`")?;
+
     let mut async_inv = pin!(async_inv);
+    let mut reset_inv = pin!(reset_inv);
     let mut sync_inv = pin!(sync_inv);
 
+    join!(
+        async {
+            info!("receiving `test.reset` parameters");
+            reset_inv
+                .try_next()
+                .await
+                .expect("failed to accept invocation")
+                .expect("unexpected end of stream");
+            info!("receiving `test.reset` parameters");
+            reset_inv
+                .try_next()
+                .await
+                .expect("failed to accept invocation")
+                .expect("unexpected end of stream");
+            anyhow::Ok(())
+        }
+        .instrument(info_span!("server")),
+        async {
+            info!("invoking `test.reset`");
+            clt.invoke_values_blocking::<_, _, (String,)>(
+                C::default(),
+                "test",
+                "reset",
+                ("arg",),
+                &[[]; 0],
+            )
+            .await
+            .expect_err("`test.reset` should have failed");
+            info!("invoking `test.reset`");
+            clt.invoke_values_blocking::<_, _, (String,)>(
+                C::default(),
+                "test",
+                "reset",
+                ("arg",),
+                &[[]; 0],
+            )
+            .await
+            .expect_err("`test.reset` should have failed");
+        }
+        .instrument(info_span!("client")),
+    );
+
     join!(
         async {
             info!("receiving `test.sync` parameters");
@@ -730,7 +778,6 @@ where
             assert_eq!(m, "test");
             assert_eq!(n, [[b"foo"]]);
             assert_eq!(o, [Some(vec![Ok(Some(String::from("bar")))])]);
-            info!("finishing `test.sync` session");
         }
         .instrument(info_span!("client")),
     );
@@ -867,6 +914,7 @@ where
         }
         .instrument(info_span!("client")),
     );
+
     Ok(())
 }
 
@@ -942,7 +990,7 @@ async fn rust_dynamic_quic() -> anyhow::Result<()> {
     use core::pin::pin;
 
     common::with_quic(
-        &["sync.test", "async.test"],
+        &["sync.test", "async.test", "reset.test"],
         |port, clt_ep, srv_ep| async move {
             let clt = wrpc_transport_quic::Client::new(clt_ep, (Ipv6Addr::LOCALHOST, port));
             let srv = wrpc_transport_quic::Server::default();