diff --git a/macros/src/span.rs b/macros/src/span.rs index f4ff8f862..377ecd40b 100644 --- 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 index 7f8a67ae8..ae2cccf3c 100644 --- 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();