diff --git a/CHANGELOG.md b/CHANGELOG.md index 81e16b5f..54c23e3d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,8 @@ 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. +* Scale broadcast backlog reclamation with the number of messages released rather than the number of subscriptions, removing the per-receive cursor scan and the lingering cost of receivers dropped after a peak. +* Avoid heap allocation when waking up to 32 waiters in Barrier, broadcast, condvar, event, MPSC, phaser, and watch notifications; larger waiter sets spill to a single heap allocation. ## v0.7.2 (2026-09-11) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index acf4d4d6..b3a0637e 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -55,6 +55,7 @@ use std::task::Poll; 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; @@ -188,9 +189,10 @@ impl Barrier { if state.arrived == self.n { state.arrived = 0; state.generation += 1; - let wakers = state.waiters.drain(); + let mut wakers = WakerBatch::new(); + state.waiters.drain_into(&mut wakers); drop(state); - wake_all(wakers); + wake_all(&mut wakers); return BarrierWaitResult(true); } diff --git a/asyncband/src/broadcast/mpmc/bounded/mod.rs b/asyncband/src/broadcast/mpmc/bounded/mod.rs index cadb8037..5cc64edd 100644 --- a/asyncband/src/broadcast/mpmc/bounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/bounded/mod.rs @@ -120,6 +120,7 @@ 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)] @@ -195,7 +196,7 @@ impl Shared { /// 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_consumed`, which is + /// `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 @@ -438,7 +439,8 @@ impl BoundedSender { /// observe an empty buffer and park after this message became visible. fn publish

(&self, payload: P, into_msg: impl FnOnce(P) -> Arc) -> Result<(), P> { let mut discarded = None; - let wakers = { + let mut wakers = WakerBatch::new(); + { let mut inner = self.shared.inner.lock(); if !inner.log.has_receivers() { @@ -453,10 +455,10 @@ impl BoundedSender { inner.log.publish_retained(into_msg(payload)); } - inner.waiters.drain() - }; + inner.waiters.drain_into(&mut wakers); + } - wake_all(wakers); + wake_all(&mut wakers); drop(discarded); Ok(()) } diff --git a/asyncband/src/broadcast/mpmc/common.rs b/asyncband/src/broadcast/mpmc/common.rs index 94f0b1eb..dced0992 100644 --- a/asyncband/src/broadcast/mpmc/common.rs +++ b/asyncband/src/broadcast/mpmc/common.rs @@ -94,6 +94,17 @@ enum Retention { Fixed, } +/// A retained message together with the number of receiver cursors positioned at its version. +/// +/// While a cursor sits on a version the message stays readable, so once the count reaches zero the +/// head can advance past the slot. Tracking the count per slot is what lets reclaim release the +/// invisible prefix directly instead of scanning every subscription for the slowest cursor. +struct Slot { + msg: Arc, + /// The number of receivers whose next read is this message. + cursors: usize, +} + /// The committed backlog: every message whose version falls in `[head, tail)`, plus one cursor for /// each active subscription. /// @@ -102,17 +113,20 @@ enum Retention { /// is placed in `buffer` in the same critical section, so a later publication can never become /// visible ahead of an earlier one. pub struct Backlog { - /// Messages whose versions are in the range `[head, tail)`. + /// Messages whose versions are in the range `[head, tail)`, each with its cursor count. /// /// Each message is held behind an `Arc` so the receive path can move the payload out of the /// critical section. Cloning the `Arc` under the lock keeps `T::clone` — and, for reclaimed /// messages, `T::drop` — outside it, which matters because both are arbitrary user code that /// may call back into this channel. - buffer: VecDeque>, + buffer: VecDeque>, /// The version of the first message in `buffer`. head: u64, - /// The number of active receivers whose cursor equals `head`. - head_receivers: usize, + /// The number of receivers whose cursor equals `tail`. + /// + /// Caught-up cursors have no buffered slot to count against, so they are tallied here. Every + /// cursor is counted exactly once: either in one slot's `cursors` or in `at_tail`. + at_tail: usize, /// The next message version to assign. tail: u64, /// Cursor for each active receiver. @@ -131,11 +145,11 @@ impl Backlog { Self::new(VecDeque::with_capacity(capacity), Retention::Fixed) } - fn new(buffer: VecDeque>, retention: Retention) -> Self { + fn new(buffer: VecDeque>, retention: Retention) -> Self { Self { buffer, head: 0, - head_receivers: 0, + at_tail: 0, tail: 0, receivers: Arena::new(), retention, @@ -187,7 +201,7 @@ impl Backlog { pub fn publish_discarded(&mut self) { debug_assert!(!self.has_receivers()); debug_assert!(self.buffer.is_empty()); - debug_assert_eq!(self.head_receivers, 0); + debug_assert_eq!(self.at_tail, 0); self.advance_tail(); self.head = self.tail; } @@ -222,50 +236,41 @@ impl Backlog { pub fn publish_retained(&mut self, msg: Arc) { debug_assert!(self.has_receivers()); self.advance_tail(); - self.buffer.push_back(msg); + // Every cursor that was caught up now has this message as its next read. + let cursors = mem::take(&mut self.at_tail); + self.buffer.push_back(Slot { msg, cursors }); if let Retention::Elastic { peak_len } = &mut self.retention { *peak_len = (*peak_len).max(self.buffer.len()); } } - fn insert_receiver(&mut self, head: u64) -> SlotId { - if head == self.head { - self.head_receivers += 1; - } - - self.receivers.insert(head) - } - /// Registers a new subscription at the committed tail. /// /// A new cursor never lowers `retained()`, so this can never release capacity. pub fn subscribe(&mut self) -> SlotId { - let head = self.tail; - self.insert_receiver(head) + self.at_tail += 1; + self.receivers.insert(self.tail) } pub fn remove_receiver(&mut self, key: SlotId) -> Reclaimed { - let head = self.receivers.remove(key); - - if head == self.head { - self.release_head_receiver() - } else { - Reclaimed::empty() + let cursor = self.receivers.remove(key); + if cursor == self.tail { + self.at_tail -= 1; + return Reclaimed::empty(); } - } - - fn release_head_receiver(&mut self) -> Reclaimed { - self.head_receivers -= 1; - if self.head_receivers == 0 { - self.reclaim_consumed() + let offset = (cursor - self.head) as usize; + let slot = &mut self.buffer[offset]; + slot.cursors -= 1; + if offset == 0 && slot.cursors == 0 { + self.reclaim_vacated() } else { Reclaimed::empty() } } pub fn receive(&mut self, key: SlotId) -> Option> { - let head = { + let version = { let cursor = self .receivers .get_mut(key) @@ -273,22 +278,32 @@ impl Backlog { if *cursor >= self.tail { return None; } - let head = *cursor; + let version = *cursor; *cursor += 1; - head + version }; - debug_assert!(head >= self.head); - let offset = (head - self.head) as usize; - let msg = self.buffer[offset].clone(); - let reclaimed = if head == self.head { - self.release_head_receiver() + debug_assert!(version >= self.head); + let offset = (version - self.head) as usize; + // Count the cursor at its next version before it leaves this one, so a reclaim triggered + // by leaving `head` stops at the message this receiver reads next. + if version + 1 == self.tail { + self.at_tail += 1; + } else { + self.buffer[offset + 1].cursors += 1; + } + + let slot = &mut self.buffer[offset]; + let msg = slot.msg.clone(); + slot.cursors -= 1; + let reclaimed = if offset == 0 && slot.cursors == 0 { + self.reclaim_vacated() } else { Reclaimed::empty() }; // A reclaim triggered by this receive always begins with this receiver's own message: the - // reclaim path runs only for a cursor sitting at `head`, so the first slot drained is - // `msg`. `take_msg` relies on this to recognize that it owns the payload. + // reclaim path runs only for a cursor leaving `head`, so the first slot drained is `msg`. + // `take_msg` relies on this to recognize that it owns the payload. debug_assert!( reclaimed .first() @@ -297,47 +312,46 @@ impl Backlog { Some((msg, reclaimed)) } - /// Advances `head` to the slowest active cursor and hands the released prefix to the caller. + /// Releases the prefix no cursor can still read and hands it to the caller. + /// + /// The head slot's cursor count has just reached zero, so every message up to the next counted + /// slot is invisible: each cursor is counted in exactly one place, and none of those places is + /// a slot being popped. Popping the zero-count prefix replaces the old scan over every + /// 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 + /// it leaves `head`, so the zero-count prefix ends there. Only removing a lagging subscription + /// can release more. /// /// `buffer` shrinks here and grows only in [`Backlog::publish`], so this is the one place /// `retained()` can fall. A bounded channel therefore accounts for released capacity at /// exactly the two call sites that reach this: [`Backlog::receive`] and /// [`Backlog::remove_receiver`]. - fn reclaim_consumed(&mut self) -> Reclaimed { - let mut next_head = self.tail; - let mut head_receivers = 0; - - for head in self.receivers.values() { - if *head < next_head { - next_head = *head; - head_receivers = 1; - } else if *head == next_head { - head_receivers += 1; - } - } + fn reclaim_vacated(&mut self) -> Reclaimed { + debug_assert!(self.buffer.front().is_some_and(|slot| slot.cursors == 0)); - debug_assert!(next_head >= self.head); - let consumed = usize::try_from(next_head - self.head) - .expect("retained broadcast message count exceeds usize"); // Move reclaimed messages out so their Drop impls run after the channel is unlocked. Keep - // the first one separate so the usual one-message reclaim does not allocate another buffer. - let first = if consumed == 0 { - None - } else { - self.buffer.pop_front() - }; - // Reclaiming exactly one message is the overwhelmingly common case — a cursor advances by - // one at a time — so skip building a `Drain` that would yield nothing. - let rest = if consumed > 1 { - self.buffer.drain(..consumed - 1).collect() + // the first one separate so the usual one-message reclaim does not allocate, and skip + // building a `Drain` that would yield nothing: even an empty one costs a few nanoseconds + // on every receive. A bulk reclaim counts the zero-count prefix up front so it moves out + // in one drain instead of growing a vector geometrically, which measured 15% slower for a + // 32-message backlog. + let first = self.buffer.pop_front().map(|slot| slot.msg); + let extra = self + .buffer + .iter() + .take_while(|slot| slot.cursors == 0) + .count(); + let rest = if extra == 0 { + Vec::new() } else { - vec![] + self.buffer.drain(..extra).map(|slot| slot.msg).collect() }; let reclaimed = Reclaimed { first, rest }; - debug_assert_eq!(reclaimed.len(), consumed); - self.head = next_head; - self.head_receivers = head_receivers; + self.head += reclaimed.len() as u64; + debug_assert!(self.head <= self.tail); self.shrink_buffer(); reclaimed } diff --git a/asyncband/src/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/broadcast/mpmc/unbounded/mod.rs index 934374d3..663e5a6b 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -71,6 +71,7 @@ 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)] @@ -176,16 +177,17 @@ impl UnboundedSender { // 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 (unretained, wakers) = { + let mut wakers = WakerBatch::new(); + let unretained = { let mut inner = self.shared.inner.lock(); let unretained = inner.log.publish(msg); - let wakers = inner.waiters.drain(); - (unretained, wakers) + inner.waiters.drain_into(&mut wakers); + unretained }; // Notify all waiting receivers. An unsent message is dropped here too, once the lock is // released. - wake_all(wakers); + wake_all(&mut wakers); drop(unretained); } diff --git a/asyncband/src/condvar/mod.rs b/asyncband/src/condvar/mod.rs index 8dcc643e..06a8550a 100644 --- a/asyncband/src/condvar/mod.rs +++ b/asyncband/src/condvar/mod.rs @@ -157,9 +157,9 @@ impl Condvar { /// If no task is currently waiting, this call has no effect. Notifications are not buffered for /// future calls to [`wait`](Self::wait) or [`wait_owned`](Self::wait_owned). pub fn notify_all(&self) { - let wakers = { + let mut wakers = WakerBatch::new(); + { let mut waiters = self.waiters.lock(); - let mut wakers = WakerBatch::new(); while waiters .unlink_first_waiter(|node| { @@ -173,11 +173,9 @@ impl Condvar { }) .is_some() {} + } - wakers - }; - - wake_all(wakers.into_iter()); + wake_all(&mut wakers); } /// Waits for a notification, atomically releasing and then reacquiring the mutex. diff --git a/asyncband/src/event/manual_reset.rs b/asyncband/src/event/manual_reset.rs index d51b16a4..074dea4f 100644 --- a/asyncband/src/event/manual_reset.rs +++ b/asyncband/src/event/manual_reset.rs @@ -103,7 +103,8 @@ impl ManualResetEvent { /// Panics if waking a selected task panics. The event remains set, and waking is still /// attempted for every other selected task before the panic resumes. pub fn set(&self) { - let wakers = { + let mut wakers = WakerBatch::new(); + { let mut state = self.state.lock(); if state.is_set { return; @@ -112,7 +113,6 @@ 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 = WakerBatch::new(); while let Some((_id, waiter)) = state.waiters.unlink_first_waiter(|waiter| { waiter.notified = true; true @@ -121,10 +121,9 @@ impl ManualResetEvent { wakers.push(waker); } } - wakers - }; + } - wake_all(wakers.into_iter()); + wake_all(&mut wakers); } /// Clears the set state. diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 7998e25a..7974b26b 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -140,8 +140,8 @@ pub(crate) mod waitlist; feature = "waitgroup", feature = "watch", ))] -// Waker-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. +// Only the semaphore refills a batch and asks whether it is full, so other feature subsets leave +// that method unused. #[allow(dead_code)] pub(crate) mod waker_batch; diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index ca4c3019..c4631b34 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -25,9 +25,7 @@ // https://github.com/tokio-rs/tokio/blob/bb9d57017e100985f86d8ca41ac105ee9140423e/tokio/src/sync/batch_semaphore.rs use std::future::Future; -use std::mem::MaybeUninit; use std::pin::Pin; -use std::ptr; use std::sync::MutexGuard; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -39,6 +37,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)] @@ -56,61 +55,6 @@ struct WaitNode { waker: Option, } -const WAKE_BATCH_SIZE: usize = 32; - -/// The initialized entries in `wakers` are exactly `start..end`. -struct WakeBatch { - wakers: [MaybeUninit; WAKE_BATCH_SIZE], - start: usize, - end: usize, -} - -impl WakeBatch { - fn new() -> Self { - const UNINIT: MaybeUninit = MaybeUninit::uninit(); - Self { - wakers: [UNINIT; WAKE_BATCH_SIZE], - start: 0, - end: 0, - } - } - - fn push(&mut self, waker: Waker) { - debug_assert_eq!(self.start, 0); - debug_assert!(self.end < WAKE_BATCH_SIZE); - self.wakers[self.end].write(waker); - self.end += 1; - } - - fn is_full(&self) -> bool { - self.end == WAKE_BATCH_SIZE - } - - fn take_next(&mut self) -> Option { - if self.start == self.end { - self.start = 0; - self.end = 0; - return None; - } - - let index = self.start; - self.start += 1; - // SAFETY: `index` was within the initialized range before advancing `start`. - Some(unsafe { self.wakers[index].assume_init_read() }) - } -} - -impl Drop for WakeBatch { - fn drop(&mut self) { - let start = self.wakers[self.start..self.end] - .as_mut_ptr() - .cast::(); - let remaining = ptr::slice_from_raw_parts_mut(start, self.end - self.start); - // SAFETY: The initialized entries are exactly `start..end`. - unsafe { ptr::drop_in_place(remaining) }; - } -} - impl Semaphore { pub const fn new(permits: usize) -> Self { Self { @@ -216,7 +160,7 @@ impl Semaphore { #[cfg(any(feature = "broadcast", feature = "mpmc"))] 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; @@ -235,7 +179,7 @@ impl Semaphore { } } drop(waiters); - wake_all(wakers.into_iter()); + wake_all(&mut wakers); } fn insert_permits_with_lock( @@ -243,14 +187,14 @@ impl Semaphore { mut rem: usize, waiters: MutexGuard<'_, WaitList>, ) { - let mut batch = WakeBatch::new(); + let mut batch = WakerBatch::new(); let mut lock = Some(waiters); // One iterator covers the entire release. If a callback panics, `wake_all` keeps pulling // batches during unwinding, so the remaining permits are still distributed and notified. wake_all(std::iter::from_fn(|| { loop { - if let Some(waker) = batch.take_next() { + if let Some(waker) = batch.next() { return Some(waker); } if rem == 0 { diff --git a/asyncband/src/internal/waker_batch.rs b/asyncband/src/internal/waker_batch.rs index 0e770e35..1898ae69 100644 --- a/asyncband/src/internal/waker_batch.rs +++ b/asyncband/src/internal/waker_batch.rs @@ -15,35 +15,62 @@ // specific language governing permissions and limitations // under the License. +use std::collections::VecDeque; +use std::mem::MaybeUninit; +use std::ptr; use std::task::Waker; -/// An owning waker collection that stores the first entry without allocating. -#[derive(Debug)] +/// Wakers kept on the stack before the batch spills to the heap. +/// +/// This is also the most wakers the semaphore collects per lock acquisition, so a drain that +/// wakes a typical waiter set never allocates; larger sets pay one allocation for the overflow. +pub const INLINE_CAPACITY: usize = 32; + +/// An owning FIFO of wakers that stores the first [`INLINE_CAPACITY`] entries without allocating. +/// +/// The batch is filled through [`WakerBatch::push`] or [`Extend`] and consumed as its own +/// iterator. Entries are written only as they are pushed, so constructing an empty or small batch +/// touches nothing beyond the two indices. Once every inline entry has been yielded the batch +/// reuses that storage, so a caller that alternates between filling and draining, as the +/// semaphore does, keeps running on the stack. pub struct WakerBatch { - first: Option, - rest: Vec, + /// The initialized entries are exactly `start..end`. + inline: [MaybeUninit; INLINE_CAPACITY], + /// The next inline entry to yield. + start: usize, + /// The next inline slot to push into. + end: usize, + /// Entries pushed while the inline storage was full, yielded after it. + /// + /// While this is non-empty every push lands here, so the batch never yields a later push + /// ahead of an earlier one. + spilled: VecDeque, } impl WakerBatch { pub const fn new() -> Self { Self { - first: None, - rest: vec![], + inline: [const { MaybeUninit::uninit() }; INLINE_CAPACITY], + start: 0, + end: 0, + spilled: VecDeque::new(), } } - pub fn with_capacity(capacity: usize) -> Self { - Self { - first: None, - rest: Vec::with_capacity(capacity.saturating_sub(1)), - } + /// Whether the next push would spill to the heap. + /// + /// The semaphore stops filling a batch here so it can release its lock and wake what it has + /// before collecting more. + pub fn is_full(&self) -> bool { + self.end == INLINE_CAPACITY || !self.spilled.is_empty() } pub fn push(&mut self, waker: Waker) { - if self.first.is_none() { - self.first = Some(waker); + if self.end < INLINE_CAPACITY && self.spilled.is_empty() { + self.inline[self.end].write(waker); + self.end += 1; } else { - self.rest.push(waker); + self.spilled.push_back(waker); } } } @@ -56,11 +83,153 @@ impl Extend for WakerBatch { } } -impl IntoIterator for WakerBatch { +impl Iterator for WakerBatch { type Item = Waker; - type IntoIter = std::iter::Chain, std::vec::IntoIter>; - fn into_iter(self) -> Self::IntoIter { - self.first.into_iter().chain(self.rest) + fn next(&mut self) -> Option { + if self.start < self.end { + let index = self.start; + self.start += 1; + // SAFETY: `index` was within the initialized range before advancing `start`. + return Some(unsafe { self.inline[index].assume_init_read() }); + } + + // Every inline entry has been yielded, so later pushes can start over from the front. + self.start = 0; + self.end = 0; + self.spilled.pop_front() + } +} + +impl Drop for WakerBatch { + fn drop(&mut self) { + let initialized = ptr::slice_from_raw_parts_mut( + self.inline[self.start..self.end] + .as_mut_ptr() + .cast::(), + self.end - self.start, + ); + // SAFETY: The initialized entries are exactly `start..end`. + unsafe { ptr::drop_in_place(initialized) }; + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::Mutex; + use std::task::Wake; + use std::task::Waker; + + use super::INLINE_CAPACITY; + use super::WakerBatch; + + /// The ids of the wakers woken so far, in order. + /// + /// Every waker holds one clone of the `Arc`, so its strong count tells how many wakers + /// are still alive. + struct Log(Mutex>); + + struct Tagged { + id: usize, + log: Arc, + } + + impl Wake for Tagged { + fn wake(self: Arc) { + self.log.0.lock().unwrap().push(self.id); + } + } + + fn log() -> Arc { + Arc::new(Log(Mutex::new(vec![]))) + } + + fn waker(log: &Arc, id: usize) -> Waker { + Waker::from(Arc::new(Tagged { + id, + log: Arc::clone(log), + })) + } + + fn wakers(log: &Arc, count: usize) -> impl Iterator + '_ { + (0..count).map(move |id| waker(log, id)) + } + + fn alive(log: &Arc) -> usize { + Arc::strong_count(log) - 1 + } + + fn woken(log: &Arc) -> Vec { + log.0.lock().unwrap().clone() + } + + #[test] + fn yields_in_push_order_across_the_spill() { + let log = log(); + let count = INLINE_CAPACITY + 8; + let mut batch = WakerBatch::new(); + batch.extend(wakers(&log, count)); + assert!(batch.is_full()); + + for waker in &mut batch { + waker.wake(); + } + + assert_eq!(woken(&log), (0..count).collect::>()); + assert!(batch.next().is_none()); + assert_eq!(alive(&log), 0); + } + + #[test] + fn drops_unconsumed_entries_exactly_once() { + let log = log(); + let count = INLINE_CAPACITY + 8; + for consumed in [0, 5, INLINE_CAPACITY, INLINE_CAPACITY + 3, count] { + let mut batch = WakerBatch::new(); + batch.extend(wakers(&log, count)); + for _ in 0..consumed { + drop(batch.next().unwrap()); + } + assert_eq!(alive(&log), count - consumed); + + drop(batch); + assert_eq!(alive(&log), 0, "after consuming {consumed}"); + } + } + + #[test] + fn reuses_inline_storage_after_draining() { + let log = log(); + let mut batch = WakerBatch::new(); + for round in 0..3 { + batch.extend(wakers(&log, INLINE_CAPACITY)); + assert!(batch.is_full(), "round {round}"); + assert_eq!(batch.by_ref().count(), INLINE_CAPACITY); + assert!(!batch.is_full(), "round {round}"); + assert_eq!(alive(&log), 0, "round {round}"); + } + } + + #[test] + fn keeps_push_order_while_spilled() { + let log = log(); + let mut batch = WakerBatch::new(); + batch.extend(wakers(&log, INLINE_CAPACITY + 1)); + // Free inline room; the spilled entry must still come out before anything pushed now. + for _ in 0..4 { + batch.next().unwrap().wake(); + } + batch.push(waker(&log, 999)); + assert!(batch.is_full()); + + for waker in &mut batch { + waker.wake(); + } + + let mut expected = (0..INLINE_CAPACITY + 1).collect::>(); + expected.push(999); + assert_eq!(woken(&log), expected); + assert_eq!(alive(&log), 0); } } diff --git a/asyncband/src/internal/wakerset.rs b/asyncband/src/internal/wakerset.rs index c49d49fd..c74a89b5 100644 --- a/asyncband/src/internal/wakerset.rs +++ b/asyncband/src/internal/wakerset.rs @@ -31,7 +31,7 @@ use crate::internal::waker_batch::WakerBatch; /// An exclusive handle to one waker slot in a [`WakerSet`]. /// /// This token deliberately does not implement `Clone` or `Copy`. Its owner must not pass it back -/// to the set after the registration has been detached by [`WakerSet::drain`] or +/// to the set after the registration has been detached by [`WakerSet::drain_into`] or /// [`WakerSet::take_all`]. #[derive(Debug)] pub struct WakerToken(SlotId); @@ -57,18 +57,15 @@ impl WakerSet { } } - /// Drains all registered wakers into an owning batch while retaining slot capacity. + /// Drains all registered wakers into `batch` while retaining slot capacity. /// - /// The caller must invalidate every outstanding token and consume or drop the iterator after + /// The batch is filled in place because its inline storage is too large to move for free: + /// returning it by value costs every publish about 6ns even when nothing is registered. The + /// caller must invalidate every outstanding token and consume or drop the batch after /// releasing the lock that protects this set. #[inline] - pub fn drain(&mut self) -> impl Iterator + 'static { - let mut wakers = WakerBatch::with_capacity(self.wakers.len()); - if self.wakers.is_empty() { - return wakers.into_iter(); - } - wakers.extend(self.wakers.drain()); - wakers.into_iter() + pub fn drain_into(&mut self, batch: &mut WakerBatch) { + batch.extend(self.wakers.drain()); } /// Takes all registered wakers together with the set's backing allocation. diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs index 8b46e059..cf7436c7 100644 --- a/asyncband/src/mpsc/bounded/receiver.rs +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -45,21 +45,21 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - let (queue, recv_waker, wakers) = { + let mut wakers = WakerBatch::new(); + let (queue, recv_waker) = { let mut state = self.shared.lock(); state.receiver = false; let queue = mem::take(&mut state.queue); let recv_waker = state.recv_waker.take(); - let mut wakers = WakerBatch::new(); while let Some((_, waiter)) = state.send_waiters.unlink_first_waiter(|_| true) { if let Some(waker) = waiter.waker.take() { wakers.push(waker); } } - (queue, recv_waker, wakers) + (queue, recv_waker) }; // Local ownership also drains the queue if a wake or waker destructor unwinds. - wake_all(wakers.into_iter()); + wake_all(&mut wakers); drop(recv_waker); drop(queue); } diff --git a/asyncband/src/phaser/mod.rs b/asyncband/src/phaser/mod.rs index ae7fbe12..72bdea1a 100644 --- a/asyncband/src/phaser/mod.rs +++ b/asyncband/src/phaser/mod.rs @@ -128,10 +128,10 @@ use std::pin::Pin; use std::sync::Arc; 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; @@ -172,13 +172,14 @@ struct State { } impl State { - fn advance_if_ready(&mut self) -> Option + 'static> { + /// Completes the phase once every participant has arrived, moving its waiters into `wakers`. + fn advance_if_ready(&mut self, wakers: &mut WakerBatch) { if self.closed || self.unarrived != 0 { - return None; + return; } self.phase = self.phase.wrapping_add(1); self.unarrived = self.registered; - Some(self.waiters.drain()) + self.waiters.drain_into(wakers); } fn completion(&self, observed: u64) -> Poll> { @@ -389,15 +390,16 @@ impl Drop for PhaserParticipants { if self.remaining == 0 { return; } - let wakers = { + let mut wakers = WakerBatch::new(); + { let mut state = self.phaser.state.lock(); // Unyielded participants have never arrived and prevent their phase from advancing. state.registered -= self.remaining; state.unarrived -= self.remaining; self.remaining = 0; - state.advance_if_ready() - }; - wake_all(wakers.into_iter().flatten()); + state.advance_if_ready(&mut wakers); + } + wake_all(&mut wakers); } } @@ -438,7 +440,8 @@ impl PhaserParticipant { /// /// Arrival and the pending observation remain committed if notifying a waker panics. pub fn arrive(&mut self) -> Result { - let (phase, wakers) = { + let mut wakers = WakerBatch::new(); + let phase = { let mut state = self.phaser.state.lock(); if state.closed { return Err(Closed(())); @@ -448,9 +451,10 @@ impl PhaserParticipant { state.unarrived -= 1; } self.pending = Some(phase); - (phase, state.advance_if_ready()) + state.advance_if_ready(&mut wakers); + phase }; - wake_all(wakers.into_iter().flatten()); + wake_all(&mut wakers); Ok(phase) } @@ -485,7 +489,8 @@ impl PhaserParticipant { } fn do_deregister(&mut self) -> Result { - let (result, wakers) = { + let mut wakers = WakerBatch::new(); + let result = { let mut state = self.phaser.state.lock(); self.registered = false; state.registered -= 1; @@ -497,9 +502,10 @@ impl PhaserParticipant { } else { Ok(state.phase) }; - (result, state.advance_if_ready()) + state.advance_if_ready(&mut wakers); + result }; - wake_all(wakers.into_iter().flatten()); + wake_all(&mut wakers); result } } diff --git a/asyncband/src/watch/mod.rs b/asyncband/src/watch/mod.rs index 92aab71c..90b71923 100644 --- a/asyncband/src/watch/mod.rs +++ b/asyncband/src/watch/mod.rs @@ -75,6 +75,7 @@ pub use self::error::RecvError; pub use self::error::SendError; 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; @@ -168,7 +169,8 @@ impl Sender { /// /// Panics if the channel has already published `u64::MAX` updates. pub fn send(&self, value: T) -> Result<(), SendError> { - let (wakers, replaced) = { + let mut wakers = WakerBatch::new(); + let replaced = { let mut state = self.shared.state.lock(); if state.receivers == 0 { return Err(SendError::new(value)); @@ -179,11 +181,11 @@ impl Sender { .expect("watch channel version counter overflowed"); let replaced = mem::replace(&mut state.value, value); state.version = version; - let wakers = state.waiters.drain(); - (wakers, replaced) + state.waiters.drain_into(&mut wakers); + replaced }; // Waker callbacks and the replaced value's destructor may reenter this channel. - wake_all(wakers); + wake_all(&mut wakers); drop(replaced); Ok(()) } @@ -197,7 +199,8 @@ impl Sender { /// /// Panics if the channel has already published `u64::MAX` updates. pub fn send_replace(&self, value: T) -> T { - let (wakers, replaced) = { + let mut wakers = WakerBatch::new(); + let replaced = { let mut state = self.shared.state.lock(); let version = state .version @@ -205,10 +208,10 @@ impl Sender { .expect("watch channel version counter overflowed"); let replaced = mem::replace(&mut state.value, value); state.version = version; - let wakers = state.waiters.drain(); - (wakers, replaced) + state.waiters.drain_into(&mut wakers); + replaced }; - wake_all(wakers); + wake_all(&mut wakers); replaced } diff --git a/benchmarks/asyncband/broadcast/mpmc/bounded.rs b/benchmarks/asyncband/broadcast/mpmc/bounded.rs index aabae2ef..5ebb1e4c 100644 --- a/benchmarks/asyncband/broadcast/mpmc/bounded.rs +++ b/benchmarks/asyncband/broadcast/mpmc/bounded.rs @@ -73,8 +73,8 @@ fn try_send_and_drain_fanout(bencher: Bencher, receiver_count: usize) { receivers.push(tx.subscribe()); } - // One message in, every receiver drains it out: the last one to read pays the reclaim scan and - // the capacity release, and the channel is empty again for the next iteration. + // One message in, every receiver drains it out: the last one to read pays the head reclaim + // and the capacity release, and the channel is empty again for the next iteration. bencher.bench_local(|| { tx.try_send(black_box(1)).unwrap(); for receiver in &mut receivers { diff --git a/benchmarks/asyncband/broadcast/mpmc/unbounded.rs b/benchmarks/asyncband/broadcast/mpmc/unbounded.rs index d4fc08e1..58605e59 100644 --- a/benchmarks/asyncband/broadcast/mpmc/unbounded.rs +++ b/benchmarks/asyncband/broadcast/mpmc/unbounded.rs @@ -34,9 +34,9 @@ const RECEIVER_COUNTS: &[usize] = &[1, 8, 32]; /// A channel that peaked at `peak` receivers and currently has `live` of them. /// -/// The two are measured separately because a dropped receiver leaves its slot behind: the reclaim -/// scan walks every slot the channel ever handed out, so a channel that shed receivers keeps -/// paying for the peak. Pairing each peak with a drained arena is what makes that visible. +/// The two are measured separately to show that a dropped receiver no longer costs the channels +/// that outlive it: reclaim releases the prefix no cursor can read, without walking the +/// subscription slots a peak left behind. #[derive(Clone, Copy)] struct Fanout { peak: usize, @@ -115,8 +115,9 @@ fn send_and_try_recv_owned_shared(bencher: Bencher) { }); } -// Measures the reclaim scan, which runs when the slowest cursor advances. Comparing a peak against -// the same peak drained down to fewer receivers shows what the slots left behind still cost. +// Measures advancing the shared backlog head: the receive that vacates the last cursor at the +// head releases the invisible prefix. Comparing a peak against the same peak drained down to +// fewer receivers confirms the slots left behind no longer add to that cost. #[divan::bench(args = RECLAIM_FANOUTS)] fn drain_with_receivers(bencher: Bencher, fanout: Fanout) { let (sender, receiver) = mpmc::unbounded();