From 0d6c7af3e43457350bdc03a6dbcafa276fab7352 Mon Sep 17 00:00:00 2001 From: ADD-SP Date: Tue, 3 Feb 2026 05:07:09 -0800 Subject: [PATCH] runtime: wake deferred tasks before entering `block_in_place` (#7879) --- .../runtime/scheduler/multi_thread/worker.rs | 4 ++ tokio/tests/rt_threaded.rs | 50 ++++++++++++++++++- 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/tokio/src/runtime/scheduler/multi_thread/worker.rs b/tokio/src/runtime/scheduler/multi_thread/worker.rs index efd2f5793..f48e6ba52 100644 --- a/tokio/src/runtime/scheduler/multi_thread/worker.rs +++ b/tokio/src/runtime/scheduler/multi_thread/worker.rs @@ -424,6 +424,10 @@ where let cx = maybe_cx.expect("no .is_some() == false cases above should lead here"); + // Since deferred tasks don't stay on `core`, make sure to wake them + // before blocking. + cx.defer.wake(); + // Get the worker core. If none is set, then blocking is fine! let mut core = match cx.core.borrow_mut().take() { Some(core) => core, diff --git a/tokio/tests/rt_threaded.rs b/tokio/tests/rt_threaded.rs index ca2a952fb..432add50f 100644 --- a/tokio/tests/rt_threaded.rs +++ b/tokio/tests/rt_threaded.rs @@ -6,7 +6,7 @@ use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tokio::net::{TcpListener, TcpStream}; use tokio::runtime; use tokio::sync::oneshot; -use tokio_test::{assert_err, assert_ok}; +use tokio_test::{assert_err, assert_ok, assert_pending}; use std::future::{poll_fn, Future}; use std::pin::Pin; @@ -692,6 +692,54 @@ fn mutex_in_block_in_place() { }) } +#[test] +/// Deferred tasks should be woken before starting the [`tokio::task::block_in_place`] +// https://github.com/tokio-rs/tokio/issues/7877 +fn wake_deferred_tasks_before_block_in_place() { + let (tx1, rx1) = oneshot::channel::<()>(); + let (tx2, rx2) = oneshot::channel::<()>(); + + let deferred_task = tokio_test::task::spawn(tokio::task::yield_now()); + let deffered_task = Arc::new(Mutex::new(deferred_task)); + + let rt = runtime::Builder::new_multi_thread() + .worker_threads(1) + .build() + .unwrap(); + + let jh = { + let deferred_task = Arc::clone(&deffered_task); + rt.spawn(async move { + { + let mut lock = deferred_task.lock().unwrap(); + assert_pending!(lock.poll()); + } + tokio::task::block_in_place(|| { + // signal that the `block_in_place` has started + tx2.send(()).unwrap(); + // wait for the shutdown signal + rx1.blocking_recv().unwrap(); + }); + }) + }; + + // wait for the `block_in_place` to start + rx2.blocking_recv().unwrap(); + + // check that the deferred task was woken before the `block_in_place` ends + let is_woken = { + let lock = deffered_task.lock().unwrap(); + lock.is_woken() + }; + + // signal the `block_in_place` to shutdown + tx1.send(()).unwrap(); + + rt.block_on(jh).unwrap(); + + assert!(is_woken); +} + // Testing the tuning logic is tricky as it is inherently timing based, and more // of a heuristic than an exact behavior. This test checks that the interval // changes over time based on load factors. There are no assertions, completion