taskdump: implement task dumps for multi-thread runtime

This PR implements task dumps on the multi-thread runtime. It
complements tokio-rs#5608, which implemented task dumps on the current-thread
runtime.

Author: jswrenn

examples/dump.rs

@@ -6,7 +6,7 @@
     target_os = "linux",
     any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
 ))]
-#[tokio::main(flavor = "current_thread")]
+#[tokio::main]
 async fn main() {
     use std::hint::black_box;
 
@@ -22,21 +22,29 @@ async fn main() {
 
     #[inline(never)]
     async fn c() {
-        black_box(tokio::task::yield_now()).await
+        loop {
+            tokio::task::yield_now().await;
+        }
     }
 
-    tokio::spawn(a());
-    tokio::spawn(b());
-    tokio::spawn(c());
+    async fn dump() {
+        let handle = tokio::runtime::Handle::current();
+        let dump = handle.dump().await;
 
-    let handle = tokio::runtime::Handle::current();
-    let dump = handle.dump();
-
-    for (i, task) in dump.tasks().iter().enumerate() {
-        let trace = task.trace();
-        println!("task {i} trace:");
-        println!("{trace}");
+        for (i, task) in dump.tasks().iter().enumerate() {
+            let trace = task.trace();
+            println!("task {i} trace:");
+            println!("{trace}\n");
+        }
     }
+
+    tokio::select!(
+        biased;
+        _ = tokio::spawn(a()) => {},
+        _ = tokio::spawn(b()) => {},
+        _ = tokio::spawn(c()) => {},
+        _ = dump() => {},
+    );
 }
 
 #[cfg(not(all(

tokio/src/runtime/handle.rs

@@ -341,15 +341,40 @@ cfg_metrics! {
 cfg_taskdump! {
     impl Handle {
         /// Capture a snapshot of this runtime's state.
-        pub fn dump(&self) -> crate::runtime::Dump {
+        pub async fn dump(&self) -> crate::runtime::Dump {
             match &self.inner {
                 scheduler::Handle::CurrentThread(handle) => handle.dump(),
                 #[cfg(all(feature = "rt-multi-thread", not(tokio_wasi)))]
-                scheduler::Handle::MultiThread(_) =>
-                    unimplemented!("taskdumps are unsupported on the multi-thread runtime"),
+                scheduler::Handle::MultiThread(handle) => {
+                    // perform the trace in a separate thread so that the
+                    // trace itself does not appear in the taskdump.
+                    let handle = handle.clone();
+                    spawn_thread(async {
+                        let handle = handle;
+                        handle.dump().await
+                    }).await
+                },
             }
         }
     }
+
+    cfg_rt_multi_thread! {
+        /// Spawn a new thread and asynchronously await on its result.
+        async fn spawn_thread<F>(f: F) -> <F as Future>::Output
+        where
+            F: Future + Send + 'static,
+            <F as Future>::Output: Send + 'static
+        {
+            let (tx, rx) = crate::sync::oneshot::channel();
+            crate::loom::thread::spawn(|| {
+                let rt = crate::runtime::Builder::new_current_thread().build().unwrap();
+                rt.block_on(async {
+                    let _ = tx.send(f.await);
+                });
+            });
+            rx.await.unwrap()
+        }
+    }
 }
 
 /// Error returned by `try_current` when no Runtime has been started

tokio/src/runtime/scheduler/multi_thread/handle.rs

@@ -95,6 +95,31 @@ cfg_metrics! {
     }
 }
 
+cfg_taskdump! {
+    impl Handle {
+        pub(crate) async fn dump(&self) -> crate::runtime::Dump {
+            let trace_status = &self.shared.trace_status;
+
+            // If a dump is in progress, block.
+            trace_status.start_trace_request(&self).await;
+
+            let result = loop {
+                if let Some(result) = trace_status.take_result() {
+                    break result;
+                } else {
+                    self.notify_all();
+                    trace_status.result_ready.notified().await;
+                }
+            };
+
+            // Allow other queued dumps to proceed.
+            trace_status.end_trace_request(&self).await;
+
+            result
+        }
+    }
+}
+
 impl fmt::Debug for Handle {
     fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
         fmt.debug_struct("multi_thread::Handle { ... }").finish()