perf(completion): make terminal waits lock-free

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
 

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<T>() -> (Completer<T>, Completion<T>) {
     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<T>() -> (Completer<T>, Completion<T>) {
 }
 
 struct Shared<T> {
-    value: OnceLock<T>,
-    state: Mutex<State>,
-}
-
-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<Option<T>>,
+    waiters: Mutex<WakerSet>,
 }
 
 /// The error returned by [`Completion::wait`] when the completer was dropped without a value.
@@ -117,9 +103,9 @@ pub struct Completer<T> {
 }
 
 // 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<T>` 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<T>` 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<T: Send> Send for Completer<T> {}
 unsafe impl<T: Send> Sync for Completer<T> {}
@@ -145,21 +131,11 @@ impl<T> Completer<T> {
             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<T> Drop for Completer<T> {
             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,15 +194,15 @@ impl<T> Completion<T> {
     /// 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
     }
 }
 
 struct Wait<'a, T> {
-    completion: &'a Completion<T>,
+    shared: &'a Shared<T>,
     token: Option<WakerToken>,
 }
 
@@ -238,33 +211,23 @@ 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 {
-            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<T> = 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<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.
-            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);
     }
 }

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::<usize>();
+    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::<usize>();
+    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::<Vec<_>>();
         for waiter in &mut waiters {
             poll_pending(waiter.as_mut(), &mut context);

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::<Cell<u8>>();
 
-    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::<usize>();
-    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::<usize>();
     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::<usize>();
-    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::<usize>();
     let tracker = Arc::new(WakeCounter::default());
     let waker = Waker::from(tracker.clone());