fix issue where we had hooks for term but not spawn with blocking tasks

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;

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<F, R>(func: F) -> JoinHandle<R>
+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<F, R>(&self, rt: &Handle, func: F) -> JoinHandle<R>
+    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<F, R>(&self, rt: &Handle, func: F) -> JoinHandle<R>
+    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<F, R>(
+        &self,
+        rt: &Handle,
+        func: F,
+        #[cfg(tokio_unstable)] run_task_hooks: bool,
+    ) -> JoinHandle<R>
     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<crate::runtime::TaskData>,
+        #[cfg(tokio_unstable)] run_task_hooks: bool,
     ) -> (JoinHandle<R>, 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

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
+            },
         }
     }
 }

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;

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));
         }
     }
 }

tokio/src/runtime/task/mod.rs

@@ -537,6 +537,15 @@ impl<S: Schedule> UnownedTask<S> {
         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);

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,
             )
         };
 

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<Mutex<Option<tokio::runtime::Handle>>> = 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]