diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index 305f9a5d0..6d6cd7c0b 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/tokio/src/runtime/scheduler/current_thread/mod.rs @@ -382,7 +382,7 @@ impl Context { // This check will fail if `before_park` spawns a task for us to run // instead of parking the thread - if core.tasks.is_empty() { + if !self.has_pending_work(&core) { // Park until the thread is signaled core.metrics.about_to_park(); core.submit_metrics(handle); @@ -414,6 +414,10 @@ impl Context { core } + fn has_pending_work(&self, core: &Core) -> bool { + !core.tasks.is_empty() || !self.defer.is_empty() + } + fn park_internal( &self, core: Box, @@ -775,7 +779,7 @@ impl CoreGuard<'_> { None => { core.metrics.end_processing_scheduled_tasks(); - core = if !context.defer.is_empty() { + core = if context.has_pending_work(&core) { context.park_yield(core, handle) } else { context.park(core, handle) diff --git a/tokio/tests/rt_basic.rs b/tokio/tests/rt_basic.rs index 665dce3c9..dc835825e 100644 --- a/tokio/tests/rt_basic.rs +++ b/tokio/tests/rt_basic.rs @@ -456,3 +456,44 @@ fn rt() -> Runtime { .build() .unwrap() } + +#[test] +fn before_park_yields() { + use futures::task::ArcWake; + use std::sync::Arc; + use tokio::runtime::Builder; + use tokio::sync::Notify; + + struct MyWaker(Notify); + + impl ArcWake for MyWaker { + fn wake_by_ref(arc_self: &Arc) { + arc_self.0.notify_one(); + } + } + + let notify = Arc::new(MyWaker(Notify::new())); + let notify2 = notify.clone(); + let waker = futures::task::waker(notify2); + let woken = Arc::new(AtomicBool::new(false)); + let woken2 = woken.clone(); + + let rt = Builder::new_current_thread() + .enable_all() + .on_thread_park(move || { + if !woken2.swap(true, Ordering::SeqCst) { + let mut cx = Context::from_waker(&waker); + // `yield_now` pushes the waker to the defer slot. + let fut = std::pin::pin!(tokio::task::yield_now()); + let _ = fut.poll(&mut cx); + } + }) + .build() + .unwrap(); + + rt.block_on(async { + notify.0.notified().await; + }); + + assert!(woken.load(Ordering::SeqCst)); +}