chore: fine tune waker wake

asyncband/src/barrier/mod.rs

@@ -52,10 +52,9 @@ use std::future::Future;
 use std::pin::Pin;
 use std::task::Context;
 use std::task::Poll;
+use std::task::Waker;
 
 use crate::internal::mutex::Mutex;
-use crate::internal::wake_all;
-use crate::internal::waker_batch::WakerBatch;
 use crate::internal::wakerset::WakerSet;
 use crate::internal::wakerset::WakerToken;
 
@@ -185,14 +184,13 @@ impl Barrier {
             state.arrived += 1;
 
             // The final arrival completes this generation. Advance the generation while holding
-            // the state lock, then wake the drained followers after releasing it.
+            // the state lock, then wake the detached followers after releasing it.
             if state.arrived == self.n {
                 state.arrived = 0;
                 state.generation += 1;
-                let mut wakers = WakerBatch::new();
-                state.waiters.drain_into(&mut wakers);
+                let mut wakers = state.waiters.take_all();
                 drop(state);
-                wake_all(&mut wakers);
+                wakers.by_ref().for_each(Waker::wake);
                 return BarrierWaitResult(true);
             }
 
@@ -239,7 +237,7 @@ impl Future for BarrierWait<'_> {
 
         let mut state = barrier.state.lock();
         if *generation < state.generation {
-            // Completion advances the generation and drains its old waiters under this same lock,
+            // Completion advances the generation and detaches its old waiters under this same lock,
             // so no registration represented by this token remains in the waker set.
             *token = None;
             return Poll::Ready(());

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

@@ -108,6 +108,7 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::task::Context;
 use std::task::Poll;
+use std::task::Waker;
 
 use super::common;
 use super::common::Backlog;
@@ -119,8 +120,6 @@ use crate::internal::arena::SlotId;
 use crate::internal::mutex::Mutex;
 use crate::internal::semaphore::Acquire;
 use crate::internal::semaphore::Semaphore;
-use crate::internal::wake_all;
-use crate::internal::waker_batch::WakerBatch;
 use crate::internal::wakerset::WakerToken;
 
 #[cfg(test)]
@@ -181,64 +180,39 @@ struct Shared<T> {
     /// and parks again if another producer took the space first. The semaphore starts empty and
     /// only ever grows when a reclaim finds someone waiting, so an idle channel accumulates none.
     tx_permits: Semaphore,
-    /// How many producers are somewhere inside the waiting path of [`BoundedSender::send`].
-    ///
-    /// An upper bound on the number of parked producers, and the only thing either release path
-    /// consults. It answers both questions a reclaim has — whether to wake anyone, and how many
-    /// permits are worth handing out — without taking the semaphore's lock. Reclaiming is far more
-    /// frequent than blocking — under fan-out every message is reclaimed, while a channel with
-    /// headroom never blocks at all — so paying an atomic load there instead of a lock acquisition
-    /// is what keeps an uncontended receive off the semaphore entirely.
+    /// Number of producers in the waiting path of [`BoundedSender::send`].
+    ///
+    /// This upper bound lets receives skip the semaphore lock when no producer is waiting.
     blocked_senders: AtomicUsize,
 }
 
 impl<T> Shared<T> {
-    /// Hands `freed` released slots back to producers parked in `send`.
-    ///
-    /// Capacity is `retained()`, which is `buffer.len()`. The buffer grows only in
-    /// `Backlog::publish_retained` and shrinks only in `Backlog::reclaim_vacated`, which is
-    /// reachable from exactly two places: a receive that vacates the last cursor at the backlog
-    /// head, and removing a subscription. Those are the only callers of this method, so no path
-    /// can free capacity without waking a producer. Subscribing cannot: a new cursor starts at the
-    /// tail and never lowers `retained()`.
+    /// Notifies blocked producers after a receive or subscription removal frees capacity.
     ///
-    /// Callers must invoke this with the channel unlocked, and — on the receive path — before
-    /// touching the payload, since `common::take_msg` runs user code that may panic.
+    /// Call with the channel unlocked, before payload cloning or destruction can panic.
     fn release_reclaimed(&self, freed: usize) {
-        // Release no more permits than there are producers to wake. A permit the semaphore cannot
-        // hand to a waiter is kept as slack, and the next producer to block has to burn it off one
-        // futile publish attempt — a channel lock apiece — at a time before it can park. Freeing a
-        // large prefix at once is not exotic: dropping a lagging subscription reclaims the whole
-        // backlog, which would otherwise leave nearly `capacity` permits behind.
-        //
-        // Capping cannot lose a wake-up, by the same argument that lets this read the count at all:
-        // a producer this load observes is one the release covers, and one it misses incremented
-        // after the load, which it does before taking the channel lock to recheck — so its recheck
-        // runs after the reclaim and finds the capacity itself.
+        // Cap permits at the number of blocked producers. Surplus permits make later sends retry
+        // a full channel instead of parking, especially after dropping a lagging subscription.
         let waiting = self.waiting_senders();
         if freed > 0 && waiting > 0 {
             self.tx_permits.release_if_nonempty(freed.min(waiting));
         }
     }
 
-    /// Wakes every parked producer, however many slots came back.
+    /// Wakes every blocked producer when the last subscription leaves.
     ///
-    /// The last subscription leaving is not a reclaim of some number of slots — it removes the
-    /// limit itself, because a channel with no receivers discards instead of retaining. Releasing
-    /// only as many permits as that final reclaim freed would strand every producer beyond that
-    /// count, so this is the one release that must be unbounded.
+    /// Sends now discard payloads without consuming capacity, so every producer can proceed.
     fn release_all(&self) {
         if self.waiting_senders() > 0 {
             self.tx_permits.notify_all();
         }
     }
 
-    /// How many producers might be waiting, answered without touching the semaphore's lock.
+    /// Returns an upper bound on parked producers without locking the semaphore.
     ///
-    /// This cannot miss a wake-up. A producer increments the count before it ever takes the
-    /// channel lock to recheck capacity, and every caller here loads it after releasing that same
-    /// lock, so the mutex orders the two: either this load observes the producer, or the
-    /// producer's recheck runs after the change and finds the capacity itself.
+    /// Producers increment before rechecking capacity under the channel lock; reclaim paths load
+    /// after releasing it. A producer is either covered by this count or rechecks capacity after
+    /// the reclaim, so skipping or capping notifications cannot strand it.
     fn waiting_senders(&self) -> usize {
         self.blocked_senders.load(Ordering::Acquire)
     }
@@ -429,36 +403,26 @@ impl<T> BoundedSender<T> {
         self.publish(msg, |msg| msg)
     }
 
-    /// The publish step both send paths share.
-    ///
-    /// `into_msg` is called only once this decides the message will actually be retained, which is
-    /// what lets `try_send` defer its allocation past the capacity check while `try_publish` hands
-    /// over an `Arc` it allocated with the channel unlocked.
+    /// Publishes a message for both send paths.
     ///
-    /// Publishing and draining the wait set share one critical section, so a receiver can never
-    /// observe an empty buffer and park after this message became visible.
+    /// Calls `into_msg` only when retaining the payload, so `try_send` can defer its allocation
+    /// until capacity is available while `try_publish` passes through its existing `Arc`.
     fn publish<P>(&self, payload: P, into_msg: impl FnOnce(P) -> Arc<T>) -> Result<(), P> {
         let mut discarded = None;
-        let mut wakers = WakerBatch::new();
-        {
-            let mut inner = self.shared.inner.lock();
-
-            if !inner.log.has_receivers() {
-                // Nothing can read this message. The payload leaves the critical section with us
-                // and is dropped below, so `T::drop` never runs under the lock.
-                inner.log.publish_discarded();
-                discarded = Some(payload);
-            } else if inner.log.retained() == self.shared.capacity {
-                // Nothing was published, so there is no wait set to drain.
-                return Err(payload);
-            } else {
-                inner.log.publish_retained(into_msg(payload));
-            }
-
-            inner.waiters.drain_into(&mut wakers);
+        let mut inner = self.shared.inner.lock();
+        if !inner.log.has_receivers() {
+            // Drop the discarded payload after unlocking: its destructor may reenter the channel.
+            inner.log.publish_discarded();
+            discarded = Some(payload);
+        } else if inner.log.retained() == self.shared.capacity {
+            // Leave waiters registered because no message was published.
+            return Err(payload);
+        } else {
+            inner.log.publish_retained(into_msg(payload));
         }
-
-        wake_all(&mut wakers);
+        let mut wakers = inner.waiters.take_all();
+        drop(inner);
+        wakers.by_ref().for_each(Waker::wake);
         drop(discarded);
         Ok(())
     }

asyncband/src/broadcast/mpmc/common.rs

@@ -37,7 +37,6 @@ use super::error::TryRecvError;
 use crate::internal::arena::Arena;
 use crate::internal::arena::SlotId;
 use crate::internal::mutex::Mutex;
-use crate::internal::wake_all;
 use crate::internal::wakerset::WakerSet;
 use crate::internal::wakerset::WakerToken;
 
@@ -320,7 +319,7 @@ impl<T> Backlog<T> {
     /// subscription for the slowest cursor, so advancing the head costs the messages released
     /// instead of the receivers subscribed.
     ///
-    /// A receive releases exactly one message: its cursor is counted at the next version before
+    /// A `receive` releases exactly one message: its cursor is counted at the next version before
     /// it leaves `head`, so the zero-count prefix ends there. Only removing a lagging subscription
     /// can release more.
     ///
@@ -406,10 +405,8 @@ impl<T> Backlog<T> {
 
 /// Buffer, receiver cursors, and parked receivers, all under one lock.
 ///
-/// The wait set lives beside the backlog so that publishing a message and draining the waiters
-/// happen in one critical section. That is what makes the park path race-free: a receiver that
-/// finds no message and then registers still holds this lock, so a concurrent send cannot slip
-/// between the two steps and skip the wake-up.
+/// Checking the backlog and registering a waker share this lock with publishing and detaching
+/// registrations, so a send cannot slip between an empty check and registration.
 pub struct Inner<T> {
     pub log: Backlog<T>,
     pub waiters: WakerSet,
@@ -432,14 +429,11 @@ impl<T> Inner<T> {
 ///
 /// Both families call this from the last sender's `Drop`.
 pub fn disconnect<T>(inner: &Mutex<Inner<T>>) {
-    let wakers = {
-        let mut inner = inner.lock();
-        inner.waiters.take_all()
-    };
-    wake_all(wakers);
+    let wakers = mem::take(&mut inner.lock().waiters);
+    wakers.wake_all();
 }
 
-/// Releases a cancelled receive's waker registration, dropping the waker unlocked.
+/// Removes the waker registration for a cancelled `receive`, dropping the waker unlocked.
 pub fn unregister<T>(
     inner: &Mutex<Inner<T>>,
     senders: &AtomicUsize,
@@ -480,9 +474,8 @@ pub fn try_receive<T>(
 
 /// The one poll step behind `recv` on both channels.
 ///
-/// Checking the backlog and registering a waker under the same lock prevents a publication from
-/// landing between those steps. Publication and disconnection detach all registrations, so their
-/// ready paths clear the token without unregistering it.
+/// Publication and disconnection detach all registrations, so their ready paths clear the token
+/// without unregistering it.
 pub fn poll_receive<T>(
     inner: &Mutex<Inner<T>>,
     senders: &AtomicUsize,
@@ -582,7 +575,7 @@ mod tests {
         assert_eq!(log.remove_receiver(b).len(), 3);
         log.assert_cursor_accounting();
 
-        // A receive that catches up to the tail releases exactly one message and moves the cursor
+        // A `receive` that catches up to the tail releases exactly one message and moves the cursor
         // back to `at_tail`.
         assert!(log.publish(Arc::new(4)).is_none());
         let (msg, reclaimed) = log.receive(a).unwrap();

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

@@ -62,6 +62,7 @@ use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering;
 use std::task::Context;
 use std::task::Poll;
+use std::task::Waker;
 
 use super::common;
 use super::common::Backlog;
@@ -70,8 +71,6 @@ use super::error::RecvError;
 use super::error::TryRecvError;
 use crate::internal::arena::SlotId;
 use crate::internal::mutex::Mutex;
-use crate::internal::wake_all;
-use crate::internal::waker_batch::WakerBatch;
 use crate::internal::wakerset::WakerToken;
 
 #[cfg(test)]
@@ -175,19 +174,13 @@ impl<T> UnboundedSender<T> {
     pub fn send(&self, msg: T) {
         let msg = Arc::new(msg);
 
-        // Publishing and draining the wait set share one critical section, so a receiver can never
-        // observe an empty buffer and park after this message became visible.
-        let mut wakers = WakerBatch::new();
-        let unretained = {
-            let mut inner = self.shared.inner.lock();
-            let unretained = inner.log.publish(msg);
-            inner.waiters.drain_into(&mut wakers);
-            unretained
-        };
+        let mut inner = self.shared.inner.lock();
+        let unretained = inner.log.publish(msg);
+        let mut wakers = inner.waiters.take_all();
+        drop(inner);
 
-        // Notify all waiting receivers. An unsent message is dropped here too, once the lock is
-        // released.
-        wake_all(&mut wakers);
+        // Wake callbacks and payload destruction may reenter the channel.
+        wakers.by_ref().for_each(Waker::wake);
         drop(unretained);
     }
 

asyncband/src/completion/mod.rs

@@ -50,6 +50,7 @@
 
 use std::fmt;
 use std::future::Future;
+use std::mem;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::sync::OnceLock;
@@ -58,7 +59,6 @@ use std::task::Context;
 use std::task::Poll;
 
 use crate::internal::mutex::Mutex;
-use crate::internal::wake_all;
 use crate::internal::wakerset::WakerSet;
 use crate::internal::wakerset::WakerToken;
 
@@ -125,17 +125,14 @@ impl<T> Completer<T> {
         let Some(shared) = self.shared.upgrade() else {
             return Err(value);
         };
-        let wakers = {
-            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.
+        let mut waiters = shared.waiters.lock();
+        // Detach registrations before publishing completion to lock-free observers.
+        let wakers = mem::take(&mut *waiters);
+        assert!(shared.result.set(Some(value)).is_ok());
+        drop(waiters);
+        // Disarm abandonment handling before invoking wake callbacks.
         self.shared = Weak::new();
-        wake_all(wakers);
+        wakers.wake_all();
         Ok(())
     }
 }
@@ -145,13 +142,11 @@ impl<T> Drop for Completer<T> {
         let Some(shared) = self.shared.upgrade() else {
             return;
         };
-        let wakers = {
-            let mut waiters = shared.waiters.lock();
-            let wakers = waiters.take_all();
-            assert!(shared.result.set(None).is_ok());
-            wakers
-        };
-        wake_all(wakers);
+        let mut waiters = shared.waiters.lock();
+        let wakers = mem::take(&mut *waiters);
+        assert!(shared.result.set(None).is_ok());
+        drop(waiters);
+        wakers.wake_all();
     }
 }
 

asyncband/src/condvar/mod.rs

@@ -68,7 +68,6 @@ use std::task::Waker;
 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;
@@ -161,21 +160,17 @@ impl Condvar {
         {
             let mut waiters = self.waiters.lock();
 
-            while waiters
-                .unlink_first_waiter(|node| {
-                    let WaitState::Waiting(waker) =
-                        mem::replace(&mut node.state, WaitState::NotifiedAll)
-                    else {
-                        unreachable!("only waiting tasks remain linked")
-                    };
-                    wakers.push(waker);
-                    true
-                })
-                .is_some()
-            {}
+            while let Some((_, node)) = waiters.unlink_first_waiter(|_| true) {
+                let WaitState::Waiting(waker) =
+                    mem::replace(&mut node.state, WaitState::NotifiedAll)
+                else {
+                    unreachable!("only waiting tasks remain linked")
+                };
+                wakers.push(waker);
+            }
         }
 
-        wake_all(&mut wakers);
+        wakers.by_ref().for_each(Waker::wake);
     }
 
     /// Waits for a notification, atomically releasing and then reacquiring the mutex.