From a3860a9f376ba4d5f8a07ec05e7e63e928306004 Mon Sep 17 00:00:00 2001 From: noah Date: Sat, 9 May 2026 19:18:10 -0500 Subject: [PATCH] hooks: overhaul task hooks This change overhauls our task hooks system, allowing users to attach user data to individual tasks, which are passed into hooks. Spawn hooks can see both parent and child metadata. --- tokio/src/runtime/blocking/pool.rs | 11 + tokio/src/runtime/blocking/schedule.rs | 19 +- tokio/src/runtime/builder.rs | 15 +- tokio/src/runtime/config.rs | 4 +- tokio/src/runtime/context.rs | 22 + tokio/src/runtime/handle.rs | 93 +++- tokio/src/runtime/mod.rs | 6 +- .../runtime/scheduler/current_thread/mod.rs | 87 +++- tokio/src/runtime/scheduler/mod.rs | 48 +- .../runtime/scheduler/multi_thread/handle.rs | 60 ++- .../runtime/scheduler/multi_thread/worker.rs | 26 - tokio/src/runtime/task/core.rs | 55 ++- tokio/src/runtime/task/harness.rs | 110 ++++- tokio/src/runtime/task/list.rs | 79 ++- tokio/src/runtime/task/mod.rs | 76 +-- tokio/src/runtime/task/raw.rs | 39 ++ tokio/src/runtime/task_hooks.rs | 192 +++++++- tokio/src/runtime/tests/mod.rs | 25 +- tokio/src/runtime/tests/task.rs | 41 +- tokio/src/task/builder.rs | 63 ++- tokio/src/task/join_set.rs | 9 + tokio/src/task/local.rs | 102 +++- tokio/src/task/spawn.rs | 30 +- tokio/tests/task_hooks.rs | 450 +++++++++++++++++- 24 files changed, 1431 insertions(+), 231 deletions(-) diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index dae98bc94..6a489810d 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -307,6 +307,8 @@ impl Spawner { Mandatory::NonMandatory, SpawnMeta::new_unnamed(fn_size), rt, + #[cfg(tokio_unstable)] + None, ) } else { self.spawn_blocking_inner( @@ -314,6 +316,8 @@ impl Spawner { Mandatory::NonMandatory, SpawnMeta::new_unnamed(fn_size), rt, + #[cfg(tokio_unstable)] + None, ) }; @@ -345,6 +349,8 @@ impl Spawner { Mandatory::Mandatory, SpawnMeta::new_unnamed(fn_size), rt, + #[cfg(tokio_unstable)] + None, ) } else { self.spawn_blocking_inner( @@ -352,6 +358,8 @@ impl Spawner { Mandatory::Mandatory, SpawnMeta::new_unnamed(fn_size), rt, + #[cfg(tokio_unstable)] + None, ) }; @@ -370,6 +378,7 @@ impl Spawner { is_mandatory: Mandatory, spawn_meta: SpawnMeta<'_>, rt: &Handle, + #[cfg(tokio_unstable)] user_data: Option, ) -> (JoinHandle, Result<(), SpawnError>) where F: FnOnce() -> R + Send + 'static, @@ -384,6 +393,8 @@ impl Spawner { BlockingSchedule::new(rt), id, task::SpawnLocation::capture(), + #[cfg(tokio_unstable)] + user_data, ); let spawned = self.spawn_task(Task::new(task, is_mandatory), rt); diff --git a/tokio/src/runtime/blocking/schedule.rs b/tokio/src/runtime/blocking/schedule.rs index 0e97c5aea..8484398c6 100644 --- a/tokio/src/runtime/blocking/schedule.rs +++ b/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::{TaskCallback, TaskMeta}; /// `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)] + task_terminate_callback: Option, } 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)] + task_terminate_callback: handle.inner.hooks().task_terminate_callback.clone(), } } } @@ -58,9 +60,10 @@ impl task::Schedule for BlockingSchedule { unreachable!(); } - fn hooks(&self) -> TaskHarnessScheduleHooks { - TaskHarnessScheduleHooks { - task_terminate_callback: self.hooks.task_terminate_callback.clone(), + #[cfg(tokio_unstable)] + fn task_terminate_callback(&self, meta: &mut TaskMeta<'_>) { + if let Some(task_terminate_callback) = &self.task_terminate_callback { + task_terminate_callback(meta); } } } diff --git a/tokio/src/runtime/builder.rs b/tokio/src/runtime/builder.rs index a37d47f18..b27656f7a 100644 --- a/tokio/src/runtime/builder.rs +++ b/tokio/src/runtime/builder.rs @@ -2,7 +2,8 @@ use crate::runtime::handle::Handle; use crate::runtime::{ - blocking, driver, Callback, HistogramBuilder, Runtime, TaskCallback, TimerFlavor, + blocking, driver, Callback, HistogramBuilder, Runtime, TaskCallback, TaskSpawnCallback, + TimerFlavor, }; #[cfg(tokio_unstable)] use crate::runtime::{metrics::HistogramConfiguration, TaskMeta}; @@ -96,7 +97,7 @@ pub struct Builder { pub(super) after_unpark: Option, /// To run before each task is spawned. - pub(super) before_spawn: Option, + pub(super) before_spawn: Option, /// To run before each poll #[cfg(tokio_unstable)] @@ -873,7 +874,7 @@ impl Builder { /// # use tokio::runtime; /// # pub fn main() { /// let runtime = runtime::Builder::new_current_thread() - /// .on_task_spawn(|_| { + /// .on_task_spawn(|_, _| { /// println!("spawning task"); /// }) /// .build() @@ -892,7 +893,7 @@ impl Builder { #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub fn on_task_spawn(&mut self, f: F) -> &mut Self where - F: Fn(&TaskMeta<'_>) + Send + Sync + 'static, + F: Fn(&mut TaskMeta<'_>, Option>) + Send + Sync + 'static, { self.before_spawn = Some(std::sync::Arc::new(f)); self @@ -939,7 +940,7 @@ impl Builder { #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub fn on_before_task_poll(&mut self, f: F) -> &mut Self where - F: Fn(&TaskMeta<'_>) + Send + Sync + 'static, + F: Fn(&mut TaskMeta<'_>) + Send + Sync + 'static, { self.before_poll = Some(std::sync::Arc::new(f)); self @@ -986,7 +987,7 @@ impl Builder { #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub fn on_after_task_poll(&mut self, f: F) -> &mut Self where - F: Fn(&TaskMeta<'_>) + Send + Sync + 'static, + F: Fn(&mut TaskMeta<'_>) + Send + Sync + 'static, { self.after_poll = Some(std::sync::Arc::new(f)); self @@ -1035,7 +1036,7 @@ impl Builder { #[cfg_attr(docsrs, doc(cfg(tokio_unstable)))] pub fn on_task_terminate(&mut self, f: F) -> &mut Self where - F: Fn(&TaskMeta<'_>) + Send + Sync + 'static, + F: Fn(&mut TaskMeta<'_>) + Send + Sync + 'static, { self.after_termination = Some(std::sync::Arc::new(f)); self diff --git a/tokio/src/runtime/config.rs b/tokio/src/runtime/config.rs index e97482fa9..ddbd5ed8c 100644 --- a/tokio/src/runtime/config.rs +++ b/tokio/src/runtime/config.rs @@ -2,7 +2,7 @@ any(not(all(tokio_unstable, feature = "full")), target_family = "wasm"), allow(dead_code) )] -use crate::runtime::{Callback, TaskCallback}; +use crate::runtime::{Callback, TaskCallback, TaskSpawnCallback}; use crate::util::RngSeedGenerator; pub(crate) struct Config { @@ -19,7 +19,7 @@ pub(crate) struct Config { pub(crate) after_unpark: Option, /// To run before each task is spawned. - pub(crate) before_spawn: Option, + pub(crate) before_spawn: Option, /// To run after each task is terminated. pub(crate) after_termination: Option, diff --git a/tokio/src/runtime/context.rs b/tokio/src/runtime/context.rs index 3bf991940..d201a06df 100644 --- a/tokio/src/runtime/context.rs +++ b/tokio/src/runtime/context.rs @@ -21,6 +21,8 @@ cfg_rt! { use crate::runtime::{scheduler, task::Id}; + #[cfg(tokio_unstable)] + use std::ptr::NonNull; use std::task::Waker; cfg_taskdump! { @@ -49,6 +51,9 @@ struct Context { #[cfg(feature = "rt")] current_task_id: Cell>, + #[cfg(all(feature = "rt", tokio_unstable))] + current_task: Cell>>, + /// 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 +97,9 @@ tokio_thread_local! { #[cfg(feature = "rt")] current_task_id: Cell::new(None), + #[cfg(all(feature = "rt", tokio_unstable))] + current_task: Cell::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 @@ -159,6 +167,20 @@ cfg_rt! { CONTEXT.try_with(|ctx| ctx.current_task_id.get()).unwrap_or(None) } + #[cfg(tokio_unstable)] + pub(crate) fn set_current_task(task: Option>) -> Option> { + CONTEXT + .try_with(|ctx| ctx.current_task.replace(task)) + .unwrap_or(None) + } + + #[cfg(tokio_unstable)] + pub(crate) fn current_task() -> Option> { + CONTEXT + .try_with(|ctx| ctx.current_task.get()) + .unwrap_or(None) + } + #[cfg(tokio_unstable)] pub(crate) fn worker_index() -> Option { with_scheduler(|ctx| ctx.and_then(|c| c.worker_index())) diff --git a/tokio/src/runtime/handle.rs b/tokio/src/runtime/handle.rs index 79fea39da..2d78dee50 100644 --- a/tokio/src/runtime/handle.rs +++ b/tokio/src/runtime/handle.rs @@ -372,6 +372,40 @@ impl Handle { #[track_caller] pub(crate) fn spawn_named(&self, future: F, meta: SpawnMeta<'_>) -> JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + self.spawn_named_inner( + future, + meta, + #[cfg(tokio_unstable)] + None, + ) + } + + #[cfg(all(tokio_unstable, feature = "tracing"))] + #[track_caller] + pub(crate) fn spawn_named_with_data( + &self, + future: F, + meta: SpawnMeta<'_>, + user_data: Option, + ) -> JoinHandle + where + F: Future + Send + 'static, + F::Output: Send + 'static, + { + self.spawn_named_inner(future, meta, user_data) + } + + #[track_caller] + fn spawn_named_inner( + &self, + future: F, + meta: SpawnMeta<'_>, + #[cfg(tokio_unstable)] user_data: Option, + ) -> JoinHandle where F: Future + Send + 'static, F::Output: Send + 'static, @@ -387,7 +421,13 @@ 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()); - self.inner.spawn(future, id, meta.spawned_at) + self.inner.spawn( + future, + id, + meta.spawned_at, + #[cfg(tokio_unstable)] + user_data, + ) } #[track_caller] @@ -401,6 +441,47 @@ impl Handle { future: F, meta: SpawnMeta<'_>, ) -> JoinHandle + where + F: Future + 'static, + F::Output: 'static, + { + unsafe { + self.spawn_local_named_inner( + future, + meta, + #[cfg(tokio_unstable)] + None, + ) + } + } + + /// # Safety + /// + /// This must only be called in `LocalRuntime` if the runtime has been verified to be owned + /// by the current thread. + #[cfg(all(tokio_unstable, feature = "tracing"))] + #[track_caller] + #[allow(dead_code)] + pub(crate) unsafe fn spawn_local_named_with_data( + &self, + future: F, + meta: SpawnMeta<'_>, + user_data: Option, + ) -> JoinHandle + where + F: Future + 'static, + F::Output: 'static, + { + unsafe { self.spawn_local_named_inner(future, meta, user_data) } + } + + #[track_caller] + unsafe fn spawn_local_named_inner( + &self, + future: F, + meta: SpawnMeta<'_>, + #[cfg(tokio_unstable)] user_data: Option, + ) -> JoinHandle where F: Future + 'static, F::Output: 'static, @@ -416,7 +497,15 @@ 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()); - unsafe { self.inner.spawn_local(future, id, meta.spawned_at) } + unsafe { + self.inner.spawn_local( + future, + id, + meta.spawned_at, + #[cfg(tokio_unstable)] + user_data, + ) + } } /// Returns the flavor of the current `Runtime`. diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index 2b3f0ad9b..2cfdd92f9 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -589,9 +589,11 @@ cfg_rt! { } mod task_hooks; - pub(crate) use task_hooks::{TaskHooks, TaskCallback}; + pub(crate) use task_hooks::{TaskCallback, TaskHooks, TaskSpawnCallback}; + #[cfg(tokio_unstable)] + pub(crate) use task_hooks::TaskData; cfg_unstable! { - pub use task_hooks::TaskMeta; + pub use task_hooks::{TaskMeta, TaskMetaRef}; } #[cfg(not(tokio_unstable))] pub(crate) use task_hooks::TaskMeta; diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index f0b072d57..95cc304e1 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/tokio/src/runtime/scheduler/current_thread/mod.rs @@ -3,7 +3,7 @@ use crate::loom::sync::Arc; use crate::runtime::driver::{self, Driver}; use crate::runtime::scheduler::{self, Defer, Inject}; use crate::runtime::task::{ - self, JoinHandle, OwnedTasks, Schedule, SpawnLocation, Task, TaskHarnessScheduleHooks, + self, JoinHandle, OwnedTasks, Schedule, SpawnLocation, Task, }; use crate::runtime::{ blocking, context, Config, MetricsBatch, SchedulerMetrics, TaskHooks, TaskMeta, WorkerMetrics, @@ -470,18 +470,35 @@ impl Handle { future: F, id: crate::runtime::task::Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> JoinHandle where F: crate::future::Future + Send + 'static, F::Output: Send + 'static, { - let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at); - - me.task_hooks.spawn(&TaskMeta { + #[cfg(tokio_unstable)] + let parent = task::current_task_meta(); + #[cfg(tokio_unstable)] + let (handle, notified) = me.shared.owned.bind_with_spawn_hook( + future, + me.clone(), id, spawned_at, - _phantom: Default::default(), - }); + user_data, + |task| { + // Safety: the task is freshly allocated and not published yet. + let mut meta = unsafe { task.task_meta() }; + me.task_hooks.spawn(&mut meta, parent); + }, + ); + #[cfg(not(tokio_unstable))] + let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at); + + #[cfg(not(tokio_unstable))] + { + let mut meta = TaskMeta::new(id, spawned_at); + me.task_hooks.spawn(&mut meta, None); + } if let Some(notified) = notified { me.schedule(notified); @@ -503,23 +520,44 @@ impl Handle { future: F, id: crate::runtime::task::Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> JoinHandle where F: crate::future::Future + 'static, F::Output: 'static, { // Safety: the caller guarantees that this is only called on a `LocalRuntime`. + #[cfg(tokio_unstable)] + let parent = task::current_task_meta(); + #[cfg(tokio_unstable)] + let before_bind = |task: &Task>| { + // Safety: the task is freshly allocated and not published yet. + let mut meta = unsafe { task.task_meta() }; + me.task_hooks.spawn(&mut meta, parent); + }; + #[cfg(tokio_unstable)] + let (handle, notified) = unsafe { + me.shared.owned.bind_local_with_spawn_hook( + future, + me.clone(), + id, + spawned_at, + user_data, + before_bind, + ) + }; + #[cfg(not(tokio_unstable))] let (handle, notified) = unsafe { me.shared .owned .bind_local(future, me.clone(), id, spawned_at) }; - me.task_hooks.spawn(&TaskMeta { - id, - spawned_at, - _phantom: Default::default(), - }); + #[cfg(not(tokio_unstable))] + { + let mut meta = TaskMeta::new(id, spawned_at); + me.task_hooks.spawn(&mut meta, None); + } if let Some(notified) = notified { me.schedule(notified); @@ -687,13 +725,19 @@ impl Schedule for Arc { }); } - fn hooks(&self) -> TaskHarnessScheduleHooks { - TaskHarnessScheduleHooks { - task_terminate_callback: self.task_hooks.task_terminate_callback.clone(), - } - } - cfg_unstable! { + fn task_terminate_callback(&self, meta: &mut TaskMeta<'_>) { + self.task_hooks.task_terminate_callback(meta); + } + + fn task_poll_start_callback(&self, meta: &mut TaskMeta<'_>) { + self.task_hooks.poll_start_callback(meta); + } + + fn task_poll_stop_callback(&self, meta: &mut TaskMeta<'_>) { + self.task_hooks.poll_stop_callback(meta); + } + fn unhandled_panic(&self) { use crate::runtime::UnhandledPanic; @@ -815,17 +859,8 @@ impl CoreGuard<'_> { let task = context.handle.shared.owned.assert_owner(task); - #[cfg(tokio_unstable)] - let task_meta = task.task_meta(); - let (c, ()) = context.run_task(core, || { - #[cfg(tokio_unstable)] - context.handle.task_hooks.poll_start_callback(&task_meta); - task.run(); - - #[cfg(tokio_unstable)] - context.handle.task_hooks.poll_stop_callback(&task_meta); }); core = c; diff --git a/tokio/src/runtime/scheduler/mod.rs b/tokio/src/runtime/scheduler/mod.rs index 8bbd110cb..881cefa2d 100644 --- a/tokio/src/runtime/scheduler/mod.rs +++ b/tokio/src/runtime/scheduler/mod.rs @@ -8,8 +8,8 @@ cfg_rt! { pub(crate) mod inject; pub(crate) use inject::Inject; + #[cfg(tokio_unstable)] use crate::runtime::TaskHooks; - use crate::runtime::WorkerMetrics; } @@ -149,16 +149,36 @@ cfg_rt! { } } - pub(crate) fn spawn(&self, future: F, id: Id, spawned_at: SpawnLocation) -> JoinHandle + pub(crate) fn spawn( + &self, + future: F, + id: Id, + spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, + ) -> JoinHandle where F: Future + Send + 'static, F::Output: Send + 'static, { match self { - Handle::CurrentThread(h) => current_thread::Handle::spawn(h, future, id, spawned_at), + Handle::CurrentThread(h) => current_thread::Handle::spawn( + h, + future, + id, + spawned_at, + #[cfg(tokio_unstable)] + user_data, + ), #[cfg(feature = "rt-multi-thread")] - Handle::MultiThread(h) => multi_thread::Handle::spawn(h, future, id, spawned_at), + Handle::MultiThread(h) => multi_thread::Handle::spawn( + h, + future, + id, + spawned_at, + #[cfg(tokio_unstable)] + user_data, + ), } } @@ -170,14 +190,29 @@ cfg_rt! { /// by the current thread. #[allow(irrefutable_let_patterns)] #[track_caller] - pub(crate) unsafe fn spawn_local(&self, future: F, id: Id, spawned_at: SpawnLocation) -> JoinHandle + pub(crate) unsafe fn spawn_local( + &self, + future: F, + id: Id, + spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, + ) -> JoinHandle where F: Future + 'static, F::Output: 'static, { if let Handle::CurrentThread(h) = self { // Safety: caller guarantees that this is a `LocalRuntime`. - unsafe { current_thread::Handle::spawn_local(h, future, id, spawned_at) } + unsafe { + current_thread::Handle::spawn_local( + h, + future, + id, + spawned_at, + #[cfg(tokio_unstable)] + user_data, + ) + } } else { panic!("Only current_thread and LocalSet have spawn_local internals implemented") } @@ -204,6 +239,7 @@ cfg_rt! { } } + #[cfg(tokio_unstable)] pub(crate) fn hooks(&self) -> &TaskHooks { match self { Handle::CurrentThread(h) => &h.task_hooks, diff --git a/tokio/src/runtime/scheduler/multi_thread/handle.rs b/tokio/src/runtime/scheduler/multi_thread/handle.rs index 657e3fe34..4f809fed4 100644 --- a/tokio/src/runtime/scheduler/multi_thread/handle.rs +++ b/tokio/src/runtime/scheduler/multi_thread/handle.rs @@ -1,7 +1,7 @@ use crate::future::Future; use crate::loom::sync::Arc; use crate::runtime::scheduler::multi_thread::worker; -use crate::runtime::task::{Notified, Task, TaskHarnessScheduleHooks}; +use crate::runtime::task::{Notified, Task}; use crate::runtime::{ blocking, driver, task::{self, JoinHandle, SpawnLocation}, @@ -57,12 +57,20 @@ impl Handle { future: F, id: task::Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> JoinHandle where 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)] + user_data, + ) } #[cfg(all(tokio_unstable, feature = "time"))] @@ -83,18 +91,35 @@ impl Handle { future: T, id: task::Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> JoinHandle 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 { + #[cfg(tokio_unstable)] + let parent = task::current_task_meta(); + #[cfg(tokio_unstable)] + let (handle, notified) = me.shared.owned.bind_with_spawn_hook( + future, + me.clone(), id, spawned_at, - _phantom: Default::default(), - }); + user_data, + |task| { + // Safety: the task is freshly allocated and not published yet. + let mut meta = unsafe { task.task_meta() }; + me.task_hooks.spawn(&mut meta, parent); + }, + ); + #[cfg(not(tokio_unstable))] + let (handle, notified) = me.shared.owned.bind(future, me.clone(), id, spawned_at); + + #[cfg(not(tokio_unstable))] + { + let mut meta = TaskMeta::new(id, spawned_at); + me.task_hooks.spawn(&mut meta, None); + } me.schedule_option_task_without_yield(notified); @@ -111,15 +136,24 @@ impl task::Schedule for Arc { self.schedule_task(task, false); } - fn hooks(&self) -> TaskHarnessScheduleHooks { - TaskHarnessScheduleHooks { - task_terminate_callback: self.task_hooks.task_terminate_callback.clone(), - } - } - fn yield_now(&self, task: Notified) { self.schedule_task(task, true); } + + #[cfg(tokio_unstable)] + fn task_terminate_callback(&self, meta: &mut TaskMeta<'_>) { + self.task_hooks.task_terminate_callback(meta); + } + + #[cfg(tokio_unstable)] + fn task_poll_start_callback(&self, meta: &mut TaskMeta<'_>) { + self.task_hooks.poll_start_callback(meta); + } + + #[cfg(tokio_unstable)] + fn task_poll_stop_callback(&self, meta: &mut TaskMeta<'_>) { + self.task_hooks.poll_stop_callback(meta); + } } impl Handle { diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index e222edbf9..82b5d83f1 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -628,9 +628,6 @@ impl Context { } fn run_task(&self, task: Notified, mut core: Box) -> RunResult { - #[cfg(tokio_unstable)] - let task_meta = task.task_meta(); - let task = self.worker.handle.shared.owned.assert_owner(task); // Make sure the worker is not in the **searching** state. This enables @@ -673,19 +670,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_meta); - task.run(); - #[cfg(tokio_unstable)] - self.worker.handle.task_hooks.poll_stop_callback(&task_meta); - let mut lifo_polls = 0; // As long as there is budget remaining and a task exists in the @@ -749,19 +735,7 @@ impl Context { *self.core.borrow_mut() = Some(core); let task = self.worker.handle.shared.owned.assert_owner(task); - #[cfg(tokio_unstable)] - let task_meta = task.task_meta(); - - #[cfg(tokio_unstable)] - self.worker - .handle - .task_hooks - .poll_start_callback(&task_meta); - task.run(); - - #[cfg(tokio_unstable)] - self.worker.handle.task_hooks.poll_stop_callback(&task_meta); } }) } diff --git a/tokio/src/runtime/task/core.rs b/tokio/src/runtime/task/core.rs index aa3f61a22..58a7592ce 100644 --- a/tokio/src/runtime/task/core.rs +++ b/tokio/src/runtime/task/core.rs @@ -23,7 +23,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::TaskData; use crate::util::linked_list; use std::num::NonZeroU64; @@ -203,9 +205,8 @@ pub(super) struct Trailer { pub(super) owned: linked_list::Pointers
, /// Consumer task waiting on completion of this task. pub(super) waker: UnsafeCell>, - /// Optional hooks needed in the harness. - #[cfg_attr(not(tokio_unstable), allow(dead_code))] //TODO: remove when hooks are stabilized - pub(super) hooks: TaskHarnessScheduleHooks, + #[cfg(tokio_unstable)] + pub(super) user_data: UnsafeCell>, } generate_addr_of_methods! { @@ -233,6 +234,7 @@ impl Cell { state: State, task_id: Id, #[cfg(tokio_unstable)] spawned_at: &'static Location<'static>, + #[cfg(tokio_unstable)] user_data: Option, ) -> Box> { // Separated into a non-generic function to reduce LLVM codegen fn new_header( @@ -254,7 +256,10 @@ impl Cell { let tracing_id = future.id(); let vtable = raw::vtable::(); let result = Box::new(Cell { - trailer: Trailer::new(scheduler.hooks()), + trailer: Trailer::new( + #[cfg(tokio_unstable)] + user_data, + ), header: new_header( state, vtable, @@ -359,7 +364,11 @@ impl Core { /// /// `self` must also be pinned. This is handled by storing the task on the /// heap. - pub(super) fn poll(&self, mut cx: Context<'_>) -> Poll { + pub(super) fn poll( + &self, + #[cfg(tokio_unstable)] header: NonNull
, + mut cx: Context<'_>, + ) -> Poll { let res = { self.stage.stage.with_mut(|ptr| { // Safety: The caller ensures mutual exclusion to the field. @@ -372,6 +381,8 @@ impl Core { let future = unsafe { Pin::new_unchecked(future) }; let _guard = TaskIdGuard::enter(self.task_id); + #[cfg(tokio_unstable)] + let _current_task = CurrentTaskGuard::enter(header); future.poll(&mut cx) }) }; @@ -430,6 +441,27 @@ impl Core { } } +#[cfg(tokio_unstable)] +pub(crate) struct CurrentTaskGuard { + parent_task: Option>, +} + +#[cfg(tokio_unstable)] +impl CurrentTaskGuard { + fn enter(header: NonNull
) -> Self { + CurrentTaskGuard { + parent_task: context::set_current_task(Some(header.cast())), + } + } +} + +#[cfg(tokio_unstable)] +impl Drop for CurrentTaskGuard { + fn drop(&mut self) { + context::set_current_task(self.parent_task); + } +} + impl Header { pub(super) unsafe fn set_next(&self, next: Option>) { self.queue_next.with_mut(|ptr| *ptr = next); @@ -537,14 +569,21 @@ impl Header { } impl Trailer { - fn new(hooks: TaskHarnessScheduleHooks) -> Self { + fn new(#[cfg(tokio_unstable)] user_data: Option) -> Self { Trailer { waker: UnsafeCell::new(None), owned: linked_list::Pointers::new(), - hooks, + #[cfg(tokio_unstable)] + user_data: UnsafeCell::new(user_data), } } + #[cfg(tokio_unstable)] + pub(super) fn user_data_ptr(&self) -> NonNull> { + self.user_data + .with_mut(|ptr| unsafe { NonNull::new_unchecked(ptr) }) + } + pub(super) unsafe fn set_waker(&self, waker: Option) { self.waker.with_mut(|ptr| { *ptr = waker; diff --git a/tokio/src/runtime/task/harness.rs b/tokio/src/runtime/task/harness.rs index 6f20d66ef..46298995e 100644 --- a/tokio/src/runtime/task/harness.rs +++ b/tokio/src/runtime/task/harness.rs @@ -48,6 +48,21 @@ where fn core(&self) -> &Core { unsafe { &self.cell.as_ref().core } } + + /// # Safety + /// + /// The returned metadata must only be used while the caller has exclusive + /// access to task hook data. + #[cfg(tokio_unstable)] + unsafe fn task_meta<'meta>(&self) -> TaskMeta<'meta> { + unsafe { + TaskMeta::new( + self.core().task_id, + self.core().spawned_at.into(), + Some(self.trailer().user_data_ptr()), + ) + } + } } /// Task operations that can be implemented without being generic over the @@ -204,10 +219,56 @@ where TransitionToIdle::Cancelled => PollFuture::Complete, } } + + #[cfg(tokio_unstable)] + { + // Safety: the task is in the RUNNING state, which excludes + // concurrent shutdown and termination metadata access. + let mut task_meta = unsafe { self.task_meta() }; + let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { + self.core() + .scheduler + .task_poll_start_callback(&mut task_meta); + })); + + if let Err(panic) = res { + poll_hook_panic(self.core(), panic); + return PollFuture::Complete; + } + + if self.state().load().is_cancelled() { + cancel_task(self.core()); + return PollFuture::Complete; + } + } + let header_ptr = self.header_ptr(); let waker_ref = waker_ref::(&header_ptr); let cx = Context::from_waker(&waker_ref); - let res = poll_future(self.core(), cx); + let res = poll_future( + self.core(), + #[cfg(tokio_unstable)] + header_ptr, + cx, + ); + + #[cfg(tokio_unstable)] + { + // Safety: the task is still in the RUNNING state, which + // excludes concurrent shutdown and termination metadata + // access. + let mut task_meta = unsafe { self.task_meta() }; + let hook_res = panic::catch_unwind(panic::AssertUnwindSafe(|| { + self.core() + .scheduler + .task_poll_stop_callback(&mut task_meta); + })); + + if let Err(panic) = hook_res { + poll_hook_panic(self.core(), panic); + return PollFuture::Complete; + } + } if res == Poll::Ready(()) { // The future completed. Move on to complete the task. @@ -255,6 +316,8 @@ where // because we are going to drop them. This only matters when running // under loom. self.trailer().waker.with_mut(|_| ()); + #[cfg(tokio_unstable)] + self.trailer().user_data.with_mut(|_| ()); self.core().stage.with_mut(|_| ()); // Safety: The caller of this method just transitioned our ref-count to @@ -369,13 +432,12 @@ 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. #[cfg(tokio_unstable)] - if let Some(f) = self.trailer().hooks.task_terminate_callback.as_ref() { + { + // Safety: completion owns the task lifecycle transition, and the + // terminate hook is invoked synchronously with exclusive metadata access. + let mut meta = unsafe { self.task_meta() }; let _ = panic::catch_unwind(panic::AssertUnwindSafe(|| { - f(&TaskMeta { - id: self.core().task_id, - spawned_at: self.core().spawned_at.into(), - _phantom: Default::default(), - }) + self.core().scheduler.task_terminate_callback(&mut meta) })); } @@ -516,9 +578,35 @@ fn panic_result_to_join_error( } } +#[cfg(tokio_unstable)] +fn poll_hook_panic( + core: &Core, + hook_panic: Box, +) { + let drop_res = panic::catch_unwind(panic::AssertUnwindSafe(|| { + core.drop_future_or_output(); + })); + let join_error = match drop_res { + Ok(()) => panic_to_error(&core.scheduler, core.task_id, hook_panic), + Err(drop_panic) => panic_to_error(&core.scheduler, core.task_id, drop_panic), + }; + + let res = panic::catch_unwind(panic::AssertUnwindSafe(|| { + core.store_output(Err(join_error)); + })); + + if res.is_err() { + core.scheduler.unhandled_panic(); + } +} + /// Polls the future. If the future completes, the output is written to the /// stage field. -fn poll_future(core: &Core, cx: Context<'_>) -> Poll<()> { +fn poll_future( + core: &Core, + #[cfg(tokio_unstable)] header: NonNull
, + cx: Context<'_>, +) -> Poll<()> { // Poll the future. let output = panic::catch_unwind(panic::AssertUnwindSafe(|| { struct Guard<'a, T: Future, S: Schedule> { @@ -532,7 +620,11 @@ fn poll_future(core: &Core, cx: Context<'_>) -> Po } } let guard = Guard { core }; - let res = guard.core.poll(cx); + let res = guard.core.poll( + #[cfg(tokio_unstable)] + header, + cx, + ); mem::forget(guard); res })); diff --git a/tokio/src/runtime/task/list.rs b/tokio/src/runtime/task/list.rs index e4efa242a..3bdd7455f 100644 --- a/tokio/src/runtime/task/list.rs +++ b/tokio/src/runtime/task/list.rs @@ -86,19 +86,50 @@ impl OwnedTasks { /// Binds the provided task to this `OwnedTasks` instance. This fails if the /// `OwnedTasks` has been closed. + #[cfg_attr(tokio_unstable, allow(dead_code))] pub(crate) fn bind( &self, task: T, scheduler: S, id: super::Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> (JoinHandle, Option>) where S: Schedule, T: Future + Send + 'static, T::Output: Send + 'static, { - let (task, notified, join) = super::new_task(task, scheduler, id, spawned_at); + let (task, notified, join) = super::new_task( + task, + scheduler, + id, + spawned_at, + #[cfg(tokio_unstable)] + user_data, + ); + let notified = unsafe { self.bind_inner(task, notified) }; + (join, notified) + } + + #[cfg(tokio_unstable)] + pub(crate) fn bind_with_spawn_hook( + &self, + task: T, + scheduler: S, + id: super::Id, + spawned_at: SpawnLocation, + user_data: Option, + before_bind: impl FnOnce(&Task), + ) -> (JoinHandle, Option>) + where + S: Schedule, + T: Future + Send + 'static, + T::Output: Send + 'static, + { + let (task, notified, join) = + super::new_task(task, scheduler, id, spawned_at, user_data); + before_bind(&task); let notified = unsafe { self.bind_inner(task, notified) }; (join, notified) } @@ -108,19 +139,53 @@ impl OwnedTasks { /// # Safety /// /// Only use this in `LocalRuntime` where the task cannot move + #[cfg_attr(tokio_unstable, allow(dead_code))] pub(crate) unsafe fn bind_local( &self, task: T, scheduler: S, id: super::Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> (JoinHandle, Option>) where S: Schedule, T: Future + 'static, T::Output: 'static, { - let (task, notified, join) = super::new_task(task, scheduler, id, spawned_at); + let (task, notified, join) = super::new_task( + task, + scheduler, + id, + spawned_at, + #[cfg(tokio_unstable)] + user_data, + ); + let notified = unsafe { self.bind_inner(task, notified) }; + (join, notified) + } + + /// # Safety + /// + /// Only use this in `LocalRuntime` where the task cannot move. + #[cfg(tokio_unstable)] + pub(crate) unsafe fn bind_local_with_spawn_hook( + &self, + task: T, + scheduler: S, + id: super::Id, + spawned_at: SpawnLocation, + user_data: Option, + before_bind: impl FnOnce(&Task), + ) -> (JoinHandle, Option>) + where + S: Schedule, + T: Future + 'static, + T::Output: 'static, + { + let (task, notified, join) = + super::new_task(task, scheduler, id, spawned_at, user_data); + before_bind(&task); let notified = unsafe { self.bind_inner(task, notified) }; (join, notified) } @@ -264,13 +329,21 @@ impl LocalOwnedTasks { scheduler: S, id: super::Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> (JoinHandle, Option>) where S: Schedule, T: Future + 'static, T::Output: 'static, { - let (task, notified, join) = super::new_task(task, scheduler, id, spawned_at); + let (task, notified, join) = super::new_task( + task, + scheduler, + id, + spawned_at, + #[cfg(tokio_unstable)] + user_data, + ); unsafe { // safety: We just created the task, so we have exclusive access diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index 6f964c1de..d83b1ac98 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -119,6 +119,12 @@ //! the `JoinHandle` exclusive access again so that it is able to drop the waker //! at a later point. //! +//! * The `user_data` field is accessed only through scoped task hook metadata. +//! The spawn hook runs before the task is scheduled, poll hooks run while +//! the task holds the RUNNING lock but outside the actual future poll, and +//! the terminate hook runs after completion. Parent task metadata exposed to +//! spawn hooks is read-only. +//! //! All other fields are immutable and can be accessed immutably without //! synchronization by anyone. //! @@ -221,7 +227,6 @@ use crate::future::Future; use crate::util::linked_list; use crate::util::sharded_list; -use crate::runtime::TaskCallback; use std::marker::PhantomData; use std::panic::Location; use std::ptr::NonNull; @@ -241,14 +246,6 @@ unsafe impl Sync for Task {} #[repr(transparent)] pub(crate) struct Notified(Task); -impl Notified { - #[cfg(all(tokio_unstable, feature = "rt-multi-thread"))] - #[inline] - pub(crate) fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { - self.0.task_meta() - } -} - // safety: This type cannot be used to touch the task without first verifying // that the value is on a thread where it is safe to poll the task. unsafe impl Send for Notified {} @@ -262,14 +259,6 @@ pub(crate) struct LocalNotified { _not_send: PhantomData<*const ()>, } -impl LocalNotified { - #[cfg(tokio_unstable)] - #[inline] - pub(crate) fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { - self.task.task_meta() - } -} - /// A task that is not owned by any `OwnedTasks`. Used for blocking tasks. /// This type holds two ref-counts. pub(crate) struct UnownedTask { @@ -284,12 +273,6 @@ unsafe impl Sync for UnownedTask {} /// Task result sent back. pub(crate) type Result = std::result::Result; -/// Hooks for scheduling tasks which are needed in the task harness. -#[derive(Clone)] -pub(crate) struct TaskHarnessScheduleHooks { - pub(crate) task_terminate_callback: Option, -} - 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 @@ -301,8 +284,6 @@ pub(crate) trait Schedule: Sync + Sized + 'static { /// Schedule the task fn schedule(&self, task: Notified); - fn hooks(&self) -> TaskHarnessScheduleHooks; - /// Schedule the task to run in the near future, yielding the thread to /// other tasks. fn yield_now(&self, task: Notified) { @@ -313,6 +294,15 @@ pub(crate) trait Schedule: Sync + Sized + 'static { fn unhandled_panic(&self) { // By default, do nothing. This maintains the 1.0 behavior. } + + #[cfg(tokio_unstable)] + fn task_terminate_callback(&self, _meta: &mut crate::runtime::TaskMeta<'_>) {} + + #[cfg(tokio_unstable)] + fn task_poll_start_callback(&self, _meta: &mut crate::runtime::TaskMeta<'_>) {} + + #[cfg(tokio_unstable)] + fn task_poll_stop_callback(&self, _meta: &mut crate::runtime::TaskMeta<'_>) {} } cfg_rt! { @@ -325,6 +315,7 @@ cfg_rt! { scheduler: S, id: Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> (Task, Notified, JoinHandle) where S: Schedule, @@ -336,6 +327,8 @@ cfg_rt! { scheduler, id, spawned_at, + #[cfg(tokio_unstable)] + user_data, ); let task = Task { raw, @@ -359,6 +352,7 @@ cfg_rt! { scheduler: S, id: Id, spawned_at: SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> (UnownedTask, JoinHandle) where S: Schedule, @@ -370,6 +364,8 @@ cfg_rt! { scheduler, id, spawned_at, + #[cfg(tokio_unstable)] + user_data, ); // This transfers the ref-count of task and notified into an UnownedTask. @@ -429,22 +425,16 @@ impl Task { unsafe { Header::get_id(self.raw.header_ptr()) } } - #[cfg(tokio_unstable)] - pub(crate) fn spawned_at(&self) -> &'static Location<'static> { - // Safety: The header pointer is valid. - unsafe { Header::get_spawn_location(self.raw.header_ptr()) } - } - // 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`. #[cfg(tokio_unstable)] - pub(crate) fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { - crate::runtime::TaskMeta { - id: self.id(), - spawned_at: self.spawned_at().into(), - _phantom: PhantomData, - } + /// # Safety + /// + /// The returned metadata must have exclusive access to hook data for as long + /// as it can expose mutable references. + pub(crate) unsafe fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { + unsafe { self.raw.task_meta() } } cfg_taskdump! { @@ -677,3 +667,15 @@ impl SpawnLocation { Self::from(Location::caller()) } } + +#[cfg(tokio_unstable)] +pub(crate) fn current_task_meta<'meta>() -> Option> { + let ptr = crate::runtime::context::current_task()?; + + // Safety: the context stores this pointer only while the referenced task is + // being polled, so the allocation is alive for the duration of this call. + let raw = unsafe { RawTask::from_raw(ptr.cast()) }; + // Safety: parent metadata is exposed read-only during synchronous spawn + // hook invocation while no mutable parent hook metadata is live. + Some(unsafe { raw.task_meta_ref() }) +} diff --git a/tokio/src/runtime/task/raw.rs b/tokio/src/runtime/task/raw.rs index a9d143f74..ca409f72e 100644 --- a/tokio/src/runtime/task/raw.rs +++ b/tokio/src/runtime/task/raw.rs @@ -11,6 +11,8 @@ use crate::future::Future; use crate::runtime::task::core::{Core, Trailer}; use crate::runtime::task::{Cell, Harness, Header, Id, Schedule, State}; #[cfg(tokio_unstable)] +use crate::runtime::TaskData; +#[cfg(tokio_unstable)] use std::panic::Location; use std::ptr::NonNull; use std::task::{Poll, Waker}; @@ -213,6 +215,7 @@ impl RawTask { scheduler: S, id: Id, _spawned_at: super::SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option, ) -> RawTask where T: Future, @@ -225,6 +228,8 @@ impl RawTask { id, #[cfg(tokio_unstable)] _spawned_at.0, + #[cfg(tokio_unstable)] + user_data, )); let ptr = unsafe { NonNull::new_unchecked(ptr.cast()) }; @@ -267,6 +272,40 @@ impl RawTask { unsafe { &*self.trailer_ptr().as_ptr() } } + #[cfg(tokio_unstable)] + /// # Safety + /// + /// The task allocation must be live, and the returned metadata must have + /// exclusive access to hook data for as long as it can expose mutable references. + pub(crate) unsafe fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { + // Safety: `self` holds a live task reference, and callers use the + // metadata only for the current hook invocation. + unsafe { + crate::runtime::TaskMeta::new( + Header::get_id(self.ptr), + Header::get_spawn_location(self.ptr).into(), + Some(self.trailer().user_data_ptr()), + ) + } + } + + #[cfg(tokio_unstable)] + /// # Safety + /// + /// The task allocation must be live, and hook data must not be mutated while + /// references exposed through the returned metadata are live. + pub(crate) unsafe fn task_meta_ref<'meta>(&self) -> crate::runtime::TaskMetaRef<'meta> { + // Safety: `self` holds a live task reference, and this only exposes + // shared access to task data. + unsafe { + crate::runtime::TaskMetaRef::new( + Header::get_id(self.ptr), + Header::get_spawn_location(self.ptr).into(), + Some(self.trailer().user_data_ptr()), + ) + } + } + /// Returns a reference to the task's state. pub(super) fn state(&self) -> &State { &self.header().state diff --git a/tokio/src/runtime/task_hooks.rs b/tokio/src/runtime/task_hooks.rs index 6df3837b5..e1cfbcac5 100644 --- a/tokio/src/runtime/task_hooks.rs +++ b/tokio/src/runtime/task_hooks.rs @@ -1,10 +1,18 @@ use super::Config; +#[cfg(tokio_unstable)] +use std::any::Any; use std::marker::PhantomData; +#[cfg(tokio_unstable)] +use std::ptr::NonNull; +use std::sync::Arc; + +#[cfg(tokio_unstable)] +pub(crate) type TaskData = Box; impl TaskHooks { - pub(crate) fn spawn(&self, meta: &TaskMeta<'_>) { + pub(crate) fn spawn(&self, meta: &mut TaskMeta<'_>, parent: Option>) { if let Some(f) = self.task_spawn_callback.as_ref() { - f(meta) + f(meta, parent) } } @@ -22,7 +30,7 @@ impl TaskHooks { #[cfg(tokio_unstable)] #[inline] - pub(crate) fn poll_start_callback(&self, meta: &TaskMeta<'_>) { + pub(crate) fn poll_start_callback(&self, meta: &mut TaskMeta<'_>) { if let Some(poll_start) = &self.before_poll_callback { (poll_start)(meta); } @@ -30,16 +38,25 @@ impl TaskHooks { #[cfg(tokio_unstable)] #[inline] - pub(crate) fn poll_stop_callback(&self, meta: &TaskMeta<'_>) { + pub(crate) fn poll_stop_callback(&self, meta: &mut TaskMeta<'_>) { if let Some(poll_stop) = &self.after_poll_callback { (poll_stop)(meta); } } + + #[cfg(tokio_unstable)] + #[inline] + pub(crate) fn task_terminate_callback(&self, meta: &mut TaskMeta<'_>) { + if let Some(task_terminate) = &self.task_terminate_callback { + (task_terminate)(meta); + } + } } #[derive(Clone)] pub(crate) struct TaskHooks { - pub(crate) task_spawn_callback: Option, + pub(crate) task_spawn_callback: Option, + #[cfg_attr(not(tokio_unstable), allow(dead_code))] pub(crate) task_terminate_callback: Option, #[cfg(tokio_unstable)] pub(crate) before_poll_callback: Option, @@ -62,10 +79,44 @@ 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, + #[cfg(tokio_unstable)] + pub(crate) user_data: Option>>, pub(crate) _phantom: PhantomData<&'a ()>, } impl<'a> TaskMeta<'a> { + #[cfg(not(tokio_unstable))] + pub(crate) fn new( + id: super::task::Id, + spawned_at: crate::runtime::task::SpawnLocation, + ) -> Self { + Self { + id, + spawned_at, + _phantom: PhantomData, + } + } + + /// # Safety + /// + /// If `user_data` is present, it must point to live task storage, and this + /// metadata value must have exclusive access to that storage while it can + /// expose mutable references to it. + #[cfg(tokio_unstable)] + pub(crate) unsafe fn new( + id: super::task::Id, + spawned_at: crate::runtime::task::SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option>>, + ) -> Self { + Self { + id, + spawned_at, + #[cfg(tokio_unstable)] + user_data, + _phantom: PhantomData, + } + } + /// Return the opaque ID of the task. #[cfg_attr(not(tokio_unstable), allow(unreachable_pub, dead_code))] pub fn id(&self) -> super::task::Id { @@ -77,7 +128,136 @@ impl<'a> TaskMeta<'a> { pub fn spawned_at(&self) -> &'static std::panic::Location<'static> { self.spawned_at.0 } + + /// Returns a shared reference to this task's user data when the stored type + /// is `T`. + #[cfg(tokio_unstable)] + pub fn data(&self) -> Option<&T> { + let user_data = self.user_data?; + + // Safety: `TaskMeta` is only constructed while the task allocation is + // known to be alive. Shared access is allowed for the duration of hook + // invocation. + unsafe { user_data.as_ref().as_ref()?.downcast_ref::() } + } + + /// Returns a mutable reference to this task's user data when the stored type + /// is `T`. + #[cfg(tokio_unstable)] + pub fn data_mut(&mut self) -> Option<&mut T> { + let mut user_data = self.user_data?; + + // Safety: mutable `TaskMeta` is only handed to hooks for the exact task + // currently being initialized, polled under the RUNNING task state, or + // terminated. Tokio does not hold this borrow across polling the future. + unsafe { user_data.as_mut().as_mut()?.downcast_mut::() } + } + + /// Replaces this task's user data. + #[cfg(tokio_unstable)] + pub fn set_data(&mut self, data: T) { + if let Some(mut user_data) = self.user_data { + // Safety: see `data_mut`. + unsafe { + *user_data.as_mut() = Some(Box::new(data)); + } + } + } + + /// Takes this task's user data when the stored type is `T`. + #[cfg(tokio_unstable)] + pub fn take_data(&mut self) -> Option> { + let mut user_data = self.user_data?; + + // Safety: see `data_mut`. + unsafe { + let user_data = user_data.as_mut(); + if !user_data.as_ref()?.is::() { + return None; + } + + user_data.take()?.downcast::().ok() + } + } + + /// Clears this task's user data, returning whether any data was present. + #[cfg(tokio_unstable)] + pub fn clear_data(&mut self) -> bool { + let Some(mut user_data) = self.user_data else { + return false; + }; + + // Safety: see `data_mut`. + unsafe { user_data.as_mut().take().is_some() } + } +} + +/// Read-only task metadata supplied to task spawn hooks for parent tasks. +/// +/// **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 TaskMetaRef<'a> { + /// The opaque ID of the task. + pub(crate) id: super::task::Id, + /// 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, + #[cfg(tokio_unstable)] + pub(crate) user_data: Option>>, + pub(crate) _phantom: PhantomData<&'a ()>, +} + +impl<'a> TaskMetaRef<'a> { + /// # Safety + /// + /// If `user_data` is present, it must point to live task storage for the + /// duration of any references exposed through this metadata value. + #[cfg(tokio_unstable)] + pub(crate) unsafe fn new( + id: super::task::Id, + spawned_at: crate::runtime::task::SpawnLocation, + #[cfg(tokio_unstable)] user_data: Option>>, + ) -> Self { + Self { + id, + spawned_at, + #[cfg(tokio_unstable)] + user_data, + _phantom: PhantomData, + } + } + + /// 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 + } + + /// Return the source code location where the task was spawned. + #[cfg(tokio_unstable)] + pub fn spawned_at(&self) -> &'static std::panic::Location<'static> { + self.spawned_at.0 + } + + /// Returns a shared reference to this task's user data when the stored type + /// is `T`. + #[cfg(tokio_unstable)] + pub fn data(&self) -> Option<&T> { + let user_data = self.user_data?; + + // Safety: `TaskMetaRef` is only constructed while the task allocation is + // known to be alive, and it does not expose mutation. + unsafe { user_data.as_ref().as_ref()?.downcast_ref::() } + } } /// Runs on specific task-related events -pub(crate) type TaskCallback = std::sync::Arc) + Send + Sync>; +pub(crate) type TaskCallback = Arc) + Send + Sync>; + +pub(crate) type TaskSpawnCallback = + Arc, Option>) + Send + Sync>; diff --git a/tokio/src/runtime/tests/mod.rs b/tokio/src/runtime/tests/mod.rs index b277cfa65..fb470f659 100644 --- a/tokio/src/runtime/tests/mod.rs +++ b/tokio/src/runtime/tests/mod.rs @@ -6,7 +6,7 @@ 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}; /// `task::Schedule` implementation that does nothing, for testing. pub(crate) struct NoopSchedule; @@ -20,11 +20,6 @@ mod noop_scheduler { unreachable!(); } - fn hooks(&self) -> TaskHarnessScheduleHooks { - TaskHarnessScheduleHooks { - task_terminate_callback: None, - } - } } } @@ -43,7 +38,14 @@ mod unowned_wrapper { let span = tracing::trace_span!("test_span"); let task = task.instrument(span); let (task, handle) = - crate::runtime::task::unowned(task, NoopSchedule, Id::next(), SpawnLocation::capture()); + crate::runtime::task::unowned( + task, + NoopSchedule, + Id::next(), + SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, + ); (task.into_notified(), handle) } @@ -55,7 +57,14 @@ mod unowned_wrapper { T::Output: Send + 'static, { let (task, handle) = - crate::runtime::task::unowned(task, NoopSchedule, Id::next(), SpawnLocation::capture()); + crate::runtime::task::unowned( + task, + NoopSchedule, + Id::next(), + SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, + ); (task.into_notified(), handle) } } diff --git a/tokio/src/runtime/tests/task.rs b/tokio/src/runtime/tests/task.rs index 7a10ac4a3..c28b046e4 100644 --- a/tokio/src/runtime/tests/task.rs +++ b/tokio/src/runtime/tests/task.rs @@ -1,6 +1,5 @@ use crate::runtime::task::{ self, unowned, Id, JoinHandle, OwnedTasks, Schedule, SpawnLocation, Task, - TaskHarnessScheduleHooks, }; use crate::runtime::tests::NoopSchedule; @@ -61,6 +60,8 @@ fn create_drop1() { NoopSchedule, Id::next(), SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, ); drop(notified); handle.assert_not_dropped(); @@ -79,6 +80,8 @@ fn create_drop2() { NoopSchedule, Id::next(), SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, ); drop(join); handle.assert_not_dropped(); @@ -97,6 +100,8 @@ fn drop_abort_handle1() { NoopSchedule, Id::next(), SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, ); let abort = join.abort_handle(); drop(join); @@ -118,6 +123,8 @@ fn drop_abort_handle2() { NoopSchedule, Id::next(), SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, ); let abort = join.abort_handle(); drop(notified); @@ -139,6 +146,8 @@ fn drop_abort_handle_clone() { NoopSchedule, Id::next(), SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, ); let abort = join.abort_handle(); let abort_clone = abort.clone(); @@ -164,6 +173,8 @@ fn create_shutdown1() { NoopSchedule, Id::next(), SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, ); drop(join); handle.assert_not_dropped(); @@ -182,6 +193,8 @@ fn create_shutdown2() { NoopSchedule, Id::next(), SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, ); handle.assert_not_dropped(); notified.shutdown(); @@ -191,7 +204,14 @@ fn create_shutdown2() { #[test] fn unowned_poll() { - let (task, _) = unowned(async {}, NoopSchedule, Id::next(), SpawnLocation::capture()); + let (task, _) = unowned( + async {}, + NoopSchedule, + Id::next(), + SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, + ); task.run(); } @@ -402,9 +422,14 @@ impl Runtime { T::Output: 'static + Send, { let (handle, notified) = - self.0 - .owned - .bind(future, self.clone(), Id::next(), SpawnLocation::capture()); + self.0.owned.bind( + future, + self.clone(), + Id::next(), + SpawnLocation::capture(), + #[cfg(tokio_unstable)] + None, + ); if let Some(notified) = notified { self.schedule(notified); @@ -460,10 +485,4 @@ impl Schedule for Runtime { fn schedule(&self, task: task::Notified) { self.0.core.try_lock().unwrap().queue.push_back(task); } - - fn hooks(&self) -> TaskHarnessScheduleHooks { - TaskHarnessScheduleHooks { - task_terminate_callback: None, - } - } } diff --git a/tokio/src/task/builder.rs b/tokio/src/task/builder.rs index c253e853e..bf5d210aa 100644 --- a/tokio/src/task/builder.rs +++ b/tokio/src/task/builder.rs @@ -4,7 +4,7 @@ use crate::{ task::{JoinHandle, LocalSet}, util::trace::SpawnMeta, }; -use std::{future::Future, io, mem}; +use std::{any::Any, fmt, future::Future, io, mem}; /// Factory which is used to configure the properties of a new task. /// @@ -14,10 +14,11 @@ use std::{future::Future, io, mem}; /// /// Methods can be chained in order to configure it. /// -/// Currently, there is only one configuration option: +/// Configuration options include: /// /// - [`name`], which specifies an associated name for /// the task +/// - [`data`], which stores user data for runtime task hooks /// /// There are three types of task that can be spawned from a Builder: /// - [`spawn_local`] for executing not [`Send`] futures @@ -55,13 +56,15 @@ use std::{future::Future, io, mem}; /// ``` /// [unstable]: crate#unstable-features /// [`name`]: Builder::name +/// [`data`]: Builder::data /// [`spawn_local`]: Builder::spawn_local /// [`spawn`]: Builder::spawn /// [`spawn_blocking`]: Builder::spawn_blocking -#[derive(Default, Debug)] +#[derive(Default)] #[cfg_attr(docsrs, doc(cfg(all(tokio_unstable, feature = "tracing"))))] pub struct Builder<'a> { name: Option<&'a str>, + data: Option, } impl<'a> Builder<'a> { @@ -71,8 +74,22 @@ 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(self, name: &'a str) -> Self { + Self { + name: Some(name), + ..self + } + } + + /// Sets task data visible to runtime task hooks. + pub fn data(self, data: T) -> Self + where + T: Any + Send + Sync + 'static, + { + Self { + data: Some(Box::new(data)), + ..self + } } /// Spawns a task with this builder's settings on the current runtime. @@ -89,11 +106,12 @@ impl<'a> Builder<'a> { Fut: Future + Send + 'static, Fut::Output: Send + 'static, { + let Builder { name, data } = self; let fut_size = mem::size_of::(); 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(name, fut_size), data) } else { - super::spawn::spawn_inner(future, SpawnMeta::new(self.name, fut_size)) + super::spawn::spawn_inner(future, SpawnMeta::new(name, fut_size), data) }) } @@ -110,11 +128,12 @@ impl<'a> Builder<'a> { Fut: Future + Send + 'static, Fut::Output: Send + 'static, { + let Builder { name, data } = self; let fut_size = mem::size_of::(); Ok(if fut_size > BOX_FUTURE_THRESHOLD { - handle.spawn_named(Box::pin(future), SpawnMeta::new(self.name, fut_size)) + handle.spawn_named_with_data(Box::pin(future), SpawnMeta::new(name, fut_size), data) } else { - handle.spawn_named(future, SpawnMeta::new(self.name, fut_size)) + handle.spawn_named_with_data(future, SpawnMeta::new(name, fut_size), data) }) } @@ -141,11 +160,12 @@ impl<'a> Builder<'a> { Fut: Future + 'static, Fut::Output: 'static, { + let Builder { name, data } = self; let fut_size = mem::size_of::(); 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(name, fut_size), data) } else { - super::local::spawn_local_inner(future, SpawnMeta::new(self.name, fut_size)) + super::local::spawn_local_inner(future, SpawnMeta::new(name, fut_size), data) }) } @@ -166,11 +186,12 @@ impl<'a> Builder<'a> { Fut: Future + 'static, Fut::Output: 'static, { + let Builder { name, data } = self; let fut_size = mem::size_of::(); Ok(if fut_size > BOX_FUTURE_THRESHOLD { - local_set.spawn_named(Box::pin(future), SpawnMeta::new(self.name, fut_size)) + local_set.spawn_named_with_data(Box::pin(future), SpawnMeta::new(name, fut_size), data) } else { - local_set.spawn_named(future, SpawnMeta::new(self.name, fut_size)) + local_set.spawn_named_with_data(future, SpawnMeta::new(name, fut_size), data) }) } @@ -212,20 +233,23 @@ impl<'a> Builder<'a> { Output: Send + 'static, { use crate::runtime::Mandatory; + let Builder { name, data } = self; let fn_size = mem::size_of::(); let (join_handle, spawn_result) = if fn_size > BOX_FUTURE_THRESHOLD { handle.inner.blocking_spawner().spawn_blocking_inner( Box::new(function), Mandatory::NonMandatory, - SpawnMeta::new(self.name, fn_size), + SpawnMeta::new(name, fn_size), handle, + data, ) } else { handle.inner.blocking_spawner().spawn_blocking_inner( function, Mandatory::NonMandatory, - SpawnMeta::new(self.name, fn_size), + SpawnMeta::new(name, fn_size), handle, + data, ) }; @@ -233,3 +257,12 @@ impl<'a> Builder<'a> { Ok(join_handle) } } + +impl fmt::Debug for Builder<'_> { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("Builder") + .field("name", &self.name) + .field("data", &self.data.as_ref().map(|_| "")) + .finish() + } +} diff --git a/tokio/src/task/join_set.rs b/tokio/src/task/join_set.rs index f280a7461..631282241 100644 --- a/tokio/src/task/join_set.rs +++ b/tokio/src/task/join_set.rs @@ -705,6 +705,15 @@ impl<'a, T: 'static> Builder<'a, T> { Self { builder, ..self } } + /// Sets task data visible to runtime task hooks. + pub fn data(self, data: D) -> Self + where + D: std::any::Any + Send + Sync + 'static, + { + let builder = self.builder.data(data); + Self { builder, ..self } + } + /// Spawn the provided task with this builder's settings and store it in the /// [`JoinSet`], returning an [`AbortHandle`] that can be used to remotely /// cancel the task. diff --git a/tokio/src/task/local.rs b/tokio/src/task/local.rs index cf75097a8..f3df25092 100644 --- a/tokio/src/task/local.rs +++ b/tokio/src/task/local.rs @@ -2,9 +2,7 @@ use crate::loom::cell::UnsafeCell; use crate::loom::sync::{Arc, Mutex}; use crate::runtime; -use crate::runtime::task::{ - self, JoinHandle, LocalOwnedTasks, SpawnLocation, Task, TaskHarnessScheduleHooks, -}; +use crate::runtime::task::{self, JoinHandle, LocalOwnedTasks, SpawnLocation, Task}; use crate::runtime::{context, ThreadId, BOX_FUTURE_THRESHOLD}; use crate::sync::AtomicWaker; use crate::util::trace::SpawnMeta; @@ -399,21 +397,37 @@ cfg_rt! { { let fut_size = std::mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { - spawn_local_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) + spawn_local_inner( + Box::pin(future), + SpawnMeta::new_unnamed(fut_size), + #[cfg(tokio_unstable)] + None, + ) } else { - spawn_local_inner(future, SpawnMeta::new_unnamed(fut_size)) + spawn_local_inner( + future, + SpawnMeta::new_unnamed(fut_size), + #[cfg(tokio_unstable)] + None, + ) } } #[track_caller] - pub(super) fn spawn_local_inner(future: F, meta: SpawnMeta<'_>) -> JoinHandle + pub(super) fn spawn_local_inner( + future: F, + meta: SpawnMeta<'_>, + #[cfg(tokio_unstable)] user_data: Option, + ) -> JoinHandle where F: Future + 'static, F::Output: 'static { use crate::runtime::{context, task}; let mut future = Some(future); + #[cfg(tokio_unstable)] + let mut user_data = Some(user_data); let res = context::with_current(|handle| { Some(if handle.is_local() { @@ -439,11 +453,24 @@ 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 - unsafe { handle.spawn_local(task, id, meta.spawned_at) } + unsafe { + handle.spawn_local( + task, + id, + meta.spawned_at, + #[cfg(tokio_unstable)] + user_data.take().unwrap(), + ) + } } else { match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) { None => panic!("`spawn_local` called from outside of a `task::LocalSet` or `runtime::LocalRuntime`"), - Some(cx) => cx.spawn(future.take().unwrap(), meta) + Some(cx) => cx.spawn( + future.take().unwrap(), + meta, + #[cfg(tokio_unstable)] + user_data.take().unwrap(), + ) } }) }); @@ -453,7 +480,12 @@ cfg_rt! { Ok(Some(join_handle)) => join_handle, Err(_) => match CURRENT.with(|LocalData { ctx, .. }| ctx.get()) { None => panic!("`spawn_local` called from outside of a `task::LocalSet` or `runtime::LocalRuntime`"), - Some(cx) => cx.spawn(future.unwrap(), meta) + Some(cx) => cx.spawn( + future.unwrap(), + meta, + #[cfg(tokio_unstable)] + user_data.unwrap(), + ) } } } @@ -729,16 +761,46 @@ impl LocalSet { F: Future + 'static, F::Output: 'static, { - self.spawn_named_inner(future, meta) + self.spawn_named_inner( + future, + meta, + #[cfg(tokio_unstable)] + None, + ) } + #[cfg(all(tokio_unstable, feature = "tracing"))] #[track_caller] - fn spawn_named_inner(&self, future: F, meta: SpawnMeta<'_>) -> JoinHandle + pub(in crate::task) fn spawn_named_with_data( + &self, + future: F, + meta: SpawnMeta<'_>, + user_data: Option, + ) -> JoinHandle where F: Future + 'static, F::Output: 'static, { - let handle = self.context.spawn(future, meta); + self.spawn_named_inner(future, meta, user_data) + } + + #[track_caller] + fn spawn_named_inner( + &self, + future: F, + meta: SpawnMeta<'_>, + #[cfg(tokio_unstable)] user_data: Option, + ) -> JoinHandle + where + F: Future + 'static, + F::Output: 'static, + { + let handle = self.context.spawn( + future, + meta, + #[cfg(tokio_unstable)] + user_data, + ); // Because a task was spawned from *outside* the `LocalSet`, wake the // `LocalSet` future to execute the new task, if it hasn't been woken. @@ -1024,7 +1086,12 @@ impl Drop for LocalSet { impl Context { #[track_caller] - fn spawn(&self, future: F, meta: SpawnMeta<'_>) -> JoinHandle + fn spawn( + &self, + future: F, + meta: SpawnMeta<'_>, + #[cfg(tokio_unstable)] user_data: Option, + ) -> JoinHandle where F: Future + 'static, F::Output: 'static, @@ -1040,6 +1107,8 @@ impl Context { self.shared.clone(), id, SpawnLocation::capture(), + #[cfg(tokio_unstable)] + user_data, ) }; @@ -1148,13 +1217,6 @@ impl task::Schedule for Arc { Shared::schedule(self, task); } - // localset does not currently support task hooks - fn hooks(&self) -> TaskHarnessScheduleHooks { - TaskHarnessScheduleHooks { - task_terminate_callback: None, - } - } - cfg_unstable! { fn unhandled_panic(&self) { use crate::runtime::UnhandledPanic; diff --git a/tokio/src/task/spawn.rs b/tokio/src/task/spawn.rs index cd26c6d76..c7d43cb2f 100644 --- a/tokio/src/task/spawn.rs +++ b/tokio/src/task/spawn.rs @@ -178,14 +178,28 @@ cfg_rt! { { let fut_size = std::mem::size_of::(); if fut_size > BOX_FUTURE_THRESHOLD { - spawn_inner(Box::pin(future), SpawnMeta::new_unnamed(fut_size)) + spawn_inner( + Box::pin(future), + SpawnMeta::new_unnamed(fut_size), + #[cfg(tokio_unstable)] + None, + ) } else { - spawn_inner(future, SpawnMeta::new_unnamed(fut_size)) + spawn_inner( + future, + SpawnMeta::new_unnamed(fut_size), + #[cfg(tokio_unstable)] + None, + ) } } #[track_caller] - pub(super) fn spawn_inner(future: T, meta: SpawnMeta<'_>) -> JoinHandle + pub(super) fn spawn_inner( + future: T, + meta: SpawnMeta<'_>, + #[cfg(tokio_unstable)] user_data: Option, + ) -> JoinHandle where T: Future + Send + 'static, T::Output: Send + 'static, @@ -207,7 +221,15 @@ cfg_rt! { let id = task::Id::next(); let task = crate::util::trace::task(future, "task", meta, id.as_u64()); - match context::with_current(|handle| handle.spawn(task, id, meta.spawned_at)) { + match context::with_current(|handle| { + handle.spawn( + task, + id, + meta.spawned_at, + #[cfg(tokio_unstable)] + user_data, + ) + }) { Ok(join_handle) => join_handle, Err(e) => panic!("{}", e), } diff --git a/tokio/tests/task_hooks.rs b/tokio/tests/task_hooks.rs index 42bb3fd94..75502d44b 100644 --- a/tokio/tests/task_hooks.rs +++ b/tokio/tests/task_hooks.rs @@ -2,8 +2,12 @@ #![cfg(all(feature = "full", tokio_unstable, target_has_atomic = "64"))] use std::collections::HashSet; +use std::future::Future; +use std::pin::Pin; use std::sync::atomic::{AtomicUsize, Ordering}; -use std::sync::{Arc, Mutex}; +use std::sync::{Arc, Condvar, Mutex}; +use std::task::{Context, Poll}; +use std::time::Duration; use tokio::runtime::Builder; @@ -19,7 +23,8 @@ fn spawn_task_hook_fires() { let ids2 = Arc::clone(&ids); let runtime = Builder::new_current_thread() - .on_task_spawn(move |data| { + .on_task_spawn(move |data, _parent| { + assert!(data.data::().is_none()); ids2.lock().unwrap().insert(data.id()); count2.fetch_add(1, Ordering::SeqCst); @@ -85,11 +90,11 @@ fn task_hook_spawn_location_current_thread() { "(current_thread) on_task_spawn", &spawns, )) - .on_before_task_poll(mk_spawn_location_hook( + .on_before_task_poll(mk_poll_location_hook( "(current_thread) on_before_task_poll", &poll_starts, )) - .on_after_task_poll(mk_spawn_location_hook( + .on_after_task_poll(mk_poll_location_hook( "(current_thread) on_after_task_poll", &poll_ends, )) @@ -136,11 +141,11 @@ fn task_hook_spawn_location_multi_thread() { "(multi_thread) on_task_spawn", &spawns, )) - .on_before_task_poll(mk_spawn_location_hook( + .on_before_task_poll(mk_poll_location_hook( "(multi_thread) on_before_task_poll", &poll_starts, )) - .on_after_task_poll(mk_spawn_location_hook( + .on_after_task_poll(mk_poll_location_hook( "(multi_thread) on_after_task_poll", &poll_ends, )) @@ -174,21 +179,430 @@ fn task_hook_spawn_location_multi_thread() { assert_eq!(poll_starts, poll_ends.fetch_add(0, Ordering::SeqCst)); } +#[derive(Debug)] +struct PollState { + before: usize, + after: usize, +} + +#[test] +fn task_data_mutates_across_current_thread_hooks() { + let terminated = Arc::new(Mutex::new(Vec::new())); + let terminated2 = Arc::clone(&terminated); + + let runtime = Builder::new_current_thread() + .on_task_spawn(|meta, parent| { + assert!(parent.is_none()); + assert!(!meta.clear_data()); + meta.set_data(PollState { + before: 0, + after: 0, + }); + }) + .on_before_task_poll(|meta| { + meta.data_mut::().unwrap().before += 1; + }) + .on_after_task_poll(|meta| { + if let Some(data) = meta.data_mut::() { + data.after += 1; + } + }) + .on_task_terminate(move |meta| { + let data = meta.take_data::().unwrap(); + assert!(meta.data::().is_none()); + terminated2.lock().unwrap().push((data.before, data.after)); + }) + .build() + .unwrap(); + + runtime.block_on(async { + tokio::spawn(async { + for _ in 0..3 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + }); + + let terminated = terminated.lock().unwrap(); + assert_eq!(terminated.len(), 1); + let (before, after) = terminated[0]; + assert!(before > 1); + assert_eq!(before, after); +} + +#[cfg_attr( + target_os = "wasi", + ignore = "WASI does not support multi-threaded runtime" +)] +#[test] +fn task_data_mutates_across_multi_thread_hooks() { + let terminated = Arc::new(Mutex::new(Vec::new())); + let terminated2 = Arc::clone(&terminated); + + let runtime = Builder::new_multi_thread() + .worker_threads(2) + .on_task_spawn(|meta, _parent| { + meta.set_data(PollState { + before: 0, + after: 0, + }); + }) + .on_before_task_poll(|meta| { + meta.data_mut::().unwrap().before += 1; + }) + .on_after_task_poll(|meta| { + if let Some(data) = meta.data_mut::() { + data.after += 1; + } + }) + .on_task_terminate(move |meta| { + let data = meta.take_data::().unwrap(); + terminated2.lock().unwrap().push((data.before, data.after)); + }) + .build() + .unwrap(); + + runtime.block_on(async { + tokio::spawn(async { + for _ in 0..3 { + tokio::task::yield_now().await; + } + }) + .await + .unwrap(); + }); + + runtime.shutdown_timeout(std::time::Duration::from_secs(60)); + + let terminated = terminated.lock().unwrap(); + assert_eq!(terminated.len(), 1); + let (before, after) = terminated[0]; + assert!(before > 1); + assert_eq!(before, after); +} + +#[cfg_attr( + target_os = "wasi", + ignore = "WASI does not support multi-threaded runtime" +)] +#[test] +fn poll_hook_data_is_not_terminated_during_multi_thread_shutdown() { + let entered = Arc::new((Mutex::new(false), Condvar::new())); + let entered2 = Arc::clone(&entered); + let release = Arc::new((Mutex::new(false), Condvar::new())); + let release2 = Arc::clone(&release); + let terminated = Arc::new(Mutex::new(Vec::new())); + let terminated2 = Arc::clone(&terminated); + + let runtime = Builder::new_multi_thread() + .worker_threads(2) + .on_task_spawn(|meta, _parent| { + meta.set_data(Vec::<&'static str>::new()); + }) + .on_before_task_poll(move |meta| { + let (lock, cvar) = &*entered2; + *lock.lock().unwrap() = true; + cvar.notify_one(); + + let (lock, cvar) = &*release2; + let mut released = lock.lock().unwrap(); + while !*released { + released = cvar.wait(released).unwrap(); + } + + meta.data_mut::>() + .unwrap() + .push("before_done"); + }) + .on_task_terminate(move |meta| { + let data = meta.take_data::>().unwrap(); + terminated2.lock().unwrap().push(*data); + }) + .build() + .unwrap(); + + drop(runtime.spawn(std::future::pending::<()>())); + + let (lock, cvar) = &*entered; + let entered_guard = lock.lock().unwrap(); + let (entered_guard, wait_result) = cvar + .wait_timeout_while(entered_guard, Duration::from_secs(5), |entered| !*entered) + .unwrap(); + assert!(*entered_guard); + assert!(!wait_result.timed_out()); + + let shutdown = std::thread::spawn(move || { + runtime.shutdown_timeout(Duration::from_secs(5)); + }); + + std::thread::sleep(Duration::from_millis(50)); + + let (lock, cvar) = &*release; + *lock.lock().unwrap() = true; + cvar.notify_one(); + + shutdown.join().unwrap(); + + assert_eq!(*terminated.lock().unwrap(), vec![vec!["before_done"]]); +} + +#[cfg_attr( + target_os = "wasi", + ignore = "WASI does not support multi-threaded runtime" +)] +#[test] +fn abort_during_before_poll_hook_does_not_poll_future() { + struct CountPolls { + polls: Arc, + } + + impl Future for CountPolls { + type Output = (); + + fn poll(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<()> { + self.polls.fetch_add(1, Ordering::SeqCst); + Poll::Pending + } + } + + let entered = Arc::new((Mutex::new(false), Condvar::new())); + let entered2 = Arc::clone(&entered); + let release = Arc::new((Mutex::new(false), Condvar::new())); + let release2 = Arc::clone(&release); + let polls = Arc::new(AtomicUsize::new(0)); + + let runtime = Builder::new_multi_thread() + .worker_threads(2) + .on_before_task_poll(move |_meta| { + let (lock, cvar) = &*entered2; + *lock.lock().unwrap() = true; + cvar.notify_one(); + + let (lock, cvar) = &*release2; + let mut released = lock.lock().unwrap(); + while !*released { + released = cvar.wait(released).unwrap(); + } + }) + .build() + .unwrap(); + + let task = runtime.spawn(CountPolls { + polls: Arc::clone(&polls), + }); + + let (lock, cvar) = &*entered; + let entered_guard = lock.lock().unwrap(); + let (entered_guard, wait_result) = cvar + .wait_timeout_while(entered_guard, Duration::from_secs(5), |entered| !*entered) + .unwrap(); + assert!(*entered_guard); + assert!(!wait_result.timed_out()); + + task.abort(); + + let (lock, cvar) = &*release; + *lock.lock().unwrap() = true; + cvar.notify_one(); + + let err = runtime.block_on(task).unwrap_err(); + assert!(err.is_cancelled()); + assert_eq!(polls.load(Ordering::SeqCst), 0); +} + +#[derive(Clone, Debug, Eq, PartialEq)] +struct Lineage { + depth: usize, +} + +#[test] +fn spawn_hook_can_inherit_parent_task_data() { + let terminated = Arc::new(Mutex::new(Vec::new())); + let terminated2 = Arc::clone(&terminated); + + let runtime = Builder::new_current_thread() + .on_task_spawn(|meta, parent| { + let depth = match parent { + Some(parent) => parent + .data::() + .map_or(0, |parent| parent.depth + 1), + None => 0, + }; + + meta.set_data(Lineage { depth }); + }) + .on_task_terminate(move |meta| { + let data = meta.take_data::().unwrap(); + terminated2.lock().unwrap().push(data.depth); + }) + .build() + .unwrap(); + + runtime.block_on(async { + tokio::spawn(async { + tokio::spawn(async {}).await.unwrap(); + }) + .await + .unwrap(); + }); + + let mut terminated = terminated.lock().unwrap().clone(); + terminated.sort_unstable(); + assert_eq!(terminated, vec![0, 1]); +} + +#[test] +fn spawn_hook_runs_before_terminate_when_current_thread_runtime_is_closed() { + struct ClosedSpawnData; + + let events = Arc::new(Mutex::new(Vec::new())); + let events2 = Arc::clone(&events); + let events3 = Arc::clone(&events); + + let runtime = Builder::new_current_thread() + .on_task_spawn(move |meta, _parent| { + meta.set_data(ClosedSpawnData); + events2.lock().unwrap().push("spawn"); + }) + .on_task_terminate(move |meta| { + if meta.take_data::().is_some() { + events3.lock().unwrap().push("terminate_with_data"); + } + }) + .build() + .unwrap(); + + let handle = runtime.handle().clone(); + drop(runtime); + + drop(handle.spawn(async {})); + + assert_eq!(*events.lock().unwrap(), ["spawn", "terminate_with_data"]); +} + +#[cfg_attr( + target_os = "wasi", + ignore = "WASI does not support multi-threaded runtime" +)] +#[test] +fn spawn_hook_runs_before_terminate_when_multi_thread_runtime_is_closed() { + struct ClosedSpawnData; + + let events = Arc::new(Mutex::new(Vec::new())); + let events2 = Arc::clone(&events); + let events3 = Arc::clone(&events); + + let runtime = Builder::new_multi_thread() + .worker_threads(1) + .on_task_spawn(move |meta, _parent| { + meta.set_data(ClosedSpawnData); + events2.lock().unwrap().push("spawn"); + }) + .on_task_terminate(move |meta| { + if meta.take_data::().is_some() { + events3.lock().unwrap().push("terminate_with_data"); + } + }) + .build() + .unwrap(); + + let handle = runtime.handle().clone(); + drop(runtime); + + drop(handle.spawn(async {})); + + assert_eq!(*events.lock().unwrap(), ["spawn", "terminate_with_data"]); +} + +#[cfg(feature = "tracing")] +#[test] +fn task_builder_data_is_visible_to_hooks() { + let terminated = Arc::new(Mutex::new(Vec::new())); + let terminated2 = Arc::clone(&terminated); + + let runtime = Builder::new_current_thread() + .on_task_spawn(|meta, _parent| { + let value = meta.data_mut::().unwrap(); + *value += 1; + }) + .on_task_terminate(move |meta| { + let value = meta.take_data::().unwrap(); + terminated2.lock().unwrap().push(*value); + }) + .build() + .unwrap(); + + runtime.block_on(async { + tokio::task::Builder::new() + .data(41usize) + .spawn(async {}) + .unwrap() + .await + .unwrap(); + }); + + assert_eq!(*terminated.lock().unwrap(), vec![42]); +} + +#[cfg(feature = "tracing")] +#[test] +fn task_builder_data_is_not_dropped_for_spawn_blocking() { + let terminated = Arc::new(Mutex::new(Vec::new())); + let terminated2 = Arc::clone(&terminated); + + let runtime = Builder::new_current_thread() + .on_task_terminate(move |meta| { + if let Some(value) = meta.take_data::() { + terminated2.lock().unwrap().push(*value); + } + }) + .build() + .unwrap(); + + runtime.block_on(async { + tokio::task::Builder::new() + .data(7usize) + .spawn_blocking(|| {}) + .unwrap() + .await + .unwrap(); + }); + + assert_eq!(*terminated.lock().unwrap(), vec![7]); +} + fn mk_spawn_location_hook( event: &'static str, count: &Arc, -) -> impl Fn(&tokio::runtime::TaskMeta<'_>) { +) -> impl Fn(&mut tokio::runtime::TaskMeta<'_>, Option>) { let count = Arc::clone(count); - move |data| { - eprintln!("{event} ({:?}): {:?}", data.id(), data.spawned_at()); - // Assert that the spawn location is in this file. - // Don't make assertions about line number/column here, as these - // may change as new code is added to the test file... - assert_eq!( - data.spawned_at().file(), - file!(), - "incorrect spawn location in {event} hook", - ); - count.fetch_add(1, Ordering::SeqCst); + move |data, _parent| { + assert_spawn_location(event, count.as_ref(), data); } } + +fn mk_poll_location_hook( + event: &'static str, + count: &Arc, +) -> impl Fn(&mut tokio::runtime::TaskMeta<'_>) { + let count = Arc::clone(count); + move |data| { + assert_spawn_location(event, count.as_ref(), data); + } +} + +fn assert_spawn_location( + event: &'static str, + count: &AtomicUsize, + data: &tokio::runtime::TaskMeta<'_>, +) { + eprintln!("{event} ({:?}): {:?}", data.id(), data.spawned_at()); + assert_eq!( + data.spawned_at().file(), + file!(), + "incorrect spawn location in {event} hook", + ); + count.fetch_add(1, Ordering::SeqCst); +}