address review comments

Author: Ariel Ben-Yehuda

tokio/src/runtime/dump.rs

@@ -227,25 +227,21 @@ impl Trace {
     ///
     /// Example usage:
     /// ```
-    /// let mut trace = None;
-    ///
     /// // some future
-    /// let mut test_future = Box::pin(async move { tokio::task::yield_now().await; 0 });
-    ///
+    /// let mut test_future = std::pin::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;@@
+    /// 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,
+    /// };
     ///
-    /// // check that there are backtraces
-    /// assert!(!trace.unwrap().resolve_backtraces().is_empty());
+    /// println!("{trace}");
     /// ```
     ///
     /// ### Nested calls

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,21 @@ 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}");
+}