diff --git a/tokio/Cargo.toml b/tokio/Cargo.toml index d6360335a..f0cd9e8b0 100644 --- a/tokio/Cargo.toml +++ b/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 } diff --git a/tokio/src/runtime/dump.rs b/tokio/src/runtime/dump.rs index 57df2ca44..40f364f91 100644 --- a/tokio/src/runtime/dump.rs +++ b/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. /// diff --git a/tokio/src/runtime/task/trace/mod.rs b/tokio/src/runtime/task/trace/mod.rs index 7b018185c..5455e1133 100644 --- a/tokio/src/runtime/task/trace/mod.rs +++ b/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>>, - /// The place to stash backtraces. - collector: Cell>, + + /// The function that is invoked at each leaf future inside of Tokio + /// + /// For example, within tokio::time:sleep, sockets. etc. + trace_leaf_fn: Cell>, } /// 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>, } @@ -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: F) -> R - where - F: FnOnce(&Cell>) -> 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: F) -> Option + where + F: FnOnce(&Cell>) -> 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 = 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: 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) { + 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| { diff --git a/tokio/src/runtime/task/trace/trace_impl.rs b/tokio/src/runtime/task/trace/trace_impl.rs new file mode 100644 index 000000000..5bf907925 --- /dev/null +++ b/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>, +} + +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: 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 = 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)); + } + }); +} diff --git a/tokio/tests/task_trace_self.rs b/tokio/tests/task_trace_self.rs index a2840dc79..1fe605668 100644 --- a/tokio/tests/task_trace_self.rs +++ b/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 { @@ -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 = 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>> = + const { std::cell::RefCell::new(vec![]) }; +} + +/// Strip the trailing `::h` 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 { + #[pin] + f: Root, + logs: Arc>>>, + } +} + +impl TaskDump { + pub fn new(f: F, logs: Arc>>>) -> Self { + TaskDump { + f: Trace::root(f), + logs, + } + } +} + +impl Future for TaskDump { + type Output = F::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + 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>>> = 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:#?}" + ); + } +}