diff --git a/CHANGELOG.md b/CHANGELOG.md index 7113dc6b..ea2f1e3c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -18,6 +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. ## v0.7.2 diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index ea34a073..00b03526 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -65,11 +65,8 @@ 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(), - state: Mutex::new(State { - status: Status::Pending, - waiters: WakerSet::new(), - }), + result: OnceLock::new(), + waiters: Mutex::new(WakerSet::new()), }); let completer = Completer { shared: Arc::downgrade(&shared), @@ -79,20 +76,9 @@ pub fn new() -> (Completer, Completion) { } struct Shared { - value: OnceLock, - state: Mutex, -} - -struct State { - status: Status, - waiters: WakerSet, -} - -#[derive(Clone, Copy, Debug, Eq, PartialEq)] -enum Status { - Pending, - Completed, - Abandoned, + // A stored None means abandonment. Initialization also publishes waiter detachment. + result: OnceLock>, + waiters: Mutex, } /// The error returned by [`Completion::wait`] when the completer was dropped without a value. @@ -117,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 {} @@ -145,21 +131,11 @@ impl Completer { return Err(value); }; let wakers = { - let mut state = shared.state.lock(); - assert_eq!( - state.status, - Status::Pending, - "a live completer must refer to a pending completion" - ); - - if let Err(value) = shared.value.set(value) { - drop(state); - 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 mut waiters = shared.waiters.lock(); + let wakers = waiters.take_all(); + // 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 // wake callbacks; the completed state no longer needs abandonment handling. @@ -175,13 +151,10 @@ impl Drop for Completer { return; }; let wakers = { - let mut state = shared.state.lock(); - if state.status != Status::Pending { - return; - } - // Publish abandonment and detach its waiters atomically with respect to registration. - state.status = Status::Abandoned; - state.waiters.take_all() + let mut waiters = shared.waiters.lock(); + let wakers = waiters.take_all(); + assert!(shared.result.set(None).is_ok()); + wakers }; wake_all(wakers); } @@ -221,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 @@ -229,7 +202,7 @@ impl Completion { } struct Wait<'a, T> { - completion: &'a Completion, + shared: &'a Shared, token: Option, } @@ -238,33 +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(); - let mut state = this.completion.shared.state.lock(); - let (poll, retired_waker) = match state.status { - Status::Pending => { - let retired = state.waiters.register(&mut this.token, cx.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 - .shared - .value - .get() - .expect("completed value must be initialized"); - (Poll::Ready(Ok(value)), None) - } - Status::Abandoned => { - // Abandonment uses the same terminal detach protocol as completion. - this.token = None; - (Poll::Ready(Err(Abandoned(()))), None) + 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(state); - drop(retired_waker); - poll + + this.token = None; + Poll::Ready(result.as_ref().ok_or(Abandoned(()))) } } @@ -274,15 +237,17 @@ 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. - self.token = None; + if self.shared.result.get().is_some() { + return; + } + + let mut waiters = self.shared.waiters.lock(); + if self.shared.result.get().is_some() { 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/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 c4d3d9af..106d36de 100644 --- a/tests-integration/tests/completion_test.rs +++ b/tests-integration/tests/completion_test.rs @@ -37,31 +37,31 @@ 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] @@ -74,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(()); @@ -114,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()); @@ -138,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] @@ -197,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());