From 89291980f15596ee9313e82c869556cda63c5ba5 Mon Sep 17 00:00:00 2001 From: noah Date: Sun, 10 May 2026 15:44:27 -0500 Subject: [PATCH] fix issue where we had hooks for term but not spawn with blocking tasks --- tokio/src/runtime/blocking/mod.rs | 2 + tokio/src/runtime/blocking/pool.rs | 77 ++++++++++- tokio/src/runtime/blocking/schedule.rs | 11 +- tokio/src/runtime/mod.rs | 2 + .../runtime/scheduler/multi_thread/worker.rs | 4 +- tokio/src/runtime/task/mod.rs | 9 ++ tokio/src/task/builder.rs | 4 + tokio/tests/task_hooks.rs | 122 ++++++++++++++++++ 8 files changed, 226 insertions(+), 5 deletions(-) diff --git a/tokio/src/runtime/blocking/mod.rs b/tokio/src/runtime/blocking/mod.rs index c42924be7..3df8070ae 100644 --- a/tokio/src/runtime/blocking/mod.rs +++ b/tokio/src/runtime/blocking/mod.rs @@ -5,6 +5,8 @@ mod pool; pub(crate) use pool::{spawn_blocking, BlockingPool, Spawner}; +#[cfg(feature = "rt-multi-thread")] +pub(crate) use pool::spawn_blocking_internal; cfg_fs! { pub(crate) use pool::spawn_mandatory_blocking; diff --git a/tokio/src/runtime/blocking/pool.rs b/tokio/src/runtime/blocking/pool.rs index 6a489810d..1a23baecb 100644 --- a/tokio/src/runtime/blocking/pool.rs +++ b/tokio/src/runtime/blocking/pool.rs @@ -185,6 +185,20 @@ where rt.spawn_blocking(func) } +/// Runs an internal runtime worker on the blocking pool without invoking task +/// lifecycle hooks. +#[track_caller] +#[cfg(feature = "rt-multi-thread")] +#[cfg_attr(target_os = "wasi", allow(dead_code))] +pub(crate) fn spawn_blocking_internal(func: F) -> JoinHandle +where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, +{ + let rt = Handle::current(); + rt.inner.blocking_spawner().spawn_blocking_internal(&rt, func) +} + cfg_fs! { #[cfg_attr(any( all(loom, not(test)), // the function is covered by loom tests @@ -296,6 +310,40 @@ impl fmt::Debug for BlockingPool { impl Spawner { #[track_caller] pub(crate) fn spawn_blocking(&self, rt: &Handle, func: F) -> JoinHandle + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + self.spawn_blocking_impl( + rt, + func, + #[cfg(tokio_unstable)] + true, + ) + } + + #[track_caller] + #[cfg(feature = "rt-multi-thread")] + pub(crate) fn spawn_blocking_internal(&self, rt: &Handle, func: F) -> JoinHandle + where + F: FnOnce() -> R + Send + 'static, + R: Send + 'static, + { + self.spawn_blocking_impl( + rt, + func, + #[cfg(tokio_unstable)] + false, + ) + } + + #[track_caller] + fn spawn_blocking_impl( + &self, + rt: &Handle, + func: F, + #[cfg(tokio_unstable)] run_task_hooks: bool, + ) -> JoinHandle where F: FnOnce() -> R + Send + 'static, R: Send + 'static, @@ -309,6 +357,8 @@ impl Spawner { rt, #[cfg(tokio_unstable)] None, + #[cfg(tokio_unstable)] + run_task_hooks, ) } else { self.spawn_blocking_inner( @@ -318,6 +368,8 @@ impl Spawner { rt, #[cfg(tokio_unstable)] None, + #[cfg(tokio_unstable)] + run_task_hooks, ) }; @@ -351,6 +403,8 @@ impl Spawner { rt, #[cfg(tokio_unstable)] None, + #[cfg(tokio_unstable)] + true, ) } else { self.spawn_blocking_inner( @@ -360,6 +414,8 @@ impl Spawner { rt, #[cfg(tokio_unstable)] None, + #[cfg(tokio_unstable)] + true, ) }; @@ -379,6 +435,7 @@ impl Spawner { spawn_meta: SpawnMeta<'_>, rt: &Handle, #[cfg(tokio_unstable)] user_data: Option, + #[cfg(tokio_unstable)] run_task_hooks: bool, ) -> (JoinHandle, Result<(), SpawnError>) where F: FnOnce() -> R + Send + 'static, @@ -390,13 +447,27 @@ impl Spawner { let (task, handle) = task::unowned( fut, - BlockingSchedule::new(rt), + BlockingSchedule::new( + rt, + #[cfg(tokio_unstable)] + run_task_hooks, + ), id, task::SpawnLocation::capture(), #[cfg(tokio_unstable)] user_data, ); + #[cfg(tokio_unstable)] + if run_task_hooks { + task::with_current_task_meta(|parent| { + // Safety: the task is freshly allocated and has not been published + // to the blocking queue yet. + let mut meta = unsafe { task.task_meta() }; + rt.inner.hooks().spawn(&mut meta, parent); + }); + } + let spawned = self.spawn_task(Task::new(task, is_mandatory), rt); (handle, spawned) } @@ -408,6 +479,10 @@ impl Spawner { // Shutdown the task: it's fine to shutdown this task (even if // mandatory) because it was scheduled after the shutdown of the // runtime began. + // + // Dropping the task can run lifecycle hooks, and those hooks are + // allowed to re-enter the blocking pool. + drop(shared); task.task.shutdown(); // no need to even push this task; it would never get picked up diff --git a/tokio/src/runtime/blocking/schedule.rs b/tokio/src/runtime/blocking/schedule.rs index 8484398c6..8678a8832 100644 --- a/tokio/src/runtime/blocking/schedule.rs +++ b/tokio/src/runtime/blocking/schedule.rs @@ -20,7 +20,10 @@ pub(crate) struct BlockingSchedule { impl BlockingSchedule { #[cfg_attr(not(feature = "test-util"), allow(unused_variables))] - pub(crate) fn new(handle: &Handle) -> Self { + pub(crate) fn new( + handle: &Handle, + #[cfg(tokio_unstable)] run_task_hooks: bool, + ) -> Self { #[cfg(feature = "test-util")] { match &handle.inner { @@ -35,7 +38,11 @@ impl BlockingSchedule { #[cfg(feature = "test-util")] handle: handle.clone(), #[cfg(tokio_unstable)] - task_terminate_callback: handle.inner.hooks().task_terminate_callback.clone(), + task_terminate_callback: if run_task_hooks { + handle.inner.hooks().task_terminate_callback.clone() + } else { + None + }, } } } diff --git a/tokio/src/runtime/mod.rs b/tokio/src/runtime/mod.rs index 2cfdd92f9..22c8fb73d 100644 --- a/tokio/src/runtime/mod.rs +++ b/tokio/src/runtime/mod.rs @@ -530,6 +530,8 @@ cfg_rt! { mod blocking; #[cfg_attr(target_os = "wasi", allow(unused_imports))] pub(crate) use blocking::spawn_blocking; + #[cfg(feature = "rt-multi-thread")] + pub(crate) use blocking::spawn_blocking_internal; cfg_trace! { pub(crate) use blocking::Mandatory; diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index 82b5d83f1..9e5cacf48 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -475,7 +475,7 @@ where // Once the blocking task is done executing, we will attempt to // steal the core back. let worker = cx.worker.clone(); - runtime::spawn_blocking(move || run(worker)); + runtime::spawn_blocking_internal(move || run(worker)); Ok(()) }); @@ -500,7 +500,7 @@ where impl Launch { pub(crate) fn launch(mut self) { for worker in self.0.drain(..) { - runtime::spawn_blocking(move || run(worker)); + runtime::spawn_blocking_internal(move || run(worker)); } } } diff --git a/tokio/src/runtime/task/mod.rs b/tokio/src/runtime/task/mod.rs index eaaaae767..d5a1de4eb 100644 --- a/tokio/src/runtime/task/mod.rs +++ b/tokio/src/runtime/task/mod.rs @@ -537,6 +537,15 @@ impl UnownedTask { task } + /// # Safety + /// + /// The returned metadata must have exclusive access to hook data for as long + /// as it can expose mutable references. + #[cfg(tokio_unstable)] + pub(crate) unsafe fn task_meta<'meta>(&self) -> crate::runtime::TaskMeta<'meta> { + unsafe { self.raw.task_meta() } + } + pub(crate) fn run(self) { let raw = self.raw; mem::forget(self); diff --git a/tokio/src/task/builder.rs b/tokio/src/task/builder.rs index bf5d210aa..d4c862c21 100644 --- a/tokio/src/task/builder.rs +++ b/tokio/src/task/builder.rs @@ -242,6 +242,8 @@ impl<'a> Builder<'a> { SpawnMeta::new(name, fn_size), handle, data, + #[cfg(tokio_unstable)] + true, ) } else { handle.inner.blocking_spawner().spawn_blocking_inner( @@ -250,6 +252,8 @@ impl<'a> Builder<'a> { SpawnMeta::new(name, fn_size), handle, data, + #[cfg(tokio_unstable)] + true, ) }; diff --git a/tokio/tests/task_hooks.rs b/tokio/tests/task_hooks.rs index dc24b7008..1d649803b 100644 --- a/tokio/tests/task_hooks.rs +++ b/tokio/tests/task_hooks.rs @@ -77,6 +77,128 @@ fn terminate_task_hook_fires() { assert_eq!(TASKS, count.load(Ordering::SeqCst)); } +#[test] +fn spawn_blocking_task_hooks_are_balanced() { + let (spawned_tx, spawned_rx) = std::sync::mpsc::channel(); + let (terminated_tx, terminated_rx) = std::sync::mpsc::channel(); + + let runtime = Builder::new_current_thread() + .on_task_spawn(move |meta, parent| { + assert!(parent.is_none()); + spawned_tx.send(meta.id()).unwrap(); + }) + .on_task_terminate(move |meta| { + terminated_tx.send(meta.id()).unwrap(); + }) + .build() + .unwrap(); + + runtime.block_on(async { + tokio::task::spawn_blocking(|| {}).await.unwrap(); + }); + + let spawned = spawned_rx + .recv_timeout(Duration::from_secs(5)) + .expect("spawn_blocking task did not fire spawn hook"); + let terminated = terminated_rx + .recv_timeout(Duration::from_secs(5)) + .expect("spawn_blocking task did not fire terminate hook"); + + assert_eq!(spawned, terminated); + assert!(matches!( + spawned_rx.try_recv(), + Err(std::sync::mpsc::TryRecvError::Empty) + )); + assert!(matches!( + terminated_rx.try_recv(), + Err(std::sync::mpsc::TryRecvError::Empty) + )); +} + +#[cfg_attr( + target_os = "wasi", + ignore = "WASI does not support multi-threaded runtime" +)] +#[test] +fn internal_runtime_blocking_tasks_do_not_fire_task_hooks() { + let spawned = Arc::new(Mutex::new(Vec::new())); + let spawned2 = Arc::clone(&spawned); + let terminated = Arc::new(Mutex::new(Vec::new())); + let terminated2 = Arc::clone(&terminated); + + let runtime = Builder::new_multi_thread() + .worker_threads(1) + .on_task_spawn(move |meta, _parent| { + spawned2.lock().unwrap().push(meta.spawned_at().file()); + }) + .on_task_terminate(move |meta| { + terminated2.lock().unwrap().push(meta.spawned_at().file()); + }) + .build() + .unwrap(); + + runtime.block_on(async { + tokio::spawn(async { + tokio::task::block_in_place(|| {}); + }) + .await + .unwrap(); + }); + + runtime.shutdown_timeout(Duration::from_secs(5)); + + let spawned = spawned.lock().unwrap(); + assert!( + !spawned.iter().any(is_multi_thread_worker_file), + "internal worker task fired spawn hook: {spawned:?}" + ); + + let terminated = terminated.lock().unwrap(); + assert!( + !terminated.iter().any(is_multi_thread_worker_file), + "internal worker task fired terminate hook: {terminated:?}" + ); +} + +fn is_multi_thread_worker_file(file: &&'static str) -> bool { + file.ends_with("runtime/scheduler/multi_thread/worker.rs") + || file.ends_with(r"runtime\scheduler\multi_thread\worker.rs") +} + +#[test] +fn spawn_blocking_after_shutdown_terminate_hook_can_reenter_pool() { + let handle: Arc>> = Arc::new(Mutex::new(None)); + let hook_handle = Arc::clone(&handle); + let terminated = Arc::new(AtomicUsize::new(0)); + let terminated2 = Arc::clone(&terminated); + + let runtime = Builder::new_current_thread() + .on_task_terminate(move |_meta| { + if terminated2.fetch_add(1, Ordering::SeqCst) == 0 { + let handle = hook_handle.lock().unwrap().clone().unwrap(); + drop(handle.spawn_blocking(|| {})); + } + }) + .build() + .unwrap(); + + let runtime_handle = runtime.handle().clone(); + *handle.lock().unwrap() = Some(runtime_handle.clone()); + runtime.shutdown_timeout(Duration::from_secs(5)); + + let (done_tx, done_rx) = std::sync::mpsc::channel(); + std::thread::spawn(move || { + drop(runtime_handle.spawn_blocking(|| {})); + done_tx.send(()).unwrap(); + }); + + done_rx + .recv_timeout(Duration::from_secs(5)) + .expect("terminate hook deadlocked while re-entering the blocking pool"); + + assert_eq!(terminated.load(Ordering::SeqCst), 2); +} + /// Test that the correct spawn location is provided to the task hooks on a /// current thread runtime. #[test]