perf(event): avoid locking manual-reset state reads and ready waits

CHANGELOG.md

@@ -32,6 +32,10 @@ All notable changes to this project will be documented in this file.
 * Prevent deadlocks in `singleflight::Group::work` and `try_work` when a duplicate key's destructor calls back into the same group.
 * Fix undefined behavior when sending or dropping a oneshot sender races with receiver completion or cancellation.
 
+### Improvements
+
+* Reduce `ManualResetEvent` state-query and already-set wait overhead, especially with concurrent readers.
+
 ### Notable changes
 
 * Executor waker operations, including cloning, waking, and dropping, are expected not to panic; recovery from panicking waker callbacks is no longer supported.

asyncband/src/event/manual_reset.rs

@@ -19,12 +19,13 @@ use std::fmt;
 use std::future::Future;
 use std::pin::Pin;
 use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
 use std::task::Context;
 use std::task::Poll;
 use std::task::Waker;
 
 use crate::internal::mutex::Mutex;
-use crate::internal::register_waker;
 use crate::internal::waitlist::WaitList;
 use crate::internal::waitlist::WaiterId;
 use crate::internal::waker_batch::WakerBatch;
@@ -71,7 +72,9 @@ use crate::internal::waker_batch::WakerBatch;
 /// # }
 /// ```
 pub struct ManualResetEvent {
-    state: Mutex<State>,
+    // Flag writes and waiter registration are serialized by the same lock.
+    is_set: AtomicBool,
+    waiters: Mutex<WaitList<Waiter>>,
 }
 
 impl ManualResetEvent {
@@ -85,10 +88,8 @@ impl ManualResetEvent {
     /// If `is_set` is `true`, waits complete immediately until the event is reset.
     pub const fn with_state(is_set: bool) -> Self {
         Self {
-            state: Mutex::new(State {
-                is_set,
-                waiters: WaitList::new(),
-            }),
+            is_set: AtomicBool::new(is_set),
+            waiters: Mutex::new(WaitList::new()),
         }
     }
 
@@ -99,15 +100,16 @@ impl ManualResetEvent {
     pub fn set(&self) {
         let mut wakers = WakerBatch::new();
         {
-            let mut state = self.state.lock();
-            if state.is_set {
+            let mut waiters = self.waiters.lock();
+            if self.is_set.load(Ordering::Relaxed) {
                 return;
             }
 
-            state.is_set = true;
+            // Publish to successful observations that do not acquire the waiter lock.
+            self.is_set.store(true, Ordering::Release);
             // Detach the complete cohort before invoking any waker. A wake callback may reset the
             // event and register a new wait, which must belong to the state current at that point.
-            while let Some((_id, waiter)) = state.waiters.unlink_first_waiter(|waiter| {
+            while let Some((_id, waiter)) = waiters.unlink_first_waiter(|waiter| {
                 waiter.notified = true;
                 true
             }) {
@@ -124,7 +126,9 @@ impl ManualResetEvent {
     /// Waits already released by a preceding [`set`](Self::set) remain ready. If the event is
     /// already unset, this has no effect.
     pub fn reset(&self) {
-        self.state.lock().is_set = false;
+        let _waiters = self.waiters.lock();
+        // Reset cannot race ahead of a set that is still detaching its waiter cohort.
+        self.is_set.store(false, Ordering::Relaxed);
     }
 
     /// Returns whether the event is currently set.
@@ -142,7 +146,7 @@ impl ManualResetEvent {
     /// assert!(event.is_set());
     /// ```
     pub fn is_set(&self) -> bool {
-        self.state.lock().is_set
+        self.is_set.load(Ordering::Acquire)
     }
 
     /// Attempts to wait without registering a waiter.
@@ -211,46 +215,69 @@ impl ManualResetEvent {
     /// event is set never enqueues. A linked waiter therefore always belongs to an unset event, so
     /// `notified` alone decides whether a registered waiter is already committed.
     fn poll_wait(&self, waiter_id: &mut Option<WaiterId>, cx: &mut Context<'_>) -> Poll<()> {
-        let (poll, retired_waker) = {
-            let mut state = self.state.lock();
+        // Only an unregistered wait can skip the lock: a released waiter still owns a node.
+        let mut prepared_waker = if waiter_id.is_none() {
+            if self.is_set() {
+                return Poll::Ready(());
+            }
+            Some(cx.waker().clone())
+        } else {
+            None
+        };
+
+        let (poll, retired_waker) = loop {
+            let mut waiters = self.waiters.lock();
             match *waiter_id {
-                Some(id) if state.waiters.waiter_mut(id).notified => {
-                    let waiter = state.remove_waiter(id);
+                Some(id) if waiters.waiter_mut(id).notified => {
+                    let waiter = waiters.remove_unlinked_waiter(id);
                     *waiter_id = None;
-                    (Poll::Ready(()), waiter.waker)
+                    break (Poll::Ready(()), waiter.waker);
                 }
                 Some(id) => {
                     debug_assert!(
-                        !state.is_set,
+                        !self.is_set.load(Ordering::Relaxed),
                         "a linked waiter must belong to an unset event"
                     );
-                    let waiter = state.waiters.waiter_mut(id);
-                    assert!(
-                        waiter.waker.is_some(),
-                        "an unnotified waiter must retain its waker"
-                    );
-                    let retired = register_waker(&mut waiter.waker, cx.waker());
-                    (Poll::Pending, retired)
+                    let waiter = waiters.waiter_mut(id);
+                    let current = waiter
+                        .waker
+                        .as_ref()
+                        .expect("an unnotified waiter must retain its waker");
+                    if current.will_wake(cx.waker()) {
+                        break (Poll::Pending, None);
+                    }
+                    if let Some(waker) = prepared_waker.take() {
+                        break (Poll::Pending, waiter.waker.replace(waker));
+                    }
+                    // Clone only a changed waker, outside the lock, then recheck commitment:
+                    // set/reset may run while preparing the replacement.
+                    drop(waiters);
+                    prepared_waker = Some(cx.waker().clone());
                 }
-                None if state.is_set => (Poll::Ready(()), None),
+                // A set after the false fast probe must either be observed here or release
+                // the node that we register while holding this same lock.
+                None if self.is_set.load(Ordering::Relaxed) => break (Poll::Ready(()), None),
                 None => {
-                    *waiter_id = Some(state.waiters.push_back(Waiter {
+                    *waiter_id = Some(waiters.push_back(Waiter {
                         notified: false,
-                        waker: Some(cx.waker().clone()),
+                        waker: prepared_waker.take(),
                     }));
-                    (Poll::Pending, None)
+                    break (Poll::Pending, None);
                 }
             }
         };
 
+        drop(prepared_waker);
         drop(retired_waker);
         poll
     }
 
     fn unregister_waiter(&self, id: WaiterId) {
         let waiter = {
-            let mut state = self.state.lock();
-            state.remove_waiter(id)
+            let mut waiters = self.waiters.lock();
+            // Released waits retain their detached nodes until repolled or cancelled.
+            waiters.unlink_waiter(id, |_| true);
+            waiters.remove_unlinked_waiter(id)
         };
         drop(waiter);
     }
@@ -264,29 +291,13 @@ impl Default for ManualResetEvent {
 
 impl fmt::Debug for ManualResetEvent {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
-        let is_set = self.state.lock().is_set;
+        let is_set = self.is_set();
         f.debug_struct("ManualResetEvent")
             .field("is_set", &is_set)
             .finish_non_exhaustive()
     }
 }
 
-#[derive(Debug)]
-struct State {
-    is_set: bool,
-    waiters: WaitList<Waiter>,
-}
-
-impl State {
-    /// Removes a waiter whether or not [`ManualResetEvent::set`] already unlinked it.
-    fn remove_waiter(&mut self, id: WaiterId) -> Waiter {
-        // Unlinking is idempotent: a waiter that `set` detached keeps its node until it is removed
-        // here, and an unconditional predicate never declines.
-        self.waiters.unlink_waiter(id, |_| true);
-        self.waiters.remove_unlinked_waiter(id)
-    }
-}
-
 #[derive(Debug)]
 struct Waiter {
     notified: bool,

benchmarks/benches/primitives/event/wait.rs

@@ -16,9 +16,14 @@
 // under the License.
 
 use std::pin::pin;
+use std::sync::Arc;
+use std::task::Context;
+use std::task::Wake;
+use std::task::Waker;
 
 use asyncband::event::ManualResetEvent;
 use benchmarks::support::bench_context;
+use benchmarks::support::defer_input_drop;
 use benchmarks::support::poll_pending;
 use benchmarks::support::poll_pinned_ready;
 use divan::Bencher;
@@ -29,6 +34,15 @@ const WAITER_COUNTS: &[usize] = &[1, 4, 16];
 const THREAD_COUNTS: &[usize] = &[1, 2, 4, 8];
 const CONTENDED_SAMPLE_SIZE: u32 = 256;
 
+// All benchmark threads read the same event without concurrent state transitions.
+#[divan::bench(args = [false, true], threads = THREAD_COUNTS, sample_size = CONTENDED_SAMPLE_SIZE)]
+fn is_set(bencher: Bencher, is_set: bool) {
+    let event = ManualResetEvent::with_state(is_set);
+
+    bencher.bench(|| black_box(black_box(&event).is_set()));
+}
+
+// Each iteration creates and polls a fresh wait against the same already-set event.
 #[divan::bench(threads = THREAD_COUNTS, sample_size = CONTENDED_SAMPLE_SIZE)]
 fn wait_already_set(bencher: Bencher) {
     let event = ManualResetEvent::with_state(true);
@@ -68,6 +82,7 @@ fn reused_waiter_handoff(bencher: Bencher) {
     });
 }
 
+// Measures the full fan-out lifecycle, including allocation, registration, set, and completion.
 #[divan::bench(args = WAITER_COUNTS)]
 fn waiter_fan_out(bencher: Bencher, waiter_count: usize) {
     let mut context = bench_context();
@@ -91,6 +106,7 @@ fn waiter_fan_out(bencher: Bencher, waiter_count: usize) {
         });
 }
 
+// A fresh event per iteration includes allocating its waiter storage.
 #[divan::bench]
 fn cancel_pending(bencher: Bencher) {
     let mut context = bench_context();
@@ -104,3 +120,71 @@ fn cancel_pending(bencher: Bencher) {
         black_box(event)
     });
 }
+
+// Reusing the event retains waiter storage capacity across registration/cancellation cycles.
+#[divan::bench]
+fn cancel_pending_reused(bencher: Bencher) {
+    let mut context = bench_context();
+    let event = ManualResetEvent::new();
+
+    bencher.bench_local(|| {
+        let mut wait = pin!(event.wait());
+        poll_pending(wait.as_mut(), &mut context);
+    });
+}
+
+// Initial registration is outside the timed loop; every repoll uses the same waker.
+#[divan::bench]
+fn repoll_pending(bencher: Bencher) {
+    let mut context = bench_context();
+    let event = ManualResetEvent::new();
+    let mut wait = pin!(event.wait());
+    poll_pending(wait.as_mut(), &mut context);
+
+    bencher.bench_local(|| poll_pending(wait.as_mut(), &mut context));
+}
+
+#[divan::bench]
+fn replace_pending_waker(bencher: Bencher) {
+    struct Task;
+    #[allow(clippy::manual_noop_waker)]
+    impl Wake for Task {
+        fn wake(self: Arc<Self>) {}
+    }
+
+    // Alternate distinct reference-counted wakers so every timed poll replaces the registration.
+    // These wakers exercise clone/drop costs but do not schedule executor tasks.
+    let wakers = [Waker::from(Arc::new(Task)), Waker::from(Arc::new(Task))];
+    let event = ManualResetEvent::new();
+    let mut wait = pin!(event.wait());
+    poll_pending(wait.as_mut(), &mut Context::from_waker(&wakers[0]));
+    let mut next = 1;
+
+    bencher.bench_local(|| {
+        poll_pending(wait.as_mut(), &mut Context::from_waker(&wakers[next]));
+        next ^= 1;
+    });
+}
+
+// Isolates set/detach/wake from registration and cleanup. The smaller sample size limits setup
+// costs for large cohorts; use --sample-count 200 without a time limit for repeated measurements.
+#[divan::bench(args = [1, 16, 256], sample_size = 64)]
+fn set_registered(bencher: Bencher, waiter_count: usize) {
+    bencher
+        .with_inputs(|| {
+            let event = Arc::new(ManualResetEvent::new());
+            let mut context = bench_context();
+            let mut waiters = (0..waiter_count)
+                .map(|_| Box::pin(event.clone().wait_owned()))
+                .collect::<Vec<_>>();
+            for waiter in &mut waiters {
+                poll_pending(waiter.as_mut(), &mut context);
+            }
+            (event, waiters)
+        })
+        .bench_local_values(|(event, waiters)| {
+            // Time cohort detachment and wake callbacks, excluding registration and cleanup.
+            event.set();
+            defer_input_drop((event, waiters), ())
+        });
+}