taskdump: implement task dumps for current-thread runtime (#5608)

Task dumps are snapshots of runtime state. Taskdumps are collected by
instrumenting Tokio's leaves to conditionally collect backtraces, which
are then coalesced per-task into execution tree traces.

This initial implementation only supports collecting taskdumps from
within the context of a current-thread runtime, and only `yield_now()`
is instrumented.

Author: jswrenn

.github/workflows/ci.yml

@@ -186,10 +186,10 @@ jobs:
     runs-on: ${{ matrix.os }}
     strategy:
       matrix:
-        os:
-          - windows-latest
-          - ubuntu-latest
-          - macos-latest
+        include:
+          - os: windows-latest
+          - os: ubuntu-latest
+          - os: macos-latest
     steps:
       - uses: actions/checkout@v3
       - name: Install Rust ${{ env.rust_stable }}
@@ -206,6 +206,31 @@ jobs:
           # in order to run doctests for unstable features, we must also pass
           # the unstable cfg to RustDoc
           RUSTDOCFLAGS: --cfg tokio_unstable
+    
+  test-unstable-taskdump:
+    name: test tokio full --unstable --taskdump
+    needs: basics
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        include:
+          - os: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: dtolnay/rust-toolchain@master
+        with:
+            toolchain: ${{ env.rust_stable }}
+      - uses: Swatinem/rust-cache@v2
+      # Run `tokio` with "unstable" and "taskdump" cfg flags.
+      - name: test tokio full --cfg unstable --cfg taskdump
+        run: cargo test --all-features
+        working-directory: tokio
+        env:
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_taskdump -Dwarnings
+          # in order to run doctests for unstable features, we must also pass
+          # the unstable cfg to RustDoc
+          RUSTDOCFLAGS: --cfg tokio_unstable --cfg tokio_taskdump
 
   miri:
     name: miri
@@ -293,9 +318,11 @@ jobs:
       matrix:
         include:
           - target: i686-unknown-linux-gnu
+            rustflags: --cfg tokio_taskdump
           - target: arm-unknown-linux-gnueabihf
           - target: armv7-unknown-linux-gnueabihf
           - target: aarch64-unknown-linux-gnu
+            rustflags: --cfg tokio_taskdump
 
           # Run a platform without AtomicU64 and no const Mutex::new
           - target: arm-unknown-linux-gnueabihf
@@ -341,15 +368,15 @@ jobs:
           target: i686-unknown-linux-gnu
       - run: cargo test -Zbuild-std --target target-specs/i686-unknown-linux-gnu.json -p tokio --all-features
         env:
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings --cfg tokio_no_atomic_u64
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_taskdump -Dwarnings --cfg tokio_no_atomic_u64
       # https://github.com/tokio-rs/tokio/pull/5356
       # https://github.com/tokio-rs/tokio/issues/5373
       - run: cargo hack build -p tokio --feature-powerset --depth 2 -Z avoid-dev-deps --keep-going
         env:
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings --cfg tokio_no_atomic_u64 --cfg tokio_no_const_mutex_new
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_taskdump -Dwarnings --cfg tokio_no_atomic_u64 --cfg tokio_no_const_mutex_new
       - run: cargo hack build -p tokio --feature-powerset --depth 2 -Z avoid-dev-deps --keep-going
         env:
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings --cfg tokio_no_atomic_u64
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_taskdump -Dwarnings --cfg tokio_no_atomic_u64
 
   features:
     name: features
@@ -372,6 +399,11 @@ jobs:
         run: cargo hack check --all --feature-powerset --depth 2 -Z avoid-dev-deps --keep-going
         env:
           RUSTFLAGS: --cfg tokio_unstable -Dwarnings
+      # Try with unstable and taskdump feature flags
+      - name: check --feature-powerset --unstable --taskdump
+        run: cargo hack check --all --feature-powerset --depth 2 -Z avoid-dev-deps --keep-going
+        env:
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_taskdump -Dwarnings
 
   minrust:
     name: minrust
@@ -424,7 +456,7 @@ jobs:
           cargo hack check --all-features --ignore-private
       - name: "check --all-features --unstable -Z minimal-versions"
         env:
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_taskdump -Dwarnings
         run: |
           # Remove dev-dependencies from Cargo.toml to prevent the next `cargo update`
           # from determining minimal versions based on dev-dependencies.
@@ -481,8 +513,8 @@ jobs:
       - name: "doc --lib --all-features"
         run: cargo doc --lib --no-deps --all-features --document-private-items
         env:
-          RUSTFLAGS: --cfg docsrs --cfg tokio_unstable
-          RUSTDOCFLAGS: --cfg docsrs --cfg tokio_unstable -Dwarnings
+          RUSTFLAGS: --cfg docsrs --cfg tokio_unstable --cfg tokio_taskdump
+          RUSTDOCFLAGS: --cfg docsrs --cfg tokio_unstable --cfg tokio_taskdump -Dwarnings
 
   loom-compile:
     name: build loom tests

examples/Cargo.toml

@@ -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,50 @@
+//! This example demonstrates tokio's experimental taskdumping functionality.
+
+#[cfg(all(
+    tokio_unstable,
+    tokio_taskdump,
+    target_os = "linux",
+    any(target_arch = "aarch64", target_arch = "x86", 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,
+    tokio_taskdump,
+    target_os = "linux",
+    any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
+)))]
+fn main() {
+    println!("task dumps are not available")
+}

tokio/Cargo.toml

@@ -115,6 +115,11 @@ socket2 = { version = "0.4.9", optional = true, features = [ "all" ] }
 [target.'cfg(tokio_unstable)'.dependencies]
 tracing = { version = "0.1.25", default-features = false, features = ["std"], optional = true } # Not in full
 
+# Currently unstable. The API exposed by these features may be broken at any time.
+# Requires `--cfg tokio_unstable` to enable.
+[target.'cfg(tokio_taskdump)'.dependencies]
+backtrace = { version = "0.3.58" }
+
 [target.'cfg(unix)'.dependencies]
 libc = { version = "0.2.42", optional = true }
 signal-hook-registry = { version = "1.1.1", optional = true }

tokio/src/lib.rs

@@ -487,6 +487,21 @@ compile_error!("Tokio's build script has incorrectly detected wasm.");
 ))]
 compile_error!("Only features sync,macros,io-util,rt,time are supported on wasm.");
 
+#[cfg(all(not(tokio_unstable), tokio_taskdump))]
+compile_error!("The `tokio_taskdump` feature requires `--cfg tokio_unstable`.");
+
+#[cfg(all(
+    tokio_taskdump,
+    not(all(
+        target_os = "linux",
+        any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
+    ))
+))]
+compile_error!(
+    "The `tokio_taskdump` feature is only currently supported on \
+linux, on `aarch64`, `x86` and `x86_64`."
+);
+
 // Includes re-exports used by macros.
 //
 // This module is not intended to be part of the public API. In general, any

tokio/src/macros/cfg.rs

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

tokio/src/runtime/context.rs

@@ -12,6 +12,10 @@ cfg_rt! {
     use std::cell::RefCell;
     use std::marker::PhantomData;
     use std::time::Duration;
+
+    cfg_taskdump! {
+        use crate::runtime::task::trace;
+    }
 }
 
 struct Context {
@@ -45,6 +49,15 @@ 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,
+        tokio_taskdump,
+        feature = "rt",
+        target_os = "linux",
+        any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
+    ))]
+    trace: trace::Context,
 }
 
 tokio_thread_local! {
@@ -75,6 +88,19 @@ tokio_thread_local! {
             rng: FastRand::new(RngSeed::new()),
 
             budget: Cell::new(coop::Budget::unconstrained()),
+
+            #[cfg(all(
+                tokio_unstable,
+                tokio_taskdump,
+                feature = "rt",
+                target_os = "linux",
+                any(
+                    target_arch = "aarch64",
+                    target_arch = "x86",
+                    target_arch = "x86_64"
+                )
+            ))]
+            trace: trace::Context::new(),
         }
     }
 }
@@ -378,6 +404,14 @@ cfg_rt! {
             matches!(self, EnterRuntime::Entered { .. })
         }
     }
+
+    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

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)
+    }
+}