From 0ba59a006ec56f935574c2ee442a54aba044906a Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 3 Oct 2026 11:31:41 +0800 Subject: [PATCH 1/2] refactor: replace waker batch iteration with wake_all Signed-off-by: tison --- asyncband/src/barrier/mod.rs | 3 +- asyncband/src/broadcast/mpmc/bounded/mod.rs | 3 +- asyncband/src/broadcast/mpmc/unbounded/mod.rs | 3 +- asyncband/src/condvar/mod.rs | 2 +- asyncband/src/event/manual_reset.rs | 2 +- asyncband/src/internal/mod.rs | 2 + asyncband/src/internal/semaphore.rs | 4 +- asyncband/src/internal/waker_batch.rs | 160 +++++++----------- asyncband/src/phaser/mod.rs | 7 +- asyncband/src/watch/mod.rs | 5 +- 10 files changed, 79 insertions(+), 112 deletions(-) diff --git a/asyncband/src/barrier/mod.rs b/asyncband/src/barrier/mod.rs index e52db7c..2a5d8bf 100644 --- a/asyncband/src/barrier/mod.rs +++ b/asyncband/src/barrier/mod.rs @@ -52,7 +52,6 @@ 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::wakerset::WakerSet; @@ -190,7 +189,7 @@ impl Barrier { state.generation += 1; let mut wakers = state.waiters.take_all(); drop(state); - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); return BarrierWaitResult(true); } diff --git a/asyncband/src/broadcast/mpmc/bounded/mod.rs b/asyncband/src/broadcast/mpmc/bounded/mod.rs index 48ec26c..f91faa2 100644 --- a/asyncband/src/broadcast/mpmc/bounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/bounded/mod.rs @@ -108,7 +108,6 @@ 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; @@ -422,7 +421,7 @@ impl BoundedSender { } let mut wakers = inner.waiters.take_all(); drop(inner); - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); drop(discarded); Ok(()) } diff --git a/asyncband/src/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/broadcast/mpmc/unbounded/mod.rs index 17f9d11..15366ef 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -62,7 +62,6 @@ 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; @@ -180,7 +179,7 @@ impl UnboundedSender { drop(inner); // Wake callbacks and payload destruction may reenter the channel. - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); drop(unretained); } diff --git a/asyncband/src/condvar/mod.rs b/asyncband/src/condvar/mod.rs index 9cc2e8d..ffae538 100644 --- a/asyncband/src/condvar/mod.rs +++ b/asyncband/src/condvar/mod.rs @@ -170,7 +170,7 @@ impl Condvar { } } - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); } /// 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 39d5bcf..45cfb33 100644 --- a/asyncband/src/event/manual_reset.rs +++ b/asyncband/src/event/manual_reset.rs @@ -116,7 +116,7 @@ impl ManualResetEvent { } } } - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); } /// Clears the set state. diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index f9dceb0..ccaa11a 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -124,6 +124,8 @@ pub(crate) mod waitlist; feature = "waitgroup", feature = "watch", ))] +// Terminal-only primitives wake WakerSet directly, leaving WakerBatch::wake_all unused. +#[allow(dead_code)] pub(crate) mod waker_batch; #[cfg(any( diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index 4ce4ad7..1a539a2 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/asyncband/src/internal/semaphore.rs @@ -172,7 +172,7 @@ impl Semaphore { } } drop(waiters); - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); } fn insert_permits_with_lock<'a>( @@ -211,7 +211,7 @@ impl Semaphore { } drop(waiters); - batch.by_ref().for_each(Waker::wake); + batch.wake_all(); if rem == 0 { return; } diff --git a/asyncband/src/internal/waker_batch.rs b/asyncband/src/internal/waker_batch.rs index f1ec29e..7cfe0cb 100644 --- a/asyncband/src/internal/waker_batch.rs +++ b/asyncband/src/internal/waker_batch.rs @@ -15,27 +15,21 @@ // specific language governing permissions and limitations // under the License. -use std::collections::VecDeque; +use std::mem; use std::mem::MaybeUninit; use std::ptr; use std::task::Waker; /// An owning FIFO of wakers that stores the first [`Self::STACK_SIZE`] entries without allocating. /// -/// Only pushed entries are initialized. Iterate by reference to avoid moving the inline storage -/// into iterator adapters and to reuse the batch, as the semaphore does between notifications. +/// Only pushed entries are initialized. [`Self::wake_all`] empties the batch in place so the +/// semaphore can reuse its storage between notifications. pub struct WakerBatch { - /// The initialized entries are exactly `start..end`. + /// The initialized entries are exactly `0..inline_len`. inline: [MaybeUninit; Self::STACK_SIZE], - /// 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, + inline_len: usize, + /// Entries pushed after the inline storage filled, woken after it. + spilled: Vec, } impl WakerBatch { @@ -47,9 +41,8 @@ impl WakerBatch { pub const fn new() -> Self { Self { inline: [const { MaybeUninit::uninit() }; Self::STACK_SIZE], - start: 0, - end: 0, - spilled: VecDeque::new(), + inline_len: 0, + spilled: Vec::new(), } } @@ -59,17 +52,30 @@ impl WakerBatch { /// before collecting more. #[inline] pub fn will_spill(&self) -> bool { - self.end == Self::STACK_SIZE || !self.spilled.is_empty() + self.inline_len == Self::STACK_SIZE } #[inline] pub fn push(&mut self, waker: Waker) { if self.will_spill() { - self.spilled.push_back(waker); + self.spilled.push(waker); } else { - self.inline[self.end].write(waker); - self.end += 1; + self.inline[self.inline_len].write(waker); + self.inline_len += 1; + } + } + + /// Wakes every entry in push order, retaining storage for the next batch. + /// + /// Call after releasing the owning primitive's state lock. + #[inline] + pub fn wake_all(&mut self) { + let len = mem::take(&mut self.inline_len); + for slot in &mut self.inline[..len] { + // SAFETY: This prefix was initialized by `push` and is no longer owned by the batch. + unsafe { slot.assume_init_read() }.wake(); } + self.spilled.drain(..).for_each(Waker::wake); } } @@ -84,35 +90,14 @@ impl FromIterator for WakerBatch { } } -impl Iterator for WakerBatch { - type Item = Waker; - - #[inline] - 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 { #[inline] 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, + self.inline.as_mut_ptr().cast::(), + self.inline_len, ); - // SAFETY: The initialized entries are exactly `start..end`. + // SAFETY: The initialized entries are exactly `0..inline_len`. unsafe { ptr::drop_in_place(initialized) }; } } @@ -171,70 +156,55 @@ mod tests { } #[test] - fn yields_in_push_order_across_the_spill() { - let log = log(); - let count = STACK_SIZE + 8; - let mut batch: WakerBatch = (0..count).map(|id| waker(&log, id)).collect(); - assert!(batch.will_spill()); - - for waker in &mut batch { - waker.wake(); + fn wakes_in_push_order_across_the_spill() { + for count in [ + 0, + 1, + STACK_SIZE - 1, + STACK_SIZE, + STACK_SIZE + 1, + 2 * STACK_SIZE + 1, + ] { + let log = log(); + let mut batch: WakerBatch = (0..count).map(|id| waker(&log, id)).collect(); + batch.wake_all(); + + assert_eq!(woken(&log), (0..count).collect::>()); + assert_eq!(alive(&log), 0); } - - 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 = STACK_SIZE + 8; - for consumed in [0, 5, STACK_SIZE, STACK_SIZE + 3, count] { - let mut batch = WakerBatch::new(); - push_wakers(&mut batch, &log, count); - for _ in 0..consumed { - drop(batch.next().unwrap()); + fn drops_unwoken_entries_exactly_once() { + for previously_woken in [0, STACK_SIZE + 1] { + for count in [0, 1, STACK_SIZE, STACK_SIZE + 1, 2 * STACK_SIZE + 1] { + let log = log(); + let mut batch = WakerBatch::new(); + push_wakers(&mut batch, &log, previously_woken); + batch.wake_all(); + push_wakers(&mut batch, &log, count); + assert_eq!(alive(&log), count); + + drop(batch); + assert_eq!(alive(&log), 0); + assert_eq!(woken(&log), (0..previously_woken).collect::>()); } - assert_eq!(alive(&log), count - consumed); - - drop(batch); - assert_eq!(alive(&log), 0, "after consuming {consumed}"); } } #[test] - fn reuses_inline_storage_after_draining() { + fn reuses_inline_storage_after_waking() { let log = log(); let mut batch = WakerBatch::new(); - for _ in 0..3 { - push_wakers(&mut batch, &log, STACK_SIZE); - assert!(batch.will_spill()); - assert_eq!(batch.by_ref().count(), STACK_SIZE); + let mut expected = vec![]; + for count in [STACK_SIZE + 8, STACK_SIZE, 1, 0, STACK_SIZE + 1] { + push_wakers(&mut batch, &log, count); + batch.wake_all(); + expected.extend(0..count); + + assert_eq!(woken(&log), expected); assert!(!batch.will_spill()); assert_eq!(alive(&log), 0); } } - - #[test] - fn keeps_push_order_while_spilled() { - let log = log(); - let mut batch = WakerBatch::new(); - push_wakers(&mut batch, &log, STACK_SIZE + 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.will_spill()); - - for waker in &mut batch { - waker.wake(); - } - - let mut expected = (0..STACK_SIZE + 1).collect::>(); - expected.push(999); - assert_eq!(woken(&log), expected); - assert_eq!(alive(&log), 0); - } } diff --git a/asyncband/src/phaser/mod.rs b/asyncband/src/phaser/mod.rs index 5d97a51..ce2c707 100644 --- a/asyncband/src/phaser/mod.rs +++ b/asyncband/src/phaser/mod.rs @@ -129,7 +129,6 @@ 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::waker_batch::WakerBatch; @@ -393,7 +392,7 @@ impl Drop for PhaserParticipants { self.remaining = 0; let mut wakers = state.advance_if_ready(); drop(state); - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); } } @@ -443,7 +442,7 @@ impl PhaserParticipant { self.pending = Some(phase); let mut wakers = state.advance_if_ready(); drop(state); - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); Ok(phase) } @@ -491,7 +490,7 @@ impl PhaserParticipant { }; let mut wakers = state.advance_if_ready(); drop(state); - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); result } } diff --git a/asyncband/src/watch/mod.rs b/asyncband/src/watch/mod.rs index c2f2072..cc8d40e 100644 --- a/asyncband/src/watch/mod.rs +++ b/asyncband/src/watch/mod.rs @@ -70,7 +70,6 @@ use std::pin::Pin; use std::sync::Arc; use std::task::Context; use std::task::Poll; -use std::task::Waker; pub use self::error::RecvError; pub use self::error::SendError; @@ -181,7 +180,7 @@ impl Sender { drop(state); // Waker callbacks and the replaced value's destructor may reenter this channel. - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); drop(replaced); Ok(()) } @@ -204,7 +203,7 @@ impl Sender { state.version = version; let mut wakers = state.waiters.take_all(); drop(state); - wakers.by_ref().for_each(Waker::wake); + wakers.wake_all(); replaced } From 58114679d0eaf79c3f920daa3c20c6f0f8aa3267 Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 3 Oct 2026 12:45:34 +0800 Subject: [PATCH 2/2] fixup Signed-off-by: tison --- asyncband/src/broadcast/mpmc/common.rs | 2 +- asyncband/src/internal/waker_batch.rs | 2 +- asyncband/src/internal/wakerset.rs | 8 +++++++- asyncband/src/phaser/mod.rs | 4 ++-- examples/src/notify_vs_event.rs | 4 ++-- examples/src/phaser_completion.rs | 2 +- examples/src/phaser_groups.rs | 2 +- examples/src/phaser_rounds.rs | 2 +- tests-integration/tests/mpmc_test/concurrency.rs | 2 +- tests-integration/tests/mpsc_test/concurrency.rs | 2 +- tests-integration/tests/mpsc_test/reservation.rs | 2 +- tests-integration/tests/spmc_test/concurrency.rs | 4 ++-- 12 files changed, 21 insertions(+), 15 deletions(-) diff --git a/asyncband/src/broadcast/mpmc/common.rs b/asyncband/src/broadcast/mpmc/common.rs index 73d96e0..ab7f0ff 100644 --- a/asyncband/src/broadcast/mpmc/common.rs +++ b/asyncband/src/broadcast/mpmc/common.rs @@ -343,7 +343,7 @@ impl Backlog { .take_while(|slot| slot.cursors == 0) .count(); let rest = if extra == 0 { - Vec::new() + vec![] } else { self.buffer.drain(..extra).map(|slot| slot.msg).collect() }; diff --git a/asyncband/src/internal/waker_batch.rs b/asyncband/src/internal/waker_batch.rs index 7cfe0cb..ff9c0e8 100644 --- a/asyncband/src/internal/waker_batch.rs +++ b/asyncband/src/internal/waker_batch.rs @@ -42,7 +42,7 @@ impl WakerBatch { Self { inline: [const { MaybeUninit::uninit() }; Self::STACK_SIZE], inline_len: 0, - spilled: Vec::new(), + spilled: vec![], } } diff --git a/asyncband/src/internal/wakerset.rs b/asyncband/src/internal/wakerset.rs index 75649cb..cbd3ff0 100644 --- a/asyncband/src/internal/wakerset.rs +++ b/asyncband/src/internal/wakerset.rs @@ -35,11 +35,17 @@ use crate::internal::waker_batch::WakerBatch; pub struct WakerToken(SlotId); /// Cancellable waker storage without an implicit lifecycle or generation. -#[derive(Default, Debug)] +#[derive(Debug)] pub struct WakerSet { wakers: Arena, } +impl Default for WakerSet { + fn default() -> Self { + Self::new() + } +} + impl WakerSet { /// Constructs an empty waker set. pub const fn new() -> Self { diff --git a/asyncband/src/phaser/mod.rs b/asyncband/src/phaser/mod.rs index ce2c707..70c524d 100644 --- a/asyncband/src/phaser/mod.rs +++ b/asyncband/src/phaser/mod.rs @@ -44,7 +44,7 @@ //! let phaser = Phaser::new(); //! let mut coordinator = phaser.register_one()?; //! let participants = phaser.register(documents.len())?; -//! let mut tasks = Vec::new(); +//! let mut tasks = vec![]; //! //! for (document, mut participant) in documents.into_iter().zip(participants) { //! let dictionary = dictionary.clone(); @@ -71,7 +71,7 @@ //! } //! coordinator.wait().await?; //! -//! let mut encoded_documents = Vec::new(); +//! let mut encoded_documents = vec![]; //! for task in tasks { //! encoded_documents.push(task.await.unwrap()?); //! } diff --git a/examples/src/notify_vs_event.rs b/examples/src/notify_vs_event.rs index 84cb529..5c6e759 100644 --- a/examples/src/notify_vs_event.rs +++ b/examples/src/notify_vs_event.rs @@ -45,7 +45,7 @@ async fn cache_worker_with_notify() { let entries = Mutex::new(VecDeque::new()); let changed = Notify::new(); let worker = async { - let mut flushed = Vec::new(); + let mut flushed = vec![]; while flushed.len() < 3 { let entry = entries.lock().unwrap().pop_front(); if let Some(entry) = entry { @@ -77,7 +77,7 @@ async fn cache_worker_with_event() { let entries = Mutex::new(VecDeque::new()); let changed = AutoResetEvent::new(); let worker = async { - let mut flushed = Vec::new(); + let mut flushed = vec![]; while flushed.len() < 3 { let entry = entries.lock().unwrap().pop_front(); if let Some(entry) = entry { diff --git a/examples/src/phaser_completion.rs b/examples/src/phaser_completion.rs index cef5d20..a1da43f 100644 --- a/examples/src/phaser_completion.rs +++ b/examples/src/phaser_completion.rs @@ -95,7 +95,7 @@ async fn publish_until_target() -> Result<(), Closed> { let mut coordinator = Member::register(&ready, &resume)?; let values = Arc::new((0..3).map(|_| AtomicU64::new(0)).collect::>()); let published = Arc::new(AtomicU64::new(0)); - let mut tasks = Vec::new(); + let mut tasks = vec![]; // Register everyone before polling any worker; the coordinator also keeps both phases open. for id in 0..3 { tasks.push(tokio::spawn(compute( diff --git a/examples/src/phaser_groups.rs b/examples/src/phaser_groups.rs index 9db8c35..3b5d265 100644 --- a/examples/src/phaser_groups.rs +++ b/examples/src/phaser_groups.rs @@ -145,7 +145,7 @@ async fn run_groups(fail_one_worker: bool) -> Result<(), Box>(), ); - let mut tasks = Vec::new(); + let mut tasks = vec![]; for group in 0..GROUPS { let ready = Phaser::new(); let resume = Phaser::new(); diff --git a/examples/src/phaser_rounds.rs b/examples/src/phaser_rounds.rs index 132ee33..b8a3542 100644 --- a/examples/src/phaser_rounds.rs +++ b/examples/src/phaser_rounds.rs @@ -40,7 +40,7 @@ async fn main() -> Result<(), Closed> { async fn start_gate() -> Result<(), Closed> { let phaser = Phaser::new(); let setup = phaser.register_one()?; - let mut tasks = Vec::new(); + let mut tasks = vec![]; for mut participant in phaser.register(3)? { tasks.push(tokio::spawn(async move { participant.wait().await?; diff --git a/tests-integration/tests/mpmc_test/concurrency.rs b/tests-integration/tests/mpmc_test/concurrency.rs index 2d9d266..ae20a74 100644 --- a/tests-integration/tests/mpmc_test/concurrency.rs +++ b/tests-integration/tests/mpmc_test/concurrency.rs @@ -42,7 +42,7 @@ where .map(|_| { let receiver = receiver.clone(); tokio::spawn(async move { - let mut values = Vec::new(); + let mut values = vec![]; while let Ok(value) = receiver.recv().await { values.push(value); } diff --git a/tests-integration/tests/mpsc_test/concurrency.rs b/tests-integration/tests/mpsc_test/concurrency.rs index 8937ec4..69968b0 100644 --- a/tests-integration/tests/mpsc_test/concurrency.rs +++ b/tests-integration/tests/mpsc_test/concurrency.rs @@ -284,7 +284,7 @@ fn unbounded_collects_from_multiple_producers() { } drop(tx); - let mut received = Vec::new(); + let mut received = vec![]; while let Ok(i) = rx.recv().await { received.push(i); } diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index 99751d4..80b44dd 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -155,7 +155,7 @@ fn concurrent_cancellation_preserves_capacity_and_message_order() { let (tx, mut rx) = mpsc::bounded(3); let mut out_of_order = 0; std::thread::scope(|scope| { - let mut workers = Vec::new(); + let mut workers = vec![]; for producer in 0..PRODUCERS { let tx = tx.clone(); workers.push(scope.spawn(move || { diff --git a/tests-integration/tests/spmc_test/concurrency.rs b/tests-integration/tests/spmc_test/concurrency.rs index 6ef9c83..d44ba4d 100644 --- a/tests-integration/tests/spmc_test/concurrency.rs +++ b/tests-integration/tests/spmc_test/concurrency.rs @@ -56,7 +56,7 @@ async fn bounded_values_are_delivered_exactly_once_to_eight_consumers() { let start = start.clone(); tokio::spawn(async move { start.wait().await; - let mut values = Vec::new(); + let mut values = vec![]; while let Ok(value) = receiver.recv().await { values.push(value); } @@ -86,7 +86,7 @@ async fn unbounded_values_are_delivered_exactly_once_to_eight_consumers() { let start = start.clone(); tokio::spawn(async move { start.wait().await; - let mut values = Vec::new(); + let mut values = vec![]; while let Ok(value) = receiver.recv().await { values.push(value); }