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 commit of a simpler 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.

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

Author: LloydW93

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

@@ -461,6 +461,8 @@ impl Handle {
         me.task_hooks.spawn(&TaskMeta {
             id,
             spawned_at,
+            #[cfg(tokio_unstable)]
+            user_data: None,
             _phantom: Default::default(),
         });
 
@@ -496,6 +498,8 @@ impl Handle {
         me.task_hooks.spawn(&TaskMeta {
             id,
             spawned_at,
+            #[cfg(tokio_unstable)]
+            user_data: None,
             _phantom: Default::default(),
         });
 

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

@@ -70,6 +70,8 @@ impl Handle {
         me.task_hooks.spawn(&TaskMeta {
             id,
             spawned_at,
+            #[cfg(tokio_unstable)]
+            user_data: None,
             _phantom: Default::default(),
         });
 

tokio/src/runtime/task/core.rs

@@ -182,6 +182,10 @@ pub(crate) struct Header {
     /// The tracing ID for this instrumented task.
     #[cfg(all(tokio_unstable, feature = "tracing"))]
     pub(super) tracing_id: Option<tracing::Id>,
+
+    /// Custom user defined metadata for this task for use in hooks.
+    #[cfg(tokio_unstable)]
+    pub(super) user_data: Option<&'static dyn std::any::Any>,
 }
 
 unsafe impl Send for Header {}
@@ -237,6 +241,8 @@ impl<T: Future, S: Schedule> Cell<T, S> {
                 owner_id: UnsafeCell::new(None),
                 #[cfg(all(tokio_unstable, feature = "tracing"))]
                 tracing_id,
+                #[cfg(tokio_unstable)]
+                user_data: None,
             }
         }
 
@@ -515,6 +521,28 @@ impl Header {
         *ptr
     }
 
+    /// Gets the user data from the task header.
+    ///
+    /// # Safety
+    ///
+    /// The provided raw pointer must point at the header of a task.
+    #[cfg(tokio_unstable)]
+    pub(super) unsafe fn get_user_data(me: NonNull<Header>) -> Option<&'static dyn std::any::Any> {
+        me.as_ref().user_data
+    }
+
+    /// Sets the user data in the task header.
+    ///
+    /// # Safety
+    ///
+    /// The provided raw pointer must point at the header of a task.
+    /// The caller must ensure exclusive access to the header.
+    #[cfg(tokio_unstable)]
+    pub(super) unsafe fn set_user_data(me: NonNull<Header>, data: Option<&'static dyn std::any::Any>) {
+        let header = &mut *me.as_ptr();
+        header.user_data = data;
+    }
+
     /// Gets the tracing id of the task containing this `Header`.
     ///
     /// # Safety

tokio/src/runtime/task/harness.rs

@@ -374,6 +374,8 @@ where
                 f(&TaskMeta {
                     id: self.core().task_id,
                     spawned_at: self.core().spawned_at.into(),
+                    #[cfg(tokio_unstable)]
+                    user_data: None,
                     _phantom: Default::default(),
                 })
             }));

tokio/src/runtime/task/mod.rs

@@ -437,6 +437,18 @@ impl<S: 'static> Task<S> {
         unsafe { Header::get_spawn_location(self.raw.header_ptr()) }
     }
 
+    #[cfg(tokio_unstable)]
+    pub(crate) fn get_user_data(&self) -> Option<&'static dyn std::any::Any> {
+        // Safety: The header pointer is valid.
+        unsafe { Header::get_user_data(self.raw.header_ptr()) }
+    }
+
+    #[cfg(tokio_unstable)]
+    pub(crate) fn set_user_data(&self, data: Option<&'static dyn std::any::Any>) {
+        // Safety: The header pointer is valid.
+        unsafe { Header::set_user_data(self.raw.header_ptr(), data) }
+    }
+
     // Explicit `'task` and `'meta` lifetimes are necessary here, as otherwise,
     // the compiler infers the lifetimes to be the same, and considers the task
     // to be borrowed for the lifetime of the returned `TaskMeta`.
@@ -445,6 +457,8 @@ impl<S: 'static> Task<S> {
         crate::runtime::TaskMeta {
             id: self.id(),
             spawned_at: self.spawned_at().into(),
+            #[cfg(tokio_unstable)]
+            user_data: self.get_user_data(),
             _phantom: PhantomData,
         }
     }

tokio/src/runtime/task_hooks.rs

@@ -1,4 +1,6 @@
 use super::Config;
+#[cfg(tokio_unstable)]
+use std::any::Any;
 use std::marker::PhantomData;
 
 impl TaskHooks {
@@ -62,6 +64,9 @@ pub struct TaskMeta<'a> {
     /// The location where the task was spawned.
     #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))]
     pub(crate) spawned_at: crate::runtime::task::SpawnLocation,
+    /// Optional user-defined metadata for the task.
+    #[cfg(tokio_unstable)]
+    pub(crate) user_data: Option<&'static dyn Any>,
     pub(crate) _phantom: PhantomData<&'a ()>,
 }
 
@@ -77,6 +82,13 @@ impl<'a> TaskMeta<'a> {
     pub fn spawned_at(&self) -> &'static std::panic::Location<'static> {
         self.spawned_at.0
     }
+
+    /// Return the user-defined metadata for this task if it is set and of the
+    /// correct type.
+    #[cfg(tokio_unstable)]
+    pub fn get_data<T: 'static>(&self) -> Option<&T> {
+        self.user_data?.downcast_ref::<T>()
+    }
 }
 
 /// Runs on specific task-related events