From 6e64f61f49457fe4f73aaa833c153b94ac5e68a7 Mon Sep 17 00:00:00 2001 From: Ariel Ben-Yehuda Date: Wed, 13 Nov 2024 16:24:26 +0000 Subject: [PATCH] make self-tracing public --- tokio/src/runtime/dump.rs | 64 ++++++++++++++++- tokio/src/runtime/task/trace/mod.rs | 3 +- tokio/tests/task_trace_self.rs | 108 ++++++++++++++++++++++++++++ 3 files changed, 173 insertions(+), 2 deletions(-) create mode 100644 tokio/tests/task_trace_self.rs diff --git a/tokio/src/runtime/dump.rs b/tokio/src/runtime/dump.rs index d262f3987cc..6039120fdd8 100644 --- a/tokio/src/runtime/dump.rs +++ b/tokio/src/runtime/dump.rs @@ -3,7 +3,9 @@ //! See [`Handle::dump`][crate::runtime::Handle::dump]. use crate::task::Id; -use std::{fmt, path::Path}; +use std::{fmt, future::Future, path::Path}; + +pub use crate::runtime::task::trace::Root; /// A snapshot of a runtime's state. /// @@ -214,6 +216,66 @@ impl Trace { }) .collect() } + + /// Runs the function `f` in tracing mode, and returns its result along with the resulting [`Trace`]. + /// + /// The passed-in function is normally a poll function. Due to the way tracing is implemented, + /// Tokio leaf futures might, instead of doing their actual work, do the equivalent of a + /// `yield_now` (returning a `Poll::Pending` and scheduling the current context for execution), + /// which means forward progress is only guaranteed if you eventually call your future outside of + /// `capture`. + /// + /// Example usage: + /// ``` + /// let mut trace = None; + /// + /// // some future + /// let mut test_future = Box::pin(async move { tokio::task::yield_now().await; 0 }); + /// + /// // trace it once + /// Trace::root(std::future::poll_fn(|cx| { + /// if trace.is_none() { + /// // make sure we only trace once + /// let (res, captured) = Trace::capture(|| test_future.as_mut().poll(cx)); + /// trace = Some(captured); + /// res + /// } else { + /// test_future.as_mut().poll(cx) + /// } + /// })).await;@@ + /// + /// // check that there are backtraces + /// assert!(!trace.unwrap().resolve_backtraces().is_empty()); + /// ``` + /// + /// ### Nested calls + /// + /// Nested calls to `capture` might return partial traces, but will not do any other undesirable behavior (for + /// example, they will not panic). + /// + /// ### Unstable Implementation Details + /// + /// When run in tracing mode, Tokio leaf futures will return `Poll::Pending` and emit a stack trace + /// that will be captured in the returning [`Trace`]. + pub fn capture(f: F) -> (R, Trace) + where + F: FnOnce() -> R, + { + let (res, trace) = super::task::trace::Trace::capture(f); + (res, Trace { inner: trace }) + } + + /// Create a root for stack traces captured using [`Trace::capture`]. Stack frames above + /// the root will not be captured. + /// + /// Nesting multiple [`Root`] futures is fine. Captures will stop at the first root. Not having + /// a [`Root`] is fine as well, but there is no guarantee on where the capture will stop. + pub fn root(f: F) -> Root + where + F: Future, + { + crate::runtime::task::trace::Trace::root(f) + } } impl Dump { diff --git a/tokio/src/runtime/task/trace/mod.rs b/tokio/src/runtime/task/trace/mod.rs index 16cb3477ffc..71aa3b22657 100644 --- a/tokio/src/runtime/task/trace/mod.rs +++ b/tokio/src/runtime/task/trace/mod.rs @@ -56,7 +56,8 @@ pub(crate) struct Trace { pin_project_lite::pin_project! { #[derive(Debug, Clone)] #[must_use = "futures do nothing unless you `.await` or poll them"] - pub(crate) struct Root { + /// A future wrapper that roots traces (captured with [`Trace::capture`]). + pub struct Root { #[pin] future: T, } diff --git a/tokio/tests/task_trace_self.rs b/tokio/tests/task_trace_self.rs new file mode 100644 index 00000000000..723252d14f7 --- /dev/null +++ b/tokio/tests/task_trace_self.rs @@ -0,0 +1,108 @@ +#![allow(unknown_lints, unexpected_cfgs)] +#![cfg(all( + tokio_unstable, + tokio_taskdump, + target_os = "linux", + any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64") +))] + +use std::{ + future::Future, + pin::Pin, + sync::{Arc, Mutex}, + task::{Context, Poll}, + time::{Duration, Instant}, +}; + +use tokio::runtime::dump::{Root, Trace}; + +pin_project_lite::pin_project! { pub struct PrettyFuture { + #[pin] + f: Root, + t_last: State, + logs: Arc>>, +} +} + +enum State { + NotStarted, + Running { since: Instant }, + Alerted, +} + +impl PrettyFuture { + pub fn pretty(f: F, logs: Arc>>) -> Self { + PrettyFuture { + f: Trace::root(f), + t_last: State::NotStarted, + logs, + } + } +} + +impl Future for PrettyFuture { + type Output = F::Output; + + fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + let mut this = self.project(); + let now = Instant::now(); + let t_last = match this.t_last { + State::Running { since } => Some(*since), + State::NotStarted => { + *this.t_last = State::Running { since: now }; + None + } + State::Alerted => { + // don't double-alert for the same future + None + } + }; + if t_last.is_some_and(|t_last| now.duration_since(t_last) > Duration::from_millis(500)) { + let (res, trace) = tokio::runtime::dump::Trace::capture(|| this.f.as_mut().poll(cx)); + this.logs.lock().unwrap().push(trace); + *this.t_last = State::Alerted; + return res; + } + this.f.poll(cx) + } +} + +#[tokio::test] +async fn task_trace_self() { + let log = Arc::new(Mutex::new(vec![])); + let log2 = Arc::new(Mutex::new(vec![])); + let mut good_line = vec![]; + let mut bad_line = vec![]; + PrettyFuture::pretty( + PrettyFuture::pretty( + async { + bad_line.push(line!() + 1); + tokio::task::yield_now().await; + bad_line.push(line!() + 1); + tokio::time::sleep(Duration::from_millis(1)).await; + for _ in 0..100 { + good_line.push(line!() + 1); + tokio::time::sleep(Duration::from_millis(10)).await; + } + }, + log.clone(), + ), + log2.clone(), + ) + .await; + for line in good_line { + let s = format!("{}:{}:", file!(), line); + assert!(log.lock().unwrap().iter().any(|x| { + eprintln!("{}", x); + format!("{}", x).contains(&s) + })); + } + for line in bad_line { + let s = format!("{}:{}:", file!(), line); + assert!(!log + .lock() + .unwrap() + .iter() + .any(|x| format!("{}", x).contains(&s))); + } +}