taskdump: instrument JoinHandle and tokio::fs

tokio/src/fs/file.rs

@@ -498,6 +498,7 @@ impl AsyncRead for File {
         cx: &mut Context<'_>,
         dst: &mut ReadBuf<'_>,
     ) -> Poll<io::Result<()>> {
+        crate::runtime::task::trace_leaf();
         let me = self.get_mut();
         let inner = me.inner.get_mut();
 
@@ -594,6 +595,7 @@ impl AsyncSeek for File {
     }
 
     fn poll_complete(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<io::Result<u64>> {
+        crate::runtime::task::trace_leaf();
         let inner = self.inner.get_mut();
 
         loop {
@@ -629,6 +631,7 @@ impl AsyncWrite for File {
         cx: &mut Context<'_>,
         src: &[u8],
     ) -> Poll<io::Result<usize>> {
+        crate::runtime::task::trace_leaf();
         let me = self.get_mut();
         let inner = me.inner.get_mut();
 
@@ -695,11 +698,13 @@ impl AsyncWrite for File {
     }
 
     fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        crate::runtime::task::trace_leaf();
         let inner = self.inner.get_mut();
         inner.poll_flush(cx)
     }
 
     fn poll_shutdown(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), io::Error>> {
+        crate::runtime::task::trace_leaf();
         self.poll_flush(cx)
     }
 }
@@ -774,8 +779,18 @@ impl Inner {
     async fn complete_inflight(&mut self) {
         use crate::future::poll_fn;
 
-        if let Err(e) = poll_fn(|cx| Pin::new(&mut *self).poll_flush(cx)).await {
-            self.last_write_err = Some(e.kind());
+        poll_fn(|cx| self.poll_complete_inflight(cx)).await
+    }
+
+    fn poll_complete_inflight(&mut self, cx: &mut Context<'_>) -> Poll<()> {
+        crate::runtime::task::trace_leaf();
+        match self.poll_flush(cx) {
+            Poll::Ready(Err(e)) => {
+                self.last_write_err = Some(e.kind());
+                Poll::Ready(())
+            }
+            Poll::Ready(Ok(())) => Poll::Ready(()),
+            Poll::Pending => Poll::Pending,
         }
     }
 

tokio/src/macros/cfg.rs

@@ -387,7 +387,15 @@ macro_rules! cfg_taskdump {
                     target_arch = "x86_64"
                 )
             ))]
-            #[cfg_attr(docsrs, doc(cfg(all(
+            $item
+        )*
+    };
+}
+
+macro_rules! cfg_not_taskdump {
+    ($($item:item)*) => {
+        $(
+            #[cfg(not(all(
                 tokio_unstable,
                 tokio_taskdump,
                 feature = "rt",
@@ -397,7 +405,7 @@ macro_rules! cfg_taskdump {
                     target_arch = "x86",
                     target_arch = "x86_64"
                 )
-            ))))]
+            )))]
             $item
         )*
     };

tokio/src/runtime/task/join.rs

@@ -313,6 +313,7 @@ impl<T> Future for JoinHandle<T> {
     type Output = super::Result<T>;
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
+        crate::runtime::task::trace_leaf();
         let mut ret = Poll::Pending;
 
         // Keep track of task budget

tokio/src/runtime/task/mod.rs

@@ -209,6 +209,13 @@ mod waker;
 
 cfg_taskdump! {
     pub(crate) mod trace;
+
+    pub(crate) use self::trace::trace_leaf;
+}
+
+cfg_not_taskdump! {
+    #[inline(always)]
+    pub(crate) fn trace_leaf() {}
 }
 
 use crate::future::Future;

tokio/src/runtime/task/trace/mod.rs

@@ -116,51 +116,50 @@ impl Trace {
     pub(crate) fn root<F>(future: F) -> Root<F> {
         Root { future }
     }
+}
 
-    /// If this is a sub-invocation of [`Trace::capture`], capture a backtrace.
-    ///
-    /// The captured backtrace will be returned by [`Trace::capture`].
-    ///
-    /// Invoking this function does nothing when it is not a sub-invocation
-    /// [`Trace::capture`].
-    // This function is marked `#[inline(never)]` to ensure that it gets a distinct `Frame` in the
-    // backtrace, below which frames should not be included in the backtrace (since they reflect the
-    // internal implementation details of this crate).
-    #[inline(never)]
-    pub(crate) fn leaf() {
-        // Safety: We don't manipulate the current context's active frame.
-        unsafe {
-            Context::with_current(|context_cell| {
-                if let Some(mut collector) = context_cell.collector.take() {
-                    let mut frames = vec![];
-                    let mut above_leaf = false;
-
-                    if let Some(active_frame) = context_cell.active_frame.get() {
-                        let active_frame = active_frame.as_ref();
-
-                        backtrace::trace(|frame| {
-                            let below_root =
-                                !ptr::eq(frame.symbol_address(), active_frame.inner_addr);
-
-                            // only capture frames above `Trace::leaf` and below
-                            // `Trace::root`.
-                            if above_leaf && below_root {
-                                frames.push(frame.to_owned().into());
-                            }
-
-                            if ptr::eq(frame.symbol_address(), Self::leaf as *const _) {
-                                above_leaf = true;
-                            }
-
-                            // only continue unwinding if we're below `Trace::root`
-                            below_root
-                        });
-                    }
-                    collector.backtraces.push(frames);
-                    context_cell.collector.set(Some(collector));
+/// If this is a sub-invocation of [`Trace::capture`], capture a backtrace.
+///
+/// The captured backtrace will be returned by [`Trace::capture`].
+///
+/// Invoking this function does nothing when it is not a sub-invocation
+/// [`Trace::capture`].
+// This function is marked `#[inline(never)]` to ensure that it gets a distinct `Frame` in the
+// backtrace, below which frames should not be included in the backtrace (since they reflect the
+// internal implementation details of this crate).
+#[inline(never)]
+pub(crate) fn trace_leaf() {
+    // Safety: We don't manipulate the current context's active frame.
+    unsafe {
+        Context::with_current(|context_cell| {
+            if let Some(mut collector) = context_cell.collector.take() {
+                let mut frames = vec![];
+                let mut above_leaf = false;
+
+                if let Some(active_frame) = context_cell.active_frame.get() {
+                    let active_frame = active_frame.as_ref();
+
+                    backtrace::trace(|frame| {
+                        let below_root = !ptr::eq(frame.symbol_address(), active_frame.inner_addr);
+
+                        // only capture frames above `Trace::leaf` and below
+                        // `Trace::root`.
+                        if above_leaf && below_root {
+                            frames.push(frame.to_owned().into());
+                        }
+
+                        if ptr::eq(frame.symbol_address(), trace_leaf as *const _) {
+                            above_leaf = true;
+                        }
+
+                        // only continue unwinding if we're below `Trace::root`
+                        below_root
+                    });
                 }
-            });
-        }
+                collector.backtraces.push(frames);
+                context_cell.collector.set(Some(collector));
+            }
+        });
     }
 }
 

tokio/src/task/yield_now.rs

@@ -46,14 +46,7 @@ pub async fn yield_now() {
         type Output = ();
 
         fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<()> {
-            #[cfg(all(
-                tokio_unstable,
-                tokio_taskdump,
-                feature = "rt",
-                target_os = "linux",
-                any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
-            ))]
-            crate::runtime::task::trace::Trace::leaf();
+            crate::runtime::task::trace_leaf();
 
             if self.yielded {
                 return Poll::Ready(());