fix(rs-nats): create a Tokio runtime if one does not exist

Signed-off-by: Roman Volosatovs <[email protected]>
bytecodealliance · Sep 25, 2024 · 34f3603 · 34f3603
1 parent 424edd4
commit 34f3603
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 4 deletions.
diff --git a/crates/transport-nats/src/lib.rs b/crates/transport-nats/src/lib.rs
@@ -17,7 +17,7 @@ use futures::{Stream, StreamExt};
 use tokio::io::{AsyncRead, AsyncWrite, ReadBuf};
 use tokio::sync::oneshot;
 use tokio::try_join;
-use tracing::{debug, instrument, trace, warn};
+use tracing::{debug, error, instrument, trace, warn};
 use wrpc_transport::Index as _;
 
 pub const PROTOCOL: &str = "wrpc.0.0.1";
@@ -472,12 +472,21 @@ impl Drop for SubjectWriter {
         if !self.shutdown {
             let nats = self.nats.clone();
             let subject = mem::replace(&mut self.tx, Subject::from_static(""));
-            tokio::spawn(async move {
+            let fut = 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")
                 }
-            });
+            };
+            match tokio::runtime::Handle::try_current() {
+                Ok(rt) => {
+                    rt.spawn(fut);
+                }
+                Err(_) => match tokio::runtime::Runtime::new() {
+                    Ok(rt) => rt.block_on(fut),
+                    Err(err) => error!(?err, "failed to create a new Tokio runtime"),
+                },
+            }
         }
     }
 }

diff --git a/tests/rust.rs b/tests/rust.rs
@@ -7,6 +7,7 @@ use core::str;
 use core::time::Duration;
 
 use std::sync::Arc;
+use std::thread;
 
 use anyhow::Context;
 use bytes::Bytes;
@@ -576,7 +577,7 @@ where
 #[instrument(skip_all, ret)]
 async fn assert_dynamic<C, I, S>(clt: Arc<I>, srv: Arc<S>) -> anyhow::Result<()>
 where
-    C: Send + Sync + Default,
+    C: Send + Sync + Default + 'static,
     I: wrpc::Invoke<Context = C>,
     S: wrpc::Serve<Context = C>,
 {
@@ -625,6 +626,12 @@ where
                 .await
                 .expect("failed to accept invocation")
                 .expect("unexpected end of stream");
+            let inv = reset_inv
+                .try_next()
+                .await
+                .expect("failed to accept invocation")
+                .expect("unexpected end of stream");
+            thread::spawn(|| inv);
             anyhow::Ok(())
         }
         .instrument(info_span!("server")),
@@ -649,6 +656,16 @@ where
             )
             .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")),
     );