Implement task dumps for current-thread runtime.

examples/Cargo.toml

@@ -7,7 +7,7 @@ edition = "2018"
 # If you copy one of the examples into a new project, you should be using
 # [dependencies] instead, and delete the **path**.
 [dev-dependencies]
-tokio = { version = "1.0.0", path = "../tokio", features = ["full", "tracing"] }
+tokio = { version = "1.0.0", path = "../tokio", features = ["full", "tracing", "taskdump"] }
 tokio-util = { version = "0.7.0", path = "../tokio-util", features = ["full"] }
 tokio-stream = { version = "0.1", path = "../tokio-stream" }
 
@@ -90,3 +90,7 @@ path = "named-pipe-ready.rs"
 [[example]]
 name = "named-pipe-multi-client"
 path = "named-pipe-multi-client.rs"
+
+[[example]]
+name = "dump"
+path = "dump.rs"

examples/dump.rs

@@ -0,0 +1,48 @@
+//! This example demonstrates tokio's experimental taskdumping functionality.
+
+#[cfg(all(
+    tokio_unstable,
+    target_os = "linux",
+    any(target_arch = "aarch64", target_arch = "i686", target_arch = "x86_64")
+))]
+#[tokio::main(flavor = "current_thread")]
+async fn main() {
+    use std::hint::black_box;
+
+    #[inline(never)]
+    async fn a() {
+        black_box(b()).await
+    }
+
+    #[inline(never)]
+    async fn b() {
+        black_box(c()).await
+    }
+
+    #[inline(never)]
+    async fn c() {
+        black_box(tokio::task::yield_now()).await
+    }
+
+    tokio::spawn(a());
+    tokio::spawn(b());
+    tokio::spawn(c());
+
+    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}");
+    }
+}
+
+#[cfg(not(all(
+    tokio_unstable,
+    target_os = "linux",
+    any(target_arch = "aarch64", target_arch = "i686", target_arch = "x86_64")
+)))]
+fn main() {
+    println!("task dumps are not available")
+}

tokio/Cargo.toml

@@ -85,6 +85,7 @@ signal = [
   "windows-sys/Win32_System_Console",
 ]
 sync = []
+taskdump = ["rt", "backtrace"]
 test-util = ["rt", "sync", "time"]
 time = []
 
@@ -114,6 +115,7 @@ socket2 = { version = "0.4.9", optional = true, features = [ "all" ] }
 # Requires `--cfg tokio_unstable` to enable.
 [target.'cfg(tokio_unstable)'.dependencies]
 tracing = { version = "0.1.25", default-features = false, features = ["std"], optional = true } # Not in full
+backtrace = { version = "0.3.58", optional = true }
 
 [target.'cfg(unix)'.dependencies]
 libc = { version = "0.2.42", optional = true }

tokio/src/macros/cfg.rs

@@ -373,6 +373,34 @@ macro_rules! cfg_not_rt_multi_thread {
     }
 }
 
+macro_rules! cfg_taskdump {
+    ($($item:item)*) => {
+        $(
+            #[cfg(all(
+                tokio_unstable,
+                feature = "taskdump",
+                target_os = "linux",
+                any(
+                    target_arch = "aarch64",
+                    target_arch = "i686",
+                    target_arch = "x86_64"
+                )
+            ))]
+            #[cfg_attr(docsrs, doc(cfg(all(
+                tokio_unstable,
+                feature = "taskdump",
+                target_os = "linux",
+                any(
+                    target_arch = "aarch64",
+                    target_arch = "i686",
+                    target_arch = "x86_64"
+                )
+            ))))]
+            $item
+        )*
+    };
+}
+
 macro_rules! cfg_test_util {
     ($($item:item)*) => {
         $(

tokio/src/runtime/context.rs

@@ -14,6 +14,10 @@ cfg_rt! {
     use std::time::Duration;
 }
 
+cfg_taskdump! {
+    use crate::runtime::task::trace;
+}
+
 struct Context {
     /// Uniquely identifies the current thread
     #[cfg(feature = "rt")]
@@ -45,6 +49,14 @@ struct Context {
     /// Tracks the amount of "work" a task may still do before yielding back to
     /// the sheduler
     budget: Cell<coop::Budget>,
+
+    #[cfg(all(
+        tokio_unstable,
+        feature = "taskdump",
+        target_os = "linux",
+        any(target_arch = "aarch64", target_arch = "i686", target_arch = "x86_64")
+    ))]
+    trace: trace::Context,
 }
 
 tokio_thread_local! {
@@ -75,6 +87,18 @@ tokio_thread_local! {
             rng: FastRand::new(RngSeed::new()),
 
             budget: Cell::new(coop::Budget::unconstrained()),
+
+            #[cfg(all(
+                tokio_unstable,
+                feature = "taskdump",
+                target_os = "linux",
+                any(
+                    target_arch = "aarch64",
+                    target_arch = "i686",
+                    target_arch = "x86_64"
+                )
+            ))]
+            trace: trace::Context::new(),
         }
     }
 }
@@ -380,6 +404,14 @@ cfg_rt! {
     }
 }
 
+cfg_taskdump! {
+    /// SAFETY: Callers of this function must ensure that trace frames always
+    /// form a valid linked list.
+    pub(crate) unsafe fn with_trace<R>(f: impl FnOnce(&trace::Context) -> R) -> R {
+        CONTEXT.with(|c| f(&c.trace))
+    }
+}
+
 // Forces the current "entered" state to be cleared while the closure
 // is executed.
 //

tokio/src/runtime/dump.rs

@@ -0,0 +1,66 @@
+//! Snapshots of runtime state.
+
+use std::fmt;
+
+/// A snapshot of a runtime's state.
+#[derive(Debug)]
+pub struct Dump {
+    tasks: Tasks,
+}
+
+/// Snapshots of tasks.
+#[derive(Debug)]
+pub struct Tasks {
+    tasks: Vec<Task>,
+}
+
+/// A snapshot of a task.
+#[derive(Debug)]
+pub struct Task {
+    trace: Trace,
+}
+
+/// An execution trace of a task's last poll.
+#[derive(Debug)]
+pub struct Trace {
+    inner: super::task::trace::Trace,
+}
+
+impl Dump {
+    pub(crate) fn new(tasks: Vec<Task>) -> Self {
+        Self {
+            tasks: Tasks { tasks },
+        }
+    }
+
+    /// Tasks in this snapshot.
+    pub fn tasks(&self) -> &Tasks {
+        &self.tasks
+    }
+}
+
+impl Tasks {
+    /// Iterate over tasks.
+    pub fn iter(&self) -> impl Iterator<Item = &Task> {
+        self.tasks.iter()
+    }
+}
+
+impl Task {
+    pub(crate) fn new(trace: super::task::trace::Trace) -> Self {
+        Self {
+            trace: Trace { inner: trace },
+        }
+    }
+
+    /// A trace of this task's state.
+    pub fn trace(&self) -> &Trace {
+        &self.trace
+    }
+}
+
+impl fmt::Display for Trace {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        self.inner.fmt(f)
+    }
+}

tokio/src/runtime/handle.rs

@@ -252,6 +252,14 @@ impl Handle {
     /// [`tokio::time`]: crate::time
     #[track_caller]
     pub fn block_on<F: Future>(&self, future: F) -> F::Output {
+        #[cfg(all(
+            tokio_unstable,
+            feature = "taskdump",
+            target_os = "linux",
+            any(target_arch = "aarch64", target_arch = "i686", target_arch = "x86_64")
+        ))]
+        let future = super::task::trace::Trace::root(future);
+
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let future =
             crate::util::trace::task(future, "block_on", None, super::task::Id::next().as_u64());
@@ -274,6 +282,13 @@ impl Handle {
         F::Output: Send + 'static,
     {
         let id = crate::runtime::task::Id::next();
+        #[cfg(all(
+            tokio_unstable,
+            feature = "taskdump",
+            target_os = "linux",
+            any(target_arch = "aarch64", target_arch = "i686", target_arch = "x86_64")
+        ))]
+        let future = super::task::trace::Trace::root(future);
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let future = crate::util::trace::task(future, "task", _name, id.as_u64());
         self.inner.spawn(future, id)
@@ -321,6 +336,20 @@ cfg_metrics! {
     }
 }
 
+cfg_taskdump! {
+    impl Handle {
+        /// Capture a snapshot of this runtime's state.
+        pub 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"),
+            }
+        }
+    }
+}
+
 /// Error returned by `try_current` when no Runtime has been started
 #[derive(Debug)]
 pub struct TryCurrentError {

tokio/src/runtime/mod.rs

@@ -233,6 +233,11 @@ cfg_rt! {
     mod defer;
     pub(crate) use defer::Defer;
 
+    cfg_taskdump! {
+        pub mod dump;
+        pub use dump::Dump;
+    }
+
     mod handle;
     pub use handle::{EnterGuard, Handle, TryCurrentError};
 

tokio/src/runtime/runtime.rs

@@ -288,6 +288,14 @@ impl Runtime {
     /// [handle]: fn@Handle::block_on
     #[track_caller]
     pub fn block_on<F: Future>(&self, future: F) -> F::Output {
+        #[cfg(all(
+            tokio_unstable,
+            feature = "taskdump",
+            target_os = "linux",
+            any(target_arch = "aarch64", target_arch = "i686", target_arch = "x86_64")
+        ))]
+        let future = super::task::trace::Trace::root(future);
+
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let future = crate::util::trace::task(
             future,

tokio/src/runtime/scheduler/current_thread.rs

@@ -377,6 +377,51 @@ impl Handle {
         handle
     }
 
+    /// Capture a snapshot of this runtime's state.
+    #[cfg(all(
+        tokio_unstable,
+        feature = "taskdump",
+        target_os = "linux",
+        any(target_arch = "aarch64", target_arch = "i686", target_arch = "x86_64")
+    ))]
+    pub(crate) fn dump(&self) -> crate::runtime::Dump {
+        use crate::runtime::dump;
+        use task::trace::trace_current_thread;
+
+        let mut traces = vec![];
+
+        // todo: how to make this work outside of a runtime context?
+        CURRENT.with(|maybe_context| {
+            // drain the local queue
+            let context = if let Some(context) = maybe_context {
+                context
+            } else {
+                return;
+            };
+            let mut maybe_core = context.core.borrow_mut();
+            let core = if let Some(core) = maybe_core.as_mut() {
+                core
+            } else {
+                return;
+            };
+            let local = &mut core.tasks;
+
+            let mut injection = self.shared.queue.lock();
+            let injection = if let Some(injection) = injection.as_mut() {
+                injection
+            } else {
+                return;
+            };
+
+            traces = trace_current_thread(&self.shared.owned, local, injection)
+                .into_iter()
+                .map(dump::Task::new)
+                .collect();
+        });
+
+        dump::Dump::new(traces)
+    }
+
     fn pop(&self) -> Option<task::Notified<Arc<Handle>>> {
         match self.shared.queue.lock().as_mut() {
             Some(queue) => queue.pop_front(),

tokio/src/runtime/task/list.rs

@@ -172,6 +172,18 @@ impl<S: 'static> OwnedTasks<S> {
     }
 }
 
+cfg_taskdump! {
+    impl<S: 'static> OwnedTasks<S> {
+        /// Locks the tasks, and calls `f` on an iterator over them.
+        pub(crate) fn for_each<F>(&self, f: F)
+        where
+            F: FnMut(&Task<S>)
+        {
+            self.inner.lock().list.for_each(f)
+        }
+    }
+}
+
 impl<S: 'static> LocalOwnedTasks<S> {
     pub(crate) fn new() -> Self {
         Self {