Merge pull request #39 from palantir/fix-async-spanned

Fix send-ness of spanned async functions
palantir · May 7, 2021 · 2fc7b84 · 2fc7b84
2 parents 0ebdbf8 + b9e6ce8
commit 2fc7b84
Show file tree

Hide file tree

Showing 3 changed files with 26 additions and 8 deletions.
diff --git a/zipkin-macros/Cargo.toml b/zipkin-macros/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "zipkin-macros"
-version = "0.1.0"
+version = "0.1.1"
 authors = ["Steven Fackler <sfackler@palantir.com>"]
 edition = "2018"
 license = "Apache-2.0"

diff --git a/zipkin-macros/src/lib.rs b/zipkin-macros/src/lib.rs
@@ -75,14 +75,20 @@ fn spanned_impl(args: AttributeArgs, mut func: ImplItemMethod) -> Result<TokenSt
 
     if func.sig.asyncness.is_some() {
         let stmts = &func.block.stmts;
-        let stmt = quote! {
-            zipkin::next_span()
+        func.block.stmts = vec![
+            syn::parse2(quote! {
+            let __macro_impl_span = zipkin::next_span()
                 .with_name(#name)
-                .detach()
-                .bind(async move { #(#stmts)* })
-                .await
-        };
-        func.block.stmts = vec![Stmt::Expr(syn::parse2(stmt).unwrap())];
+                .detach();
+            })
+            .unwrap(),
+            Stmt::Expr(
+                syn::parse2(quote! {
+                    __macro_impl_span.bind(async move { #(#stmts)* }).await
+                })
+                .unwrap(),
+            ),
+        ];
     } else {
         let stmt = quote! {
             let __macro_impl_span = zipkin::next_span().with_name(#name);

diff --git a/zipkin/src/test/macros.rs b/zipkin/src/test/macros.rs
@@ -15,6 +15,12 @@ use crate as zipkin; // hack to get the macro codegen to work in the same crate
 use crate::{spanned, test};
 use futures::executor;
 
+fn is_send<T>(_: T)
+where
+    T: Send,
+{
+}
+
 #[test]
 fn blocking_free_function() {
     #[spanned(name = "foobar")]
@@ -105,6 +111,8 @@ fn async_free_function() {
         zipkin::next_span().with_name("fizzbuzz");
     }
 
+    is_send(foo());
+
     test::init();
 
     let future = zipkin::next_span().with_name("root").detach().bind(foo());
@@ -133,6 +141,8 @@ fn async_associated_function() {
         }
     }
 
+    is_send(Foo::foo());
+
     test::init();
 
     let future = zipkin::next_span()
@@ -164,6 +174,8 @@ fn async_method() {
         }
     }
 
+    is_send(Foo.foo());
+
     test::init();
 
     let future = zipkin::next_span()