rt: add optional per-task user data exposed to hooks

In tokio-rs#7197 and tokio-rs#7306, improved capabilities of task hooks were
discussed and an initial implementation provided. However, that
involved quite wide-reaching changes, modifying every spawn site
and introducing a global map to provide the full inheritance
capabilities originally proposed.

This is the first part of a more basic version where we only use
the existing hooks and provide the capabilities for consumers to
be able to implement more complex relationships if needed, just
adding an optional user data ref to the task header.

The additional data is 2*usize, and is not enough to result in
the struct requiring more than one cache line.

A user is now able to use their own global map to build inheritance
capabilities if needed, and this would be made simpler by also
exposing the current task user data to the on_task_spawn hook,
which a followup will look to do.

Author: LloydW93

tokio/src/runtime/builder.rs

@@ -3,7 +3,10 @@
 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::runtime::{
+    metrics::HistogramConfiguration, LocalOptions, LocalRuntime, TaskMeta, TaskSpawnCallback,
+    UserData,
+};
 use crate::util::rand::{RngSeed, RngSeedGenerator};
 
 use crate::runtime::blocking::BlockingPool;
@@ -89,6 +92,9 @@ pub struct Builder {
     pub(super) after_unpark: Option<Callback>,
 
     /// To run before each task is spawned.
+    #[cfg(tokio_unstable)]
+    pub(super) before_spawn: Option<TaskSpawnCallback>,
+    #[cfg(not(tokio_unstable))]
     pub(super) before_spawn: Option<TaskCallback>,
 
     /// To run before each poll
@@ -731,8 +737,15 @@ impl Builder {
     /// Executes function `f` just before a task is spawned.
     ///
     /// `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.
+    /// [`tokio::spawn`](crate::spawn) can be called, and may result in this callback
+    /// being invoked immediately.
+    ///
+    /// `f` must return an `Option<&'static dyn Any>`. A value returned by this callback
+    /// is attached to the task and can be retrieved using [`TaskMeta::get_data`] in
+    /// subsequent calls to other hooks for this task such as
+    /// [`on_before_task_poll`](crate::runtime::Builder::on_before_task_poll),
+    /// [`on_after_task_poll`](crate::runtime::Builder::on_after_task_poll), and
+    /// [`on_task_terminate`](crate::runtime::Builder::on_task_terminate).
     ///
     /// This can be used for bookkeeping or monitoring purposes.
     ///
@@ -755,6 +768,7 @@ impl Builder {
     /// let runtime = runtime::Builder::new_current_thread()
     ///     .on_task_spawn(|_| {
     ///         println!("spawning task");
+    ///         None::<()>
     ///     })
     ///     .build()
     ///     .unwrap();
@@ -768,13 +782,70 @@ impl Builder {
     /// })
     /// # }
     /// ```
+    ///
+    /// ```
+    /// # #[cfg(not(target_family = "wasm"))]
+    /// # use tokio::runtime;
+    /// # use std::sync::atomic::{AtomicUsize, Ordering};
+    /// # pub fn main() {
+    /// struct YieldingTaskMetadata {
+    ///     pub yield_count: AtomicUsize,
+    /// }
+    /// let runtime = runtime::Builder::new_current_thread()
+    ///     .on_task_spawn(|meta| {
+    ///         println!("spawning task {}", meta.id());
+    ///         Some(YieldingTaskMetadata { yield_count: AtomicUsize::new(0) })
+    ///     })
+    ///     .on_after_task_poll(|meta| {
+    ///         if let Some(data) = meta.get_data::<YieldingTaskMetadata>() {
+    ///             println!("task {} yield count: {}", meta.id(), data.yield_count.fetch_add(1, Ordering::Relaxed));
+    ///         }
+    ///     })
+    ///     .on_task_terminate(|meta| {
+    ///         match meta.get_data::<YieldingTaskMetadata>() {
+    ///             Some(data) => {
+    ///                 let yield_count = data.yield_count.load(Ordering::Relaxed);
+    ///                 println!("task {} total yield count: {}", meta.id(), yield_count);
+    ///                 assert!(yield_count == 64);
+    ///             },
+    ///             None => panic!("task has missing or incorrect user data"),
+    ///         }
+    ///     })
+    ///     .build()
+    ///     .unwrap();
+    ///
+    /// runtime.block_on(async {
+    ///     let _ = tokio::task::spawn(async {
+    ///         for _ in 0..64 {
+    ///             println!("yielding");
+    ///             tokio::task::yield_now().await;
+    ///         }
+    ///     }).await.unwrap();
+    /// })
+    /// # }
+    /// ```
     #[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 on_task_spawn<F, T>(&mut self, f: F) -> &mut Self
     where
-        F: Fn(&TaskMeta<'_>) + Send + Sync + 'static,
+        F: Fn(&TaskMeta<'_>) -> Option<T> + Send + Sync + 'static,
+        T: 'static,
     {
-        self.before_spawn = Some(std::sync::Arc::new(f));
+        use std::any::Any;
+
+        fn wrap<F, T>(f: F) -> impl Fn(&TaskMeta<'_>) -> UserData + Send + Sync + 'static
+        where
+            F: Fn(&TaskMeta<'_>) -> Option<T> + Send + Sync + 'static,
+            T: 'static,
+        {
+            move |meta| {
+                f(meta).map(|value| {
+                    let boxed: Box<dyn Any> = Box::new(value);
+                    Box::leak(boxed) as &'static dyn Any
+                })
+            }
+        }
+        self.before_spawn = Some(std::sync::Arc::new(wrap(f)));
         self
     }
 

tokio/src/runtime/config.rs

@@ -3,6 +3,8 @@
     allow(dead_code)
 )]
 use crate::runtime::{Callback, TaskCallback};
+#[cfg(tokio_unstable)]
+use crate::runtime::TaskSpawnCallback;
 use crate::util::RngSeedGenerator;
 
 pub(crate) struct Config {
@@ -19,6 +21,9 @@ pub(crate) struct Config {
     pub(crate) after_unpark: Option<Callback>,
 
     /// To run before each task is spawned.
+    #[cfg(tokio_unstable)]
+    pub(crate) before_spawn: Option<TaskSpawnCallback>,
+    #[cfg(not(tokio_unstable))]
     pub(crate) before_spawn: Option<TaskCallback>,
 
     /// To run after each task is terminated.

tokio/src/runtime/mod.rs

@@ -391,6 +391,8 @@ cfg_rt! {
 
     mod task_hooks;
     pub(crate) use task_hooks::{TaskHooks, TaskCallback};
+    #[cfg(tokio_unstable)]
+    pub(crate) use task_hooks::{TaskSpawnCallback, UserData};
     cfg_unstable! {
         pub use task_hooks::TaskMeta;
     }

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

@@ -5,6 +5,8 @@ use crate::runtime::scheduler::{self, Defer, Inject};
 use crate::runtime::task::{
     self, JoinHandle, OwnedTasks, Schedule, SpawnLocation, Task, TaskHarnessScheduleHooks,
 };
+#[cfg(tokio_unstable)]
+use crate::runtime::UserData;
 use crate::runtime::{
     blocking, context, Config, MetricsBatch, SchedulerMetrics, TaskHooks, TaskMeta, WorkerMetrics,
 };
@@ -456,13 +458,43 @@ impl Handle {
         F: crate::future::Future + Send + 'static,
         F::Output: Send + 'static,
     {
-        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
+        Self::spawn_with_user_data(me, future, id, spawned_at, #[cfg(tokio_unstable)] None)
+    }
 
-        me.task_hooks.spawn(&TaskMeta {
+    #[track_caller]
+    pub(crate) fn spawn_with_user_data<F>(
+        me: &Arc<Self>,
+        future: F,
+        id: crate::runtime::task::Id,
+        spawned_at: SpawnLocation,
+        #[cfg(tokio_unstable)]
+        user_data: UserData,
+    ) -> JoinHandle<F::Output>
+    where
+        F: crate::future::Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        let task_meta = TaskMeta {
             id,
             spawned_at,
+            #[cfg(tokio_unstable)]
+            user_data,
             _phantom: Default::default(),
-        });
+        };
+
+        #[cfg(not(tokio_unstable))]
+        {
+            me.task_hooks.spawn(&task_meta);
+        }
+
+        let (handle, notified) = me.shared.owned.bind(
+            future,
+            me.clone(),
+            id,
+            spawned_at,
+            #[cfg(tokio_unstable)]
+            me.task_hooks.spawn(&task_meta),
+        );
 
         if let Some(notified) = notified {
             me.schedule(notified);
@@ -488,16 +520,27 @@ impl Handle {
         F: crate::future::Future + 'static,
         F::Output: 'static,
     {
-        let (handle, notified) = me
-            .shared
-            .owned
-            .bind_local(future, me.clone(), id, spawned_at);
-
-        me.task_hooks.spawn(&TaskMeta {
+        let task_meta = TaskMeta {
             id,
             spawned_at,
+            #[cfg(tokio_unstable)]
+            user_data: None,
             _phantom: Default::default(),
-        });
+        };
+
+        #[cfg(not(tokio_unstable))]
+        {
+            me.task_hooks.spawn(&task_meta);
+        }
+
+        let (handle, notified) = me.shared.owned.bind_local(
+            future,
+            me.clone(),
+            id,
+            spawned_at,
+            #[cfg(tokio_unstable)]
+            me.task_hooks.spawn(&task_meta),
+        );
 
         if let Some(notified) = notified {
             me.schedule(notified);

tokio/src/runtime/scheduler/mod.rs

@@ -68,6 +68,8 @@ impl Handle {
 cfg_rt! {
     use crate::future::Future;
     use crate::loom::sync::Arc;
+    #[cfg(tokio_unstable)]
+    use crate::runtime::UserData;
     use crate::runtime::{blocking, task::{Id, SpawnLocation}};
     use crate::runtime::context;
     use crate::task::JoinHandle;
@@ -130,6 +132,20 @@ cfg_rt! {
             }
         }
 
+        #[cfg(tokio_unstable)]
+        pub(crate) fn spawn_with_user_data<F>(&self, future: F, id: Id, spawned_at: SpawnLocation, user_data: UserData) -> JoinHandle<F::Output>
+        where
+            F: Future + Send + 'static,
+            F::Output: Send + 'static,
+        {
+            match self {
+                Handle::CurrentThread(h) => current_thread::Handle::spawn_with_user_data(h, future, id, spawned_at, user_data),
+
+                #[cfg(feature = "rt-multi-thread")]
+                Handle::MultiThread(h) => multi_thread::Handle::spawn_with_user_data(h, future, id, spawned_at, user_data),
+            }
+        }
+
         /// Spawn a local task
         ///
         /// # Safety

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

@@ -8,6 +8,8 @@ use crate::runtime::{
     TaskHooks, TaskMeta,
 };
 use crate::util::RngSeedGenerator;
+#[cfg(tokio_unstable)]
+use crate::runtime::task_hooks::UserData;
 
 use std::fmt;
 
@@ -47,7 +49,23 @@ impl Handle {
         F: crate::future::Future + Send + 'static,
         F::Output: Send + 'static,
     {
-        Self::bind_new_task(me, future, id, spawned_at)
+        Self::bind_new_task(me, future, id, spawned_at, #[cfg(tokio_unstable)] None)
+    }
+
+    /// Spawns a future with user data onto the thread pool
+    #[cfg(tokio_unstable)]
+    pub(crate) fn spawn_with_user_data<F>(
+        me: &Arc<Self>,
+        future: F,
+        id: task::Id,
+        spawned_at: SpawnLocation,
+        user_data: UserData,
+    ) -> JoinHandle<F::Output>
+    where
+        F: crate::future::Future + Send + 'static,
+        F::Output: Send + 'static,
+    {
+        Self::bind_new_task(me, future, id, spawned_at, user_data)
     }
 
     pub(crate) fn shutdown(&self) {
@@ -60,18 +78,34 @@ impl Handle {
         future: T,
         id: task::Id,
         spawned_at: SpawnLocation,
+        #[cfg(tokio_unstable)]
+        user_data: UserData,
     ) -> JoinHandle<T::Output>
     where
         T: Future + Send + 'static,
         T::Output: Send + 'static,
     {
-        let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at);
-
-        me.task_hooks.spawn(&TaskMeta {
+        let task_meta = TaskMeta {
             id,
             spawned_at,
+            #[cfg(tokio_unstable)]
+            user_data,
             _phantom: Default::default(),
-        });
+        };
+
+        #[cfg(not(tokio_unstable))]
+        {
+            me.task_hooks.spawn(&task_meta);
+        }
+
+        let (handle, notified) = me.shared.owned.bind(
+            future,
+            me.clone(),
+            id,
+            spawned_at,
+            #[cfg(tokio_unstable)]
+            me.task_hooks.spawn(&task_meta),
+        );
 
         me.schedule_option_task_without_yield(notified);