implement usdt probes for task events

fix spellcheck

fix inference errors caused by new serde dependency

split task spawn details out

vendor usdt impl

implement aarch64 support for linux

fix typos

move poll measurements to somewhere that wont monomorphise, remove semaphores from most linux probes

significant overhaul to reduce monomorphisation of probes, also expose task termination reason

Revert "fix inference errors caused by new serde dependency"

This reverts commit e3a7a09.

fix disabled polyfill

fix bad assembler expressions

fix lints

refactor stapsdt handling

include usdt cfg in CI testing

refactor macos provider

document stapsdt format

fix rt feature

revert back to feature flag again

fix CI

fix probe names on linux and inline PollGuard drop

add retained flag to notes, allowing us to inline and monomorphise again

proper fix for stapsdt note and gc-sections interaction

Author: conradludgate

.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 }}
@@ -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

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/builder.rs

@@ -940,6 +940,8 @@ impl Builder {
     /// # }
     /// ```
     pub fn build(&mut self) -> io::Result<Runtime> {
+        crate::util::usdt::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", 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;
     }