perf: reduce fan-out notification overhead

CHANGELOG.md

@@ -38,6 +38,7 @@ All notable changes to this project will be documented in this file.
 
 ### Improvements
 
+* Reduce fan-out notification overhead by avoiding heap allocation for a single waiter and transferring terminal waiter storage out of state locks.
 * Reduce MPSC receiver registration and wake latency by storing receiver wakers inline instead of allocating them on the heap.
 * Reduce semaphore and mutex hot-path overhead by avoiding wake-buffer allocation when no tasks are queued and batching queued wakes on the stack.
 * Allocate `OnceMap` and `singleflight::Group` registries lazily to reduce construction overhead.

asyncband/src/broadcast/mpmc/unbounded/mod.rs

@@ -398,7 +398,7 @@ impl<T> Drop for UnboundedSender<T> {
                 // Wake every parked receiver so it can observe the channel's disconnected state.
                 let wakers = {
                     let mut inner = self.shared.inner.lock();
-                    inner.waiters.drain()
+                    inner.waiters.take_all()
                 };
                 wake_all(wakers);
             }

asyncband/src/completion/mod.rs

@@ -160,7 +160,7 @@ impl<T> Completer<T> {
             }
             // Publish the value before making completion observable and detaching its waiters.
             state.status = Status::Completed;
-            state.waiters.drain()
+            state.waiters.take_all()
         };
         // `complete` consumes the only completer. Disarm its destructor before invoking arbitrary
         // wake callbacks; the completed state no longer needs abandonment handling.
@@ -182,7 +182,7 @@ impl<T> Drop for Completer<T> {
             }
             // Publish abandonment and detach its waiters atomically with respect to registration.
             state.status = Status::Abandoned;
-            state.waiters.drain()
+            state.waiters.take_all()
         };
         wake_all(wakers);
     }

asyncband/src/condvar/mod.rs

@@ -69,6 +69,7 @@ use crate::internal::mutex::Mutex;
 use crate::internal::waitlist::WaitList;
 use crate::internal::waitlist::WaiterId;
 use crate::internal::wake_all;
+use crate::internal::waker_batch::WakerBatch;
 use crate::mutex;
 use crate::mutex::MutexGuard;
 use crate::mutex::OwnedMutexGuard;
@@ -158,7 +159,7 @@ impl Condvar {
     pub fn notify_all(&self) {
         let wakers = {
             let mut waiters = self.waiters.lock();
-            let mut wakers = vec![];
+            let mut wakers = WakerBatch::new();
 
             while waiters
                 .unlink_first_waiter(|node| {

asyncband/src/event/mod.rs

@@ -63,6 +63,7 @@ use crate::internal::mutex::Mutex;
 use crate::internal::waitlist::WaitList;
 use crate::internal::waitlist::WaiterId;
 use crate::internal::wake_all;
+use crate::internal::waker_batch::WakerBatch;
 
 /// A reusable event that remains set until explicitly reset.
 ///
@@ -171,7 +172,7 @@ impl ManualResetEvent {
             state.is_set = true;
             // 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.
-            let mut wakers = vec![];
+            let mut wakers = WakerBatch::new();
             while let Some((_id, waiter)) = state.waiters.unlink_first_waiter(|waiter| {
                 waiter.notified = true;
                 true

asyncband/src/internal/arena.rs

@@ -59,26 +59,6 @@ pub struct Arena<T> {
     len: usize,
 }
 
-/// Values extracted from an [`Arena`], storing the common single-value case inline.
-///
-/// This is the specialized subset of a small-vector abstraction needed here: keep one value inline,
-/// store additional values in a `Vec`, and support consuming iteration. Keeping that representation
-/// focused avoids a general unsafe collection implementation for a single internal operation.
-#[derive(Debug)]
-struct ArenaValues<T> {
-    first: Option<T>,
-    rest: Vec<T>,
-}
-
-impl<T> IntoIterator for ArenaValues<T> {
-    type Item = T;
-    type IntoIter = std::iter::Chain<std::option::IntoIter<T>, std::vec::IntoIter<T>>;
-
-    fn into_iter(self) -> Self::IntoIter {
-        self.first.into_iter().chain(self.rest)
-    }
-}
-
 #[derive(Debug)]
 enum Slot<T> {
     Occupied(T),
@@ -181,40 +161,32 @@ impl<T> Arena<T> {
         value
     }
 
-    /// Takes every occupied value in slot order while retaining the allocation for reuse.
+    /// Drains every occupied value in slot order while retaining the allocation for reuse.
     ///
-    /// After a non-empty take, every previously issued slot ID becomes invalid, including IDs for
+    /// After a non-empty drain, every previously issued slot ID becomes invalid, including IDs for
     /// slots that were already vacant. Consumers that retain IDs across this operation must supply
     /// their own epoch check.
     #[inline]
-    pub fn take_all(&mut self) -> impl Iterator<Item = T> + use<T> {
-        let len = self.len;
-        let mut values = ArenaValues {
-            first: None,
-            rest: vec![],
-        };
-        if len == 0 {
-            // Individually removed values leave vacant slots behind. Keep their free list intact
-            // instead of scanning the arena's historical high-water mark to drain no values.
-            return values.into_iter();
-        }
-
-        for slot in self.slots.drain(..) {
-            if let Slot::Occupied(value) = slot {
-                if values.first.is_none() {
-                    values.first = Some(value);
-                } else {
-                    if values.rest.is_empty() {
-                        values.rest.reserve(len - 1);
-                    }
-                    values.rest.push(value);
-                }
-            }
-        }
+    pub fn drain(&mut self) -> impl Iterator<Item = T> + '_ {
+        self.next_vacant = 0;
+        self.len = 0;
+        self.slots.drain(..).filter_map(|slot| match slot {
+            Slot::Occupied(value) => Some(value),
+            Slot::Vacant { .. } => None,
+        })
+    }
 
+    /// Takes every occupied value and the backing allocation in slot order.
+    #[inline]
+    pub fn take_all(&mut self) -> impl Iterator<Item = T> + use<T> {
         self.next_vacant = 0;
         self.len = 0;
-        values.into_iter()
+        mem::take(&mut self.slots)
+            .into_iter()
+            .filter_map(|slot| match slot {
+                Slot::Occupied(value) => Some(value),
+                Slot::Vacant { .. } => None,
+            })
     }
 }
 
@@ -242,19 +214,32 @@ mod tests {
     }
 
     #[test]
-    fn take_all_restarts_slot_id_allocation() {
+    fn drain_restarts_slot_id_allocation() {
         let mut arena = Arena::with_capacity(3);
         let first = arena.insert(1);
         let second = arena.insert(2);
         let third = arena.insert(3);
         let capacity = arena.slots.capacity();
         arena.remove(second);
 
-        assert_eq!(arena.take_all().collect::<Vec<_>>(), vec![1, 3]);
+        assert_eq!(arena.drain().collect::<Vec<_>>(), vec![1, 3]);
         assert_eq!(arena.len(), 0);
         assert_eq!(arena.slots.capacity(), capacity);
 
         let slot_ids = [arena.insert(4), arena.insert(5), arena.insert(6)];
         assert_eq!(slot_ids, [first, second, third]);
     }
+
+    #[test]
+    fn take_all_releases_the_backing_allocation() {
+        let mut arena = Arena::new();
+        arena.insert(1);
+        let removed = arena.insert(2);
+        arena.insert(3);
+        arena.remove(removed);
+
+        let values = arena.take_all();
+        assert_eq!(arena.slots.capacity(), 0);
+        assert_eq!(values.collect::<Vec<_>>(), vec![1, 3]);
+    }
 }

asyncband/src/internal/countdown.rs

@@ -57,7 +57,7 @@ impl CountdownState {
     pub fn wake_all(&self) {
         let wakers = {
             let mut waiters = self.waiters.lock();
-            waiters.drain()
+            waiters.take_all()
         };
 
         wake_all(wakers);

asyncband/src/internal/mod.rs

@@ -121,6 +121,25 @@ pub(crate) mod semaphore;
 #[allow(dead_code)]
 pub(crate) mod waitlist;
 
+#[cfg(any(
+    feature = "barrier",
+    feature = "broadcast",
+    feature = "event",
+    feature = "completion",
+    feature = "latch",
+    feature = "mpsc",
+    feature = "mutex",
+    feature = "once",
+    feature = "rwlock",
+    feature = "semaphore",
+    feature = "waitgroup",
+    feature = "watch",
+))]
+// Wait-set primitives know the exact batch capacity, while linked-list primitives use the
+// allocation-free constructor. Each constructor is therefore unused in some feature subsets.
+#[allow(dead_code)]
+pub(crate) mod waker_batch;
+
 #[cfg(any(
     feature = "barrier",
     feature = "broadcast",

asyncband/src/internal/semaphore.rs

@@ -30,6 +30,7 @@ use crate::internal::mutex::Mutex;
 use crate::internal::waitlist::WaitList;
 use crate::internal::waitlist::WaiterId;
 use crate::internal::wake_all;
+use crate::internal::waker_batch::WakerBatch;
 
 /// The internal semaphore that provides low-level async primitives.
 #[derive(Debug)]
@@ -207,7 +208,7 @@ impl Semaphore {
     /// Adds as many permits until there is no waiter.
     pub fn notify_all(&self) {
         let mut waiters = self.waiters.lock();
-        let mut wakers = vec![];
+        let mut wakers = WakerBatch::new();
         loop {
             match waiters.unlink_first_waiter(|node| {
                 node.permits = 0;

asyncband/src/internal/waitset.rs

@@ -26,6 +26,7 @@ use std::task::Waker;
 
 use crate::internal::arena::Arena;
 use crate::internal::arena::SlotId;
+use crate::internal::waker_batch::WakerBatch;
 
 /// An exclusive handle to one waiter slot in a [`WaitSet`].
 ///
@@ -68,8 +69,24 @@ impl WaitSet {
     /// releasing the lock that protects this wait set.
     #[inline]
     pub fn drain(&mut self) -> impl Iterator<Item = Waker> + 'static {
+        let mut wakers = WakerBatch::with_capacity(self.waiters.len());
+        if self.waiters.is_empty() {
+            return wakers.into_iter();
+        }
+        self.advance_epoch();
+        wakers.extend(self.waiters.drain());
+        wakers.into_iter()
+    }
+
+    /// Takes all registered wakers and the wait set's backing allocation without waking them.
+    ///
+    /// This is intended for terminal state transitions where the current allocation cannot be
+    /// reused. The caller must consume or drop the iterator after releasing the lock that protects
+    /// this wait set.
+    #[inline]
+    pub fn take_all(&mut self) -> impl Iterator<Item = Waker> + 'static {
         if !self.waiters.is_empty() {
-            self.epoch = self.epoch.checked_add(1).expect("wait set epoch overflow");
+            self.advance_epoch();
         }
         self.waiters.take_all()
     }
@@ -121,6 +138,10 @@ impl WaitSet {
         )
     }
 
+    fn advance_epoch(&mut self) {
+        self.epoch = self.epoch.checked_add(1).expect("wait set epoch overflow");
+    }
+
     #[cfg(test)]
     fn registered_len(&self) -> usize {
         self.waiters.len()
@@ -210,6 +231,29 @@ mod tests {
         assert_eq!(second_task.0.load(Ordering::Relaxed), 1);
     }
 
+    #[test]
+    fn take_all_invalidates_existing_tokens() {
+        let mut waiters = WaitSet::new();
+        let first_task = Arc::new(TrackWake(AtomicUsize::new(0)));
+        let first_waker = Waker::from(first_task.clone());
+        let second_task = Arc::new(TrackWake(AtomicUsize::new(0)));
+        let second_waker = Waker::from(second_task.clone());
+        let mut first_token = None;
+        let mut second_token = None;
+
+        drop(waiters.register(&mut first_token, &first_waker));
+        assert_eq!(waiters.take_all().count(), 1);
+
+        drop(waiters.register(&mut second_token, &second_waker));
+        drop(waiters.register(&mut first_token, &first_waker));
+
+        let registered = waiters.take_all().collect::<Vec<_>>();
+        assert_eq!(registered.len(), 2);
+        wake_all(registered.into_iter());
+        assert_eq!(first_task.0.load(Ordering::Relaxed), 1);
+        assert_eq!(second_task.0.load(Ordering::Relaxed), 1);
+    }
+
     #[test]
     fn wake_all_notifies_remaining_waiters_after_a_panic() {
         let mut waiters = WaitSet::new();

asyncband/src/internal/waker_batch.rs

@@ -0,0 +1,66 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::task::Waker;
+
+/// An owning waker collection that stores the first entry without allocating.
+#[derive(Debug)]
+pub struct WakerBatch {
+    first: Option<Waker>,
+    rest: Vec<Waker>,
+}
+
+impl WakerBatch {
+    pub const fn new() -> Self {
+        Self {
+            first: None,
+            rest: vec![],
+        }
+    }
+
+    pub fn with_capacity(capacity: usize) -> Self {
+        Self {
+            first: None,
+            rest: Vec::with_capacity(capacity.saturating_sub(1)),
+        }
+    }
+
+    pub fn push(&mut self, waker: Waker) {
+        if self.first.is_none() {
+            self.first = Some(waker);
+        } else {
+            self.rest.push(waker);
+        }
+    }
+}
+
+impl Extend<Waker> for WakerBatch {
+    fn extend<I: IntoIterator<Item = Waker>>(&mut self, iter: I) {
+        for waker in iter {
+            self.push(waker);
+        }
+    }
+}
+
+impl IntoIterator for WakerBatch {
+    type Item = Waker;
+    type IntoIter = std::iter::Chain<std::option::IntoIter<Waker>, std::vec::IntoIter<Waker>>;
+
+    fn into_iter(self) -> Self::IntoIter {
+        self.first.into_iter().chain(self.rest)
+    }
+}

asyncband/src/watch/mod.rs

@@ -153,7 +153,7 @@ impl<T> Drop for Sender<T> {
             if state.senders != 0 {
                 return;
             }
-            state.waiters.drain()
+            state.waiters.take_all()
         };
         wake_all(wakers);
     }