diff --git a/tokio/src/sync/task/atomic_waker.rs b/tokio/src/sync/task/atomic_waker.rs index d06498738..a13e27eae 100644 --- a/tokio/src/sync/task/atomic_waker.rs +++ b/tokio/src/sync/task/atomic_waker.rs @@ -316,11 +316,12 @@ impl AtomicWaker { // memory the `AtomicWaker` is associated with. match self.state.fetch_or(WAKING, AcqRel) { WAITING => { - // The waking lock has been acquired. + // SAFETY: the waking lock has been acquired. let waker = unsafe { self.waker.with_mut(|t| (*t).take()) }; - // Release the lock - self.state.fetch_and(!WAKING, Release); + // Release the lock. + let old_state = self.state.swap(WAITING, Release); + debug_assert!(old_state == WAKING); waker } diff --git a/tokio/src/sync/tests/atomic_waker.rs b/tokio/src/sync/tests/atomic_waker.rs index d89b44fa4..b182574f3 100644 --- a/tokio/src/sync/tests/atomic_waker.rs +++ b/tokio/src/sync/tests/atomic_waker.rs @@ -39,6 +39,38 @@ fn wake_without_register() { assert!(!waker.is_woken()); } +#[test] +#[cfg_attr(target_family = "wasm", ignore)] // threads not supported +fn failed_wake_synchronizes() { + for _ in 0..1000 { + failed_wake_synchronizes_inner(); + } +} + +fn failed_wake_synchronizes_inner() { + use futures::task::noop_waker_ref; + use std::sync::atomic::{AtomicBool, Ordering}; + static DID_SYNCHRONIZE: AtomicBool = AtomicBool::new(false); + DID_SYNCHRONIZE.store(false, Ordering::Relaxed); + + let waker = AtomicWaker::new(); + waker.register_by_ref(noop_waker_ref()); + + std::thread::scope(|s| { + let jh = s.spawn(|| { + DID_SYNCHRONIZE.store(true, Ordering::Relaxed); + waker.take_waker() + }); + + waker.take_waker(); + waker.register_by_ref(noop_waker_ref()); + + let did_synchronize = DID_SYNCHRONIZE.load(Ordering::Relaxed); + let did_take = jh.join().unwrap().is_some(); + assert!(did_synchronize || did_take); + }); +} + #[cfg(panic = "unwind")] #[test] #[cfg(not(target_family = "wasm"))] // wasm currently doesn't support unwinding