From 7010f81cf2e3d2d3a7e9ee71a58f18f2c25d32ae Mon Sep 17 00:00:00 2001 From: ryux1 Date: Sun, 6 Sep 2026 08:40:51 +0200 Subject: [PATCH] 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(