From 7010f81cf2e3d2d3a7e9ee71a58f18f2c25d32ae Mon Sep 17 00:00:00 2001 From: ryux1 Date: Sun, 6 Sep 2026 08:40:51 +0200 Subject: [PATCH 1/2] perf(completion): make terminal waits lock-free --- CHANGELOG.md | 1 + asyncband/src/completion/mod.rs | 103 ++++++++++++++------- asyncband/src/internal/wakerset.rs | 26 ++++++ tests-integration/tests/completion_test.rs | 61 ++++++++++++ 4 files changed, 156 insertions(+), 35 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 43ffd8e0..d2a51e67 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -15,6 +15,7 @@ All notable changes to this project will be documented in this file. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Reduce bounded MPSC contention when senders or the receiver are not waiting, improving throughput without changing capacity or cancellation semantics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. +* Make completed and abandoned `Completion` waits lock-free while preserving cancellable pending registration and unlocked waker callbacks. ## v0.7.2 diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index ea34a073..cd199e34 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -54,6 +54,8 @@ use std::pin::Pin; use std::sync::Arc; use std::sync::OnceLock; use std::sync::Weak; +use std::sync::atomic::AtomicU8; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; @@ -66,10 +68,8 @@ use crate::internal::wakerset::WakerToken; pub fn new() -> (Completer, Completion) { let shared = Arc::new(Shared { value: OnceLock::new(), - state: Mutex::new(State { - status: Status::Pending, - waiters: WakerSet::new(), - }), + status: AtomicU8::new(Status::Pending as u8), + waiters: Mutex::new(WakerSet::new()), }); let completer = Completer { shared: Arc::downgrade(&shared), @@ -80,14 +80,11 @@ pub fn new() -> (Completer, Completion) { struct Shared { value: OnceLock, - state: Mutex, -} - -struct State { - status: Status, - waiters: WakerSet, + status: AtomicU8, + waiters: Mutex, } +#[repr(u8)] #[derive(Clone, Copy, Debug, Eq, PartialEq)] enum Status { Pending, @@ -95,6 +92,17 @@ enum Status { Abandoned, } +impl Status { + fn load(status: &AtomicU8) -> Self { + match status.load(Ordering::Acquire) { + value if value == Self::Pending as u8 => Self::Pending, + value if value == Self::Completed as u8 => Self::Completed, + value if value == Self::Abandoned as u8 => Self::Abandoned, + _ => unreachable!("completion status must be valid"), + } + } +} + /// The error returned by [`Completion::wait`] when the completer was dropped without a value. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct Abandoned(()); @@ -145,21 +153,24 @@ impl Completer { return Err(value); }; let wakers = { - let mut state = shared.state.lock(); + let mut waiters = shared.waiters.lock(); assert_eq!( - state.status, + Status::load(&shared.status), Status::Pending, "a live completer must refer to a pending completion" ); if let Err(value) = shared.value.set(value) { - drop(state); + drop(waiters); drop(value); panic!("pending completion value must be unset"); } - // Publish the value before making completion observable and detaching its waiters. - state.status = Status::Completed; - state.waiters.take_all() + let wakers = waiters.take_all(); + // Release publishes both the value and the detached waiter cohort to lock-free polls. + shared + .status + .store(Status::Completed as u8, Ordering::Release); + wakers }; // `complete` consumes the only completer. Disarm its destructor before invoking arbitrary // wake callbacks; the completed state no longer needs abandonment handling. @@ -175,13 +186,15 @@ impl Drop for Completer { return; }; let wakers = { - let mut state = shared.state.lock(); - if state.status != Status::Pending { + let mut waiters = shared.waiters.lock(); + if Status::load(&shared.status) != Status::Pending { return; } - // Publish abandonment and detach its waiters atomically with respect to registration. - state.status = Status::Abandoned; - state.waiters.take_all() + let wakers = waiters.take_all(); + shared + .status + .store(Status::Abandoned as u8, Ordering::Release); + wakers }; wake_all(wakers); } @@ -238,15 +251,32 @@ impl<'a, T> Future for Wait<'a, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); - let mut state = this.completion.shared.state.lock(); - let (poll, retired_waker) = match state.status { + match Status::load(&this.completion.shared.status) { + Status::Completed => { + this.token = None; + return Poll::Ready(Ok(this + .completion + .shared + .value + .get() + .expect("completed value must be initialized"))); + } + Status::Abandoned => { + this.token = None; + return Poll::Ready(Err(Abandoned(()))); + } + Status::Pending => {} + } + + // Cloning a RawWaker can execute arbitrary user code, so do it before taking the lock. + let waker = cx.waker().clone(); + let mut waiters = this.completion.shared.waiters.lock(); + let (poll, retired_waker) = match Status::load(&this.completion.shared.status) { Status::Pending => { - let retired = state.waiters.register(&mut this.token, cx.waker()); + let retired = waiters.register_owned(&mut this.token, waker); (Poll::Pending, retired) } Status::Completed => { - // Completion detaches every registration under this same lock before another poll - // can observe the terminal status. this.token = None; let completion: &'a Completion = this.completion; let value = completion @@ -254,15 +284,14 @@ impl<'a, T> Future for Wait<'a, T> { .value .get() .expect("completed value must be initialized"); - (Poll::Ready(Ok(value)), None) + (Poll::Ready(Ok(value)), Some(waker)) } Status::Abandoned => { - // Abandonment uses the same terminal detach protocol as completion. this.token = None; - (Poll::Ready(Err(Abandoned(()))), None) + (Poll::Ready(Err(Abandoned(()))), Some(waker)) } }; - drop(state); + drop(waiters); drop(retired_waker); poll } @@ -274,15 +303,19 @@ impl Drop for Wait<'_, T> { return; } - let mut state = self.completion.shared.state.lock(); - if state.status != Status::Pending { - // The terminal transition already detached this registration. + if Status::load(&self.completion.shared.status) != Status::Pending { + self.token = None; + return; + } + + let mut waiters = self.completion.shared.waiters.lock(); + if Status::load(&self.completion.shared.status) != Status::Pending { self.token = None; return; } - let waker = state.waiters.unregister(&mut self.token); - drop(state); + let waker = waiters.unregister(&mut self.token); + drop(waiters); drop(waker); } } diff --git a/asyncband/src/internal/wakerset.rs b/asyncband/src/internal/wakerset.rs index eb11366e..d95fabae 100644 --- a/asyncband/src/internal/wakerset.rs +++ b/asyncband/src/internal/wakerset.rs @@ -102,6 +102,32 @@ impl WakerSet { None } + /// Registers or updates an already-owned waker. + /// + /// If the supplied waker is unused or replaces an existing waker, the retired waker is + /// returned so the caller can drop it after releasing the lock that protects this set. + #[inline] + #[must_use = "drop the returned waker after releasing the waker set's state lock"] + pub fn register_owned( + &mut self, + token: &mut Option, + waker: Waker, + ) -> Option { + if let Some(current) = token.as_ref().map(|token| { + self.wakers + .get_mut(token.0) + .expect("waker token must refer to an occupied slot") + }) { + if current.will_wake(&waker) { + return Some(waker); + } + return Some(mem::replace(current, waker)); + } + + *token = Some(WakerToken(self.wakers.insert(waker))); + None + } + /// Removes the waker identified by `token`. /// /// The owner must clear stale tokens without calling this method after detaching the set. The diff --git a/tests-integration/tests/completion_test.rs b/tests-integration/tests/completion_test.rs index 0bfd5b01..85c95e07 100644 --- a/tests-integration/tests/completion_test.rs +++ b/tests-integration/tests/completion_test.rs @@ -17,6 +17,7 @@ use std::cell::Cell; use std::future::Future; +use std::mem::ManuallyDrop; use std::pin::Pin; use std::sync::Arc; use std::sync::Barrier; @@ -25,6 +26,8 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use std::task::RawWaker; +use std::task::RawWakerVTable; use std::task::Wake; use std::task::Waker; use std::thread; @@ -63,6 +66,8 @@ impl Wake for WakeCallback { struct DropCallbackWake(Mutex>>); +struct CloneCallbackWake(Mutex>>); + // This test needs a custom waker whose final `Arc` drop is observable. #[allow(clippy::manual_noop_waker)] impl Wake for DropCallbackWake { @@ -77,6 +82,45 @@ impl Drop for DropCallbackWake { } } +unsafe fn clone_callback_waker(data: *const ()) -> RawWaker { + // SAFETY: Every pointer using this vtable comes from `Arc::into_raw`. `ManuallyDrop` keeps the + // original waker's strong reference alive while its clone callback borrows the allocation. + let state = ManuallyDrop::new(unsafe { Arc::::from_raw(data.cast()) }); + if let Some(callback) = state.0.lock().unwrap().take() { + callback(); + } + RawWaker::new( + Arc::into_raw(Arc::clone(&state)).cast(), + &CLONE_CALLBACK_VTABLE, + ) +} + +unsafe fn wake_clone_callback_waker(data: *const ()) { + // SAFETY: `wake` consumes the raw waker's strong reference exactly once. + drop(unsafe { Arc::::from_raw(data.cast()) }); +} + +unsafe fn wake_clone_callback_waker_by_ref(_data: *const ()) {} + +unsafe fn drop_clone_callback_waker(data: *const ()) { + // SAFETY: `drop` consumes the raw waker's strong reference exactly once. + drop(unsafe { Arc::::from_raw(data.cast()) }); +} + +static CLONE_CALLBACK_VTABLE: RawWakerVTable = RawWakerVTable::new( + clone_callback_waker, + wake_clone_callback_waker, + wake_clone_callback_waker_by_ref, + drop_clone_callback_waker, +); + +fn waker_with_clone_callback(callback: impl FnOnce() + Send + 'static) -> Waker { + let state = Arc::new(CloneCallbackWake(Mutex::new(Some(Box::new(callback))))); + let raw = RawWaker::new(Arc::into_raw(state).cast(), &CLONE_CALLBACK_VTABLE); + // SAFETY: The vtable preserves the Arc strong count and all callbacks are thread safe. + unsafe { Waker::from_raw(raw) } +} + fn poll_with(future: Pin<&mut F>, waker: &Waker) -> Poll { future.poll(&mut Context::from_waker(waker)) } @@ -351,6 +395,23 @@ fn wake_callbacks_run_outside_the_completion_lock() { ); } +#[test] +fn waker_clone_callbacks_run_outside_the_completion_lock() { + assert_completes_without_deadlock( + "waker clone callback deadlocked against the completion lock", + || { + let (completer, completion) = completion::new::(); + let waker = waker_with_clone_callback(move || drop(completer)); + let mut wait = Box::pin(completion.wait()); + + assert!(matches!( + poll_with(wait.as_mut(), &waker), + Poll::Ready(Err(_)) + )); + }, + ); +} + #[test] fn replaced_wakers_are_dropped_outside_the_completion_lock() { assert_completes_without_deadlock( From fd855622af3ad7943730d242ae0bcaf1969c09c8 Mon Sep 17 00:00:00 2001 From: tison Date: Fri, 11 Sep 2026 02:00:48 +0800 Subject: [PATCH 2/2] perf(completion): consolidate terminal state in OnceLock --- CHANGELOG.md | 2 +- asyncband/src/completion/mod.rs | 126 +++++---------------- asyncband/src/internal/wakerset.rs | 26 ----- benchmarks/asyncband/completion/mod.rs | 33 ++++-- tests-integration/tests/completion_test.rs | 123 +++----------------- 5 files changed, 66 insertions(+), 244 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2e8fa9f2..ea2f1e3c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,7 +18,7 @@ All notable changes to this project will be documented in this file. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. -* Make completed and abandoned `Completion` waits lock-free while preserving cancellable pending registration and unlocked waker callbacks. +* Make completed and abandoned `Completion` waits lock-free while preserving cancellable pending registration. ## v0.7.2 diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index cd199e34..00b03526 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -54,8 +54,6 @@ use std::pin::Pin; use std::sync::Arc; use std::sync::OnceLock; use std::sync::Weak; -use std::sync::atomic::AtomicU8; -use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; @@ -67,8 +65,7 @@ use crate::internal::wakerset::WakerToken; /// Creates a single-use [`Completer`] and a cloneable [`Completion`] observer. pub fn new() -> (Completer, Completion) { let shared = Arc::new(Shared { - value: OnceLock::new(), - status: AtomicU8::new(Status::Pending as u8), + result: OnceLock::new(), waiters: Mutex::new(WakerSet::new()), }); let completer = Completer { @@ -79,30 +76,11 @@ pub fn new() -> (Completer, Completion) { } struct Shared { - value: OnceLock, - status: AtomicU8, + // A stored None means abandonment. Initialization also publishes waiter detachment. + result: OnceLock>, waiters: Mutex, } -#[repr(u8)] -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum Status { - Pending, - Completed, - Abandoned, -} - -impl Status { - fn load(status: &AtomicU8) -> Self { - match status.load(Ordering::Acquire) { - value if value == Self::Pending as u8 => Self::Pending, - value if value == Self::Completed as u8 => Self::Completed, - value if value == Self::Abandoned as u8 => Self::Abandoned, - _ => unreachable!("completion status must be valid"), - } - } -} - /// The error returned by [`Completion::wait`] when the completer was dropped without a value. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct Abandoned(()); @@ -125,9 +103,9 @@ pub struct Completer { } // SAFETY: The completer can only move an owned `T` into the shared `OnceLock` while holding the -// state mutex; it never exposes or accesses the stored value afterward. `Completion` retains its -// ordinary auto traits, so observers cannot cross threads unless `T` can be shared. `T: Send` also -// permits the shared allocation and its value to be destroyed by the completing thread if its +// waiter mutex; it never exposes or accesses the stored value afterward. `Completion` retains +// its ordinary auto traits, so observers cannot cross threads unless `T` can be shared. `T: Send` +// also permits the shared allocation and its value to be destroyed by the completing thread if its // temporary strong reference is the last one. unsafe impl Send for Completer {} unsafe impl Sync for Completer {} @@ -154,22 +132,9 @@ impl Completer { }; let wakers = { let mut waiters = shared.waiters.lock(); - assert_eq!( - Status::load(&shared.status), - Status::Pending, - "a live completer must refer to a pending completion" - ); - - if let Err(value) = shared.value.set(value) { - drop(waiters); - drop(value); - panic!("pending completion value must be unset"); - } let wakers = waiters.take_all(); - // Release publishes both the value and the detached waiter cohort to lock-free polls. - shared - .status - .store(Status::Completed as u8, Ordering::Release); + // The single completer publishes only after every waiter token has been invalidated. + assert!(shared.result.set(Some(value)).is_ok()); wakers }; // `complete` consumes the only completer. Disarm its destructor before invoking arbitrary @@ -187,13 +152,8 @@ impl Drop for Completer { }; let wakers = { let mut waiters = shared.waiters.lock(); - if Status::load(&shared.status) != Status::Pending { - return; - } let wakers = waiters.take_all(); - shared - .status - .store(Status::Abandoned as u8, Ordering::Release); + assert!(shared.result.set(None).is_ok()); wakers }; wake_all(wakers); @@ -234,7 +194,7 @@ impl Completion { /// not affect this observer, another wait, or the eventual result. pub async fn wait(&self) -> Result<&T, Abandoned> { Wait { - completion: self, + shared: &self.shared, token: None, } .await @@ -242,7 +202,7 @@ impl Completion { } struct Wait<'a, T> { - completion: &'a Completion, + shared: &'a Shared, token: Option, } @@ -251,49 +211,23 @@ impl<'a, T> Future for Wait<'a, T> { fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let this = self.get_mut(); - match Status::load(&this.completion.shared.status) { - Status::Completed => { - this.token = None; - return Poll::Ready(Ok(this - .completion - .shared - .value - .get() - .expect("completed value must be initialized"))); - } - Status::Abandoned => { - this.token = None; - return Poll::Ready(Err(Abandoned(()))); - } - Status::Pending => {} - } - - // Cloning a RawWaker can execute arbitrary user code, so do it before taking the lock. - let waker = cx.waker().clone(); - let mut waiters = this.completion.shared.waiters.lock(); - let (poll, retired_waker) = match Status::load(&this.completion.shared.status) { - Status::Pending => { - let retired = waiters.register_owned(&mut this.token, waker); - (Poll::Pending, retired) - } - Status::Completed => { - this.token = None; - let completion: &'a Completion = this.completion; - let value = completion - .shared - .value - .get() - .expect("completed value must be initialized"); - (Poll::Ready(Ok(value)), Some(waker)) - } - Status::Abandoned => { - this.token = None; - (Poll::Ready(Err(Abandoned(()))), Some(waker)) + let shared = this.shared; + let result = if let Some(result) = shared.result.get() { + result + } else { + let mut waiters = shared.waiters.lock(); + if let Some(result) = shared.result.get() { + result + } else { + let retired = waiters.register(&mut this.token, cx.waker()); + drop(waiters); + drop(retired); + return Poll::Pending; } }; - drop(waiters); - drop(retired_waker); - poll + + this.token = None; + Poll::Ready(result.as_ref().ok_or(Abandoned(()))) } } @@ -303,14 +237,12 @@ impl Drop for Wait<'_, T> { return; } - if Status::load(&self.completion.shared.status) != Status::Pending { - self.token = None; + if self.shared.result.get().is_some() { return; } - let mut waiters = self.completion.shared.waiters.lock(); - if Status::load(&self.completion.shared.status) != Status::Pending { - self.token = None; + let mut waiters = self.shared.waiters.lock(); + if self.shared.result.get().is_some() { return; } diff --git a/asyncband/src/internal/wakerset.rs b/asyncband/src/internal/wakerset.rs index e9f12b8d..c49d49fd 100644 --- a/asyncband/src/internal/wakerset.rs +++ b/asyncband/src/internal/wakerset.rs @@ -102,32 +102,6 @@ impl WakerSet { None } - /// Registers or updates an already-owned waker. - /// - /// If the supplied waker is unused or replaces an existing waker, the retired waker is - /// returned so the caller can drop it after releasing the lock that protects this set. - #[inline] - #[must_use = "drop the returned waker after releasing the waker set's state lock"] - pub fn register_owned( - &mut self, - token: &mut Option, - waker: Waker, - ) -> Option { - if let Some(current) = token.as_ref().map(|token| { - self.wakers - .get_mut(token.0) - .expect("waker token must refer to an occupied slot") - }) { - if current.will_wake(&waker) { - return Some(waker); - } - return Some(mem::replace(current, waker)); - } - - *token = Some(WakerToken(self.wakers.insert(waker))); - None - } - /// Removes the waker identified by `token`. /// /// The owner must clear stale tokens without calling this method after detaching the set. The diff --git a/benchmarks/asyncband/completion/mod.rs b/benchmarks/asyncband/completion/mod.rs index 2676a97d..576af453 100644 --- a/benchmarks/asyncband/completion/mod.rs +++ b/benchmarks/asyncband/completion/mod.rs @@ -26,7 +26,7 @@ use crate::support::poll_pending; use crate::support::poll_pinned_ready; use crate::support::poll_ready; -const OBSERVER_COUNTS: &[usize] = &[1, 2, 4, 8, 32]; +const OBSERVER_COUNTS: &[usize] = &[2, 8, 32]; #[divan::bench] fn ready_wait(bencher: Bencher) { @@ -37,6 +37,25 @@ fn ready_wait(bencher: Bencher) { bencher.bench_local(|| black_box(*poll_ready(completion.wait(), &mut context).unwrap())); } +#[divan::bench] +fn abandoned_wait(bencher: Bencher) { + let mut context = bench_context(); + let (completer, completion) = completion::new::(); + drop(completer); + + bencher.bench_local(|| black_box(poll_ready(completion.wait(), &mut context).unwrap_err())); +} + +#[divan::bench] +fn repoll_pending(bencher: Bencher) { + let mut context = bench_context(); + let (_completer, completion) = completion::new::(); + let mut wait = pin!(completion.wait()); + poll_pending(wait.as_mut(), &mut context); + + bencher.bench_local(|| poll_pending(wait.as_mut(), &mut context)); +} + #[divan::bench] fn cancel_pending(bencher: Bencher) { let mut context = bench_context(); @@ -79,15 +98,9 @@ fn notify_pending_fanout(bencher: Bencher, observer_count: usize) { let mut context = bench_context(); bencher.bench_local(|| { - let (completer, first) = black_box(completion::new()); - let mut observers = Vec::with_capacity(observer_count); - observers.push(first); - for _ in 1..observer_count { - observers.push(observers[0].clone()); - } - let mut waiters = observers - .iter() - .map(|observer| Box::pin(observer.wait())) + let (completer, completion) = black_box(completion::new()); + let mut waiters = (0..observer_count) + .map(|_| Box::pin(completion.wait())) .collect::>(); for waiter in &mut waiters { poll_pending(waiter.as_mut(), &mut context); diff --git a/tests-integration/tests/completion_test.rs b/tests-integration/tests/completion_test.rs index 2a82acf6..106d36de 100644 --- a/tests-integration/tests/completion_test.rs +++ b/tests-integration/tests/completion_test.rs @@ -16,13 +16,9 @@ // under the License. use std::cell::Cell; -use std::mem::ManuallyDrop; use std::sync::Arc; use std::sync::Barrier; -use std::sync::Mutex; use std::task::Poll; -use std::task::RawWaker; -use std::task::RawWakerVTable; use std::task::Waker; use std::thread; @@ -36,77 +32,36 @@ use tests_integration::waker_on_wake; struct NotClone(String); -struct CloneCallbackWake(Mutex>>); - -unsafe fn clone_callback_waker(data: *const ()) -> RawWaker { - // SAFETY: Every pointer using this vtable comes from `Arc::into_raw`. `ManuallyDrop` keeps the - // original waker's strong reference alive while its clone callback borrows the allocation. - let state = ManuallyDrop::new(unsafe { Arc::::from_raw(data.cast()) }); - if let Some(callback) = state.0.lock().unwrap().take() { - callback(); - } - RawWaker::new( - Arc::into_raw(Arc::clone(&state)).cast(), - &CLONE_CALLBACK_VTABLE, - ) -} - -unsafe fn wake_clone_callback_waker(data: *const ()) { - // SAFETY: `wake` consumes the raw waker's strong reference exactly once. - drop(unsafe { Arc::::from_raw(data.cast()) }); -} - -unsafe fn wake_clone_callback_waker_by_ref(_data: *const ()) {} - -unsafe fn drop_clone_callback_waker(data: *const ()) { - // SAFETY: `drop` consumes the raw waker's strong reference exactly once. - drop(unsafe { Arc::::from_raw(data.cast()) }); -} - -static CLONE_CALLBACK_VTABLE: RawWakerVTable = RawWakerVTable::new( - clone_callback_waker, - wake_clone_callback_waker, - wake_clone_callback_waker_by_ref, - drop_clone_callback_waker, -); - -fn waker_with_clone_callback(callback: impl FnOnce() + Send + 'static) -> Waker { - let state = Arc::new(CloneCallbackWake(Mutex::new(Some(Box::new(callback))))); - let raw = RawWaker::new(Arc::into_raw(state).cast(), &CLONE_CALLBACK_VTABLE); - // SAFETY: The vtable preserves the Arc strong count and all callbacks are thread safe. - unsafe { Waker::from_raw(raw) } -} - #[test] fn all_observers_borrow_the_same_non_clone_value() { let (completer, completion) = completion::new(); let first = completion.clone(); let second = completion.clone(); + drop(completion); assert!(completer.complete(NotClone(String::from("ready"))).is_ok()); let first_value = pollster::block_on(first.wait()).unwrap(); let second_value = pollster::block_on(second.wait()).unwrap(); let repeated = pollster::block_on(first.wait()).unwrap(); - let late = completion.clone(); - let late_value = pollster::block_on(late.wait()).unwrap(); - assert_eq!(first_value.0.as_str(), "ready"); assert!(std::ptr::eq(first_value, second_value)); assert!(std::ptr::eq(first_value, repeated)); - assert!(std::ptr::eq(first_value, late_value)); + + let late = first.clone(); + drop(first); + let late_value = pollster::block_on(late.wait()).unwrap(); + assert!(std::ptr::eq(second_value, late_value)); } #[test] fn completer_transfers_a_send_only_value_between_threads() { let (completer, completion) = completion::new::>(); - thread::spawn(move || completer.complete(Cell::new(7))) - .join() - .unwrap() - .unwrap(); + let worker = thread::spawn(move || completer.complete(Cell::new(7))); assert_eq!(pollster::block_on(completion.wait()).unwrap().get(), 7); + worker.join().unwrap().unwrap(); } #[test] @@ -119,33 +74,6 @@ fn complete_returns_the_value_when_no_observers_remain() { ); } -#[test] -fn dropping_the_completer_abandons_every_observer() { - let (completer, first) = completion::new::(); - let second = first.clone(); - drop(completer); - - assert!(pollster::block_on(first.wait()).is_err()); - assert!(pollster::block_on(second.wait()).is_err()); -} - -#[test] -fn dropping_one_observer_before_or_after_completion_does_not_affect_another() { - let (completer, first) = completion::new(); - let second = first.clone(); - drop(first); - - completer.complete(5).unwrap(); - assert_eq!(pollster::block_on(second.wait()), Ok(&5)); - - let (completer, first) = completion::new(); - let second = first.clone(); - completer.complete(6).unwrap(); - drop(first); - - assert_eq!(pollster::block_on(second.wait()), Ok(&6)); -} - #[test] fn completed_payload_is_released_with_the_last_observer() { let payload = Arc::new(()); @@ -159,7 +87,7 @@ fn completed_payload_is_released_with_the_last_observer() { } #[test] -fn abandonment_wakes_all_registered_waits() { +fn abandonment_wakes_registered_waits_and_is_visible_to_late_observers() { let (completer, first) = completion::new::(); let second = first.clone(); let first_tracker = Arc::new(WakeCounter::default()); @@ -183,6 +111,9 @@ fn abandonment_wakes_all_registered_waits() { poll_with(second_wait.as_mut(), &second_waker), Poll::Ready(Err(_)) )); + + let late = first.clone(); + assert!(pollster::block_on(late.wait()).is_err()); } #[test] @@ -242,21 +173,7 @@ fn cancelling_after_wake_does_not_consume_the_shared_result() { } #[test] -fn cancellation_and_completer_drop_have_clean_orderings() { - let (completer, completion) = completion::new::(); - let tracker = Arc::new(WakeCounter::default()); - let waker = Waker::from(tracker.clone()); - let baseline = Arc::strong_count(&tracker); - let mut wait = Box::pin(completion.wait()); - - assert!(poll_with(wait.as_mut(), &waker).is_pending()); - assert_eq!(Arc::strong_count(&tracker), baseline + 1); - drop(wait); - assert_eq!(Arc::strong_count(&tracker), baseline); - drop(completer); - assert_eq!(tracker.count(), 0); - assert!(pollster::block_on(completion.wait()).is_err()); - +fn cancelling_after_abandonment_does_not_retain_the_waker() { let (completer, completion) = completion::new::(); let tracker = Arc::new(WakeCounter::default()); let waker = Waker::from(tracker.clone()); @@ -324,20 +241,6 @@ fn wake_callbacks_run_outside_the_completion_lock() { }); } -#[test] -fn waker_clone_callbacks_run_outside_the_completion_lock() { - assert_completes_without_deadlock(|| { - let (completer, completion) = completion::new::(); - let waker = waker_with_clone_callback(move || drop(completer)); - let mut wait = Box::pin(completion.wait()); - - assert!(matches!( - poll_with(wait.as_mut(), &waker), - Poll::Ready(Err(_)) - )); - }); -} - #[test] fn replaced_wakers_are_dropped_outside_the_completion_lock() { assert_completes_without_deadlock(|| {