feat: USDT probes for tokio task events

.github/workflows/ci.yml

@@ -284,6 +284,41 @@ jobs:
           # the unstable cfg to RustDoc
           RUSTDOCFLAGS: --cfg tokio_unstable
 
+  test-unstable-usdt:
+    name: test tokio full --unstable --usdt
+    needs: basics
+    runs-on: ${{ matrix.os }}
+    strategy:
+      matrix:
+        include:
+          - { os: ubuntu-latest }
+          - { os: macos-latest }
+    steps:
+      - uses: actions/checkout@v5
+      - name: Install Rust ${{ env.rust_stable }}
+        uses: dtolnay/rust-toolchain@stable
+        with:
+            toolchain: ${{ env.rust_stable }}
+
+      - name: Install cargo-nextest
+        uses: taiki-e/install-action@v2
+        with:
+          tool: cargo-nextest
+
+      - uses: Swatinem/rust-cache@v2
+      # Run `tokio` with "unstable" cfg flag.
+      - name: test tokio full --cfg unstable --usdt
+        run: |
+          set -euxo pipefail
+          cargo nextest run --features $TOKIO_STABLE_FEATURES,${{ matrix.extra_features }}
+          cargo test --doc --features $TOKIO_STABLE_FEATURES,${{ matrix.extra_features }}
+        working-directory: tokio
+        env:
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_usdt -Dwarnings
+          # in order to run doctests for unstable features, we must also pass
+          # the unstable cfg to RustDoc
+          RUSTDOCFLAGS: --cfg tokio_unstable --cfg tokio_usdt
+
   test-unstable-taskdump:
     name: test tokio full --unstable --taskdump
     needs: basics
@@ -528,6 +563,7 @@ jobs:
           - target: aarch64-unknown-linux-gnu
             os: ubuntu-24.04-arm
             extra_features: "io-uring,taskdump"
+            extra_flags: "--cfg tokio_usdt"
           - target: aarch64-pc-windows-msvc
             os: windows-11-arm
     steps:
@@ -559,8 +595,8 @@ jobs:
           cargo test --doc -p tokio --features $TOKIO_STABLE_FEATURES,${{ matrix.extra_features }} --target ${{ matrix.target }}
         env:
           RUST_TEST_THREADS: 1
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings --cfg tokio_no_tuning_tests
-          RUSTDOCFLAGS: --cfg tokio_unstable -Dwarnings
+          RUSTFLAGS: --cfg tokio_unstable ${{ matrix.extra_flags }} -Dwarnings --cfg tokio_no_tuning_tests
+          RUSTDOCFLAGS: --cfg tokio_unstable ${{ matrix.extra_flags }} -Dwarnings
 
   cross-test-without-parking_lot:
     needs: basics
@@ -578,6 +614,7 @@ jobs:
           - target: aarch64-unknown-linux-gnu
             os: ubuntu-24.04-arm
             extra_features: "io-uring,taskdump"
+            extra_flags: "--cfg tokio_usdt"
           - target: aarch64-pc-windows-msvc
             os: windows-11-arm
     steps:
@@ -613,8 +650,8 @@ jobs:
           cargo test --doc -p tokio --features $TOKIO_STABLE_FEATURES,${{ matrix.extra_features }} --target ${{ matrix.target }}
         env:
           RUST_TEST_THREADS: 1
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings --cfg tokio_no_parking_lot --cfg tokio_no_tuning_tests ${{ matrix.rustflags }}
-          RUSTDOCFLAGS: --cfg tokio_unstable -Dwarnings
+          RUSTFLAGS: --cfg tokio_unstable ${{ matrix.extra_flags }} -Dwarnings --cfg tokio_no_parking_lot --cfg tokio_no_tuning_tests ${{ matrix.rustflags }}
+          RUSTDOCFLAGS: --cfg tokio_unstable ${{ matrix.extra_flags }} -Dwarnings
 
   # See https://github.com/tokio-rs/tokio/issues/5187
   no-atomic-u64-test:
@@ -688,8 +725,8 @@ jobs:
           - name: "--unstable"
             rustflags: "--cfg tokio_unstable -Dwarnings"
             exclude_features: "io-uring,taskdump"
-          - name: "--unstable io-uring,taskdump"
-            rustflags: "--cfg tokio_unstable -Dwarnings"
+          - name: "--unstable io-uring,taskdump,usdt"
+            rustflags: "--cfg tokio_unstable --cfg tokio_usdt -Dwarnings"
             exclude_features: ""
     steps:
       - uses: actions/checkout@v5
@@ -756,7 +793,7 @@ jobs:
           cargo hack check -p tokio-macros -p tokio-stream -p tokio-util -p tokio-test --all-features --ignore-private
       - name: "check --all-features --unstable -Z minimal-versions"
         env:
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_usdt -Dwarnings
         run: |
           # Remove dev-dependencies from Cargo.toml to prevent the next `cargo update`
           # from determining minimal versions based on dev-dependencies.
@@ -802,7 +839,7 @@ jobs:
       - name: "clippy --all --all-features --unstable"
         run: cargo clippy --all --tests --no-deps --all-features
         env:
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings
+          RUSTFLAGS: --cfg tokio_unstable --cfg tokio_usdt -Dwarnings
 
   docs:
     name: docs
@@ -1091,10 +1128,11 @@ jobs:
       matrix:
         include:
           - os: windows-latest
-            # Windows neither supports io-uring nor taskdump.
+            # Windows neither supports io-uring nor taskdump nor usdt.
             extra_features: "tracing"
           - os: ubuntu-latest
             # includes all unstable features.
+            extra_flags: "--cfg tokio_usdt"
             extra_features: "tracing,io-uring,taskdump"
     steps:
       - uses: actions/checkout@v5
@@ -1111,8 +1149,8 @@ jobs:
           tool: [email protected]
       - name: check-external-types
         env:
-          RUSTFLAGS: --cfg tokio_unstable -Dwarnings
-          RUSTDOCFLAGS: --cfg tokio_unstable
+          RUSTFLAGS: --cfg tokio_unstable ${{ matrix.extra_flags }} -Dwarnings
+          RUSTDOCFLAGS: --cfg tokio_unstable ${{ matrix.extra_flags }}
         run: cargo check-external-types --features $TOKIO_STABLE_FEATURES,${{ matrix.extra_features }}
         working-directory: tokio
 

Cargo.toml

@@ -28,5 +28,6 @@ unexpected_cfgs = { level = "warn", check-cfg = [
   'cfg(tokio_no_parking_lot)',
   'cfg(tokio_no_tuning_tests)',
   'cfg(tokio_unstable)',
+  'cfg(tokio_usdt)',
   'cfg(target_os, values("cygwin"))',
 ] }

spellcheck.dic

@@ -1,4 +1,4 @@
-308
+310
 &
 +
 <
@@ -293,6 +293,8 @@ Unsets
 unsynchronized
 untrusted
 uring
+usdt
+USDT
 usecases
 Valgrind
 Varghese

tokio/src/lib.rs

@@ -499,6 +499,9 @@ compile_error!(
 linux, on `aarch64`, `x86` and `x86_64`."
 );
 
+#[cfg(all(not(tokio_unstable), tokio_usdt))]
+compile_error!("The `tokio_usdt` feature requires `--cfg tokio_unstable`.");
+
 // 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

@@ -593,6 +593,35 @@ macro_rules! cfg_not_trace {
     }
 }
 
+macro_rules! cfg_usdt {
+    ($($item:item)*) => {
+        $(
+            #[cfg(all(tokio_unstable, tokio_usdt))]
+            #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, tokio_usdt))))]
+            $item
+        )*
+    };
+}
+
+macro_rules! cfg_not_usdt {
+    ($($item:item)*) => {
+        $(
+            #[cfg(any(not(tokio_unstable), not(tokio_usdt)))]
+            $item
+        )*
+    }
+}
+
+macro_rules! cfg_trace_or_usdt {
+    ($($item:item)*) => {
+        $(
+            #[cfg(all(tokio_unstable, any(feature = "tracing", tokio_usdt)))]
+            #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, any(feature = "tracing", tokio_usdt)))))]
+            $item
+        )*
+    };
+}
+
 macro_rules! cfg_coop {
     ($($item:item)*) => {
         $(

tokio/src/runtime/blocking/mod.rs

@@ -10,7 +10,7 @@ cfg_fs! {
     pub(crate) use pool::spawn_mandatory_blocking;
 }
 
-cfg_trace! {
+cfg_trace_or_usdt! {
     pub(crate) use pool::Mandatory;
 }
 

tokio/src/runtime/builder.rs

@@ -940,6 +940,9 @@ impl Builder {
     /// # }
     /// ```
     pub fn build(&mut self) -> io::Result<Runtime> {
+        #[cfg(all(tokio_unstable, tokio_usdt))]
+        crate::util::usdt::usdt_impl::register_probes()?;
+
         match &self.kind {
             Kind::CurrentThread => self.build_current_thread_runtime(),
             #[cfg(feature = "rt-multi-thread")]

tokio/src/runtime/handle.rs

@@ -19,6 +19,7 @@ use crate::runtime::task::JoinHandle;
 use crate::runtime::BOX_FUTURE_THRESHOLD;
 use crate::util::error::{CONTEXT_MISSING_ERROR, THREAD_LOCAL_DESTROYED_ERROR};
 use crate::util::trace::SpawnMeta;
+use crate::util::usdt;
 
 use std::future::Future;
 use std::marker::PhantomData;
@@ -359,9 +360,12 @@ impl Handle {
         ))]
         let future = super::task::trace::Trace::root(future);
 
+        #[cfg(all(tokio_unstable, any(feature = "tracing", tokio_usdt)))]
+        let id = super::task::Id::next();
+        #[cfg(all(tokio_unstable, tokio_usdt))]
+        let future = usdt::block_on(future, _meta, id);
         #[cfg(all(tokio_unstable, feature = "tracing"))]
-        let future =
-            crate::util::trace::task(future, "block_on", _meta, super::task::Id::next().as_u64());
+        let future = crate::util::trace::task(future, "block_on", _meta, id.as_u64());
 
         // Enter the runtime context. This sets the current driver handles and
         // prevents blocking an existing runtime.
@@ -385,6 +389,7 @@ impl Handle {
             any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
         ))]
         let future = super::task::trace::Trace::root(future);
+        usdt::start_task(usdt::TaskKind::Spawn, meta, id, std::mem::size_of::<F>());
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let future = crate::util::trace::task(future, "task", meta, id.as_u64());
         self.inner.spawn(future, id, meta.spawned_at)
@@ -410,6 +415,12 @@ impl Handle {
             any(target_arch = "aarch64", target_arch = "x86", target_arch = "x86_64")
         ))]
         let future = super::task::trace::Trace::root(future);
+        usdt::start_task(
+            usdt::TaskKind::SpawnLocal,
+            meta,
+            id,
+            std::mem::size_of::<F>(),
+        );
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let future = crate::util::trace::task(future, "task", meta, id.as_u64());
         self.inner.spawn_local(future, id, meta.spawned_at)

tokio/src/runtime/local_runtime/runtime.rs

@@ -6,6 +6,8 @@ use crate::runtime::{context, Builder, EnterGuard, Handle, BOX_FUTURE_THRESHOLD}
 use crate::task::JoinHandle;
 
 use crate::util::trace::SpawnMeta;
+#[cfg(all(tokio_unstable, tokio_usdt))]
+use crate::util::usdt;
 use std::future::Future;
 use std::marker::PhantomData;
 use std::mem;
@@ -239,13 +241,12 @@ impl LocalRuntime {
         ))]
         let future = crate::runtime::task::trace::Trace::root(future);
 
+        #[cfg(all(tokio_unstable, any(feature = "tracing", tokio_usdt)))]
+        let id = crate::runtime::task::Id::next();
+        #[cfg(all(tokio_unstable, tokio_usdt))]
+        let future = usdt::block_on(future, _meta, id);
         #[cfg(all(tokio_unstable, feature = "tracing"))]
-        let future = crate::util::trace::task(
-            future,
-            "block_on",
-            _meta,
-            crate::runtime::task::Id::next().as_u64(),
-        );
+        let future = crate::util::trace::task(future, "block_on", _meta, id.as_u64());
 
         let _enter = self.enter();
 

tokio/src/runtime/mod.rs

@@ -407,7 +407,7 @@ cfg_rt! {
     #[cfg_attr(target_os = "wasi", allow(unused_imports))]
     pub(crate) use blocking::spawn_blocking;
 
-    cfg_trace! {
+    cfg_trace_or_usdt! {
         pub(crate) use blocking::Mandatory;
     }
 

tokio/src/runtime/runtime.rs

@@ -4,6 +4,8 @@ use crate::runtime::scheduler::CurrentThread;
 use crate::runtime::{context, EnterGuard, Handle};
 use crate::task::JoinHandle;
 use crate::util::trace::SpawnMeta;
+#[cfg(all(tokio_unstable, tokio_usdt))]
+use crate::util::usdt;
 
 use std::future::Future;
 use std::mem;
@@ -354,13 +356,12 @@ impl Runtime {
         ))]
         let future = super::task::trace::Trace::root(future);
 
+        #[cfg(all(tokio_unstable, any(feature = "tracing", tokio_usdt)))]
+        let id = crate::runtime::task::Id::next();
+        #[cfg(all(tokio_unstable, tokio_usdt))]
+        let future = usdt::block_on(future, _meta, id);
         #[cfg(all(tokio_unstable, feature = "tracing"))]
-        let future = crate::util::trace::task(
-            future,
-            "block_on",
-            _meta,
-            crate::runtime::task::Id::next().as_u64(),
-        );
+        let future = crate::util::trace::task(future, "block_on", _meta, id.as_u64());
 
         let _enter = self.enter();
 

tokio/src/runtime/task/core.rs

@@ -15,7 +15,7 @@ use crate::runtime::context;
 use crate::runtime::task::raw::{self, Vtable};
 use crate::runtime::task::state::State;
 use crate::runtime::task::{Id, Schedule, TaskHarnessScheduleHooks};
-use crate::util::linked_list;
+use crate::util::{linked_list, usdt};
 
 use std::num::NonZeroU64;
 #[cfg(tokio_unstable)]
@@ -362,11 +362,13 @@ impl<T: Future, S: Schedule> Core<T, S> {
                 let future = unsafe { Pin::new_unchecked(future) };
 
                 let _guard = TaskIdGuard::enter(self.task_id);
+                let _usdt_guard = crate::util::usdt::PollGuard::new(self.task_id);
                 future.poll(&mut cx)
             })
         };
 
         if res.is_ready() {
+            usdt::finish_task(self.task_id, usdt::TerminateKind::Success);
             self.drop_future_or_output();
         }
 

tokio/src/runtime/task/harness.rs

@@ -3,6 +3,7 @@ use crate::runtime::task::core::{Cell, Core, Header, Trailer};
 use crate::runtime::task::state::{Snapshot, State};
 use crate::runtime::task::waker::waker_ref;
 use crate::runtime::task::{Id, JoinError, Notified, RawTask, Schedule, Task};
+use crate::util::usdt;
 
 #[cfg(tokio_unstable)]
 use crate::runtime::TaskMeta;
@@ -498,6 +499,8 @@ enum PollFuture {
 
 /// Cancels the task and store the appropriate error in the stage field.
 fn cancel_task<T: Future, S: Schedule>(core: &Core<T, S>) {
+    usdt::finish_task(core.task_id, usdt::TerminateKind::Cancelled);
+
     // Drop the future from a panic guard.
     let res = panic::catch_unwind(panic::AssertUnwindSafe(|| {
         core.drop_future_or_output();
@@ -526,6 +529,8 @@ fn poll_future<T: Future, S: Schedule>(core: &Core<T, S>, cx: Context<'_>) -> Po
         }
         impl<'a, T: Future, S: Schedule> Drop for Guard<'a, T, S> {
             fn drop(&mut self) {
+                usdt::finish_task(self.core.task_id, usdt::TerminateKind::Panicked);
+
                 // If the future panics on poll, we drop it inside the panic
                 // guard.
                 self.core.drop_future_or_output();