add a return param to Ctxt::with_current

core/src/ctxt.rs

@@ -7,17 +7,12 @@ pub trait Ctxt {
     fn open_root<P: Props>(&self, props: P) -> Self::Frame;
 
     fn open_push<P: Props>(&self, props: P) -> Self::Frame {
-        let mut frame = None;
-        self.with_current(|current| {
-            frame = Some(self.open_root(props.chain(current)));
-        });
-
-        frame.expect("`with` closure not called")
+        self.with_current(|current| self.open_root(props.chain(current)))
     }
 
     fn enter(&self, local: &mut Self::Frame);
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F);
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R;
 
     fn exit(&self, local: &mut Self::Frame);
 
@@ -44,7 +39,7 @@ impl<'a, C: Ctxt + ?Sized> Ctxt for &'a C {
         (**self).enter(frame)
     }
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         (**self).with_current(with)
     }
 
@@ -61,7 +56,7 @@ impl<C: Ctxt> Ctxt for Option<C> {
     type Current = Option<internal::Slot<C::Current>>;
     type Frame = Option<C::Frame>;
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         match self {
             Some(ctxt) => {
                 ctxt.with_current(|props| unsafe { with(&Some(internal::Slot::new(props))) })
@@ -102,7 +97,7 @@ impl<'a, C: Ctxt + ?Sized + 'a> Ctxt for alloc::boxed::Box<C> {
     type Current = C::Current;
     type Frame = C::Frame;
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         (**self).with_current(with)
     }
 
@@ -146,7 +141,7 @@ impl<'a, T: Ctxt + ?Sized> Ctxt for ByRef<'a, T> {
         self.0.enter(frame)
     }
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         self.0.with_current(with)
     }
 
@@ -163,7 +158,7 @@ impl Ctxt for Empty {
     type Current = Empty;
     type Frame = Empty;
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         with(&Empty)
     }
 
@@ -327,12 +322,15 @@ mod alloc_support {
         type Current = internal::ErasedCurrent;
         type Frame = ErasedFrame;
 
-        fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+        fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
             let mut f = Some(with);
+            let mut r = None;
 
             self.erase_ctxt().0.dispatch_with_current(&mut |props| {
-                f.take().expect("called multiple times")(&props)
+                r = Some(f.take().expect("called multiple times")(&props));
             });
+
+            r.expect("ctxt didn't call `with`")
         }
 
         fn open_root<P: Props>(&self, props: P) -> Self::Frame {
@@ -360,7 +358,7 @@ mod alloc_support {
         type Current = <dyn ErasedCtxt + 'a as Ctxt>::Current;
         type Frame = <dyn ErasedCtxt + 'a as Ctxt>::Frame;
 
-        fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+        fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
             (self as &(dyn ErasedCtxt + 'a)).with_current(with)
         }
 

core/src/runtime.rs

@@ -231,7 +231,7 @@ impl<TEmitter: Emitter, TFilter: Filter, TCtxt: Ctxt, TClock: Clock, TRng: Rng>
         self.ctxt.enter(scope)
     }
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         self.ctxt.with_current(with)
     }
 
@@ -294,7 +294,7 @@ impl<T: Ctxt> Ctxt for AssertInternal<T> {
         self.0.enter(local)
     }
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         self.0.with_current(with)
     }
 

src/frame.rs

@@ -35,7 +35,7 @@ impl<C: Ctxt> Frame<C> {
     }
 
     #[track_caller]
-    pub fn with(&mut self, with: impl FnOnce(&C::Current)) {
+    pub fn with<R>(&mut self, with: impl FnOnce(&C::Current) -> R) -> R {
         self.enter().with(with)
     }
 
@@ -72,7 +72,7 @@ pub struct EnterGuard<'a, C: Ctxt> {
 
 impl<'a, C: Ctxt> EnterGuard<'a, C> {
     #[track_caller]
-    pub fn with(&mut self, with: impl FnOnce(&C::Current)) {
+    pub fn with<R>(&mut self, with: impl FnOnce(&C::Current) -> R) -> R {
         self.scope.ctxt.with_current(with)
     }
 }

src/macro_hooks.rs

@@ -561,12 +561,11 @@ pub fn __private_push_span_ctxt<'a, 'b, E: Emitter, F: Filter, C: Ctxt, T: Clock
         }
     }
 
-    let mut trace_id = None;
-    let mut span_parent = None;
-
-    rt.with_current(|current| {
-        trace_id = current.pull::<TraceId, _>(KEY_TRACE_ID);
-        span_parent = current.pull::<SpanId, _>(KEY_SPAN_ID);
+    let (mut trace_id, span_parent) = rt.with_current(|current| {
+        (
+            current.pull::<TraceId, _>(KEY_TRACE_ID),
+            current.pull::<SpanId, _>(KEY_SPAN_ID),
+        )
     });
 
     trace_id = trace_id.or_else(|| TraceId::random(rt));

src/platform/thread_local_ctxt.rs

@@ -44,7 +44,7 @@ impl Ctxt for ThreadLocalCtxt {
     type Current = ThreadLocalSpan;
     type Frame = ThreadLocalSpan;
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         ACTIVE.with(|span| with(&*span.borrow()))
     }
 

targets/otlp/src/client/http.rs

@@ -154,12 +154,11 @@ async fn send_request(
             }
 
             // Propagate traceparent for the batch
-            let mut trace_id = None;
-            let mut span_id = None;
-
-            rt.ctxt().with_current(|props| {
-                trace_id = props.pull::<emit::trace::TraceId, _>(KEY_TRACE_ID);
-                span_id = props.pull::<emit::trace::SpanId, _>(KEY_SPAN_ID);
+            let (trace_id, span_id) = rt.ctxt().with_current(|props| {
+                (
+                    props.pull::<emit::trace::TraceId, _>(KEY_TRACE_ID),
+                    props.pull::<emit::trace::SpanId, _>(KEY_SPAN_ID),
+                )
             });
 
             req = if let (Some(trace_id), Some(span_id)) = (trace_id, span_id) {

targets/tracing/src/lib.rs

@@ -126,7 +126,7 @@ impl<C: emit::Ctxt, S: tracing::Subscriber> emit::Ctxt for TracingCtxt<C, S> {
         self.0.enter(&mut frame.1)
     }
 
-    fn with_current<F: FnOnce(&Self::Current)>(&self, with: F) {
+    fn with_current<R, F: FnOnce(&Self::Current) -> R>(&self, with: F) -> R {
         self.0.with_current(with)
     }