diff --git a/tokio/src/runtime/dump.rs b/tokio/src/runtime/dump.rs index aea2381127b..c4fb913c387 100644 --- a/tokio/src/runtime/dump.rs +++ b/tokio/src/runtime/dump.rs @@ -3,7 +3,14 @@ //! See [Handle::dump][crate::runtime::Handle::dump]. use crate::task::Id; -use std::fmt; +use std::{fmt, future::Future}; + +pub use crate::runtime::task::trace::Root; + +/// missing dok +pub fn root(f: F) -> Root { + crate::runtime::task::trace::Trace::root(f) +} /// A snapshot of a runtime's state. /// @@ -38,6 +45,25 @@ pub struct Trace { inner: super::task::trace::Trace, } +impl Trace { + /// document + pub fn capture(f: F) -> (R, Trace) + where + F: FnOnce() -> R, + { + let (res, trace) = super::task::trace::Trace::capture(f); + (res, Trace { inner: trace }) + } + + /// doc doc doc + pub fn root(f: F) -> Root + where + F: Future, + { + crate::runtime::task::trace::Trace::root(f) + } +} + impl Dump { pub(crate) fn new(tasks: Vec) -> Self { Self { diff --git a/tokio/src/runtime/task/trace/mod.rs b/tokio/src/runtime/task/trace/mod.rs index bb411f42d72..8b6db4e6db3 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 { + /// Dok: roots a trace + 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..fdab28328da --- /dev/null +++ b/tokio/tests/task_trace_self.rs @@ -0,0 +1,84 @@ +#![cfg(all(tokio_unstable, tokio_taskdump, feature = "full", not(target_os = "wasi"), not(miri)))] // Wasi doesn't support bind + +use std::{ + future::Future, + mem, + 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: Option, + logs: Arc>>, +} +} + +impl PrettyFuture { + pub fn pretty(f: F, logs: Arc>>) -> Self { + PrettyFuture { + f: Trace::root(f), + t_last: None, + 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 = mem::replace(this.t_last, Some(now)); + if t_last.is_some_and(|t_last| now.duration_since(t_last) > Duration::from_millis(10)) { + let (res, trace) = tokio::runtime::dump::Trace::capture(|| this.f.as_mut().poll(cx)); + this.logs.lock().unwrap().push(trace); + return res; + } + this.f.as_mut().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; + good_line.push(line!() + 1); + tokio::time::sleep(Duration::from_millis(1000)).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))); + } +}