8040: Allow impl FnMut() in taskdumps instead of just fn()

tokio/src/runtime/task/trace/mod.rs

@@ -34,7 +34,8 @@ pub(crate) struct Context {
     /// The function that is invoked at each leaf future inside of Tokio
     ///
     /// For example, within tokio::time:sleep, sockets. etc.
-    trace_leaf_fn: Cell<Option<fn(&TraceMeta)>>,
+    #[allow(clippy::type_complexity)]
+    trace_leaf_fn: Cell<Option<NonNull<dyn FnMut(&TraceMeta)>>>,
 }
 
 /// A [`Frame`] in an intrusive, doubly-linked tree of [`Frame`]s.
@@ -114,19 +115,38 @@ impl Context {
         }
     }
 
+    /// Calls the provided closure if we are being traced.
     fn try_with_current_trace_leaf_fn<F, R>(f: F) -> Option<R>
     where
-        F: FnOnce(&Cell<Option<fn(&TraceMeta)>>) -> R,
+        F: for<'a> FnOnce(&'a mut dyn FnMut(&TraceMeta)) -> R,
     {
+        let mut ret = None;
+
+        let inner = |context: &Context| {
+            if let Some(mut trace_leaf_fn) = context.trace_leaf_fn.get() {
+                context.trace_leaf_fn.set(None);
+                let _restore = defer(move || {
+                    context.trace_leaf_fn.set(Some(trace_leaf_fn));
+                });
+
+                // SAFETY: The trace leaf fn is valid for the duration in which it's stored in the
+                // context. Furthermore, re-entrant calls are not possible since we store `None` for
+                // the duration in which we hold a mutable reference, so access is exclusive for that
+                // duration.
+                ret = Some(f(unsafe { trace_leaf_fn.as_mut() }));
+            }
+        };
+
         // SAFETY: This call can only access the trace_leaf_fn field, so it cannot
         // break the trace frame linked list.
-        unsafe { Self::try_with_current(|context| f(&context.trace_leaf_fn)) }
+        unsafe { Self::try_with_current(inner) };
+
+        ret
     }
 
     /// Produces `true` if the current task is being traced; otherwise false.
     pub(crate) fn is_tracing() -> bool {
-        Self::try_with_current_trace_leaf_fn(|maybe_trace_leaf| maybe_trace_leaf.get().is_some())
-            .unwrap_or(false)
+        Self::try_with_current_trace_leaf_fn(|_| ()).is_some()
     }
 }
 
@@ -189,20 +209,46 @@ pub struct TraceMeta {
 /// assert!(count > 0);
 /// # }
 /// ```
-pub fn trace_with<F, R>(f: F, trace_leaf: fn(&TraceMeta)) -> R
+pub fn trace_with<FN, FT, R>(f: FN, mut trace_leaf: FT) -> R
 where
-    F: FnOnce() -> R,
+    FN: FnOnce() -> R,
+    FT: FnMut(&TraceMeta),
 {
-    // store our new trace_leaf function
-    let previous =
-        Context::try_with_current_trace_leaf_fn(|current| current.replace(Some(trace_leaf)));
+    let trace_leaf_dyn = (&mut trace_leaf) as &mut (dyn FnMut(&TraceMeta) + '_);
+    // SAFETY: The raw pointer is removed from the thread local before `trace_leaf` is dropped, so
+    // this transmute cannot lead to the violation of any lifetime requirements.
+    let trace_leaf_dyn = unsafe {
+        std::mem::transmute::<
+            *mut (dyn FnMut(&TraceMeta) + '_),
+            *mut (dyn FnMut(&TraceMeta) + 'static),
+        >(trace_leaf_dyn)
+    };
+    // SAFETY: Pointer comes from reference, so not null.
+    let trace_leaf_dyn = unsafe { NonNull::new_unchecked(trace_leaf_dyn) };
+
+    let mut old_trace_leaf_fn = None;
+
+    // Even if this access fails, that's okay. In that case, we still call the closure without
+    // actually performing any tracing.
+    //
+    // SAFETY: This call can only access the trace_leaf_fn field, so it cannot
+    // break the trace frame linked list.
+    unsafe {
+        Context::try_with_current(|ctx| {
+            old_trace_leaf_fn = ctx.trace_leaf_fn.replace(Some(trace_leaf_dyn));
+        })
+    };
 
-    // restore previous on drop. This is ensures state remains consistent
-    // even if the trace_leaf function panics
     let _restore = defer(move || {
-        if let Some(previous) = previous {
-            Context::try_with_current_trace_leaf_fn(|current| current.set(previous));
-        }
+        // This ensures that `trace_leaf_fn` cannot be accessed after this call returns.
+        //
+        // SAFETY: This call can only access the trace_leaf_fn field, so it cannot
+        // break the trace frame linked list.
+        unsafe {
+            Context::try_with_current(|ctx| {
+                ctx.trace_leaf_fn.set(old_trace_leaf_fn);
+            })
+        };
     });
 
     f()
@@ -249,13 +295,14 @@ impl Trace {
 // internal implementation details of this crate).
 #[inline(never)]
 pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> {
-    let trace_leaf_fn = Context::try_with_current_trace_leaf_fn(|cell| cell.get()).flatten();
-    if let Some(trace_leaf_fn) = trace_leaf_fn {
+    let root_addr = Context::current_frame_addr();
+
+    let ret = Context::try_with_current_trace_leaf_fn(|leaf_fn| {
         let meta = TraceMeta {
-            root_addr: Context::current_frame_addr(),
+            root_addr,
             trace_leaf_addr: trace_leaf as *const c_void,
         };
-        trace_leaf_fn(&meta);
+        leaf_fn(&meta);
 
         // Use the same logic that `yield_now` uses to send out wakeups after
         // the task yields.
@@ -268,10 +315,11 @@ pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> {
                 }
             }
         });
+    });
 
-        Poll::Pending
-    } else {
-        Poll::Ready(())
+    match ret {
+        Some(()) => Poll::Pending,
+        None => Poll::Ready(()),
     }
 }
 

tokio/src/runtime/task/trace/trace_impl.rs

@@ -2,73 +2,42 @@
 //!
 //! This implementation may eventually be extracted into a separate `tokio-taskdump` crate.
 
-use std::{cell::Cell, ptr};
+use std::ptr;
 
 use crate::runtime::task::trace::{trace_with, Trace, TraceMeta};
 
-use super::defer;
-
-/// Thread local state used to communicate between calling the trace and the interior `trace_leaf` function
-struct TraceContext {
-    collector: Cell<Option<Trace>>,
-}
-
-thread_local! {
-    static TRACE_CONTEXT: TraceContext = const {
-        TraceContext {
-            collector: Cell::new(None),
-        }
-    };
-}
-
 /// Capture using the default `backtrace::trace`-based implementation.
 #[inline(never)]
 pub(super) fn capture<F, R>(f: F) -> (R, Trace)
 where
     F: FnOnce() -> R,
 {
-    let collector = Trace::empty();
-
-    let previous = TRACE_CONTEXT.with(|state| state.collector.replace(Some(collector)));
+    let mut trace = Trace::empty();
 
-    // restore previous collector on drop even if the callback panics
-    let _restore = defer(move || {
-        TRACE_CONTEXT.with(|state| state.collector.set(previous));
-    });
+    let result = trace_with(f, |meta| trace_leaf(meta, &mut trace));
 
-    let result = trace_with(f, trace_leaf);
-
-    // take the collector before _restore runs
-    let collector = TRACE_CONTEXT.with(|state| state.collector.take()).unwrap();
-
-    (result, collector)
+    (result, trace)
 }
 
-/// Capture a backtrace via `backtrace::trace` and collect it into `STATE`
-#[inline(never)]
-pub(crate) fn trace_leaf(meta: &TraceMeta) {
-    TRACE_CONTEXT.with(|state| {
-        if let Some(mut collector) = state.collector.take() {
-            let mut frames: Vec<backtrace::BacktraceFrame> = vec![];
-            let mut above_leaf = false;
-
-            if let Some(root_addr) = meta.root_addr {
-                backtrace::trace(|frame| {
-                    let below_root = !ptr::eq(frame.symbol_address(), root_addr);
+/// Capture a backtrace via `backtrace::trace` and collect it into `trace`.
+pub(crate) fn trace_leaf(meta: &TraceMeta, trace: &mut Trace) {
+    let mut frames: Vec<backtrace::BacktraceFrame> = vec![];
+    let mut above_leaf = false;
 
-                    if above_leaf && below_root {
-                        frames.push(frame.to_owned().into());
-                    }
+    if let Some(root_addr) = meta.root_addr {
+        backtrace::trace(|frame| {
+            let below_root = !ptr::eq(frame.symbol_address(), root_addr);
 
-                    if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) {
-                        above_leaf = true;
-                    }
+            if above_leaf && below_root {
+                frames.push(frame.to_owned().into());
+            }
 
-                    below_root
-                });
+            if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) {
+                above_leaf = true;
             }
-            collector.push_backtrace(frames);
-            state.collector.set(Some(collector));
-        }
-    });
+
+            below_root
+        });
+    }
+    trace.push_backtrace(frames);
 }

tokio/tests/task_trace_self.rs

@@ -109,48 +109,41 @@ async fn task_trace_self() {
 
 /// Collect frames between `trace_leaf_for_test` and `root_addr` using
 /// `backtrace::trace`, resolve them, and store pretty-printed symbol names
-/// (with compiler hashes stripped) into `TRACE_WITH_LOG`.
+/// (with compiler hashes stripped) into `logs`.
 #[inline(never)]
-fn trace_leaf_for_test(meta: &TraceMeta) {
-    TRACE_WITH_LOG.with(|log| {
-        let mut frames: Vec<backtrace::BacktraceFrame> = vec![];
-        let mut above_leaf = false;
+fn trace_leaf_for_test(meta: &TraceMeta, log: &mut Vec<Vec<String>>) {
+    let mut frames: Vec<backtrace::BacktraceFrame> = vec![];
+    let mut above_leaf = false;
 
-        if let Some(root_addr) = meta.root_addr {
-            backtrace::trace(|frame| {
-                let below_root = !ptr::eq(frame.symbol_address(), root_addr);
+    if let Some(root_addr) = meta.root_addr {
+        backtrace::trace(|frame| {
+            let below_root = !ptr::eq(frame.symbol_address(), root_addr);
 
-                if above_leaf && below_root {
-                    frames.push(frame.to_owned().into());
-                }
+            if above_leaf && below_root {
+                frames.push(frame.to_owned().into());
+            }
 
-                if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) {
-                    above_leaf = true;
-                }
+            if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) {
+                above_leaf = true;
+            }
 
-                below_root
-            });
-        }
+            below_root
+        });
+    }
 
-        // Resolve frames into human-readable symbol names with hashes stripped.
-        let mut bt = backtrace::Backtrace::from(frames);
-        bt.resolve();
-        let mut names = vec![];
-        for frame in bt.frames() {
-            for symbol in frame.symbols() {
-                if let Some(name) = symbol.name() {
-                    names.push(strip_symbol_hash(&format!("{name}")).to_owned());
-                }
+    // Resolve frames into human-readable symbol names with hashes stripped.
+    let mut bt = backtrace::Backtrace::from(frames);
+    bt.resolve();
+    let mut names = vec![];
+    for frame in bt.frames() {
+        for symbol in frame.symbols() {
+            if let Some(name) = symbol.name() {
+                names.push(strip_symbol_hash(&format!("{name}")).to_owned());
             }
         }
+    }
 
-        log.borrow_mut().push(names);
-    });
-}
-
-thread_local! {
-    static TRACE_WITH_LOG: std::cell::RefCell<Vec<Vec<String>>> =
-        const { std::cell::RefCell::new(vec![]) };
+    log.push(names);
 }
 
 /// Strip the trailing `::h<hex>` hash that rustc appends to symbol names.
@@ -197,25 +190,25 @@ impl<F: Future> Future for TaskDump<F> {
             Poll::Pending => {}
         };
 
+        let mut logs = Vec::new();
+
         // Tracing poll with a noop waker. If the future is at a yield
         // point, trace_leaf fires our callback and returns Pending. We discard
         // the result — this poll is purely for capturing the backtrace.
         let noop = futures::task::noop_waker();
         let mut noop_cx = Context::from_waker(&noop);
-        let logs = this.logs.clone();
-        let trace_poll = trace_with(|| this.f.as_mut().poll(&mut noop_cx), trace_leaf_for_test);
+        let trace_poll = trace_with(
+            || this.f.as_mut().poll(&mut noop_cx),
+            |meta| trace_leaf_for_test(meta, &mut logs),
+        );
         // trace should always produce poll pending
         assert!(
             matches!(trace_poll, Poll::Pending),
             "expected trace to produce Poll::Pending but it was ready"
         );
 
         // Drain any frames captured by trace_leaf_for_test into our log.
-        TRACE_WITH_LOG.with(|tl| {
-            let mut tl = tl.borrow_mut();
-            let mut dest = logs.lock().unwrap();
-            dest.append(&mut tl);
-        });
+        this.logs.lock().unwrap().extend(logs);
         Poll::Pending
     }
 }