taskdump: add trace_with for customized task dumps (#8025)

provided trace function.

tokio/Cargo.toml

@@ -116,6 +116,9 @@ mio = { version = "1.2.0", default-features = false, features = ["os-poll", "os-
 slab = { version = "0.4.9", optional = true }
 backtrace = { version = "0.3.58", optional = true }
 
+[target.'cfg(all(tokio_unstable, target_os = "linux"))'.dev-dependencies]
+backtrace = "0.3.58"
+
 [target.'cfg(unix)'.dependencies]
 libc = { version = "0.2.168", optional = true }
 signal-hook-registry = { version = "1.1.1", optional = true }

tokio/src/runtime/dump.rs

@@ -5,7 +5,7 @@
 use crate::task::Id;
 use std::{fmt, future::Future, path::Path};
 
-pub use crate::runtime::task::trace::Root;
+pub use crate::runtime::task::trace::{trace_with, Root, TraceMeta};
 
 /// A snapshot of a runtime's state.
 ///

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

@@ -10,10 +10,11 @@ use std::ffi::c_void;
 use std::fmt;
 use std::future::Future;
 use std::pin::Pin;
-use std::ptr::{self, NonNull};
+use std::ptr::NonNull;
 use std::task::{self, Poll};
 
 mod symbol;
+mod trace_impl;
 mod tree;
 
 use symbol::Symbol;
@@ -29,8 +30,11 @@ pub(crate) struct Context {
     /// The address of [`Trace::root`] establishes an upper unwinding bound on
     /// the backtraces in `Trace`.
     active_frame: Cell<Option<NonNull<Frame>>>,
-    /// The place to stash backtraces.
-    collector: Cell<Option<Trace>>,
+
+    /// 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)>>,
 }
 
 /// A [`Frame`] in an intrusive, doubly-linked tree of [`Frame`]s.
@@ -39,6 +43,8 @@ struct Frame {
     inner_addr: *const c_void,
 
     /// The parent frame, if any.
+    ///
+    /// Tracking parent allows nested `Root` futures to correctly manage their boundaries
     parent: Option<NonNull<Frame>>,
 }
 
@@ -72,7 +78,7 @@ impl Context {
     pub(crate) const fn new() -> Self {
         Context {
             active_frame: Cell::new(None),
-            collector: Cell::new(None),
+            trace_leaf_fn: Cell::new(None),
         }
     }
 
@@ -96,28 +102,112 @@ impl Context {
         }
     }
 
-    fn with_current_collector<F, R>(f: F) -> R
-    where
-        F: FnOnce(&Cell<Option<Trace>>) -> R,
-    {
-        // SAFETY: This call can only access the collector field, so it cannot
-        // break the trace frame linked list.
+    fn current_frame_addr() -> Option<*const c_void> {
+        // SAFETY: This call does not modify the linked list structure
         unsafe {
-            Self::try_with_current(|context| f(&context.collector)).expect(FAIL_NO_THREAD_LOCAL)
+            Context::try_with_current(|ctx| {
+                ctx.active_frame
+                    .get()
+                    .map(|frame| frame.as_ref().inner_addr)
+            })
+            .flatten()
         }
     }
 
+    fn try_with_current_trace_leaf_fn<F, R>(f: F) -> Option<R>
+    where
+        F: FnOnce(&Cell<Option<fn(&TraceMeta)>>) -> R,
+    {
+        // 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)) }
+    }
+
     /// Produces `true` if the current task is being traced; otherwise false.
     pub(crate) fn is_tracing() -> bool {
-        Self::with_current_collector(|maybe_collector| {
-            let collector = maybe_collector.take();
-            let result = collector.is_some();
-            maybe_collector.set(collector);
-            result
-        })
+        Self::try_with_current_trace_leaf_fn(|maybe_trace_leaf| maybe_trace_leaf.get().is_some())
+            .unwrap_or(false)
     }
 }
 
+/// Metadata passed into the `trace_leaf` callback for [`trace_with`]
+#[non_exhaustive]
+#[derive(Debug)]
+pub struct TraceMeta {
+    /// The root boundary address set by [`Root::poll`] if any.
+    ///
+    /// When using unwinding the stack, this is the address at which
+    /// stack walking should stop. It corresponds to the `Root::poll` function pointer.
+    pub root_addr: Option<*const c_void>,
+
+    /// The address of the internal `trace_leaf` function that triggered this callback.
+    ///
+    /// When capturing a backtrace, use this as the lower bound — frames at or below
+    /// this address are internal implementation details and should be excluded.
+    pub trace_leaf_addr: *const c_void,
+}
+
+/// Runs `f`. If `f` hits a Tokio yield point `trace_leaf` will be invoked.
+///
+/// This allows taking a task dump with caller-provided task dump machinery. If `f` is the poll function of a future
+/// and that future returns `Poll::Pending`, then `trace_leaf` will be invoked. `trace_leaf` can then take a backtrace
+/// to determine exactly where the yield occurred.
+///
+/// `trace_leaf` is a function pointer (`fn`) rather than a closure (`Fn`) because it must be stored
+/// in thread-local state via a `Cell`. Use thread-locals to communicate between the callback and
+/// calling code (see example below).
+///
+/// # Example
+///
+/// ```
+/// use std::future::Future;
+/// use std::task::Poll;
+/// use tokio::runtime::dump::{trace_with, Trace, TraceMeta};
+///
+/// // Thread-local storage for the custom trace function.
+/// std::thread_local! {
+///     static LEAF_COUNT: std::cell::Cell<u32> = const { std::cell::Cell::new(0) };
+/// }
+///
+/// fn my_trace_leaf(_meta: &TraceMeta) {
+///     LEAF_COUNT.with(|c| c.set(c.get() + 1));
+/// }
+///
+/// # async fn example() {
+/// let mut fut = std::pin::pin!(async {
+///     tokio::task::yield_now().await;
+/// });
+///
+/// LEAF_COUNT.with(|c| c.set(0));
+///
+/// Trace::root(std::future::poll_fn(|cx| {
+///     trace_with(|| { let _ = fut.as_mut().poll(cx); }, my_trace_leaf);
+///     Poll::Ready(())
+/// })).await;
+///
+/// let count = LEAF_COUNT.with(|c| c.get());
+/// assert!(count > 0);
+/// # }
+/// ```
+pub fn trace_with<F, R>(f: F, trace_leaf: fn(&TraceMeta)) -> R
+where
+    F: FnOnce() -> R,
+{
+    // store our new trace_leaf function
+    let previous =
+        Context::try_with_current_trace_leaf_fn(|current| current.replace(Some(trace_leaf)));
+
+    // 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));
+        }
+    });
+
+    f()
+}
+
 impl Trace {
     /// Invokes `f`, returning both its result and the collection of backtraces
     /// captured at each sub-invocation of [`trace_leaf`].
@@ -126,16 +216,15 @@ impl Trace {
     where
         F: FnOnce() -> R,
     {
-        let collector = Trace { backtraces: vec![] };
+        trace_impl::capture(f)
+    }
 
-        let previous = Context::with_current_collector(|current| current.replace(Some(collector)));
+    pub(crate) fn empty() -> Self {
+        Self { backtraces: vec![] }
+    }
 
-        let result = f();
-
-        let collector =
-            Context::with_current_collector(|current| current.replace(previous)).unwrap();
-
-        (result, collector)
+    fn push_backtrace(&mut self, bt: Vec<BacktraceFrame>) {
+        self.backtraces.push(bt);
     }
 
     /// The root of a trace.
@@ -160,44 +249,14 @@ impl Trace {
 // internal implementation details of this crate).
 #[inline(never)]
 pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> {
-    // Safety: We don't manipulate the current context's active frame.
-    let did_trace = unsafe {
-        Context::try_with_current(|context_cell| {
-            if let Some(mut collector) = context_cell.collector.take() {
-                let mut frames = vec![];
-                let mut above_leaf = false;
+    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 meta = TraceMeta {
+            root_addr: Context::current_frame_addr(),
+            trace_leaf_addr: trace_leaf as *const c_void,
+        };
+        trace_leaf_fn(&meta);
 
-                if let Some(active_frame) = context_cell.active_frame.get() {
-                    let active_frame = active_frame.as_ref();
-
-                    backtrace::trace(|frame| {
-                        let below_root = !ptr::eq(frame.symbol_address(), active_frame.inner_addr);
-
-                        // only capture frames above `Trace::leaf` and below
-                        // `Trace::root`.
-                        if above_leaf && below_root {
-                            frames.push(frame.to_owned().into());
-                        }
-
-                        if ptr::eq(frame.symbol_address(), trace_leaf as *const _) {
-                            above_leaf = true;
-                        }
-
-                        // only continue unwinding if we're below `Trace::root`
-                        below_root
-                    });
-                }
-                collector.backtraces.push(frames);
-                context_cell.collector.set(Some(collector));
-                true
-            } else {
-                false
-            }
-        })
-        .unwrap_or(false)
-    };
-
-    if did_trace {
         // Use the same logic that `yield_now` uses to send out wakeups after
         // the task yields.
         context::with_scheduler(|scheduler| {

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

@@ -0,0 +1,74 @@
+//! Current `backtrace::trace` + collector based backtrace implementation
+//!
+//! This implementation may eventually be extracted into a separate `tokio-taskdump` crate.
+
+use std::{cell::Cell, 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)));
+
+    // 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, trace_leaf);
+
+    // take the collector before _restore runs
+    let collector = TRACE_CONTEXT.with(|state| state.collector.take()).unwrap();
+
+    (result, collector)
+}
+
+/// 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);
+
+                    if above_leaf && below_root {
+                        frames.push(frame.to_owned().into());
+                    }
+
+                    if ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) {
+                        above_leaf = true;
+                    }
+
+                    below_root
+                });
+            }
+            collector.push_backtrace(frames);
+            state.collector.set(Some(collector));
+        }
+    });
+}

tokio/tests/task_trace_self.rs

@@ -8,11 +8,12 @@
 
 use std::future::Future;
 use std::pin::Pin;
+use std::ptr;
 use std::sync::{Arc, Mutex};
 use std::task::{Context, Poll};
 use std::time::{Duration, Instant};
 
-use tokio::runtime::dump::{Root, Trace};
+use tokio::runtime::dump::{trace_with, Root, Trace, TraceMeta};
 
 pin_project_lite::pin_project! {
     pub struct PrettyFuture<F: Future> {
@@ -105,3 +106,177 @@ async fn task_trace_self() {
             .any(|x| format!("{x}").contains(&s)));
     }
 }
+
+/// 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`.
+#[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;
+
+        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 ptr::eq(frame.symbol_address(), meta.trace_leaf_addr) {
+                    above_leaf = true;
+                }
+
+                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());
+                }
+            }
+        }
+
+        log.borrow_mut().push(names);
+    });
+}
+
+thread_local! {
+    static TRACE_WITH_LOG: std::cell::RefCell<Vec<Vec<String>>> =
+        const { std::cell::RefCell::new(vec![]) };
+}
+
+/// Strip the trailing `::h<hex>` hash that rustc appends to symbol names.
+fn strip_symbol_hash(s: &str) -> &str {
+    // Symbols end with "::h" followed by hex digits. Find the last "::h".
+    if let Some(pos) = s.rfind("::h") {
+        let suffix = &s[pos + 3..];
+        if !suffix.is_empty() && suffix.chars().all(|c| c.is_ascii_hexdigit()) {
+            return &s[..pos];
+        }
+    }
+    s
+}
+
+pin_project_lite::pin_project! {
+    /// A future wrapper that uses `trace_with` to capture a backtrace on every
+    /// poll.
+    /// The captured backtraces are stored in `logs`.
+    pub struct TaskDump<F: Future> {
+        #[pin]
+        f: Root<F>,
+        logs: Arc<Mutex<Vec<Vec<String>>>>,
+    }
+}
+
+impl<F: Future> TaskDump<F> {
+    pub fn new(f: F, logs: Arc<Mutex<Vec<Vec<String>>>>) -> Self {
+        TaskDump {
+            f: Trace::root(f),
+            logs,
+        }
+    }
+}
+
+impl<F: Future> Future for TaskDump<F> {
+    type Output = F::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
+        let mut this = self.project();
+
+        // Poll the future with a real waker. if it returns Ready, exit immediately
+        match this.f.as_mut().poll(cx) {
+            Poll::Ready(result) => return Poll::Ready(result),
+            Poll::Pending => {}
+        };
+
+        // 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);
+        // 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);
+        });
+        Poll::Pending
+    }
+}
+
+#[inline(never)]
+async fn inner_yield_point() {
+    tokio::task::yield_now().await;
+}
+
+/// Validates that `trace_with` (via the `TaskDump` wrapper):
+/// 1. Invokes the trace-leaf callback when the wrapped future is at a yield point.
+/// 2. Produces a backtrace (via `backtrace::trace`) that contains the expected
+///    intermediate symbols between the root and the leaf.
+/// 3. Does not produce spurious callbacks when the future returns `Ready`.
+#[tokio::test]
+async fn trace_with_callback_and_backtrace() {
+    let logs: Arc<Mutex<Vec<Vec<String>>>> = Arc::new(Mutex::new(vec![]));
+
+    let result = TaskDump::new(
+        async {
+            inner_yield_point().await;
+            42
+        },
+        logs.clone(),
+    )
+    .await;
+
+    assert_eq!(result, 42);
+
+    let captured = logs.lock().unwrap();
+
+    assert_eq!(
+        captured.len(),
+        1,
+        "expected exactly 1 traces, got {:#?}",
+        *captured
+    );
+
+    // These symbols should appear in the trace in this exact order (substring match).
+    // trace_leaf itself should NOT appear — it's the boundary frame.
+    let expected_in_order = [
+        "tokio::task::yield_now::yield_now",
+        "core::future::poll_fn::PollFn",
+        "tokio::task::yield_now::yield_now",
+        "task_trace_self::inner_yield_point",
+        "task_trace_self::trace_with_callback_and_backtrace",
+    ];
+    let trace = &captured[0];
+
+    assert_eq!(
+        expected_in_order.len(),
+        trace.len(),
+        "expected {} frames but got {}:\n{trace:#?}",
+        expected_in_order.len(),
+        trace.len()
+    );
+
+    for (expected, actual) in expected_in_order.iter().zip(trace.iter()) {
+        assert!(
+            actual.contains(expected),
+            "expected frame containing {expected:?}, got {actual:?}\nfull trace:\n{trace:#?}"
+        );
+    }
+}