fix up drop ordering

emit-rs · Dec 9, 2024 · 645f91b · 645f91b
1 parent 58fade5
commit 645f91b
Show file tree

Hide file tree

Showing 2 changed files with 49 additions and 19 deletions.
diff --git a/macros/src/span.rs b/macros/src/span.rs
@@ -286,7 +286,7 @@ fn inject_sync(
     Ok(quote!({
         #setup_tokens
 
-        let (mut __ctxt, mut __span_guard) = emit::__private::__private_begin_span(
+        let (__ctxt, __span_guard) = emit::__private::__private_begin_span(
             #rt_tokens,
             #mdl_tokens,
             #span_name_tokens,
@@ -296,17 +296,13 @@ fn inject_sync(
             #ctxt_props_tokens,
             #default_completion_tokens,
         );
-        let __ctxt_guard = __ctxt.enter();
 
-        let #span_guard = &mut __span_guard;
-
-        let __r = #body_tokens;
+        __ctxt.call(move || {
+            let mut __span_guard = __span_guard;
+            let #span_guard = &mut __span_guard;
 
-        #[allow(unreachable_code)]
-        {
-            __span_guard.complete();
-            __r
-        }
+            #body_tokens
+        })
     }))
 }
 
@@ -378,15 +374,10 @@ fn inject_async(
         );
 
         __ctxt.in_future(async move {
-            let #span_guard = &mut __span_guard;
-
-            let __r = #body_tokens;
-
-            #[allow(unreachable_code)]
-            {
-                __span_guard.complete();
-                __r
-            }
+            async move {
+                let #span_guard = &mut __span_guard;
+                #body_tokens
+            }.await
         }).await
     }))
 }

diff --git a/test/ui/src/span.rs b/test/ui/src/span.rs
@@ -387,6 +387,45 @@ fn span_explicit_ids() {
     assert!(CALLED.was_called());
 }
 
+#[tokio::test]
+async fn async_span_explicit_ids() {
+    static CALLED: StaticCalled = StaticCalled::new();
+    static RT: StaticRuntime = static_runtime(
+        |evt| {
+            assert_eq!(emit::TraceId::from_u128(1), evt.props().pull("trace_id"));
+            assert_eq!(emit::SpanId::from_u64(2), evt.props().pull("span_parent"));
+            assert_eq!(emit::SpanId::from_u64(3), evt.props().pull("span_id"));
+
+            CALLED.record();
+        },
+        |evt| {
+            assert_eq!(emit::TraceId::from_u128(1), evt.props().pull("trace_id"));
+            assert_eq!(emit::SpanId::from_u64(2), evt.props().pull("span_parent"));
+            assert_eq!(emit::SpanId::from_u64(3), evt.props().pull("span_id"));
+
+            true
+        },
+    );
+
+    #[emit::span(rt: RT, "test", trace_id, span_parent, span_id)]
+    async fn exec(trace_id: &str, span_parent: &str, span_id: &str) {
+        let ctxt = emit::SpanCtxt::current(RT.ctxt());
+
+        assert_eq!(emit::TraceId::from_u128(1), ctxt.trace_id().copied());
+        assert_eq!(emit::SpanId::from_u64(2), ctxt.span_parent().copied());
+        assert_eq!(emit::SpanId::from_u64(3), ctxt.span_id().copied());
+    }
+
+    exec(
+        "00000000000000000000000000000001",
+        "0000000000000002",
+        "0000000000000003",
+    )
+    .await;
+
+    assert!(CALLED.was_called());
+}
+
 #[test]
 fn span_explicit_ids_ctxt() {
     static CALLED: StaticCalled = StaticCalled::new();