address review comments

tokio-rs · Jan 9, 2025 · 49e2cbd · 49e2cbd
1 parent 6e64f61
commit 49e2cbd
Show file tree

Hide file tree

Showing 2 changed files with 45 additions and 28 deletions.
diff --git a/tokio/src/runtime/dump.rs b/tokio/src/runtime/dump.rs
@@ -227,25 +227,21 @@ impl Trace {
     //github.com/
     //github.com/ Example usage:
     //github.com/ ```
-    //github.com/ let mut trace = None;
-    //github.com/
     //github.com/ // some future
-    //github.com/ let mut test_future = Box::pin(async move { tokio::task::yield_now().await; 0 });
+    //github.com/ let mut test_future = std::pin::pin!(async move { tokio::task::yield_now().await; 0 });
     //github.com/
     //github.com/ // trace it once
-    //github.com/ Trace::root(std::future::poll_fn(|cx| {
-    //github.com/     if trace.is_none() {
-    //github.com/         // make sure we only trace once
-    //github.com/         let (res, captured) = Trace::capture(|| test_future.as_mut().poll(cx));
-    //github.com/         trace = Some(captured);
-    //github.com/         res
-    //github.com/     } else {
-    //github.com/         test_future.as_mut().poll(cx)
-    //github.com/     }
-    //github.com/ })).await;@@
+    //github.com/ let (trace, res) = Trace::root(std::future::poll_fn(|cx| {
+    //github.com/     let (res, trace) = Trace::capture(|| test_future.as_mut().poll(cx));
+    //github.com/     Poll::Ready((trace, res))
+    //github.com/ })).await;
+    //github.com/
+    //github.com/ let output = match res {
+    //github.com/    Poll::Ready(output) => output,
+    //github.com/    Poll::Pending => test_future.await,
+    //github.com/ };
     //github.com/
-    //github.com/ // check that there are backtraces
-    //github.com/ assert!(!trace.unwrap().resolve_backtraces().is_empty());
+    //github.com/ println!("{trace}");
     //github.com/ ```
     //github.com/
     //github.com/ ### Nested calls

diff --git a/tokio/tests/task_trace_self.rs b/tokio/tests/task_trace_self.rs
@@ -6,22 +6,21 @@
     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 std::future::Future;
+use std::pin::Pin;
+use std::sync::{Arc, Mutex};
+use std::task::{Context, Poll};
+use std::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: State,
-    logs: Arc<Mutex<Vec<Trace>>>,
-}
+pin_project_lite::pin_project! {
+    pub struct PrettyFuture<F: Future> {
+        #[pin]
+        f: Root<F>,
+        t_last: State,
+        logs: Arc<Mutex<Vec<Trace>>>,
+    }
 }
 
 enum State {
@@ -106,3 +105,25 @@ async fn task_trace_self() {
             .any(|x| format!("{}", x).contains(&s)));
     }
 }
+
+async fn test() {
+    // some future
+    let mut test_future = std::pin::pin!(async move {
+        tokio::task::yield_now().await;
+        0
+    });
+
+    // trace it once
+    let (trace, res) = Trace::root(std::future::poll_fn(|cx| {
+        let (res, trace) = Trace::capture(|| test_future.as_mut().poll(cx));
+        Poll::Ready((trace, res))
+    }))
+    .await;
+
+    let output = match res {
+        Poll::Ready(output) => output,
+        Poll::Pending => test_future.await,
+    };
+
+    println!("{trace}");
+}