add actions

This change overhauls the entire task hooks system so that users can propagate arbitrary information between task hook invocations and pass context data between the hook "harnesses" for parent and child tasks at time of spawn.

This is intended to be significantly more extensible and long-term maintainable than the current task hooks system, and should ultimately be much easier to stabilize.

.github/workflows/ci.yml

@@ -1,8 +1,8 @@
 on:
   push:
-    branches: ["master", "tokio-*.x"]
+    branches: [ "master", "tokio-*.x" ]
   pull_request:
-    branches: ["master", "tokio-*.x"]
+    branches: [ "master", "tokio-*.x" ]
 
 name: CI
 
@@ -107,7 +107,7 @@ jobs:
       - name: Install Rust ${{ env.rust_stable }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: ${{ env.rust_stable }}
+          toolchain: ${{ env.rust_stable }}
       - name: Install cargo-nextest
         uses: taiki-e/install-action@v2
         with:
@@ -139,7 +139,7 @@ jobs:
       - name: Install Rust ${{ env.rust_stable }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: ${{ env.rust_stable }}
+          toolchain: ${{ env.rust_stable }}
       - name: Install cargo-nextest
         uses: taiki-e/install-action@v2
         with:
@@ -169,7 +169,7 @@ jobs:
       - name: Install Rust ${{ env.rust_nightly }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: ${{ env.rust_nightly }}
+          toolchain: ${{ env.rust_nightly }}
       - name: Install cargo-nextest
         uses: taiki-e/install-action@v2
         with:
@@ -197,7 +197,7 @@ jobs:
       - name: Install Rust ${{ env.rust_stable }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: ${{ env.rust_stable }}
+          toolchain: ${{ env.rust_stable }}
       - name: Install cargo-hack
         uses: taiki-e/install-action@v2
         with:
@@ -237,7 +237,7 @@ jobs:
       - name: Install Rust ${{ env.rust_stable }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: ${{ env.rust_stable }}
+          toolchain: ${{ env.rust_stable }}
 
       - name: Enable parking_lot send_guard feature
         # Inserts the line "plsend = ["parking_lot/send_guard"]" right after [features]
@@ -256,7 +256,7 @@ jobs:
       - name: Install Rust ${{ env.rust_stable }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: 1.82
+          toolchain: 1.82
 
       - name: Install Valgrind
         uses: taiki-e/install-action@valgrind
@@ -295,7 +295,7 @@ jobs:
       - name: Install Rust ${{ env.rust_stable }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: ${{ env.rust_stable }}
+          toolchain: ${{ env.rust_stable }}
 
       - name: Install cargo-nextest
         uses: taiki-e/install-action@v2
@@ -329,7 +329,7 @@ jobs:
       - name: Install Rust ${{ env.rust_stable }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: ${{ env.rust_stable }}
+          toolchain: ${{ env.rust_stable }}
 
       - name: Install cargo-nextest
         uses: taiki-e/install-action@v2
@@ -363,7 +363,7 @@ jobs:
       - name: Install Rust ${{ env.rust_stable }}
         uses: dtolnay/rust-toolchain@stable
         with:
-            toolchain: ${{ env.rust_stable }}
+          toolchain: ${{ env.rust_stable }}
       - name: Install cargo-nextest
         uses: taiki-e/install-action@v2
         with:
@@ -842,10 +842,10 @@ jobs:
           toolchain: ${{ env.rust_stable }}
       - uses: Swatinem/rust-cache@v2
       - name: build --cfg loom
-        run: cargo test --no-run --lib --features full
+        run: cargo test --no-run --lib --release --features full
         working-directory: tokio
         env:
-          RUSTFLAGS: --cfg loom --cfg tokio_unstable -Dwarnings
+          RUSTFLAGS: --cfg loom --cfg tokio_unstable -Dwarnings -Cdebug-assertions
 
   check-readme:
     name: Check README

tokio-util/src/compat.rs

@@ -1,5 +1,110 @@
 //! Compatibility between the `tokio::io` and `futures-io` versions of the
 //! `AsyncRead` and `AsyncWrite` traits.
+//!
+//! ## Bridging Tokio and Futures I/O with `compat()`
+//!
+//! The [`compat()`] function provides a compatibility layer that allows types implementing
+//! [`tokio::io::AsyncRead`] or [`tokio::io::AsyncWrite`] to be used as their
+//! [`futures::io::AsyncRead`] or [`futures::io::AsyncWrite`] counterparts — and vice versa.
+//!
+//! This is especially useful when working with libraries that expect I/O types from one ecosystem
+//! (usually `futures`) but you are using types from the other (usually `tokio`).
+//!
+//! ## Compatibility Overview
+//!
+//! | Inner Type Implements...    | `Compat<T>` Implements...   |
+//! |-----------------------------|-----------------------------|
+//! | [`tokio::io::AsyncRead`]    | [`futures::io::AsyncRead`]  |
+//! | [`futures::io::AsyncRead`]  | [`tokio::io::AsyncRead`]    |
+//! | [`tokio::io::AsyncWrite`]   | [`futures::io::AsyncWrite`] |
+//! | [`futures::io::AsyncWrite`] | [`tokio::io::AsyncWrite`]   |
+//!
+//! ## Feature Flag
+//!
+//! This functionality is available through the `compat` feature flag:
+//!
+//! ```toml
+//! tokio-util = { version = "...", features = ["compat"] }
+//! ```
+//!
+//! ## Example 1: Tokio -> Futures (`AsyncRead`)
+//!
+//! This example demonstrates sending data over a [`tokio::net::TcpStream`] and using
+//! [`futures::io::AsyncReadExt::read`] from the `futures` crate to read it after adapting the
+//! stream via [`compat()`].
+//!
+//! ```no_run
+//! use tokio::net::{TcpListener, TcpStream};
+//! use tokio::io::AsyncWriteExt;
+//! use tokio_util::compat::TokioAsyncReadCompatExt;
+//! use futures::io::AsyncReadExt;
+//!
+//! #[tokio::main]
+//! async fn main() -> std::io::Result<()> {
+//!     let listener = TcpListener::bind("127.0.0.1:8081").await?;
+//!
+//!     tokio::spawn(async {
+//!         let mut client = TcpStream::connect("127.0.0.1:8081").await.unwrap();
+//!         client.write_all(b"Hello World").await.unwrap();
+//!     });
+//!
+//!     let (stream, _) = listener.accept().await?;
+//!
+//!     // Adapt `tokio::TcpStream` to be used with `futures::io::AsyncReadExt`
+//!     let mut compat_stream = stream.compat();
+//!     let mut buffer = [0; 20];
+//!     let n = compat_stream.read(&mut buffer).await?;
+//!     println!("Received: {}", String::from_utf8_lossy(&buffer[..n]));
+//!
+//!     Ok(())
+//! }
+//! ```
+//!
+//! ## Example 2: Futures -> Tokio (`AsyncRead`)
+//!
+//! The reverse is also possible: you can take a [`futures::io::AsyncRead`] (e.g. a cursor) and
+//! adapt it to be used with [`tokio::io::AsyncReadExt::read_to_end`]
+//!
+//! ```
+//! use futures::io::Cursor;
+//! use tokio_util::compat::FuturesAsyncReadCompatExt;
+//! use tokio::io::AsyncReadExt;
+//!
+//! fn main() {
+//!     let future = async {
+//!         let reader = Cursor::new(b"Hello from futures");
+//!         let mut compat_reader = reader.compat();
+//!         let mut buf = Vec::new();
+//!         compat_reader.read_to_end(&mut buf).await.unwrap();
+//!         assert_eq!(&buf, b"Hello from futures");
+//!     };
+//!
+//!     // Run the future inside a Tokio runtime
+//!     tokio::runtime::Runtime::new().unwrap().block_on(future);
+//! }
+//! ```
+//!
+//! ## Common Use Cases
+//!
+//! - Using `tokio` sockets with `async-tungstenite`, `async-compression`, or `futures-rs`-based
+//!   libraries.
+//! - Bridging I/O interfaces between mixed-ecosystem libraries.
+//! - Avoiding rewrites or duplication of I/O code in async environments.
+//!
+//! ## See Also
+//!
+//! - [`Compat`] type
+//! - [`TokioAsyncReadCompatExt`]
+//! - [`FuturesAsyncReadCompatExt`]
+//! - [`tokio::io`]
+//! - [`futures::io`]
+//!
+//! [`futures::io`]: https://docs.rs/futures/latest/futures/io/
+//! [`futures::io::AsyncRead`]: https://docs.rs/futures/latest/futures/io/trait.AsyncRead.html
+//! [`futures::io::AsyncWrite`]: https://docs.rs/futures/latest/futures/io/trait.AsyncWrite.html
+//! [`futures::io::AsyncReadExt::read`]: https://docs.rs/futures/latest/futures/io/trait.AsyncReadExt.html#method.read
+//! [`compat()`]: TokioAsyncReadCompatExt::compat
+
 use pin_project_lite::pin_project;
 use std::io;
 use std::pin::Pin;

tokio-util/src/lib.rs

@@ -16,8 +16,6 @@
 //! This crate is not versioned in lockstep with the core
 //! [`tokio`] crate. However, `tokio-util` _will_ respect Rust's
 //! semantic versioning policy, especially with regard to breaking changes.
-//!
-//! [`tokio`]: https://docs.rs/tokio
 
 #[macro_use]
 mod cfg;

tokio/src/lib.rs

@@ -351,10 +351,7 @@
 //! - [`task::Builder`]
 //! - Some methods on [`task::JoinSet`]
 //! - [`runtime::RuntimeMetrics`]
-//! - [`runtime::Builder::on_task_spawn`]
-//! - [`runtime::Builder::on_task_terminate`]
 //! - [`runtime::Builder::unhandled_panic`]
-//! - [`runtime::TaskMeta`]
 //!
 //! This flag enables **unstable** features. The public API of these features
 //! may break in 1.x releases. To enable these features, the `--cfg

tokio/src/runtime/blocking/pool.rs

@@ -375,10 +375,15 @@ impl Spawner {
         F: FnOnce() -> R + Send + 'static,
         R: Send + 'static,
     {
+        // let parent = with_c
         let id = task::Id::next();
         let fut =
             blocking_task::<F, BlockingTask<F>>(BlockingTask::new(func), spawn_meta, id.as_u64());
 
+        #[cfg(tokio_unstable)]
+        let (task, handle) = task::unowned(fut, BlockingSchedule::new(rt), id, None);
+
+        #[cfg(not(tokio_unstable))]
         let (task, handle) = task::unowned(fut, BlockingSchedule::new(rt), id);
 
         let spawned = self.spawn_task(Task::new(task, is_mandatory), rt);

tokio/src/runtime/blocking/schedule.rs

@@ -1,7 +1,9 @@
 #[cfg(feature = "test-util")]
 use crate::runtime::scheduler;
-use crate::runtime::task::{self, Task, TaskHarnessScheduleHooks};
+use crate::runtime::task::{self, Task};
 use crate::runtime::Handle;
+#[cfg(tokio_unstable)]
+use crate::runtime::{OptionalTaskHooksFactory, OptionalTaskHooksFactoryRef};
 
 /// `task::Schedule` implementation that does nothing (except some bookkeeping
 /// in test-util builds). This is unique to the blocking scheduler as tasks
@@ -12,7 +14,8 @@ use crate::runtime::Handle;
 pub(crate) struct BlockingSchedule {
     #[cfg(feature = "test-util")]
     handle: Handle,
-    hooks: TaskHarnessScheduleHooks,
+    #[cfg(tokio_unstable)]
+    hooks_factory: OptionalTaskHooksFactory,
 }
 
 impl BlockingSchedule {
@@ -31,9 +34,8 @@ impl BlockingSchedule {
         BlockingSchedule {
             #[cfg(feature = "test-util")]
             handle: handle.clone(),
-            hooks: TaskHarnessScheduleHooks {
-                task_terminate_callback: handle.inner.hooks().task_terminate_callback.clone(),
-            },
+            #[cfg(tokio_unstable)]
+            hooks_factory: handle.inner.hooks_factory(),
         }
     }
 }
@@ -58,9 +60,13 @@ impl task::Schedule for BlockingSchedule {
         unreachable!();
     }
 
-    fn hooks(&self) -> TaskHarnessScheduleHooks {
-        TaskHarnessScheduleHooks {
-            task_terminate_callback: self.hooks.task_terminate_callback.clone(),
-        }
+    #[cfg(tokio_unstable)]
+    fn hooks_factory(&self) -> OptionalTaskHooksFactory {
+        self.hooks_factory.clone()
+    }
+
+    #[cfg(tokio_unstable)]
+    fn hooks_factory_ref(&self) -> OptionalTaskHooksFactoryRef<'_> {
+        self.hooks_factory.as_ref().map(AsRef::as_ref)
     }
 }

tokio/src/runtime/builder.rs

@@ -1,15 +1,19 @@
 #![cfg_attr(loom, allow(unused_imports))]
 
-use crate::runtime::handle::Handle;
-use crate::runtime::{blocking, driver, Callback, HistogramBuilder, Runtime, TaskCallback};
-#[cfg(tokio_unstable)]
-use crate::runtime::{metrics::HistogramConfiguration, LocalOptions, LocalRuntime, TaskMeta};
-use crate::util::rand::{RngSeed, RngSeedGenerator};
-
 use crate::runtime::blocking::BlockingPool;
+use crate::runtime::handle::Handle;
 use crate::runtime::scheduler::CurrentThread;
+use crate::runtime::{blocking, driver, Callback, HistogramBuilder, Runtime};
+#[cfg(tokio_unstable)]
+use crate::runtime::{
+    metrics::HistogramConfiguration, LocalOptions, LocalRuntime, OptionalTaskHooksFactory,
+    TaskHookHarnessFactory,
+};
+use crate::util::rand::{RngSeed, RngSeedGenerator};
 use std::fmt;
 use std::io;
+#[cfg(tokio_unstable)]
+use std::sync::Arc;
 use std::thread::ThreadId;
 use std::time::Duration;
 
@@ -85,19 +89,8 @@ pub struct Builder {
     /// To run after each thread is unparked.
     pub(super) after_unpark: Option<Callback>,
 
-    /// To run before each task is spawned.
-    pub(super) before_spawn: Option<TaskCallback>,
-
-    /// To run before each poll
     #[cfg(tokio_unstable)]
-    pub(super) before_poll: Option<TaskCallback>,
-
-    /// To run after each poll
-    #[cfg(tokio_unstable)]
-    pub(super) after_poll: Option<TaskCallback>,
-
-    /// To run after each task is terminated.
-    pub(super) after_termination: Option<TaskCallback>,
+    pub(super) task_hook_harness_factory: OptionalTaskHooksFactory,
 
     /// Customizable keep alive timeout for `BlockingPool`
     pub(super) keep_alive: Option<Duration>,
@@ -287,13 +280,8 @@ impl Builder {
             before_park: None,
             after_unpark: None,
 
-            before_spawn: None,
-            after_termination: None,
-
             #[cfg(tokio_unstable)]
-            before_poll: None,
-            #[cfg(tokio_unstable)]
-            after_poll: None,
+            task_hook_harness_factory: None,
 
             keep_alive: None,
 
@@ -685,188 +673,19 @@ impl Builder {
         self
     }
 
-    /// Executes function `f` just before a task is spawned.
+    /// Factory method for producing "fallback" task hook harnesses.
     ///
-    /// `f` is called within the Tokio context, so functions like
-    /// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being
-    /// invoked immediately.
-    ///
-    /// This can be used for bookkeeping or monitoring purposes.
-    ///
-    /// Note: There can only be one spawn callback for a runtime; calling this function more
-    /// than once replaces the last callback defined, rather than adding to it.
-    ///
-    /// This *does not* support [`LocalSet`](crate::task::LocalSet) at this time.
-    ///
-    /// **Note**: This is an [unstable API][unstable]. The public API of this type
-    /// may break in 1.x releases. See [the documentation on unstable
-    /// features][unstable] for details.
-    ///
-    /// [unstable]: crate#unstable-features
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// # use tokio::runtime;
-    /// # pub fn main() {
-    /// let runtime = runtime::Builder::new_current_thread()
-    ///     .on_task_spawn(|_| {
-    ///         println!("spawning task");
-    ///     })
-    ///     .build()
-    ///     .unwrap();
-    ///
-    /// runtime.block_on(async {
-    ///     tokio::task::spawn(std::future::ready(()));
-    ///
-    ///     for _ in 0..64 {
-    ///         tokio::task::yield_now().await;
-    ///     }
-    /// })
-    /// # }
-    /// ```
+    /// The order of operations for assigning the hook harness for a task are as follows:
+    /// 1. [`crate::task::spawn_with_hooks`], if used.
+    /// 2. [`crate::runtime::task_hooks::TaskHookHarnessFactory`], if it returns something other than [Option::None].
+    /// 3. This function.
     #[cfg(all(not(loom), tokio_unstable))]
     #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
-    pub fn on_task_spawn<F>(&mut self, f: F) -> &mut Self
+    pub fn hook_harness_factory<T>(&mut self, hooks: T) -> &mut Self
     where
-        F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
+        T: TaskHookHarnessFactory + Send + Sync + 'static,
     {
-        self.before_spawn = Some(std::sync::Arc::new(f));
-        self
-    }
-
-    /// Executes function `f` just before a task is polled
-    ///
-    /// `f` is called within the Tokio context, so functions like
-    /// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being
-    /// invoked immediately.
-    ///
-    /// **Note**: This is an [unstable API][unstable]. The public API of this type
-    /// may break in 1.x releases. See [the documentation on unstable
-    /// features][unstable] for details.
-    ///
-    /// [unstable]: crate#unstable-features
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// # use std::sync::{atomic::AtomicUsize, Arc};
-    /// # use tokio::task::yield_now;
-    /// # pub fn main() {
-    /// let poll_start_counter = Arc::new(AtomicUsize::new(0));
-    /// let poll_start = poll_start_counter.clone();
-    /// let rt = tokio::runtime::Builder::new_multi_thread()
-    ///     .enable_all()
-    ///     .on_before_task_poll(move |meta| {
-    ///         println!("task {} is about to be polled", meta.id())
-    ///     })
-    ///     .build()
-    ///     .unwrap();
-    /// let task = rt.spawn(async {
-    ///     yield_now().await;
-    /// });
-    /// let _ = rt.block_on(task);
-    ///
-    /// # }
-    /// ```
-    #[cfg(tokio_unstable)]
-    pub fn on_before_task_poll<F>(&mut self, f: F) -> &mut Self
-    where
-        F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
-    {
-        self.before_poll = Some(std::sync::Arc::new(f));
-        self
-    }
-
-    /// Executes function `f` just after a task is polled
-    ///
-    /// `f` is called within the Tokio context, so functions like
-    /// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback being
-    /// invoked immediately.
-    ///
-    /// **Note**: This is an [unstable API][unstable]. The public API of this type
-    /// may break in 1.x releases. See [the documentation on unstable
-    /// features][unstable] for details.
-    ///
-    /// [unstable]: crate#unstable-features
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// # use std::sync::{atomic::AtomicUsize, Arc};
-    /// # use tokio::task::yield_now;
-    /// # pub fn main() {
-    /// let poll_stop_counter = Arc::new(AtomicUsize::new(0));
-    /// let poll_stop = poll_stop_counter.clone();
-    /// let rt = tokio::runtime::Builder::new_multi_thread()
-    ///     .enable_all()
-    ///     .on_after_task_poll(move |meta| {
-    ///         println!("task {} completed polling", meta.id());
-    ///     })
-    ///     .build()
-    ///     .unwrap();
-    /// let task = rt.spawn(async {
-    ///     yield_now().await;
-    /// });
-    /// let _ = rt.block_on(task);
-    ///
-    /// # }
-    /// ```
-    #[cfg(tokio_unstable)]
-    pub fn on_after_task_poll<F>(&mut self, f: F) -> &mut Self
-    where
-        F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
-    {
-        self.after_poll = Some(std::sync::Arc::new(f));
-        self
-    }
-
-    /// Executes function `f` just after a task is terminated.
-    ///
-    /// `f` is called within the Tokio context, so functions like
-    /// [`tokio::spawn`](crate::spawn) can be called.
-    ///
-    /// This can be used for bookkeeping or monitoring purposes.
-    ///
-    /// Note: There can only be one task termination callback for a runtime; calling this
-    /// function more than once replaces the last callback defined, rather than adding to it.
-    ///
-    /// This *does not* support [`LocalSet`](crate::task::LocalSet) at this time.
-    ///
-    /// **Note**: This is an [unstable API][unstable]. The public API of this type
-    /// may break in 1.x releases. See [the documentation on unstable
-    /// features][unstable] for details.
-    ///
-    /// [unstable]: crate#unstable-features
-    ///
-    /// # Examples
-    ///
-    /// ```
-    /// # use tokio::runtime;
-    /// # pub fn main() {
-    /// let runtime = runtime::Builder::new_current_thread()
-    ///     .on_task_terminate(|_| {
-    ///         println!("killing task");
-    ///     })
-    ///     .build()
-    ///     .unwrap();
-    ///
-    /// runtime.block_on(async {
-    ///     tokio::task::spawn(std::future::ready(()));
-    ///
-    ///     for _ in 0..64 {
-    ///         tokio::task::yield_now().await;
-    ///     }
-    /// })
-    /// # }
-    /// ```
-    #[cfg(all(not(loom), tokio_unstable))]
-    #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))]
-    pub fn on_task_terminate<F>(&mut self, f: F) -> &mut Self
-    where
-        F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
-    {
-        self.after_termination = Some(std::sync::Arc::new(f));
+        self.task_hook_harness_factory = Some(Arc::new(hooks));
         self
     }
 
@@ -1475,12 +1294,8 @@ impl Builder {
             Config {
                 before_park: self.before_park.clone(),
                 after_unpark: self.after_unpark.clone(),
-                before_spawn: self.before_spawn.clone(),
                 #[cfg(tokio_unstable)]
-                before_poll: self.before_poll.clone(),
-                #[cfg(tokio_unstable)]
-                after_poll: self.after_poll.clone(),
-                after_termination: self.after_termination.clone(),
+                task_hook_factory: self.task_hook_harness_factory.clone(),
                 global_queue_interval: self.global_queue_interval,
                 event_interval: self.event_interval,
                 #[cfg(tokio_unstable)]
@@ -1628,12 +1443,8 @@ cfg_rt_multi_thread! {
                 Config {
                     before_park: self.before_park.clone(),
                     after_unpark: self.after_unpark.clone(),
-                    before_spawn: self.before_spawn.clone(),
                     #[cfg(tokio_unstable)]
-                    before_poll: self.before_poll.clone(),
-                    #[cfg(tokio_unstable)]
-                    after_poll: self.after_poll.clone(),
-                    after_termination: self.after_termination.clone(),
+                    task_hook_factory: self.task_hook_harness_factory.clone(),
                     global_queue_interval: self.global_queue_interval,
                     event_interval: self.event_interval,
                     #[cfg(tokio_unstable)]

tokio/src/runtime/config.rs

@@ -2,7 +2,10 @@
     any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"),
     allow(dead_code)
 )]
-use crate::runtime::{Callback, TaskCallback};
+
+use crate::runtime::Callback;
+#[cfg(tokio_unstable)]
+use crate::runtime::OptionalTaskHooksFactory;
 use crate::util::RngSeedGenerator;
 
 pub(crate) struct Config {
@@ -18,19 +21,9 @@ pub(crate) struct Config {
     /// Callback for a worker unparking itself
     pub(crate) after_unpark: Option<Callback>,
 
-    /// To run before each task is spawned.
-    pub(crate) before_spawn: Option<TaskCallback>,
-
-    /// To run after each task is terminated.
-    pub(crate) after_termination: Option<TaskCallback>,
-
-    /// To run before each poll
+    /// Called on task spawn to generate the attached task hook harness.
     #[cfg(tokio_unstable)]
-    pub(crate) before_poll: Option<TaskCallback>,
-
-    /// To run after each poll
-    #[cfg(tokio_unstable)]
-    pub(crate) after_poll: Option<TaskCallback>,
+    pub(crate) task_hook_factory: OptionalTaskHooksFactory,
 
     /// The multi-threaded scheduler includes a per-worker LIFO slot used to
     /// store the last scheduled task. This can improve certain usage patterns,

tokio/src/runtime/context.rs

@@ -1,10 +1,14 @@
+#[cfg(all(feature = "rt", tokio_unstable))]
+use crate::loom::cell::UnsafeCell;
 use crate::loom::thread::AccessError;
+#[cfg(all(feature = "rt", tokio_unstable))]
+use crate::runtime::{OptionalTaskHooksMut, OptionalTaskHooksWeak, TaskHookHarness};
 use crate::task::coop;
-
-use std::cell::Cell;
-
 #[cfg(any(feature = "rt", feature = "macros", feature = "time"))]
 use crate::util::rand::FastRand;
+use std::cell::Cell;
+#[cfg(all(feature = "rt", tokio_unstable))]
+use std::ptr::NonNull;
 
 cfg_rt! {
     mod blocking;
@@ -49,6 +53,10 @@ struct Context {
     #[cfg(feature = "rt")]
     current_task_id: Cell<Option<Id>>,
 
+    /// Tracks the current set of task hooks,
+    #[cfg(all(feature = "rt", tokio_unstable))]
+    current_task_hooks: OptionalTaskHooksWeak,
+
     /// Tracks if the current thread is currently driving a runtime.
     /// Note, that if this is set to "entered", the current scheduler
     /// handle may not reference the runtime currently executing. This
@@ -92,6 +100,9 @@ tokio_thread_local! {
             #[cfg(feature = "rt")]
             current_task_id: Cell::new(None),
 
+            #[cfg(all(feature = "rt", tokio_unstable))]
+            current_task_hooks: UnsafeCell::new(None),
+
             // Tracks if the current thread is currently driving a runtime.
             // Note, that if this is set to "entered", the current scheduler
             // handle may not reference the runtime currently executing. This
@@ -139,6 +150,16 @@ pub(crate) fn budget<R>(f: impl FnOnce(&Cell<coop::Budget>) -> R) -> Result<R, A
     CONTEXT.try_with(|ctx| f(&ctx.budget))
 }
 
+#[cfg(all(feature = "rt", tokio_unstable))]
+pub(crate) struct SetTaskHooksGuard;
+
+#[cfg(all(feature = "rt", tokio_unstable))]
+impl Drop for SetTaskHooksGuard {
+    fn drop(&mut self) {
+        let _ = clear_task_hooks();
+    }
+}
+
 cfg_rt! {
     use crate::runtime::ThreadId;
 
@@ -163,6 +184,47 @@ cfg_rt! {
         CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None)
     }
 
+    #[track_caller]
+    #[cfg(tokio_unstable)]
+    pub(super) fn set_task_hooks(hooks: Option<NonNull<dyn TaskHookHarness + Send + Sync + 'static>>) -> Result<SetTaskHooksGuard, AccessError> {
+        CONTEXT.try_with(|ctx| {
+            ctx.current_task_hooks.with_mut(|x| {
+                unsafe {
+                    *x = hooks;
+                }
+            })
+        })?;
+
+        Ok(SetTaskHooksGuard)
+    }
+
+    #[track_caller]
+    #[cfg(tokio_unstable)]
+    pub(super) fn clear_task_hooks() -> Result<(), AccessError> {
+        CONTEXT.try_with(|ctx| {
+            ctx.current_task_hooks.with_mut(|x| {
+                unsafe {
+                    *x = None;
+                }
+            })
+        })?;
+
+        Ok(())
+    }
+
+    #[track_caller]
+    #[cfg(tokio_unstable)]
+    pub(super) fn with_task_hooks<R>(f: impl FnOnce(OptionalTaskHooksMut<'_>) -> R) -> Result<R, AccessError> {
+        CONTEXT.try_with(|ctx| {
+            ctx.current_task_hooks.with_mut(|ptr| {
+                let hooks = unsafe { &mut *ptr };
+                unsafe {
+                    f(hooks.as_mut().map(|x| x.as_mut()))
+                }
+            })
+        })
+    }
+
     #[track_caller]
     pub(crate) fn defer(waker: &Waker) {
         with_scheduler(|maybe_scheduler| {

tokio/src/runtime/handle.rs

@@ -1,5 +1,5 @@
 #[cfg(tokio_unstable)]
-use crate::runtime;
+use crate::runtime::{self, OptionalTaskHooks};
 use crate::runtime::{context, scheduler, RuntimeFlavor, RuntimeMetrics};
 
 /// Handle to the runtime.
@@ -191,6 +191,13 @@ impl Handle {
         F::Output: Send + 'static,
     {
         let fut_size = mem::size_of::<F>();
+        #[cfg(tokio_unstable)]
+        return if fut_size > BOX_FUTURE_THRESHOLD {
+            self.spawn_named(Box::pin(future), SpawnMeta::new_unnamed(fut_size), None)
+        } else {
+            self.spawn_named(future, SpawnMeta::new_unnamed(fut_size), None)
+        };
+        #[cfg(not(tokio_unstable))]
         if fut_size > BOX_FUTURE_THRESHOLD {
             self.spawn_named(Box::pin(future), SpawnMeta::new_unnamed(fut_size))
         } else {
@@ -329,7 +336,12 @@ impl Handle {
     }
 
     #[track_caller]
-    pub(crate) fn spawn_named<F>(&self, future: F, _meta: SpawnMeta<'_>) -> JoinHandle<F::Output>
+    pub(crate) fn spawn_named<F>(
+        &self,
+        future: F,
+        _meta: SpawnMeta<'_>,
+        #[cfg(tokio_unstable)] parent: OptionalTaskHooks,
+    ) -> JoinHandle<F::Output>
     where
         F: Future + Send + 'static,
         F::Output: Send + 'static,
@@ -345,6 +357,9 @@ impl Handle {
         let future = super::task::trace::Trace::root(future);
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let future = crate::util::trace::task(future, "task", _meta, id.as_u64());
+        #[cfg(tokio_unstable)]
+        return self.inner.spawn(future, id, parent);
+        #[cfg(not(tokio_unstable))]
         self.inner.spawn(future, id)
     }
 
@@ -354,6 +369,7 @@ impl Handle {
         &self,
         future: F,
         _meta: SpawnMeta<'_>,
+        #[cfg(tokio_unstable)] hooks_override: OptionalTaskHooks,
     ) -> JoinHandle<F::Output>
     where
         F: Future + 'static,
@@ -370,6 +386,9 @@ impl Handle {
         let future = super::task::trace::Trace::root(future);
         #[cfg(all(tokio_unstable, feature = "tracing"))]
         let future = crate::util::trace::task(future, "task", _meta, id.as_u64());
+        #[cfg(tokio_unstable)]
+        return self.inner.spawn_local(future, id, hooks_override);
+        #[cfg(not(tokio_unstable))]
         self.inner.spawn_local(future, id)
     }
 

tokio/src/runtime/local_runtime/runtime.rs

@@ -155,9 +155,9 @@ impl LocalRuntime {
         // safety: spawn_local can only be called from `LocalRuntime`, which this is
         unsafe {
             if std::mem::size_of::<F>() > BOX_FUTURE_THRESHOLD {
-                self.handle.spawn_local_named(Box::pin(future), meta)
+                self.handle.spawn_local_named(Box::pin(future), meta, None)
             } else {
-                self.handle.spawn_local_named(future, meta)
+                self.handle.spawn_local_named(future, meta, None)
             }
         }
     }

tokio/src/runtime/metrics/batch.rs

@@ -14,11 +14,9 @@ pub(crate) struct MetricsBatch {
     /// Instant at which work last resumed (continued after park).
     processing_scheduled_tasks_started_at: Instant,
 
-    #[cfg(tokio_unstable)]
     /// Number of times the worker parked.
     park_count: u64,
 
-    #[cfg(tokio_unstable)]
     /// Number of times the worker parked and unparked.
     park_unpark_count: u64,
 
@@ -80,6 +78,8 @@ impl MetricsBatch {
                 MetricsBatch {
                     busy_duration_total: 0,
                     processing_scheduled_tasks_started_at: now,
+                    park_count: 0,
+                    park_unpark_count: 0,
                 }
             }
         },
@@ -120,7 +120,12 @@ impl MetricsBatch {
     cfg_metrics_variant! {
         stable: {
             #[inline(always)]
-            fn submit_unstable(&mut self, _worker: &WorkerMetrics, _mean_poll_time: u64) {}
+            fn submit_unstable(&mut self, worker: &WorkerMetrics, _mean_poll_time: u64) {
+                worker.park_count.store(self.park_count, Relaxed);
+                worker
+                    .park_unpark_count
+                    .store(self.park_unpark_count, Relaxed);
+            }
         },
         unstable: {
             #[inline(always)]
@@ -153,7 +158,10 @@ impl MetricsBatch {
     cfg_metrics_variant! {
         stable: {
             /// The worker is about to park.
-            pub(crate) fn about_to_park(&mut self) {}
+            pub(crate) fn about_to_park(&mut self) {
+                self.park_count += 1;
+                self.park_unpark_count += 1;
+            }
         },
         unstable: {
             /// The worker is about to park.
@@ -171,18 +179,9 @@ impl MetricsBatch {
             }
         }
     }
-
-    cfg_metrics_variant! {
-        stable: {
-            /// The worker was unparked.
-            pub(crate) fn unparked(&mut self) {}
-        },
-        unstable: {
-            /// The worker was unparked.
-            pub(crate) fn unparked(&mut self) {
-                self.park_unpark_count += 1;
-            }
-        }
+    /// The worker was unparked.
+    pub(crate) fn unparked(&mut self) {
+        self.park_unpark_count += 1;
     }
 
     /// Start processing a batch of tasks

tokio/src/runtime/metrics/runtime.rs

@@ -143,6 +143,103 @@ impl RuntimeMetrics {
                 .load(Relaxed);
             Duration::from_nanos(nanos)
         }
+
+        /// Returns the total number of times the given worker thread has parked.
+        ///
+        /// The worker park count starts at zero when the runtime is created and
+        /// increases by one each time the worker parks the thread waiting for new
+        /// inbound events to process. This usually means the worker has processed
+        /// all pending work and is currently idle.
+        ///
+        /// The counter is monotonically increasing. It is never decremented or
+        /// reset to zero.
+        ///
+        /// # Arguments
+        ///
+        /// `worker` is the index of the worker being queried. The given value must
+        /// be between 0 and `num_workers()`. The index uniquely identifies a single
+        /// worker and will continue to identify the worker throughout the lifetime
+        /// of the runtime instance.
+        ///
+        /// # Panics
+        ///
+        /// The method panics when `worker` represents an invalid worker, i.e. is
+        /// greater than or equal to `num_workers()`.
+        ///
+        /// # Examples
+        ///
+        /// ```
+        /// use tokio::runtime::Handle;
+        ///
+        /// #[tokio::main]
+        /// async fn main() {
+        ///     let metrics = Handle::current().metrics();
+        ///
+        ///     let n = metrics.worker_park_count(0);
+        ///     println!("worker 0 parked {} times", n);
+        /// }
+        /// ```
+        pub fn worker_park_count(&self, worker: usize) -> u64 {
+            self.handle
+                .inner
+                .worker_metrics(worker)
+                .park_count
+                .load(Relaxed)
+        }
+
+        /// Returns the total number of times the given worker thread has parked
+        /// and unparked.
+        ///
+        /// The worker park/unpark count starts at zero when the runtime is created
+        /// and increases by one each time the worker parks the thread waiting for
+        /// new inbound events to process. This usually means the worker has processed
+        /// all pending work and is currently idle. When new work becomes available,
+        /// the worker is unparked and the park/unpark count is again increased by one.
+        ///
+        /// An odd count means that the worker is currently parked.
+        /// An even count means that the worker is currently active.
+        ///
+        /// The counter is monotonically increasing. It is never decremented or
+        /// reset to zero.
+        ///
+        /// # Arguments
+        ///
+        /// `worker` is the index of the worker being queried. The given value must
+        /// be between 0 and `num_workers()`. The index uniquely identifies a single
+        /// worker and will continue to identify the worker throughout the lifetime
+        /// of the runtime instance.
+        ///
+        /// # Panics
+        ///
+        /// The method panics when `worker` represents an invalid worker, i.e. is
+        /// greater than or equal to `num_workers()`.
+        ///
+        /// # Examples
+        ///
+        /// ```
+        /// use tokio::runtime::Handle;
+        ///
+        /// #[tokio::main]
+        /// async fn main() {
+        ///     let metrics = Handle::current().metrics();
+        ///     let n = metrics.worker_park_unpark_count(0);
+        ///
+        ///     println!("worker 0 parked and unparked {} times", n);
+        ///
+        ///     if n % 2 == 0 {
+        ///         println!("worker 0 is active");
+        ///     } else {
+        ///         println!("worker 0 is parked");
+        ///     }
+        /// }
+        /// ```
+        pub fn worker_park_unpark_count(&self, worker: usize) -> u64 {
+            self.handle
+                .inner
+                .worker_metrics(worker)
+                .park_unpark_count
+                .load(Relaxed)
+        }
     }
 
     cfg_unstable_metrics! {
@@ -318,104 +415,6 @@ impl RuntimeMetrics {
                     .load(Relaxed)
             }
 
-            /// Returns the total number of times the given worker thread has parked.
-            ///
-            /// The worker park count starts at zero when the runtime is created and
-            /// increases by one each time the worker parks the thread waiting for new
-            /// inbound events to process. This usually means the worker has processed
-            /// all pending work and is currently idle.
-            ///
-            /// The counter is monotonically increasing. It is never decremented or
-            /// reset to zero.
-            ///
-            /// # Arguments
-            ///
-            /// `worker` is the index of the worker being queried. The given value must
-            /// be between 0 and `num_workers()`. The index uniquely identifies a single
-            /// worker and will continue to identify the worker throughout the lifetime
-            /// of the runtime instance.
-            ///
-            /// # Panics
-            ///
-            /// The method panics when `worker` represents an invalid worker, i.e. is
-            /// greater than or equal to `num_workers()`.
-            ///
-            /// # Examples
-            ///
-            /// ```
-            /// use tokio::runtime::Handle;
-            ///
-            /// #[tokio::main]
-            /// async fn main() {
-            ///     let metrics = Handle::current().metrics();
-            ///
-            ///     let n = metrics.worker_park_count(0);
-            ///     println!("worker 0 parked {} times", n);
-            /// }
-            /// ```
-            pub fn worker_park_count(&self, worker: usize) -> u64 {
-                self.handle
-                    .inner
-                    .worker_metrics(worker)
-                    .park_count
-                    .load(Relaxed)
-            }
-
-            /// Returns the total number of times the given worker thread has parked
-            /// and unparked.
-            ///
-            /// The worker park/unpark count starts at zero when the runtime is created
-            /// and increases by one each time the worker parks the thread waiting for
-            /// new inbound events to process. This usually means the worker has processed
-            /// all pending work and is currently idle. When new work becomes available,
-            /// the worker is unparked and the park/unpark count is again increased by one.
-            ///
-            /// An odd count means that the worker is currently parked.
-            /// An even count means that the worker is currently active.
-            ///
-            /// The counter is monotonically increasing. It is never decremented or
-            /// reset to zero.
-            ///
-            /// # Arguments
-            ///
-            /// `worker` is the index of the worker being queried. The given value must
-            /// be between 0 and `num_workers()`. The index uniquely identifies a single
-            /// worker and will continue to identify the worker throughout the lifetime
-            /// of the runtime instance.
-            ///
-            /// # Panics
-            ///
-            /// The method panics when `worker` represents an invalid worker, i.e. is
-            /// greater than or equal to `num_workers()`.
-            ///
-            /// # Examples
-            ///
-            /// ```
-            /// use tokio::runtime::Handle;
-            ///
-            /// #[tokio::main]
-            /// async fn main() {
-            ///     let metrics = Handle::current().metrics();
-            ///     let n = metrics.worker_park_unpark_count(0);
-            ///
-            ///     println!("worker 0 parked and unparked {} times", n);
-            ///
-            ///     if n % 2 == 0 {
-            ///         println!("worker 0 is active");
-            ///     } else {
-            ///         println!("worker 0 is parked");
-            ///     }
-            /// }
-            /// ```
-            pub fn worker_park_unpark_count(&self, worker: usize) -> u64 {
-                self.handle
-                    .inner
-                    .worker_metrics(worker)
-                    .park_unpark_count
-                    .load(Relaxed)
-            }
-
-
             /// Returns the number of times the given worker thread unparked but
             /// performed no work before parking again.
             ///

tokio/src/runtime/metrics/worker.rs

@@ -28,11 +28,9 @@ pub(crate) struct WorkerMetrics {
     /// Thread id of worker thread.
     thread_id: Mutex<Option<ThreadId>>,
 
-    #[cfg(tokio_unstable)]
     ///  Number of times the worker parked.
     pub(crate) park_count: MetricAtomicU64,
 
-    #[cfg(tokio_unstable)]
     ///  Number of times the worker parked and unparked.
     pub(crate) park_unpark_count: MetricAtomicU64,
 

tokio/src/runtime/mod.rs

@@ -380,13 +380,10 @@ cfg_rt! {
         pub use dump::Dump;
     }
 
-    mod task_hooks;
-    pub(crate) use task_hooks::{TaskHooks, TaskCallback};
     cfg_unstable! {
-        pub use task_hooks::TaskMeta;
+        mod task_hooks;
+        pub use task_hooks::*;
     }
-    #[cfg(not(tokio_unstable))]
-    pub(crate) use task_hooks::TaskMeta;
 
     mod handle;
     pub use handle::{EnterGuard, Handle, TryCurrentError};

tokio/src/runtime/runtime.rs

@@ -233,6 +233,15 @@ impl Runtime {
         F::Output: Send + 'static,
     {
         let fut_size = mem::size_of::<F>();
+        #[cfg(tokio_unstable)]
+        return if fut_size > BOX_FUTURE_THRESHOLD {
+            self.handle
+                .spawn_named(Box::pin(future), SpawnMeta::new_unnamed(fut_size), None)
+        } else {
+            self.handle
+                .spawn_named(future, SpawnMeta::new_unnamed(fut_size), None)
+        };
+        #[cfg(not(tokio_unstable))]
         if fut_size > BOX_FUTURE_THRESHOLD {
             self.handle
                 .spawn_named(Box::pin(future), SpawnMeta::new_unnamed(fut_size))

tokio/src/runtime/scheduler/current_thread/mod.rs

@@ -1,17 +1,19 @@
 use crate::loom::sync::atomic::AtomicBool;
 use crate::loom::sync::Arc;
+#[cfg(tokio_unstable)]
+use crate::runtime::context::with_task_hooks;
 use crate::runtime::driver::{self, Driver};
 use crate::runtime::scheduler::{self, Defer, Inject};
-use crate::runtime::task::{
-    self, JoinHandle, OwnedTasks, Schedule, Task, TaskHarnessScheduleHooks,
-};
+use crate::runtime::task::{self, JoinHandle, OwnedTasks, Schedule, Task};
+use crate::runtime::{blocking, context, Config, MetricsBatch, SchedulerMetrics, WorkerMetrics};
+#[cfg(tokio_unstable)]
 use crate::runtime::{
-    blocking, context, Config, MetricsBatch, SchedulerMetrics, TaskHooks, TaskMeta, WorkerMetrics,
+    OnChildTaskSpawnContext, OnTopLevelTaskSpawnContext, OptionalTaskHooks,
+    OptionalTaskHooksFactory, OptionalTaskHooksFactoryRef,
 };
 use crate::sync::notify::Notify;
 use crate::util::atomic_cell::AtomicCell;
 use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef};
-
 use std::cell::RefCell;
 use std::collections::VecDeque;
 use std::future::{poll_fn, Future};
@@ -20,7 +22,7 @@ use std::task::Poll::{Pending, Ready};
 use std::task::Waker;
 use std::thread::ThreadId;
 use std::time::Duration;
-use std::{fmt, thread};
+use std::{fmt, panic, thread};
 
 /// Executes tasks on the current thread
 pub(crate) struct CurrentThread {
@@ -47,7 +49,8 @@ pub(crate) struct Handle {
     pub(crate) seed_generator: RngSeedGenerator,
 
     /// User-supplied hooks to invoke for things
-    pub(crate) task_hooks: TaskHooks,
+    #[cfg(tokio_unstable)]
+    pub(crate) task_hooks: OptionalTaskHooksFactory,
 
     /// If this is a `LocalRuntime`, flags the owning thread ID.
     pub(crate) local_tid: Option<ThreadId>,
@@ -142,14 +145,8 @@ impl CurrentThread {
             .unwrap_or(DEFAULT_GLOBAL_QUEUE_INTERVAL);
 
         let handle = Arc::new(Handle {
-            task_hooks: TaskHooks {
-                task_spawn_callback: config.before_spawn.clone(),
-                task_terminate_callback: config.after_termination.clone(),
-                #[cfg(tokio_unstable)]
-                before_poll_callback: config.before_poll.clone(),
-                #[cfg(tokio_unstable)]
-                after_poll_callback: config.after_poll.clone(),
-            },
+            #[cfg(tokio_unstable)]
+            task_hooks: config.task_hook_factory.clone(),
             shared: Shared {
                 inject: Inject::new(),
                 owned: OwnedTasks::new(1),
@@ -448,19 +445,65 @@ impl Handle {
     pub(crate) fn spawn<F>(
         me: &Arc<Self>,
         future: F,
-        id: crate::runtime::task::Id,
+        id: task::Id,
+        #[cfg(tokio_unstable)] hooks_override: OptionalTaskHooks,
     ) -> JoinHandle<F::Output>
     where
         F: crate::future::Future + Send + 'static,
         F::Output: Send + 'static,
     {
-        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id);
-
-        me.task_hooks.spawn(&TaskMeta {
-            id,
-            _phantom: Default::default(),
+        // preference order for hook selection:
+        // 1. "hook override" - comes from builder
+        // 2. parent task's hook
+        // 3. runtime hook factory
+        #[cfg(tokio_unstable)]
+        let hooks = hooks_override.or_else(|| {
+            with_task_hooks(|parent| {
+                parent
+                    .map(|parent| {
+                        if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                            parent
+                                .on_child_spawn(&mut OnChildTaskSpawnContext {
+                                    id,
+                                    _phantom: Default::default(),
+                                })
+                                .hooks
+                        })) {
+                            r
+                        } else {
+                            None
+                        }
+                    })
+                    .flatten()
+            })
+            .ok()
+            .flatten()
+            .or_else(|| {
+                if let Some(hooks) = me.hooks_factory_ref() {
+                    if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                        hooks
+                            .on_top_level_spawn(&mut OnTopLevelTaskSpawnContext {
+                                id,
+                                _phantom: Default::default(),
+                            })
+                            .hooks
+                    })) {
+                        r
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            })
         });
 
+        #[cfg(tokio_unstable)]
+        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, hooks);
+
+        #[cfg(not(tokio_unstable))]
+        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id);
+
         if let Some(notified) = notified {
             me.schedule(notified);
         }
@@ -477,19 +520,67 @@ impl Handle {
     pub(crate) unsafe fn spawn_local<F>(
         me: &Arc<Self>,
         future: F,
-        id: crate::runtime::task::Id,
+        id: task::Id,
+        #[cfg(tokio_unstable)] hooks_override: OptionalTaskHooks,
     ) -> JoinHandle<F::Output>
     where
         F: crate::future::Future + 'static,
         F::Output: 'static,
     {
-        let (handle, notified) = me.shared.owned.bind_local(future, me.clone(), id);
-
-        me.task_hooks.spawn(&TaskMeta {
-            id,
-            _phantom: Default::default(),
+        // preference order for hook selection:
+        // 1. "hook override" - comes from builder
+        // 2. parent task's hook
+        // 3. runtime hook factory
+        #[cfg(tokio_unstable)]
+        let hooks = hooks_override.or_else(|| {
+            with_task_hooks(|parent| {
+                parent
+                    .map(|parent| {
+                        if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                            parent
+                                .on_child_spawn(&mut OnChildTaskSpawnContext {
+                                    id,
+                                    _phantom: Default::default(),
+                                })
+                                .hooks
+                        })) {
+                            r
+                        } else {
+                            None
+                        }
+                    })
+                    .flatten()
+            })
+            .ok()
+            .flatten()
+            .or_else(|| {
+                if let Some(hooks) = me.hooks_factory_ref() {
+                    if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                        hooks
+                            .on_top_level_spawn(&mut OnTopLevelTaskSpawnContext {
+                                id,
+                                _phantom: Default::default(),
+                            })
+                            .hooks
+                    })) {
+                        r
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            })
         });
 
+        let (handle, notified) = me.shared.owned.bind_local(
+            future,
+            me.clone(),
+            id,
+            #[cfg(tokio_unstable)]
+            hooks,
+        );
+
         if let Some(notified) = notified {
             me.schedule(notified);
         }
@@ -654,10 +745,14 @@ impl Schedule for Arc<Handle> {
         });
     }
 
-    fn hooks(&self) -> TaskHarnessScheduleHooks {
-        TaskHarnessScheduleHooks {
-            task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
-        }
+    #[cfg(tokio_unstable)]
+    fn hooks_factory(&self) -> OptionalTaskHooksFactory {
+        self.task_hooks.clone()
+    }
+
+    #[cfg(tokio_unstable)]
+    fn hooks_factory_ref(&self) -> OptionalTaskHooksFactoryRef<'_> {
+        self.task_hooks.as_ref().map(AsRef::as_ref)
     }
 
     cfg_unstable! {
@@ -770,17 +865,8 @@ impl CoreGuard<'_> {
 
                     let task = context.handle.shared.owned.assert_owner(task);
 
-                    #[cfg(tokio_unstable)]
-                    let task_id = task.task_id();
-
                     let (c, ()) = context.run_task(core, || {
-                        #[cfg(tokio_unstable)]
-                        context.handle.task_hooks.poll_start_callback(task_id);
-
                         task.run();
-
-                        #[cfg(tokio_unstable)]
-                        context.handle.task_hooks.poll_stop_callback(task_id);
                     });
 
                     core = c;

tokio/src/runtime/scheduler/mod.rs

@@ -8,8 +8,6 @@ cfg_rt! {
     pub(crate) mod inject;
     pub(crate) use inject::Inject;
 
-    use crate::runtime::TaskHooks;
-
     use crate::runtime::WorkerMetrics;
 }
 
@@ -25,6 +23,10 @@ cfg_rt_multi_thread! {
 }
 
 use crate::runtime::driver;
+#[cfg(all(feature = "rt", tokio_unstable))]
+use crate::runtime::task::Schedule;
+#[cfg(all(feature = "rt", tokio_unstable))]
+use crate::runtime::{OptionalTaskHooks, OptionalTaskHooksFactory};
 
 #[derive(Debug, Clone)]
 pub(crate) enum Handle {
@@ -117,11 +119,24 @@ cfg_rt! {
             }
         }
 
-        pub(crate) fn spawn<F>(&self, future: F, id: Id) -> JoinHandle<F::Output>
+        pub(crate) fn spawn<F>(&self,
+            future: F,
+            id: Id,
+            #[cfg(tokio_unstable)]
+            hooks_override: OptionalTaskHooks
+        ) -> JoinHandle<F::Output>
         where
             F: Future + Send + 'static,
             F::Output: Send + 'static,
         {
+            #[cfg(tokio_unstable)]
+            return match self {
+                Handle::CurrentThread(h) => current_thread::Handle::spawn(h, future, id, hooks_override),
+
+                #[cfg(feature = "rt-multi-thread")]
+                Handle::MultiThread(h) => multi_thread::Handle::spawn(h, future, id, hooks_override),
+            };
+            #[cfg(not(tokio_unstable))]
             match self {
                 Handle::CurrentThread(h) => current_thread::Handle::spawn(h, future, id),
 
@@ -136,12 +151,15 @@ cfg_rt! {
         /// This should only be called in `LocalRuntime` if the runtime has been verified to be owned
         /// by the current thread.
         #[allow(irrefutable_let_patterns)]
-        pub(crate) unsafe fn spawn_local<F>(&self, future: F, id: Id) -> JoinHandle<F::Output>
+        pub(crate) unsafe fn spawn_local<F>(&self, future: F, id: Id, #[cfg(tokio_unstable)] hooks_override: OptionalTaskHooks) -> JoinHandle<F::Output>
         where
             F: Future + 'static,
             F::Output: 'static,
         {
             if let Handle::CurrentThread(h) = self {
+                #[cfg(tokio_unstable)]
+                return current_thread::Handle::spawn_local(h, future, id, hooks_override);
+                #[cfg(not(tokio_unstable))]
                 current_thread::Handle::spawn_local(h, future, id)
             } else {
                 panic!("Only current_thread and LocalSet have spawn_local internals implemented")
@@ -169,12 +187,9 @@ cfg_rt! {
             }
         }
 
-        pub(crate) fn hooks(&self) -> &TaskHooks {
-            match self {
-                Handle::CurrentThread(h) => &h.task_hooks,
-                #[cfg(feature = "rt-multi-thread")]
-                Handle::MultiThread(h) => &h.task_hooks,
-            }
+        #[cfg(tokio_unstable)]
+        pub(crate) fn hooks_factory(&self) -> OptionalTaskHooksFactory {
+            match_flavor!(self, Handle(h) => h.hooks_factory())
         }
     }
 

tokio/src/runtime/scheduler/multi_thread/handle.rs

@@ -1,14 +1,20 @@
 use crate::future::Future;
 use crate::loom::sync::Arc;
+#[cfg(tokio_unstable)]
+use crate::runtime::context::with_task_hooks;
 use crate::runtime::scheduler::multi_thread::worker;
+#[cfg(tokio_unstable)]
+use crate::runtime::task::Schedule;
 use crate::runtime::{
     blocking, driver,
     task::{self, JoinHandle},
-    TaskHooks, TaskMeta,
 };
+#[cfg(tokio_unstable)]
+use crate::runtime::{OnChildTaskSpawnContext, OnTopLevelTaskSpawnContext, OptionalTaskHooks};
 use crate::util::RngSeedGenerator;
-
 use std::fmt;
+#[cfg(tokio_unstable)]
+use std::panic;
 
 mod metrics;
 
@@ -29,18 +35,24 @@ pub(crate) struct Handle {
 
     /// Current random number generator seed
     pub(crate) seed_generator: RngSeedGenerator,
-
-    /// User-supplied hooks to invoke for things
-    pub(crate) task_hooks: TaskHooks,
 }
 
 impl Handle {
     /// Spawns a future onto the thread pool
-    pub(crate) fn spawn<F>(me: &Arc<Self>, future: F, id: task::Id) -> JoinHandle<F::Output>
+    pub(crate) fn spawn<F>(
+        me: &Arc<Self>,
+        future: F,
+        id: task::Id,
+        #[cfg(tokio_unstable)] hooks_override: OptionalTaskHooks,
+    ) -> JoinHandle<F::Output>
     where
         F: crate::future::Future + Send + 'static,
         F::Output: Send + 'static,
     {
+        #[cfg(tokio_unstable)]
+        return Self::bind_new_task(me, future, id, hooks_override);
+
+        #[cfg(not(tokio_unstable))]
         Self::bind_new_task(me, future, id)
     }
 
@@ -48,18 +60,70 @@ impl Handle {
         self.close();
     }
 
-    pub(super) fn bind_new_task<T>(me: &Arc<Self>, future: T, id: task::Id) -> JoinHandle<T::Output>
+    pub(super) fn bind_new_task<T>(
+        me: &Arc<Self>,
+        future: T,
+        id: task::Id,
+        #[cfg(tokio_unstable)] hooks_override: OptionalTaskHooks,
+    ) -> JoinHandle<T::Output>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id);
-
-        me.task_hooks.spawn(&TaskMeta {
-            id,
-            _phantom: Default::default(),
+        // preference order for hook selection:
+        // 1. "hook override" - comes from builder
+        // 2. parent task's hook
+        // 3. runtime hook factory
+        #[cfg(tokio_unstable)]
+        let hooks = hooks_override.or_else(|| {
+            with_task_hooks(|parent| {
+                parent
+                    .map(|parent| {
+                        if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                            parent
+                                .on_child_spawn(&mut OnChildTaskSpawnContext {
+                                    id,
+                                    _phantom: Default::default(),
+                                })
+                                .hooks
+                        })) {
+                            r
+                        } else {
+                            None
+                        }
+                    })
+                    .flatten()
+            })
+            .ok()
+            .flatten()
+            .or_else(|| {
+                if let Some(hooks) = me.hooks_factory_ref() {
+                    if let Ok(r) = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                        hooks
+                            .on_top_level_spawn(&mut OnTopLevelTaskSpawnContext {
+                                id,
+                                _phantom: Default::default(),
+                            })
+                            .hooks
+                    })) {
+                        r
+                    } else {
+                        None
+                    }
+                } else {
+                    None
+                }
+            })
         });
 
+        let (handle, notified) = me.shared.owned.bind(
+            future,
+            me.clone(),
+            id,
+            #[cfg(tokio_unstable)]
+            hooks,
+        );
+
         me.schedule_option_task_without_yield(notified);
 
         handle

tokio/src/runtime/scheduler/multi_thread/worker.rs

@@ -58,13 +58,15 @@
 
 use crate::loom::sync::{Arc, Mutex};
 use crate::runtime;
+use crate::runtime::context;
 use crate::runtime::scheduler::multi_thread::{
     idle, queue, Counters, Handle, Idle, Overflow, Parker, Stats, TraceStatus, Unparker,
 };
 use crate::runtime::scheduler::{inject, Defer, Lock};
-use crate::runtime::task::{OwnedTasks, TaskHarnessScheduleHooks};
+use crate::runtime::task::OwnedTasks;
 use crate::runtime::{blocking, driver, scheduler, task, Config, SchedulerMetrics, WorkerMetrics};
-use crate::runtime::{context, TaskHooks};
+#[cfg(tokio_unstable)]
+use crate::runtime::{OptionalTaskHooksFactory, OptionalTaskHooksFactoryRef};
 use crate::task::coop;
 use crate::util::atomic_cell::AtomicCell;
 use crate::util::rand::{FastRand, RngSeedGenerator};
@@ -281,7 +283,6 @@ pub(super) fn create(
 
     let remotes_len = remotes.len();
     let handle = Arc::new(Handle {
-        task_hooks: TaskHooks::from_config(&config),
         shared: Shared {
             remotes: remotes.into_boxed_slice(),
             inject,
@@ -570,9 +571,6 @@ impl Context {
     }
 
     fn run_task(&self, task: Notified, mut core: Box<Core>) -> RunResult {
-        #[cfg(tokio_unstable)]
-        let task_id = task.task_id();
-
         let task = self.worker.handle.shared.owned.assert_owner(task);
 
         // Make sure the worker is not in the **searching** state. This enables
@@ -592,16 +590,8 @@ impl Context {
 
         // Run the task
         coop::budget(|| {
-            // Unlike the poll time above, poll start callback is attached to the task id,
-            // so it is tightly associated with the actual poll invocation.
-            #[cfg(tokio_unstable)]
-            self.worker.handle.task_hooks.poll_start_callback(task_id);
-
             task.run();
 
-            #[cfg(tokio_unstable)]
-            self.worker.handle.task_hooks.poll_stop_callback(task_id);
-
             let mut lifo_polls = 0;
 
             // As long as there is budget remaining and a task exists in the
@@ -665,16 +655,7 @@ impl Context {
                 *self.core.borrow_mut() = Some(core);
                 let task = self.worker.handle.shared.owned.assert_owner(task);
 
-                #[cfg(tokio_unstable)]
-                let task_id = task.task_id();
-
-                #[cfg(tokio_unstable)]
-                self.worker.handle.task_hooks.poll_start_callback(task_id);
-
                 task.run();
-
-                #[cfg(tokio_unstable)]
-                self.worker.handle.task_hooks.poll_stop_callback(task_id);
             }
         })
     }
@@ -1063,10 +1044,18 @@ impl task::Schedule for Arc<Handle> {
         self.schedule_task(task, false);
     }
 
-    fn hooks(&self) -> TaskHarnessScheduleHooks {
-        TaskHarnessScheduleHooks {
-            task_terminate_callback: self.task_hooks.task_terminate_callback.clone(),
-        }
+    #[cfg(tokio_unstable)]
+    fn hooks_factory(&self) -> OptionalTaskHooksFactory {
+        self.shared.config.task_hook_factory.clone()
+    }
+
+    #[cfg(tokio_unstable)]
+    fn hooks_factory_ref(&self) -> OptionalTaskHooksFactoryRef<'_> {
+        self.shared
+            .config
+            .task_hook_factory
+            .as_ref()
+            .map(AsRef::as_ref)
     }
 
     fn yield_now(&self, task: Notified) {

tokio/src/runtime/task/core.rs

@@ -14,7 +14,9 @@ use crate::loom::cell::UnsafeCell;
 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::runtime::task::{Id, Schedule};
+#[cfg(tokio_unstable)]
+use crate::runtime::OptionalTaskHooks;
 use crate::util::linked_list;
 
 use std::num::NonZeroU64;
@@ -186,7 +188,8 @@ pub(super) struct Trailer {
     /// Consumer task waiting on completion of this task.
     pub(super) waker: UnsafeCell<Option<Waker>>,
     /// Optional hooks needed in the harness.
-    pub(super) hooks: TaskHarnessScheduleHooks,
+    #[cfg(tokio_unstable)]
+    pub(super) hooks: UnsafeCell<OptionalTaskHooks>,
 }
 
 generate_addr_of_methods! {
@@ -208,7 +211,13 @@ pub(super) enum Stage<T: Future> {
 impl<T: Future, S: Schedule> Cell<T, S> {
     /// Allocates a new task cell, containing the header, trailer, and core
     /// structures.
-    pub(super) fn new(future: T, scheduler: S, state: State, task_id: Id) -> Box<Cell<T, S>> {
+    pub(super) fn new(
+        future: T,
+        scheduler: S,
+        state: State,
+        task_id: Id,
+        #[cfg(tokio_unstable)] hooks: OptionalTaskHooks,
+    ) -> Box<Cell<T, S>> {
         // Separated into a non-generic function to reduce LLVM codegen
         fn new_header(
             state: State,
@@ -229,7 +238,13 @@ impl<T: Future, S: Schedule> Cell<T, S> {
         let tracing_id = future.id();
         let vtable = raw::vtable::<T, S>();
         let result = Box::new(Cell {
-            trailer: Trailer::new(scheduler.hooks()),
+            #[cfg(tokio_unstable)]
+            trailer: Trailer::new(
+                #[cfg(tokio_unstable)]
+                hooks,
+            ),
+            #[cfg(not(tokio_unstable))]
+            trailer: Trailer::new(),
             header: new_header(
                 state,
                 vtable,
@@ -462,11 +477,12 @@ impl Header {
 }
 
 impl Trailer {
-    fn new(hooks: TaskHarnessScheduleHooks) -> Self {
+    fn new(#[cfg(tokio_unstable)] hooks: OptionalTaskHooks) -> Self {
         Trailer {
             waker: UnsafeCell::new(None),
             owned: linked_list::Pointers::new(),
-            hooks,
+            #[cfg(tokio_unstable)]
+            hooks: UnsafeCell::new(hooks),
         }
     }
 

tokio/src/runtime/task/harness.rs

@@ -1,10 +1,12 @@
 use crate::future::Future;
+#[cfg(tokio_unstable)]
+use crate::runtime::context::with_task_hooks;
 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::runtime::TaskMeta;
+#[cfg(tokio_unstable)]
+use crate::runtime::{AfterTaskPollContext, OnTaskTerminateContext};
 use std::any::Any;
 use std::mem;
 use std::mem::ManuallyDrop;
@@ -150,8 +152,21 @@ where
     /// All necessary state checks and transitions are performed.
     /// Panics raised while polling the future are handled.
     pub(super) fn poll(self) {
+        let res = self.poll_inner();
+
+        #[cfg(tokio_unstable)]
+        let _ = with_task_hooks(|t| {
+            if let Some(hooks) = t {
+                let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                    hooks.after_poll(&mut AfterTaskPollContext {
+                        _phantom: Default::default(),
+                    })
+                }));
+            }
+        });
+
         // We pass our ref-count to `poll_inner`.
-        match self.poll_inner() {
+        match res {
             PollFuture::Notified => {
                 // The `poll_inner` call has given us two ref-counts back.
                 // We give one of them to a new task and call `yield_now`.
@@ -367,14 +382,16 @@ where
         //
         // We call this in a separate block so that it runs after the task appears to have
         // completed and will still run if the destructor panics.
-        if let Some(f) = self.trailer().hooks.task_terminate_callback.as_ref() {
-            let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
-                f(&TaskMeta {
-                    id: self.core().task_id,
-                    _phantom: Default::default(),
-                })
-            }));
-        }
+        #[cfg(tokio_unstable)]
+        let _ = with_task_hooks(|t| {
+            if let Some(hooks) = t {
+                let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                    hooks.on_task_terminate(&mut OnTaskTerminateContext {
+                        _phantom: Default::default(),
+                    })
+                }));
+            }
+        });
 
         // The task has completed execution and will no longer be scheduled.
         let num_release = self.release();

tokio/src/runtime/task/list.rs

@@ -13,9 +13,10 @@ use crate::util::linked_list::{Link, LinkedList};
 use crate::util::sharded_list;
 
 use crate::loom::sync::atomic::{AtomicBool, Ordering};
+#[cfg(tokio_unstable)]
+use crate::runtime::OptionalTaskHooks;
 use std::marker::PhantomData;
 use std::num::NonZeroU64;
-
 // The id from the module below is used to verify whether a given task is stored
 // in this OwnedTasks, or some other task. The counter starts at one so we can
 // use `None` for tasks not owned by any list.
@@ -91,13 +92,20 @@ impl<S: 'static> OwnedTasks<S> {
         task: T,
         scheduler: S,
         id: super::Id,
+        #[cfg(tokio_unstable)] hooks: OptionalTaskHooks,
     ) -> (JoinHandle<T::Output>, Option<Notified<S>>)
     where
         S: Schedule,
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        let (task, notified, join) = super::new_task(task, scheduler, id);
+        let (task, notified, join) = super::new_task(
+            task,
+            scheduler,
+            id,
+            #[cfg(tokio_unstable)]
+            hooks,
+        );
         let notified = unsafe { self.bind_inner(task, notified) };
         (join, notified)
     }
@@ -111,13 +119,20 @@ impl<S: 'static> OwnedTasks<S> {
         task: T,
         scheduler: S,
         id: super::Id,
+        #[cfg(tokio_unstable)] parent: OptionalTaskHooks,
     ) -> (JoinHandle<T::Output>, Option<Notified<S>>)
     where
         S: Schedule,
         T: Future + 'static,
         T::Output: 'static,
     {
-        let (task, notified, join) = super::new_task(task, scheduler, id);
+        let (task, notified, join) = super::new_task(
+            task,
+            scheduler,
+            id,
+            #[cfg(tokio_unstable)]
+            parent,
+        );
         let notified = unsafe { self.bind_inner(task, notified) };
         (join, notified)
     }
@@ -258,12 +273,16 @@ impl<S: 'static> LocalOwnedTasks<S> {
         task: T,
         scheduler: S,
         id: super::Id,
+        #[cfg(tokio_unstable)] parent: OptionalTaskHooks,
     ) -> (JoinHandle<T::Output>, Option<Notified<S>>)
     where
         S: Schedule,
         T: Future + 'static,
         T::Output: 'static,
     {
+        #[cfg(tokio_unstable)]
+        let (task, notified, join) = super::new_task(task, scheduler, id, parent);
+        #[cfg(not(tokio_unstable))]
         let (task, notified, join) = super::new_task(task, scheduler, id);
 
         unsafe {

tokio/src/runtime/task/mod.rs

@@ -221,10 +221,10 @@ cfg_taskdump! {
 }
 
 use crate::future::Future;
+#[cfg(tokio_unstable)]
+use crate::runtime::{OptionalTaskHooks, OptionalTaskHooksFactory, OptionalTaskHooksFactoryRef};
 use crate::util::linked_list;
 use crate::util::sharded_list;
-
-use crate::runtime::TaskCallback;
 use std::marker::PhantomData;
 use std::ptr::NonNull;
 use std::{fmt, mem};
@@ -256,13 +256,6 @@ pub(crate) struct LocalNotified<S: 'static> {
     _not_send: PhantomData<*const ()>,
 }
 
-impl<S> LocalNotified<S> {
-    #[cfg(tokio_unstable)]
-    pub(crate) fn task_id(&self) -> Id {
-        self.task.id()
-    }
-}
-
 /// A task that is not owned by any `OwnedTasks`. Used for blocking tasks.
 /// This type holds two ref-counts.
 pub(crate) struct UnownedTask<S: 'static> {
@@ -277,12 +270,6 @@ unsafe impl<S> Sync for UnownedTask<S> {}
 /// Task result sent back.
 pub(crate) type Result<T> = std::result::Result<T, JoinError>;
 
-/// Hooks for scheduling tasks which are needed in the task harness.
-#[derive(Clone)]
-pub(crate) struct TaskHarnessScheduleHooks {
-    pub(crate) task_terminate_callback: Option<TaskCallback>,
-}
-
 pub(crate) trait Schedule: Sync + Sized + 'static {
     /// The task has completed work and is ready to be released. The scheduler
     /// should release it immediately and return it. The task module will batch
@@ -294,7 +281,11 @@ pub(crate) trait Schedule: Sync + Sized + 'static {
     /// Schedule the task
     fn schedule(&self, task: Notified<Self>);
 
-    fn hooks(&self) -> TaskHarnessScheduleHooks;
+    #[cfg(tokio_unstable)]
+    fn hooks_factory(&self) -> OptionalTaskHooksFactory;
+
+    #[cfg(tokio_unstable)]
+    fn hooks_factory_ref(&self) -> OptionalTaskHooksFactoryRef<'_>;
 
     /// Schedule the task to run in the near future, yielding the thread to
     /// other tasks.
@@ -317,13 +308,19 @@ cfg_rt! {
         task: T,
         scheduler: S,
         id: Id,
+        #[cfg(tokio_unstable)]
+        hooks: OptionalTaskHooks
     ) -> (Task<S>, Notified<S>, JoinHandle<T::Output>)
     where
         S: Schedule,
         T: Future + 'static,
         T::Output: 'static,
     {
+        #[cfg(tokio_unstable)]
+        let raw = RawTask::new::<T, S>(task, scheduler, id, hooks);
+        #[cfg(not(tokio_unstable))]
         let raw = RawTask::new::<T, S>(task, scheduler, id);
+
         let task = Task {
             raw,
             _p: PhantomData,
@@ -341,12 +338,16 @@ cfg_rt! {
     /// only when the task is not going to be stored in an `OwnedTasks` list.
     ///
     /// Currently only blocking tasks use this method.
-    pub(crate) fn unowned<T, S>(task: T, scheduler: S, id: Id) -> (UnownedTask<S>, JoinHandle<T::Output>)
+    pub(crate) fn unowned<T, S>(task: T, scheduler: S, id: Id, #[cfg(tokio_unstable)] hooks: OptionalTaskHooks) -> (UnownedTask<S>, JoinHandle<T::Output>)
     where
         S: Schedule,
         T: Send + Future + 'static,
         T::Output: Send + 'static,
     {
+        #[cfg(tokio_unstable)]
+        let (task, notified, join) = new_task(task, scheduler, id, hooks);
+
+        #[cfg(not(tokio_unstable))]
         let (task, notified, join) = new_task(task, scheduler, id);
 
         // This transfers the ref-count of task and notified into an UnownedTask.
@@ -459,6 +460,7 @@ impl<S: Schedule> LocalNotified<S> {
     /// Runs the task.
     pub(crate) fn run(self) {
         let raw = self.task.raw;
+
         mem::forget(self);
         raw.poll();
     }

tokio/src/runtime/task/raw.rs

@@ -1,7 +1,12 @@
 use crate::future::Future;
+#[cfg(tokio_unstable)]
+use crate::runtime::context::set_task_hooks;
 use crate::runtime::task::core::{Core, Trailer};
 use crate::runtime::task::{Cell, Harness, Header, Id, Schedule, State};
-
+#[cfg(tokio_unstable)]
+use crate::runtime::{BeforeTaskPollContext, OptionalTaskHooks, TaskHookHarness};
+#[cfg(tokio_unstable)]
+use std::panic;
 use std::ptr::NonNull;
 use std::task::{Poll, Waker};
 
@@ -157,12 +162,24 @@ const fn get_id_offset(
 }
 
 impl RawTask {
-    pub(super) fn new<T, S>(task: T, scheduler: S, id: Id) -> RawTask
+    pub(super) fn new<T, S>(
+        task: T,
+        scheduler: S,
+        id: Id,
+        #[cfg(tokio_unstable)] hooks: OptionalTaskHooks,
+    ) -> RawTask
     where
         T: Future,
         S: Schedule,
     {
-        let ptr = Box::into_raw(Cell::<_, S>::new(task, scheduler, State::new(), id));
+        let ptr = Box::into_raw(Cell::<_, S>::new(
+            task,
+            scheduler,
+            State::new(),
+            id,
+            #[cfg(tokio_unstable)]
+            hooks,
+        ));
         let ptr = unsafe { NonNull::new_unchecked(ptr.cast()) };
 
         RawTask { ptr }
@@ -197,8 +214,27 @@ impl RawTask {
 
     /// Safety: mutual exclusion is required to call this function.
     pub(crate) fn poll(self) {
-        let vtable = self.header().vtable;
-        unsafe { (vtable.poll)(self.ptr) }
+        #[cfg(tokio_unstable)]
+        let _guard = self.trailer().hooks.with_mut(|ptr| unsafe {
+            ptr.as_mut().and_then(|x| {
+                x.as_mut().map(|x| {
+                    let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| {
+                        x.before_poll(&mut BeforeTaskPollContext {
+                            _phantom: Default::default(),
+                        })
+                    }));
+
+                    set_task_hooks(NonNull::new(
+                        (&mut **x) as *mut (dyn TaskHookHarness + Send + Sync + 'static),
+                    ))
+                })
+            })
+        });
+
+        unsafe {
+            let vtable = self.header().vtable;
+            (vtable.poll)(self.ptr);
+        }
     }
 
     pub(super) fn schedule(self) {

tokio/src/runtime/task_hooks.rs

@@ -1,81 +0,0 @@
-use std::marker::PhantomData;
-
-use super::Config;
-
-impl TaskHooks {
-    pub(crate) fn spawn(&self, meta: &TaskMeta<'_>) {
-        if let Some(f) = self.task_spawn_callback.as_ref() {
-            f(meta)
-        }
-    }
-
-    #[allow(dead_code)]
-    pub(crate) fn from_config(config: &Config) -> Self {
-        Self {
-            task_spawn_callback: config.before_spawn.clone(),
-            task_terminate_callback: config.after_termination.clone(),
-            #[cfg(tokio_unstable)]
-            before_poll_callback: config.before_poll.clone(),
-            #[cfg(tokio_unstable)]
-            after_poll_callback: config.after_poll.clone(),
-        }
-    }
-
-    #[cfg(tokio_unstable)]
-    #[inline]
-    pub(crate) fn poll_start_callback(&self, id: super::task::Id) {
-        if let Some(poll_start) = &self.before_poll_callback {
-            (poll_start)(&TaskMeta {
-                id,
-                _phantom: std::marker::PhantomData,
-            })
-        }
-    }
-
-    #[cfg(tokio_unstable)]
-    #[inline]
-    pub(crate) fn poll_stop_callback(&self, id: super::task::Id) {
-        if let Some(poll_stop) = &self.after_poll_callback {
-            (poll_stop)(&TaskMeta {
-                id,
-                _phantom: std::marker::PhantomData,
-            })
-        }
-    }
-}
-
-#[derive(Clone)]
-pub(crate) struct TaskHooks {
-    pub(crate) task_spawn_callback: Option<TaskCallback>,
-    pub(crate) task_terminate_callback: Option<TaskCallback>,
-    #[cfg(tokio_unstable)]
-    pub(crate) before_poll_callback: Option<TaskCallback>,
-    #[cfg(tokio_unstable)]
-    pub(crate) after_poll_callback: Option<TaskCallback>,
-}
-
-/// Task metadata supplied to user-provided hooks for task events.
-///
-/// **Note**: This is an [unstable API][unstable]. The public API of this type
-/// may break in 1.x releases. See [the documentation on unstable
-/// features][unstable] for details.
-///
-/// [unstable]: crate#unstable-features
-#[allow(missing_debug_implementations)]
-#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
-pub struct TaskMeta<'a> {
-    /// The opaque ID of the task.
-    pub(crate) id: super::task::Id,
-    pub(crate) _phantom: PhantomData<&'a ()>,
-}
-
-impl<'a> TaskMeta<'a> {
-    /// Return the opaque ID of the task.
-    #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))]
-    pub fn id(&self) -> super::task::Id {
-        self.id
-    }
-}
-
-/// Runs on specific task-related events
-pub(crate) type TaskCallback = std::sync::Arc<dyn Fn(&TaskMeta<'_>) + Send + Sync>;

tokio/src/runtime/task_hooks/mod.rs

@@ -0,0 +1,156 @@
+use super::task;
+use crate::loom::cell::UnsafeCell;
+use std::marker::PhantomData;
+use std::ptr::NonNull;
+use std::sync::Arc;
+
+/// A factory which produces new [`TaskHookHarness`] objects for tasks which either have been
+/// spawned in "detached mode" via [`crate::task::spawn_with_hooks`], or which were spawned from outside the runtime or
+/// from another context where no [`TaskHookHarness`] was present.
+pub trait TaskHookHarnessFactory {
+    /// Runs a hook which may produce a new [`TaskHookHarness`] object which the runtime will attach to a given task.
+    fn on_top_level_spawn(&self, ctx: &mut OnTopLevelTaskSpawnContext<'_>)
+        -> OnTopLevelSpawnAction;
+}
+
+/// Trait for user-provided "harness" objects which are attached to tasks and provide hook
+/// implementations.
+#[allow(unused_variables)]
+pub trait TaskHookHarness {
+    /// Pre-poll task hook which runs arbitrary user logic.
+    fn before_poll(&mut self, ctx: &mut BeforeTaskPollContext<'_>) -> BeforeTaskPollAction {
+        BeforeTaskPollAction::default()
+    }
+
+    /// Post-poll task hook which runs arbitrary user logic.
+    fn after_poll(&mut self, ctx: &mut AfterTaskPollContext<'_>) -> AfterTaskPollAction {
+        AfterTaskPollAction::default()
+    }
+
+    /// Task hook which runs when this task spawns a child, unless that child is explicitly spawned
+    /// detached from the parent.
+    ///
+    /// This hook creates a harness for the child, or detaches the child from any instrumentation.
+    fn on_child_spawn(&mut self, ctx: &mut OnChildTaskSpawnContext<'_>) -> OnChildSpawnAction {
+        OnChildSpawnAction::default()
+    }
+
+    /// Task hook which runs on task termination.
+    fn on_task_terminate(&mut self, ctx: &mut OnTaskTerminateContext<'_>) -> OnTaskTerminateAction {
+        OnTaskTerminateAction::default()
+    }
+}
+
+pub(crate) type OptionalTaskHooksFactory =
+    Option<Arc<dyn TaskHookHarnessFactory + Send + Sync + 'static>>;
+pub(crate) type OptionalTaskHooks = Option<Box<dyn TaskHookHarness + Send + Sync + 'static>>;
+
+pub(crate) type OptionalTaskHooksWeak =
+    UnsafeCell<Option<NonNull<dyn TaskHookHarness + Send + Sync + 'static>>>;
+
+pub(crate) type OptionalTaskHooksMut<'a> =
+    Option<&'a mut (dyn TaskHookHarness + Send + Sync + 'static)>;
+pub(crate) type OptionalTaskHooksFactoryRef<'a> =
+    Option<&'a (dyn TaskHookHarnessFactory + Send + Sync + 'static)>;
+
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+pub struct OnTopLevelTaskSpawnContext<'a> {
+    pub(crate) id: task::Id,
+    pub(crate) _phantom: PhantomData<&'a ()>,
+}
+
+impl<'a> OnTopLevelTaskSpawnContext<'a> {
+    /// Returns the ID of the task.
+    pub fn id(&self) -> task::Id {
+        self.id
+    }
+}
+
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+pub struct OnChildTaskSpawnContext<'a> {
+    pub(crate) id: task::Id,
+    pub(crate) _phantom: PhantomData<&'a ()>,
+}
+
+impl<'a> OnChildTaskSpawnContext<'a> {
+    /// Returns the ID of the task.
+    pub fn id(&self) -> task::Id {
+        self.id
+    }
+}
+
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+pub struct OnTaskTerminateContext<'a> {
+    pub(crate) _phantom: PhantomData<&'a ()>,
+}
+
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+pub struct BeforeTaskPollContext<'a> {
+    pub(crate) _phantom: PhantomData<&'a ()>,
+}
+
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+pub struct AfterTaskPollContext<'a> {
+    pub(crate) _phantom: PhantomData<&'a ()>,
+}
+
+#[derive(Default)]
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+#[non_exhaustive]
+pub struct OnTopLevelSpawnAction {
+    pub(crate) hooks: Option<Box<dyn TaskHookHarness + Send + Sync + 'static>>,
+}
+
+impl OnTopLevelSpawnAction {
+    /// Pass in a set of task hooks for the task.
+    pub fn set_hooks<T>(&mut self, hooks: T) -> &mut Self
+    where
+        T: TaskHookHarness + Send + Sync + 'static,
+    {
+        self.hooks = Some(Box::new(hooks));
+        self
+    }
+}
+
+#[derive(Default)]
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+#[non_exhaustive]
+pub struct OnChildSpawnAction {
+    pub(crate) hooks: Option<Box<dyn TaskHookHarness + Send + Sync + 'static>>,
+}
+
+impl OnChildSpawnAction {
+    /// Pass in a set of task hooks for the child task.
+    pub fn set_hooks<T>(&mut self, hooks: T) -> &mut Self
+    where
+        T: TaskHookHarness + Send + Sync + 'static,
+    {
+        self.hooks = Some(Box::new(hooks));
+        self
+    }
+}
+
+#[derive(Default)]
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+#[non_exhaustive]
+pub struct OnTaskTerminateAction {}
+
+#[derive(Default)]
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+#[non_exhaustive]
+pub struct BeforeTaskPollAction {}
+
+#[derive(Default)]
+#[allow(missing_debug_implementations, missing_docs)]
+#[cfg_attr(not(tokio_unstable), allow(unreachable_pub))]
+#[non_exhaustive]
+pub struct AfterTaskPollAction {}

tokio/src/runtime/tests/mod.rs

@@ -6,7 +6,9 @@ use self::noop_scheduler::NoopSchedule;
 use self::unowned_wrapper::unowned;
 
 mod noop_scheduler {
-    use crate::runtime::task::{self, Task, TaskHarnessScheduleHooks};
+    use crate::runtime::task::{self, Task};
+    #[cfg(tokio_unstable)]
+    use crate::runtime::{OptionalTaskHooksFactory, OptionalTaskHooksFactoryRef};
 
     /// `task::Schedule` implementation that does nothing, for testing.
     pub(crate) struct NoopSchedule;
@@ -20,10 +22,14 @@ mod noop_scheduler {
             unreachable!();
         }
 
-        fn hooks(&self) -> TaskHarnessScheduleHooks {
-            TaskHarnessScheduleHooks {
-                task_terminate_callback: None,
-            }
+        #[cfg(tokio_unstable)]
+        fn hooks_factory(&self) -> OptionalTaskHooksFactory {
+            None
+        }
+
+        #[cfg(tokio_unstable)]
+        fn hooks_factory_ref(&self) -> OptionalTaskHooksFactoryRef<'_> {
+            None
         }
     }
 }
@@ -41,6 +47,9 @@ mod unowned_wrapper {
         use tracing::Instrument;
         let span = tracing::trace_span!("test_span");
         let task = task.instrument(span);
+        #[cfg(tokio_unstable)]
+        let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next(), None);
+        #[cfg(not(tokio_unstable))]
         let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next());
         (task.into_notified(), handle)
     }
@@ -51,6 +60,9 @@ mod unowned_wrapper {
         T: std::future::Future + Send + 'static,
         T::Output: Send + 'static,
     {
+        #[cfg(tokio_unstable)]
+        let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next(), None);
+        #[cfg(not(tokio_unstable))]
         let (task, handle) = crate::runtime::task::unowned(task, NoopSchedule, Id::next());
         (task.into_notified(), handle)
     }

tokio/src/runtime/tests/queue.rs

@@ -1,5 +1,4 @@
 use crate::runtime::scheduler::multi_thread::{queue, Stats};
-
 use std::cell::RefCell;
 use std::thread;
 use std::time::Duration;

tokio/src/runtime/tests/task.rs

@@ -1,8 +1,7 @@
-use crate::runtime::task::{
-    self, unowned, Id, JoinHandle, OwnedTasks, Schedule, Task, TaskHarnessScheduleHooks,
-};
+use crate::runtime::task::{self, unowned, Id, JoinHandle, OwnedTasks, Schedule, Task};
 use crate::runtime::tests::NoopSchedule;
-
+#[cfg(tokio_unstable)]
+use crate::runtime::{OptionalTaskHooksFactory, OptionalTaskHooksFactoryRef};
 use std::collections::VecDeque;
 use std::future::Future;
 use std::sync::atomic::{AtomicBool, Ordering};
@@ -447,9 +446,13 @@ impl Schedule for Runtime {
         self.0.core.try_lock().unwrap().queue.push_back(task);
     }
 
-    fn hooks(&self) -> TaskHarnessScheduleHooks {
-        TaskHarnessScheduleHooks {
-            task_terminate_callback: None,
-        }
+    #[cfg(tokio_unstable)]
+    fn hooks_factory(&self) -> OptionalTaskHooksFactory {
+        None
+    }
+
+    #[cfg(tokio_unstable)]
+    fn hooks_factory_ref(&self) -> OptionalTaskHooksFactoryRef<'_> {
+        None
     }
 }

tokio/src/task/builder.rs

@@ -44,8 +44,12 @@ use std::{future::Future, io, mem};
 ///     loop {
 ///         let (socket, _) = listener.accept().await?;
 ///
-///         tokio::task::Builder::new()
-///             .name("tcp connection handler")
+///         let mut builder = tokio::task::Builder::new();
+///
+///         builder
+///             .name("tcp connection handler");
+///
+///         builder
 ///             .spawn(async move {
 ///                 // Process each socket concurrently.
 ///                 process(socket).await
@@ -71,8 +75,9 @@ impl<'a> Builder<'a> {
     }
 
     /// Assigns a name to the task which will be spawned.
-    pub fn name(&self, name: &'a str) -> Self {
-        Self { name: Some(name) }
+    pub fn name(&mut self, name: &'a str) -> &mut Self {
+        self.name = Some(name);
+        self
     }
 
     /// Spawns a task with this builder's settings on the current runtime.
@@ -91,9 +96,9 @@ impl<'a> Builder<'a> {
     {
         let fut_size = mem::size_of::<Fut>();
         Ok(if fut_size > BOX_FUTURE_THRESHOLD {
-            super::spawn::spawn_inner(Box::pin(future), SpawnMeta::new(self.name, fut_size))
+            super::spawn::spawn_inner(Box::pin(future), SpawnMeta::new(self.name, fut_size), None)
         } else {
-            super::spawn::spawn_inner(future, SpawnMeta::new(self.name, fut_size))
+            super::spawn::spawn_inner(future, SpawnMeta::new(self.name, fut_size), None)
         })
     }
 
@@ -112,9 +117,9 @@ impl<'a> Builder<'a> {
     {
         let fut_size = mem::size_of::<Fut>();
         Ok(if fut_size > BOX_FUTURE_THRESHOLD {
-            handle.spawn_named(Box::pin(future), SpawnMeta::new(self.name, fut_size))
+            handle.spawn_named(Box::pin(future), SpawnMeta::new(self.name, fut_size), None)
         } else {
-            handle.spawn_named(future, SpawnMeta::new(self.name, fut_size))
+            handle.spawn_named(future, SpawnMeta::new(self.name, fut_size), None)
         })
     }
 
@@ -140,9 +145,13 @@ impl<'a> Builder<'a> {
     {
         let fut_size = mem::size_of::<Fut>();
         Ok(if fut_size > BOX_FUTURE_THRESHOLD {
-            super::local::spawn_local_inner(Box::pin(future), SpawnMeta::new(self.name, fut_size))
+            super::local::spawn_local_inner(
+                Box::pin(future),
+                SpawnMeta::new(self.name, fut_size),
+                None,
+            )
         } else {
-            super::local::spawn_local_inner(future, SpawnMeta::new(self.name, fut_size))
+            super::local::spawn_local_inner(future, SpawnMeta::new(self.name, fut_size), None)
         })
     }
 

tokio/src/task/join_set.rs

@@ -641,9 +641,13 @@ where
 #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))]
 impl<'a, T: 'static> Builder<'a, T> {
     /// Assigns a name to the task which will be spawned.
-    pub fn name(self, name: &'a str) -> Self {
-        let builder = self.builder.name(name);
-        Self { builder, ..self }
+    pub fn name(mut self, name: &'a str) -> Self {
+        self.builder.name(name);
+
+        Self {
+            builder: self.builder,
+            ..self
+        }
     }
 
     /// Spawn the provided task with this builder's settings and store it in the

tokio/src/task/local.rs

@@ -1,9 +1,11 @@
 //! Runs `!Send` futures on the current thread.
 use crate::loom::cell::UnsafeCell;
 use crate::loom::sync::{Arc, Mutex};
+use crate::runtime::task::{self, JoinHandle, LocalOwnedTasks, Task};
 #[cfg(tokio_unstable)]
-use crate::runtime;
-use crate::runtime::task::{self, JoinHandle, LocalOwnedTasks, Task, TaskHarnessScheduleHooks};
+use crate::runtime::{
+    self, OptionalTaskHooks, OptionalTaskHooksFactory, OptionalTaskHooksFactoryRef,
+};
 use crate::runtime::{context, ThreadId, BOX_FUTURE_THRESHOLD};
 use crate::sync::AtomicWaker;
 use crate::util::trace::SpawnMeta;
@@ -371,6 +373,13 @@ cfg_rt! {
         F::Output: 'static,
     {
         let fut_size = std::mem::size_of::<F>();
+        #[cfg(tokio_unstable)]
+        if fut_size > BOX_FUTURE_THRESHOLD {
+            spawn_local_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size), None)
+        } else {
+            spawn_local_inner(future, SpawnMeta::new_unnamed(fut_size), None)
+        }
+        #[cfg(not(tokio_unstable))]
         if fut_size > BOX_FUTURE_THRESHOLD {
             spawn_local_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size))
         } else {
@@ -380,7 +389,7 @@ cfg_rt! {
 
 
     #[track_caller]
-    pub(super) fn spawn_local_inner<F>(future: F, meta: SpawnMeta<'_>) -> JoinHandle<F::Output>
+    pub(super) fn spawn_local_inner<F>(future: F, meta: SpawnMeta<'_>, #[cfg(tokio_unstable)] hooks_override: OptionalTaskHooks) -> JoinHandle<F::Output>
     where F: Future + 'static,
           F::Output: 'static
     {
@@ -412,6 +421,9 @@ cfg_rt! {
                 let task = crate::util::trace::task(future, "task", meta, id.as_u64());
 
                 // safety: we have verified that this is a `LocalRuntime` owned by the current thread
+                #[cfg(tokio_unstable)]
+                unsafe { handle.spawn_local(task, id, hooks_override) }
+                #[cfg(not(tokio_unstable))]
                 unsafe { handle.spawn_local(task, id) }
             } else {
                 match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) {
@@ -1004,6 +1016,15 @@ impl Context {
         let future = crate::util::trace::task(future, "local", meta, id.as_u64());
 
         // Safety: called from the thread that owns the `LocalSet`
+        #[cfg(tokio_unstable)]
+        let (handle, notified) = {
+            self.shared.local_state.assert_called_from_owner_thread();
+            self.shared
+                .local_state
+                .owned
+                .bind(future, self.shared.clone(), id, None)
+        };
+        #[cfg(not(tokio_unstable))]
         let (handle, notified) = {
             self.shared.local_state.assert_called_from_owner_thread();
             self.shared
@@ -1117,11 +1138,15 @@ impl task::Schedule for Arc<Shared> {
         Shared::schedule(self, task);
     }
 
-    // localset does not currently support task hooks
-    fn hooks(&self) -> TaskHarnessScheduleHooks {
-        TaskHarnessScheduleHooks {
-            task_terminate_callback: None,
-        }
+    #[cfg(tokio_unstable)]
+    fn hooks_factory(&self) -> OptionalTaskHooksFactory {
+        None
+    }
+
+    // localset does not support task hooks
+    #[cfg(tokio_unstable)]
+    fn hooks_factory_ref(&self) -> OptionalTaskHooksFactoryRef<'_> {
+        None
     }
 
     cfg_unstable! {

tokio/src/task/mod.rs

@@ -311,6 +311,10 @@ cfg_rt! {
 
     pub use crate::runtime::task::{Id, id, try_id};
 
+    cfg_unstable! {
+        pub use spawn::spawn_with_hooks;
+    }
+
     cfg_trace! {
         mod builder;
         pub use builder::Builder;

tokio/src/task/spawn.rs

@@ -1,4 +1,6 @@
 use crate::runtime::BOX_FUTURE_THRESHOLD;
+#[cfg(tokio_unstable)]
+use crate::runtime::{OptionalTaskHooks, TaskHookHarness};
 use crate::task::JoinHandle;
 use crate::util::trace::SpawnMeta;
 
@@ -169,6 +171,13 @@ cfg_rt! {
         F::Output: Send + 'static,
     {
         let fut_size = std::mem::size_of::<F>();
+        #[cfg(tokio_unstable)]
+        if fut_size > BOX_FUTURE_THRESHOLD {
+            spawn_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size), None)
+        } else {
+            spawn_inner(future, SpawnMeta::new_unnamed(fut_size), None)
+        }
+        #[cfg(not(tokio_unstable))]
         if fut_size > BOX_FUTURE_THRESHOLD {
             spawn_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size))
         } else {
@@ -176,8 +185,26 @@ cfg_rt! {
         }
     }
 
+    /// Spawn a future with a custom set of task hooks
     #[track_caller]
-    pub(super) fn spawn_inner<T>(future: T, meta: SpawnMeta<'_>) -> JoinHandle<T::Output>
+    #[cfg(tokio_unstable)]
+    pub fn spawn_with_hooks<F, T>(future: F, hooks: T) -> JoinHandle<F::Output>
+    where
+        F: Future + Send + 'static,
+        F::Output: Send + 'static,
+        T: TaskHookHarness + Send + Sync + 'static,
+    {
+        let fut_size = std::mem::size_of::<F>();
+
+        if fut_size > BOX_FUTURE_THRESHOLD {
+            spawn_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size), Some(Box::new(hooks)))
+        } else {
+            spawn_inner(future, SpawnMeta::new_unnamed(fut_size), Some(Box::new(hooks)))
+        }
+    }
+
+    #[track_caller]
+    pub(super) fn spawn_inner<T>(future: T, meta: SpawnMeta<'_>, #[cfg(tokio_unstable)] hooks_override: OptionalTaskHooks) -> JoinHandle<T::Output>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
@@ -199,6 +226,13 @@ cfg_rt! {
         let id = task::Id::next();
         let task = crate::util::trace::task(future, "task", meta, id.as_u64());
 
+        #[cfg(tokio_unstable)]
+        return match context::with_current(|handle| handle.spawn(task, id, hooks_override)) {
+            Ok(join_handle) => join_handle,
+            Err(e) => panic!("{}", e),
+        };
+
+        #[cfg(not(tokio_unstable))]
         match context::with_current(|handle| handle.spawn(task, id)) {
             Ok(join_handle) => join_handle,
             Err(e) => panic!("{}", e),

tokio/tests/rt_metrics.rs

@@ -4,6 +4,7 @@
 use std::sync::mpsc;
 use std::time::Duration;
 use tokio::runtime::Runtime;
+use tokio::time;
 
 #[test]
 fn num_workers() {
@@ -125,6 +126,64 @@ fn worker_total_busy_duration() {
     }
 }
 
+#[test]
+fn worker_park_count() {
+    let rt = current_thread();
+    let metrics = rt.metrics();
+    rt.block_on(async {
+        time::sleep(Duration::from_millis(1)).await;
+    });
+    drop(rt);
+    assert!(1 <= metrics.worker_park_count(0));
+
+    let rt = threaded();
+    let metrics = rt.metrics();
+    rt.block_on(async {
+        time::sleep(Duration::from_millis(1)).await;
+    });
+    drop(rt);
+    assert!(1 <= metrics.worker_park_count(0));
+    assert!(1 <= metrics.worker_park_count(1));
+}
+
+#[test]
+fn worker_park_unpark_count() {
+    let rt = current_thread();
+    let metrics = rt.metrics();
+    rt.block_on(rt.spawn(async {})).unwrap();
+    drop(rt);
+    assert!(2 <= metrics.worker_park_unpark_count(0));
+
+    let rt = threaded();
+    let metrics = rt.metrics();
+
+    // Wait for workers to be parked after runtime startup.
+    for _ in 0..100 {
+        if 1 <= metrics.worker_park_unpark_count(0) && 1 <= metrics.worker_park_unpark_count(1) {
+            break;
+        }
+        std::thread::sleep(std::time::Duration::from_millis(100));
+    }
+    assert_eq!(1, metrics.worker_park_unpark_count(0));
+    assert_eq!(1, metrics.worker_park_unpark_count(1));
+
+    // Spawn a task to unpark and then park a worker.
+    rt.block_on(rt.spawn(async {})).unwrap();
+    for _ in 0..100 {
+        if 3 <= metrics.worker_park_unpark_count(0) || 3 <= metrics.worker_park_unpark_count(1) {
+            break;
+        }
+        std::thread::sleep(std::time::Duration::from_millis(100));
+    }
+    assert!(3 <= metrics.worker_park_unpark_count(0) || 3 <= metrics.worker_park_unpark_count(1));
+
+    // Both threads unpark for runtime shutdown.
+    drop(rt);
+    assert_eq!(0, metrics.worker_park_unpark_count(0) % 2);
+    assert_eq!(0, metrics.worker_park_unpark_count(1) % 2);
+    assert!(4 <= metrics.worker_park_unpark_count(0) || 4 <= metrics.worker_park_unpark_count(1));
+}
+
 fn try_block_threaded(rt: &Runtime) -> Result<Vec<mpsc::Sender<()>>, mpsc::RecvTimeoutError> {
     let (tx, rx) = mpsc::channel();
 

tokio/tests/rt_poll_callbacks.rs

@@ -1,128 +0,0 @@
-#![allow(unknown_lints, unexpected_cfgs)]
-#![cfg(tokio_unstable)]
-
-use std::sync::{atomic::AtomicUsize, Arc, Mutex};
-
-use tokio::task::yield_now;
-
-#[cfg(not(target_os = "wasi"))]
-#[test]
-fn callbacks_fire_multi_thread() {
-    let poll_start_counter = Arc::new(AtomicUsize::new(0));
-    let poll_stop_counter = Arc::new(AtomicUsize::new(0));
-    let poll_start = poll_start_counter.clone();
-    let poll_stop = poll_stop_counter.clone();
-
-    let before_task_poll_callback_task_id: Arc<Mutex<Option<tokio::task::Id>>> =
-        Arc::new(Mutex::new(None));
-    let after_task_poll_callback_task_id: Arc<Mutex<Option<tokio::task::Id>>> =
-        Arc::new(Mutex::new(None));
-
-    let before_task_poll_callback_task_id_ref = Arc::clone(&before_task_poll_callback_task_id);
-    let after_task_poll_callback_task_id_ref = Arc::clone(&after_task_poll_callback_task_id);
-    let rt = tokio::runtime::Builder::new_multi_thread()
-        .enable_all()
-        .on_before_task_poll(move |task_meta| {
-            before_task_poll_callback_task_id_ref
-                .lock()
-                .unwrap()
-                .replace(task_meta.id());
-            poll_start_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-        })
-        .on_after_task_poll(move |task_meta| {
-            after_task_poll_callback_task_id_ref
-                .lock()
-                .unwrap()
-                .replace(task_meta.id());
-            poll_stop_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-        })
-        .build()
-        .unwrap();
-    let task = rt.spawn(async {
-        yield_now().await;
-        yield_now().await;
-        yield_now().await;
-    });
-
-    let spawned_task_id = task.id();
-
-    rt.block_on(task).expect("task should succeed");
-    // We need to drop the runtime to guarantee the workers have exited (and thus called the callback)
-    drop(rt);
-
-    assert_eq!(
-        before_task_poll_callback_task_id.lock().unwrap().unwrap(),
-        spawned_task_id
-    );
-    assert_eq!(
-        after_task_poll_callback_task_id.lock().unwrap().unwrap(),
-        spawned_task_id
-    );
-    let actual_count = 4;
-    assert_eq!(
-        poll_start.load(std::sync::atomic::Ordering::Relaxed),
-        actual_count,
-        "unexpected number of poll starts"
-    );
-    assert_eq!(
-        poll_stop.load(std::sync::atomic::Ordering::Relaxed),
-        actual_count,
-        "unexpected number of poll stops"
-    );
-}
-
-#[test]
-fn callbacks_fire_current_thread() {
-    let poll_start_counter = Arc::new(AtomicUsize::new(0));
-    let poll_stop_counter = Arc::new(AtomicUsize::new(0));
-    let poll_start = poll_start_counter.clone();
-    let poll_stop = poll_stop_counter.clone();
-
-    let before_task_poll_callback_task_id: Arc<Mutex<Option<tokio::task::Id>>> =
-        Arc::new(Mutex::new(None));
-    let after_task_poll_callback_task_id: Arc<Mutex<Option<tokio::task::Id>>> =
-        Arc::new(Mutex::new(None));
-
-    let before_task_poll_callback_task_id_ref = Arc::clone(&before_task_poll_callback_task_id);
-    let after_task_poll_callback_task_id_ref = Arc::clone(&after_task_poll_callback_task_id);
-    let rt = tokio::runtime::Builder::new_current_thread()
-        .enable_all()
-        .on_before_task_poll(move |task_meta| {
-            before_task_poll_callback_task_id_ref
-                .lock()
-                .unwrap()
-                .replace(task_meta.id());
-            poll_start_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-        })
-        .on_after_task_poll(move |task_meta| {
-            after_task_poll_callback_task_id_ref
-                .lock()
-                .unwrap()
-                .replace(task_meta.id());
-            poll_stop_counter.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
-        })
-        .build()
-        .unwrap();
-
-    let task = rt.spawn(async {
-        yield_now().await;
-        yield_now().await;
-        yield_now().await;
-    });
-
-    let spawned_task_id = task.id();
-
-    let _ = rt.block_on(task);
-    drop(rt);
-
-    assert_eq!(
-        before_task_poll_callback_task_id.lock().unwrap().unwrap(),
-        spawned_task_id
-    );
-    assert_eq!(
-        after_task_poll_callback_task_id.lock().unwrap().unwrap(),
-        spawned_task_id
-    );
-    assert_eq!(poll_start.load(std::sync::atomic::Ordering::Relaxed), 4);
-    assert_eq!(poll_stop.load(std::sync::atomic::Ordering::Relaxed), 4);
-}

tokio/tests/rt_unstable_metrics.rs

@@ -212,64 +212,6 @@ fn worker_thread_id_threaded() {
     .unwrap()
 }
 
-#[test]
-fn worker_park_count() {
-    let rt = current_thread();
-    let metrics = rt.metrics();
-    rt.block_on(async {
-        time::sleep(Duration::from_millis(1)).await;
-    });
-    drop(rt);
-    assert!(1 <= metrics.worker_park_count(0));
-
-    let rt = threaded();
-    let metrics = rt.metrics();
-    rt.block_on(async {
-        time::sleep(Duration::from_millis(1)).await;
-    });
-    drop(rt);
-    assert!(1 <= metrics.worker_park_count(0));
-    assert!(1 <= metrics.worker_park_count(1));
-}
-
-#[test]
-fn worker_park_unpark_count() {
-    let rt = current_thread();
-    let metrics = rt.metrics();
-    rt.block_on(rt.spawn(async {})).unwrap();
-    drop(rt);
-    assert!(2 <= metrics.worker_park_unpark_count(0));
-
-    let rt = threaded();
-    let metrics = rt.metrics();
-
-    // Wait for workers to be parked after runtime startup.
-    for _ in 0..100 {
-        if 1 <= metrics.worker_park_unpark_count(0) && 1 <= metrics.worker_park_unpark_count(1) {
-            break;
-        }
-        std::thread::sleep(std::time::Duration::from_millis(100));
-    }
-    assert_eq!(1, metrics.worker_park_unpark_count(0));
-    assert_eq!(1, metrics.worker_park_unpark_count(1));
-
-    // Spawn a task to unpark and then park a worker.
-    rt.block_on(rt.spawn(async {})).unwrap();
-    for _ in 0..100 {
-        if 3 <= metrics.worker_park_unpark_count(0) || 3 <= metrics.worker_park_unpark_count(1) {
-            break;
-        }
-        std::thread::sleep(std::time::Duration::from_millis(100));
-    }
-    assert!(3 <= metrics.worker_park_unpark_count(0) || 3 <= metrics.worker_park_unpark_count(1));
-
-    // Both threads unpark for runtime shutdown.
-    drop(rt);
-    assert_eq!(0, metrics.worker_park_unpark_count(0) % 2);
-    assert_eq!(0, metrics.worker_park_unpark_count(1) % 2);
-    assert!(4 <= metrics.worker_park_unpark_count(0) || 4 <= metrics.worker_park_unpark_count(1));
-}
-
 #[test]
 fn worker_noop_count() {
     // There isn't really a great way to generate no-op parks as they happen as

tokio/tests/task_builder.rs

@@ -8,22 +8,22 @@ use tokio::{
 
 #[test]
 async fn spawn_with_name() {
-    let result = Builder::new()
-        .name("name")
-        .spawn(async { "task executed" })
-        .unwrap()
-        .await;
+    let mut b = Builder::new();
+
+    b.name("name");
+
+    let result = b.spawn(async { "task executed" }).unwrap().await;
 
     assert_eq!(result.unwrap(), "task executed");
 }
 
 #[test]
 async fn spawn_blocking_with_name() {
-    let result = Builder::new()
-        .name("name")
-        .spawn_blocking(|| "task executed")
-        .unwrap()
-        .await;
+    let mut b = Builder::new();
+
+    b.name("name");
+
+    let result = b.spawn_blocking(|| "task executed").unwrap().await;
 
     assert_eq!(result.unwrap(), "task executed");
 }
@@ -33,11 +33,11 @@ async fn spawn_local_with_name() {
     let unsend_data = Rc::new("task executed");
     let result = LocalSet::new()
         .run_until(async move {
-            Builder::new()
-                .name("name")
-                .spawn_local(async move { unsend_data })
-                .unwrap()
-                .await
+            let mut b = Builder::new();
+
+            b.name("name");
+
+            b.spawn_local(async move { unsend_data }).unwrap().await
         })
         .await;
 

tokio/tests/task_hooks.rs

@@ -1,75 +1,471 @@
-#![warn(rust_2018_idioms)]
-#![cfg(all(feature = "full", tokio_unstable, target_has_atomic = "64"))]
+#![cfg(all(
+    feature = "full",
+    tokio_unstable,
+    target_has_atomic = "64",
+    not(target_arch = "wasm32")
+))]
 
-use std::collections::HashSet;
 use std::sync::atomic::{AtomicUsize, Ordering};
-use std::sync::{Arc, Mutex};
+use std::sync::Arc;
+use tokio::runtime;
+use tokio::runtime::{
+    AfterTaskPollAction, AfterTaskPollContext, BeforeTaskPollAction, BeforeTaskPollContext,
+    OnChildSpawnAction, OnChildTaskSpawnContext, OnTaskTerminateAction, OnTaskTerminateContext,
+    OnTopLevelSpawnAction, OnTopLevelTaskSpawnContext, TaskHookHarness, TaskHookHarnessFactory,
+};
 
-use tokio::runtime::Builder;
-
-const TASKS: usize = 8;
-const ITERATIONS: usize = 64;
-/// Assert that the spawn task hook always fires when set.
 #[test]
-fn spawn_task_hook_fires() {
-    let count = Arc::new(AtomicUsize::new(0));
-    let count2 = Arc::clone(&count);
+fn runtime_default_factory() {
+    let ct = runtime::Builder::new_current_thread();
+    let mt = runtime::Builder::new_multi_thread();
 
-    let ids = Arc::new(Mutex::new(HashSet::new()));
-    let ids2 = Arc::clone(&ids);
-
-    let runtime = Builder::new_current_thread()
-        .on_task_spawn(move |data| {
-            ids2.lock().unwrap().insert(data.id());
-
-            count2.fetch_add(1, Ordering::SeqCst);
-        })
-        .build()
-        .unwrap();
-
-    for _ in 0..TASKS {
-        runtime.spawn(std::future::pending::<()>());
-    }
-
-    let count_realized = count.load(Ordering::SeqCst);
-    assert_eq!(
-        TASKS, count_realized,
-        "Total number of spawned task hook invocations was incorrect, expected {TASKS}, got {}",
-        count_realized
-    );
-
-    let count_ids_realized = ids.lock().unwrap().len();
-
-    assert_eq!(
-        TASKS, count_ids_realized,
-        "Total number of spawned task hook invocations was incorrect, expected {TASKS}, got {}",
-        count_realized
-    );
+    run_runtime_default_factory(ct);
+    run_runtime_default_factory(mt);
 }
 
-/// Assert that the terminate task hook always fires when set.
 #[test]
-fn terminate_task_hook_fires() {
-    let count = Arc::new(AtomicUsize::new(0));
-    let count2 = Arc::clone(&count);
+fn parent_child_chaining() {
+    let ct = runtime::Builder::new_current_thread();
+    let mt = runtime::Builder::new_multi_thread();
 
-    let runtime = Builder::new_current_thread()
-        .on_task_terminate(move |_data| {
-            count2.fetch_add(1, Ordering::SeqCst);
+    run_parent_child_chaining(ct);
+    run_parent_child_chaining(mt);
+}
+
+#[test]
+fn before_poll() {
+    let ct = runtime::Builder::new_current_thread();
+    let mt = runtime::Builder::new_multi_thread();
+
+    run_before_poll(ct);
+    run_before_poll(mt);
+}
+
+#[test]
+fn after_poll() {
+    let ct = runtime::Builder::new_current_thread();
+    let mt = runtime::Builder::new_multi_thread();
+
+    run_after_poll(ct);
+    run_after_poll(mt);
+}
+
+#[test]
+fn terminate() {
+    let ct = runtime::Builder::new_current_thread();
+
+    run_terminate(ct);
+}
+
+#[test]
+fn hook_switching() {
+    let ct = runtime::Builder::new_current_thread();
+    let mt = runtime::Builder::new_multi_thread();
+
+    run_hook_switching(ct);
+    run_hook_switching(mt);
+}
+
+#[test]
+fn override_hooks() {
+    let ct = runtime::Builder::new_current_thread();
+    let mt = runtime::Builder::new_multi_thread();
+
+    run_override(ct);
+    run_override(mt);
+}
+
+fn run_runtime_default_factory(mut builder: runtime::Builder) {
+    struct TestFactory {
+        counter: Arc<AtomicUsize>,
+    }
+
+    impl TaskHookHarnessFactory for TestFactory {
+        fn on_top_level_spawn(
+            &self,
+            _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
+        ) -> OnTopLevelSpawnAction {
+            self.counter.fetch_add(1, Ordering::SeqCst);
+
+            Default::default()
+        }
+    }
+
+    let counter = Arc::new(AtomicUsize::new(0));
+
+    let rt = builder
+        .hook_harness_factory(TestFactory {
+            counter: counter.clone(),
         })
         .build()
         .unwrap();
 
-    for _ in 0..TASKS {
-        runtime.spawn(std::future::ready(()));
+    rt.spawn(async {});
+
+    assert_eq!(counter.load(Ordering::SeqCst), 1);
+
+    let handle = rt.handle();
+
+    handle.spawn(async {});
+
+    assert_eq!(counter.load(Ordering::SeqCst), 2);
+
+    rt.block_on(async {});
+
+    assert_eq!(counter.load(Ordering::SeqCst), 2);
+
+    rt.block_on(async { tokio::spawn(async {}) });
+
+    assert_eq!(counter.load(Ordering::SeqCst), 3);
+
+    // block on a future which spawns a future and waits for it, which in turn spawns another future
+    //
+    // this checks that stuff works from on-worker within a multithreaded runtime
+    let _ = rt.block_on(async { tokio::spawn(async { tokio::spawn(async {}) }).await });
+
+    assert_eq!(counter.load(Ordering::SeqCst), 5);
+}
+
+fn run_parent_child_chaining(mut builder: runtime::Builder) {
+    struct TestFactory {
+        parent_spawns: Arc<AtomicUsize>,
+        child_spawns: Arc<AtomicUsize>,
     }
 
-    runtime.block_on(async {
-        // tick the runtime a bunch to close out tasks
-        for _ in 0..ITERATIONS {
-            tokio::task::yield_now().await;
+    struct TestHooks {
+        spawns: Arc<AtomicUsize>,
+    }
+
+    impl TaskHookHarnessFactory for TestFactory {
+        fn on_top_level_spawn(
+            &self,
+            _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
+        ) -> OnTopLevelSpawnAction {
+            self.parent_spawns.fetch_add(1, Ordering::SeqCst);
+
+            let mut a = OnTopLevelSpawnAction::default();
+
+            a.set_hooks(TestHooks {
+                spawns: self.child_spawns.clone(),
+            });
+
+            a
         }
+    }
+
+    impl TaskHookHarness for TestHooks {
+        fn on_child_spawn(&mut self, _ctx: &mut OnChildTaskSpawnContext<'_>) -> OnChildSpawnAction {
+            self.spawns.fetch_add(1, Ordering::SeqCst);
+
+            let mut a = OnChildSpawnAction::default();
+
+            a.set_hooks(Self {
+                spawns: self.spawns.clone(),
+            });
+
+            a
+        }
+    }
+
+    let parent_spawns = Arc::new(AtomicUsize::new(0));
+    let child_spawns = Arc::new(AtomicUsize::new(0));
+
+    let rt = builder
+        .hook_harness_factory(TestFactory {
+            parent_spawns: parent_spawns.clone(),
+            child_spawns: child_spawns.clone(),
+        })
+        .build()
+        .unwrap();
+
+    rt.spawn(async {});
+
+    assert_eq!(parent_spawns.load(Ordering::SeqCst), 1);
+    assert_eq!(child_spawns.load(Ordering::SeqCst), 0);
+
+    let _ = rt.block_on(async { tokio::spawn(async { tokio::spawn(async {}) }).await });
+
+    assert_eq!(parent_spawns.load(Ordering::SeqCst), 2);
+    assert_eq!(child_spawns.load(Ordering::SeqCst), 1);
+}
+
+fn run_before_poll(mut builder: runtime::Builder) {
+    struct TestFactory {
+        polls: Arc<AtomicUsize>,
+    }
+
+    struct TestHooks {
+        polls: Arc<AtomicUsize>,
+    }
+
+    impl TaskHookHarnessFactory for TestFactory {
+        fn on_top_level_spawn(
+            &self,
+            _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
+        ) -> OnTopLevelSpawnAction {
+            let mut a = OnTopLevelSpawnAction::default();
+
+            a.set_hooks(TestHooks {
+                polls: self.polls.clone(),
+            });
+
+            a
+        }
+    }
+
+    impl TaskHookHarness for TestHooks {
+        fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) -> BeforeTaskPollAction {
+            self.polls.fetch_add(1, Ordering::SeqCst);
+
+            Default::default()
+        }
+    }
+
+    let polls = Arc::new(AtomicUsize::new(0));
+
+    let rt = builder
+        .hook_harness_factory(TestFactory {
+            polls: polls.clone(),
+        })
+        .build()
+        .unwrap();
+
+    rt.block_on(async {});
+    assert_eq!(polls.load(Ordering::SeqCst), 0);
+
+    let _ = rt.block_on(async { tokio::spawn(async {}).await });
+    assert_eq!(polls.load(Ordering::SeqCst), 1);
+
+    let _ = rt.block_on(async { tokio::spawn(async { tokio::spawn(async {}).await }).await });
+    assert_eq!(polls.load(Ordering::SeqCst), 4);
+}
+
+fn run_after_poll(mut builder: runtime::Builder) {
+    struct TestFactory {
+        polls: Arc<AtomicUsize>,
+    }
+
+    struct TestHooks {
+        polls: Arc<AtomicUsize>,
+    }
+
+    impl TaskHookHarnessFactory for TestFactory {
+        fn on_top_level_spawn(
+            &self,
+            _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
+        ) -> OnTopLevelSpawnAction {
+            let mut a = OnTopLevelSpawnAction::default();
+
+            a.set_hooks(TestHooks {
+                polls: self.polls.clone(),
+            });
+
+            a
+        }
+    }
+
+    impl TaskHookHarness for TestHooks {
+        fn after_poll(&mut self, _ctx: &mut AfterTaskPollContext<'_>) -> AfterTaskPollAction {
+            self.polls.fetch_add(1, Ordering::SeqCst);
+
+            Default::default()
+        }
+    }
+
+    let polls = Arc::new(AtomicUsize::new(0));
+
+    let rt = builder
+        .hook_harness_factory(TestFactory {
+            polls: polls.clone(),
+        })
+        .build()
+        .unwrap();
+
+    rt.block_on(async {});
+    assert_eq!(polls.load(Ordering::SeqCst), 0);
+
+    let _ = rt.block_on(async { tokio::spawn(async {}).await });
+    assert_eq!(polls.load(Ordering::SeqCst), 1);
+
+    let _ = rt.block_on(async { tokio::spawn(async { tokio::spawn(async {}).await }).await });
+    assert_eq!(polls.load(Ordering::SeqCst), 4);
+}
+
+fn run_terminate(mut builder: runtime::Builder) {
+    struct TestFactory {
+        terminations: Arc<AtomicUsize>,
+    }
+
+    struct TestHooks {
+        terminations: Arc<AtomicUsize>,
+    }
+
+    impl TaskHookHarnessFactory for TestFactory {
+        fn on_top_level_spawn(
+            &self,
+            _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
+        ) -> OnTopLevelSpawnAction {
+            let mut a = OnTopLevelSpawnAction::default();
+
+            a.set_hooks(TestHooks {
+                terminations: self.terminations.clone(),
+            });
+
+            a
+        }
+    }
+
+    impl TaskHookHarness for TestHooks {
+        fn on_task_terminate(
+            &mut self,
+            _ctx: &mut OnTaskTerminateContext<'_>,
+        ) -> OnTaskTerminateAction {
+            self.terminations.fetch_add(1, Ordering::SeqCst);
+
+            Default::default()
+        }
+    }
+
+    let terminations = Arc::new(AtomicUsize::new(0));
+
+    let rt = builder
+        .hook_harness_factory(TestFactory {
+            terminations: terminations.clone(),
+        })
+        .build()
+        .unwrap();
+
+    let _ = rt.block_on(async { tokio::spawn(async { tokio::spawn(async {}).await }).await });
+
+    assert_eq!(terminations.load(Ordering::SeqCst), 2);
+}
+
+fn run_hook_switching(mut builder: runtime::Builder) {
+    struct TestFactory {
+        next_id: Arc<AtomicUsize>,
+        flag: Arc<AtomicUsize>,
+    }
+
+    struct TestHooks {
+        id: usize,
+        flag: Arc<AtomicUsize>,
+    }
+
+    impl TaskHookHarnessFactory for TestFactory {
+        fn on_top_level_spawn(
+            &self,
+            _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
+        ) -> OnTopLevelSpawnAction {
+            let mut a = OnTopLevelSpawnAction::default();
+
+            a.set_hooks(TestHooks {
+                id: self.next_id.fetch_add(1, Ordering::SeqCst),
+                flag: self.flag.clone(),
+            });
+
+            a
+        }
+    }
+
+    impl TaskHookHarness for TestHooks {
+        fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) -> BeforeTaskPollAction {
+            self.flag.store(self.id, Ordering::SeqCst);
+
+            Default::default()
+        }
+    }
+
+    let polls = Arc::new(AtomicUsize::new(0));
+
+    let rt = builder
+        .hook_harness_factory(TestFactory {
+            next_id: Arc::new(Default::default()),
+            flag: polls.clone(),
+        })
+        .build()
+        .unwrap();
+
+    let _ = rt.block_on(async { tokio::spawn(async {}).await });
+    assert_eq!(polls.load(Ordering::SeqCst), 0);
+
+    let _ = rt.block_on(async { tokio::spawn(async { tokio::spawn(async {}).await }).await });
+    assert_eq!(polls.load(Ordering::SeqCst), 1);
+
+    let _ = rt.block_on(async { tokio::spawn(async {}).await });
+    assert_eq!(polls.load(Ordering::SeqCst), 3);
+}
+
+fn run_override(mut builder: runtime::Builder) {
+    struct TestFactory {
+        counter: Arc<AtomicUsize>,
+    }
+
+    struct TestHooks {
+        counter: Arc<AtomicUsize>,
+    }
+
+    impl TaskHookHarness for TestHooks {
+        fn before_poll(&mut self, _ctx: &mut BeforeTaskPollContext<'_>) -> BeforeTaskPollAction {
+            self.counter.fetch_add(1, Ordering::SeqCst);
+
+            Default::default()
+        }
+
+        fn on_child_spawn(&mut self, _ctx: &mut OnChildTaskSpawnContext<'_>) -> OnChildSpawnAction {
+            let mut a = OnChildSpawnAction::default();
+
+            a.set_hooks(Self {
+                counter: self.counter.clone(),
+            });
+
+            a
+        }
+    }
+
+    impl TaskHookHarnessFactory for TestFactory {
+        fn on_top_level_spawn(
+            &self,
+            _ctx: &mut OnTopLevelTaskSpawnContext<'_>,
+        ) -> OnTopLevelSpawnAction {
+            self.counter.fetch_add(1, Ordering::SeqCst);
+
+            Default::default()
+        }
+    }
+
+    let factory_counter = Arc::new(AtomicUsize::new(0));
+    let builder_counter = Arc::new(AtomicUsize::new(0));
+
+    let rt = builder
+        .hook_harness_factory(TestFactory {
+            counter: factory_counter.clone(),
+        })
+        .build()
+        .unwrap();
+
+    rt.spawn(async {});
+
+    assert_eq!(factory_counter.load(Ordering::SeqCst), 1);
+
+    let _ = rt.block_on(async {
+        tokio::task::spawn_with_hooks(
+            async {},
+            TestHooks {
+                counter: builder_counter.clone(),
+            },
+        )
+        .await
     });
 
-    assert_eq!(TASKS, count.load(Ordering::SeqCst));
+    assert_eq!(factory_counter.load(Ordering::SeqCst), 1);
+    assert_eq!(builder_counter.load(Ordering::SeqCst), 1);
+
+    let _ = rt.block_on(async {
+        let counter = builder_counter.clone();
+        tokio::spawn(async { tokio::task::spawn_with_hooks(async {}, TestHooks { counter }).await })
+            .await
+    });
+
+    assert_eq!(factory_counter.load(Ordering::SeqCst), 2);
+    assert_eq!(builder_counter.load(Ordering::SeqCst), 2);
 }

tokio/tests/tracing_task.rs

@@ -64,9 +64,11 @@ async fn task_builder_name_recorded() {
 
     {
         let _guard = tracing::subscriber::set_default(subscriber);
-        task::Builder::new()
-            .name("test-task")
-            .spawn(futures::future::ready(()))
+        let mut b = task::Builder::new();
+
+        b.name("test-task");
+
+        b.spawn(futures::future::ready(()))
             .unwrap()
             .await
             .expect("failed to await join handle");