perf(completion): make terminal waits lock-free

CHANGELOG.md

@@ -11,6 +11,7 @@ All notable changes to this project will be documented in this file.
 ### Improvements
 
 * Reduce unbounded MPSC synchronization overhead by transferring messages in batches and coordinating receiver notifications with queued messages; release large empty batch allocations while retaining small buffers for reuse.
+* Make completed and abandoned `Completion` waits lock-free while preserving cancellable pending registration and unlocked waker callbacks.
 
 ## v0.7.2
 

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<T>() -> (Completer<T>, Completion<T>) {
     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,21 +80,29 @@ pub fn new<T>() -> (Completer<T>, Completion<T>) {
 
 struct Shared<T> {
     value: OnceLock<T>,
-    state: Mutex<State>,
-}
-
-struct State {
-    status: Status,
-    waiters: WakerSet,
+    status: AtomicU8,
+    waiters: Mutex<WakerSet>,
 }
 
+#[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(());
@@ -145,21 +153,24 @@ impl<T> Completer<T> {
             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<T> Drop for Completer<T> {
             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,31 +251,47 @@ impl<'a, T> Future for Wait<'a, T> {
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         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<T> = this.completion;
                 let value = completion
                     .shared
                     .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<T> 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);
     }
 }

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<WakerToken>,
+        waker: Waker,
+    ) -> Option<Waker> {
+        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

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<Option<Box<dyn FnOnce() + Send>>>);
 
+struct CloneCallbackWake(Mutex<Option<Box<dyn FnOnce() + Send>>>);
+
 // 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::<CloneCallbackWake>::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::<CloneCallbackWake>::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::<CloneCallbackWake>::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<F: Future>(future: Pin<&mut F>, waker: &Waker) -> Poll<F::Output> {
     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::<usize>();
+            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(