diff --git a/tokio/src/runtime/scheduler/current_thread/mod.rs b/tokio/src/runtime/scheduler/current_thread/mod.rs index 6d6cd7c0b..68ab17f14 100644 --- a/tokio/src/runtime/scheduler/current_thread/mod.rs +++ b/tokio/src/runtime/scheduler/current_thread/mod.rs @@ -15,7 +15,7 @@ use crate::util::{waker_ref, RngSeedGenerator, Wake, WakerRef}; use std::cell::RefCell; use std::collections::VecDeque; use std::future::{poll_fn, Future}; -use std::sync::atomic::Ordering::{AcqRel, Release}; +use std::sync::atomic::Ordering::{AcqRel, Acquire, Release}; use std::task::Poll::{Pending, Ready}; use std::task::Waker; use std::thread::ThreadId; @@ -380,8 +380,8 @@ impl Context { core = c; } - // This check will fail if `before_park` spawns a task for us to run - // instead of parking the thread + // If `before_park` spawns a task (or otherwise schedules work for us), then we should not + // park the thread. if !self.has_pending_work(&core) { // Park until the thread is signaled core.metrics.about_to_park(); @@ -415,7 +415,7 @@ impl Context { } fn has_pending_work(&self, core: &Core) -> bool { - !core.tasks.is_empty() || !self.defer.is_empty() + !core.tasks.is_empty() || !self.defer.is_empty() || self.handle.shared.woken.load(Acquire) } fn park_internal( @@ -724,8 +724,20 @@ impl Wake for Handle { /// Wake by reference fn wake_by_ref(arc_self: &Arc) { - arc_self.shared.woken.store(true, Release); - arc_self.driver.unpark(); + let already_woken = arc_self.shared.woken.swap(true, Release); + + if !already_woken { + use scheduler::Context::CurrentThread; + + // If we are already running on the runtime, then it's not required to wake up the + // runtime. + context::with_scheduler(|maybe_cx| match maybe_cx { + Some(CurrentThread(cx)) if Arc::ptr_eq(arc_self, &cx.handle) => {} + _ => { + arc_self.driver.unpark(); + } + }); + } } } diff --git a/tokio/tests/rt_common_before_park.rs b/tokio/tests/rt_common_before_park.rs new file mode 100644 index 000000000..cb0dc152b --- /dev/null +++ b/tokio/tests/rt_common_before_park.rs @@ -0,0 +1,92 @@ +#![warn(rust_2018_idioms)] +#![cfg(feature = "full")] +#![cfg(not(target_os = "wasi"))] // Wasi doesn't support threads + +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use tokio::runtime::Builder; +use tokio::sync::Notify; + +#[test] +fn before_park_wakes_block_on_task() { + let notify = Arc::new(Notify::new()); + let notify2 = notify.clone(); + let woken = Arc::new(AtomicBool::new(false)); + let woken2 = woken.clone(); + + let rt = Builder::new_current_thread() + .enable_all() + .on_thread_park(move || { + // Only wake once to avoid busy loop if something goes wrong, + // though in this test we expect it to unpark immediately. + if !woken2.swap(true, Ordering::SeqCst) { + notify2.notify_one(); + } + }) + .build() + .unwrap(); + + rt.block_on(async { + // This will block until `notify` is notified. + // `before_park` should run when the runtime is about to park. + // It will notify `notify`, which should wake this task. + // The runtime should then see the task is woken and NOT park. + notify.notified().await; + }); + + assert!(woken.load(Ordering::SeqCst)); +} + +#[test] +fn before_park_spawns_task() { + let notify = Arc::new(Notify::new()); + let notify2 = notify.clone(); + 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 notify = notify2.clone(); + tokio::spawn(async move { + notify.notify_one(); + }); + } + }) + .build() + .unwrap(); + + rt.block_on(async { + // This will block until `notify` is notified. + // `before_park` should run when the runtime is about to park. + // It will spawn a task that notifies `notify`. + // The runtime should see the new task and NOT park. + // If it parks, it will deadlock. + notify.notified().await; + }); + + assert!(woken.load(Ordering::SeqCst)); +} + +#[test] +fn wake_from_other_thread_block_on() { + let rt = Builder::new_current_thread().enable_all().build().unwrap(); + let handle = rt.handle().clone(); + let notify = Arc::new(Notify::new()); + let notify2 = notify.clone(); + + let th = std::thread::spawn(move || { + // Give the main thread time to park + std::thread::sleep(std::time::Duration::from_millis(5)); + handle.block_on(async move { + notify2.notify_one(); + }); + }); + + rt.block_on(async { + notify.notified().await; + }); + + th.join().unwrap(); +} diff --git a/tokio/tests/rt_metrics.rs b/tokio/tests/rt_metrics.rs index 2b0ac933f..b242c07e2 100644 --- a/tokio/tests/rt_metrics.rs +++ b/tokio/tests/rt_metrics.rs @@ -152,7 +152,7 @@ fn worker_park_unpark_count() { let metrics = rt.metrics(); rt.block_on(rt.spawn(async {})).unwrap(); drop(rt); - assert!(2 <= metrics.worker_park_unpark_count(0)); + assert_eq!(0, metrics.worker_park_unpark_count(0) % 2); let rt = threaded(); let metrics = rt.metrics();