feat: USDT probes for tokio task events

.github/workflows/ci.yml

@@ -258,6 +258,9 @@ jobs:
           # only Linux supports io_uring
           - { os: ubuntu-latest, extra_features: io-uring }
           - { os: macos-latest, extra_features: "" }
+          # only Linux and macOS support usdt
+          - { os: ubuntu-latest, extra_features: usdt }
+          - { os: macos-latest, extra_features: usdt }
     steps:
       - uses: actions/checkout@v5
       - name: Install Rust ${{ env.rust_stable }}
@@ -493,9 +496,9 @@ jobs:
       matrix:
         target:
           - name: x86_64-unknown-haiku
-            exclude_features: "taskdump" # taskdump is only available on Linux
+            exclude_features: "taskdump,usdt" # taskdump,usdt are only available on Linux
 #         - name: armv7-sony-vita-newlibeabihf
-#           exclude_features: "process,signal,rt-process-signal,full,taskdump"
+#           exclude_features: "process,signal,rt-process-signal,full,taskdump,usdt"
     steps:
       - uses: actions/checkout@v5
       - name: Install Rust ${{ env.rust_nightly }}
@@ -527,7 +530,7 @@ jobs:
             os: ubuntu-24.04-arm
           - target: aarch64-unknown-linux-gnu
             os: ubuntu-24.04-arm
-            extra_features: "io-uring,taskdump"
+            extra_features: "io-uring,taskdump,usdt"
           - target: aarch64-pc-windows-msvc
             os: windows-11-arm
     steps:
@@ -577,7 +580,7 @@ jobs:
             os: ubuntu-24.04-arm
           - target: aarch64-unknown-linux-gnu
             os: ubuntu-24.04-arm
-            extra_features: "io-uring,taskdump"
+            extra_features: "io-uring,taskdump,usdt"
           - target: aarch64-pc-windows-msvc
             os: windows-11-arm
     steps:
@@ -671,7 +674,8 @@ jobs:
       # https://github.com/tokio-rs/tokio/issues/5373
       - name: Check
         # We use `--skip io-uring` since io-uring crate doesn't provide a binding for the i686 target.
-        run: cargo hack check -Zbuild-std --target target-specs/i686-unknown-linux-gnu.json -p tokio --feature-powerset --skip io-uring --depth 2 --keep-going
+        # We use `--skip usdt` since usdt is not supported on the i686 target.
+        run: cargo hack check -Zbuild-std --target target-specs/i686-unknown-linux-gnu.json -p tokio --feature-powerset --skip io-uring --skip usdt --depth 2 --keep-going
         env:
           RUSTFLAGS: --cfg tokio_unstable -Dwarnings
 
@@ -684,11 +688,11 @@ jobs:
         include:
           - name: ""
             rustflags: ""
-            exclude_features: "io-uring,taskdump"
+            exclude_features: "io-uring,taskdump,usdt"
           - name: "--unstable"
             rustflags: "--cfg tokio_unstable -Dwarnings"
-            exclude_features: "io-uring,taskdump"
-          - name: "--unstable io-uring,taskdump"
+            exclude_features: "io-uring,taskdump,usdt"
+          - name: "--unstable io-uring,taskdump,usdt"
             rustflags: "--cfg tokio_unstable -Dwarnings"
             exclude_features: ""
     steps:
@@ -1091,11 +1095,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_features: "tracing,io-uring,taskdump"
+            extra_features: "tracing,io-uring,taskdump,usdt"
     steps:
       - uses: actions/checkout@v5
       - name: Install Rust ${{ matrix.rust }}
@@ -1112,7 +1116,7 @@ jobs:
       - name: check-external-types
         env:
           RUSTFLAGS: --cfg tokio_unstable -Dwarnings
-          RUSTDOCFLAGS: --cfg tokio_unstable
+          RUSTDOCFLAGS: --cfg tokio_unstable ${{ matrix.extra_flags }}
         run: cargo check-external-types --features $TOKIO_STABLE_FEATURES,${{ matrix.extra_features }}
         working-directory: tokio
 

spellcheck.dic

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

tokio/Cargo.toml

@@ -88,6 +88,8 @@ time = []
 io-uring = ["dep:io-uring", "libc", "mio/os-poll", "mio/os-ext", "dep:slab"]
 # Unstable feature. Requires `--cfg tokio_unstable` to enable.
 taskdump = ["dep:backtrace"]
+# Unstable feature. Requires `--cfg tokio_unstable` to enable.
+usdt = []
 
 [dependencies]
 tokio-macros = { version = "~2.6.0", path = "../tokio-macros", optional = true }

tokio/src/lib.rs

@@ -351,6 +351,7 @@
 //! Some feature flags are only available when specifying the `tokio_unstable` flag:
 //!
 //! - `tracing`: Enables tracing events.
+//! - `usdt`: Enables USDT probes.
 //!
 //! Likewise, some parts of the API are only available with the same flag:
 //!
@@ -499,6 +500,9 @@ compile_error!(
 linux, on `aarch64`, `x86` and `x86_64`."
 );
 
+#[cfg(all(not(tokio_unstable), feature = "usdt"))]
+compile_error!("The `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
@@ -577,6 +581,7 @@ mod trace {
         #[inline(always)]
         #[allow(dead_code)]
         pub(crate) fn trace_leaf(_: &mut std::task::Context<'_>) -> std::task::Poll<()> {
+            crate::util::usdt::trace_root();
             std::task::Poll::Ready(())
         }
     }

tokio/src/macros/cfg.rs

@@ -593,6 +593,35 @@ macro_rules! cfg_not_trace {
     }
 }
 
+macro_rules! cfg_usdt {
+    ($($item:item)*) => {
+        $(
+            #[cfg(all(tokio_unstable, feature = "usdt"))]
+            #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "usdt"))))]
+            $item
+        )*
+    };
+}
+
+macro_rules! cfg_not_usdt {
+    ($($item:item)*) => {
+        $(
+            #[cfg(any(not(tokio_unstable), not(feature = "usdt")))]
+            $item
+        )*
+    }
+}
+
+macro_rules! cfg_trace_or_usdt {
+    ($($item:item)*) => {
+        $(
+            #[cfg(all(tokio_unstable, any(feature = "tracing", feature = "usdt")))]
+            #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, any(feature = "tracing", feature = "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/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", feature = "usdt")))]
+        let id = super::task::Id::next();
+        #[cfg(all(tokio_unstable, feature = "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)
@@ -414,6 +419,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());
         unsafe { 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, feature = "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", feature = "usdt")))]
+        let id = crate::runtime::task::Id::next();
+        #[cfg(all(tokio_unstable, feature = "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, feature = "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", feature = "usdt")))]
+        let id = crate::runtime::task::Id::next();
+        #[cfg(all(tokio_unstable, feature = "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

@@ -24,7 +24,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)]
@@ -371,11 +371,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();

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

@@ -160,6 +160,8 @@ impl Trace {
 // internal implementation details of this crate).
 #[inline(never)]
 pub(crate) fn trace_leaf(cx: &mut task::Context<'_>) -> Poll<()> {
+    crate::util::usdt::trace_root();
+
     // Safety: We don't manipulate the current context's active frame.
     let did_trace = unsafe {
         Context::try_with_current(|context_cell| {