From 171624845432505353bdf25e674a1528c43300ef Mon Sep 17 00:00:00 2001
From: Ariel Ben-Yehuda <arielby@amazon.com>
Date: Wed, 13 Nov 2024 16:24:26 +0000
Subject: [PATCH] make self-tracing public

---
 tokio/src/runtime/dump.rs           | 28 ++++++++-
 tokio/src/runtime/task/trace/mod.rs |  3 +-
 tokio/tests/task_trace_self.rs      | 91 +++++++++++++++++++++++++++++
 3 files changed, 120 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 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: Future>(f: F) -> Root<F> {
+    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, R>(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: F) -> Root<F>
+    where
+        F: Future,
+    {
+        crate::runtime::task::trace::Trace::root(f)
+    }
+}
+
 impl Dump {
     pub(crate) fn new(tasks: Vec<Task>) -> 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<T> {
+    /// Dok: roots a trace
+    pub struct Root<T> {
         #[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..82af2a81905
--- /dev/null
+++ b/tokio/tests/task_trace_self.rs
@@ -0,0 +1,91 @@
+#![allow(unexpected_cfgs)]
+#![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<F: Future> {
+    #[pin]
+    f: Root<F>,
+    t_last: Option<Instant>,
+    logs: Arc<Mutex<Vec<Trace>>>,
+}
+}
+
+impl<F: Future> PrettyFuture<F> {
+    pub fn pretty(f: F, logs: Arc<Mutex<Vec<Trace>>>) -> Self {
+        PrettyFuture {
+            f: Trace::root(f),
+            t_last: None,
+            logs,
+        }
+    }
+}
+
+impl<F: Future> Future for PrettyFuture<F> {
+    type Output = F::Output;
+
+    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<F::Output> {
+        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)));
+    }
+}