diff --git a/CHANGELOG.md b/CHANGELOG.md index d469a841..c51ffff9 100644 --- a/CHANGELOG.md +++ b/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. diff --git a/asyncband/src/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/broadcast/mpmc/unbounded/mod.rs index 1feb1053..4556acf4 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -398,7 +398,7 @@ impl Drop for UnboundedSender { // 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); } diff --git a/asyncband/src/completion/mod.rs b/asyncband/src/completion/mod.rs index c8ab8207..ab21326e 100644 --- a/asyncband/src/completion/mod.rs +++ b/asyncband/src/completion/mod.rs @@ -160,7 +160,7 @@ impl Completer { } // 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 Drop for Completer { } // 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); } diff --git a/asyncband/src/condvar/mod.rs b/asyncband/src/condvar/mod.rs index 18d3a648..985fc7ea 100644 --- a/asyncband/src/condvar/mod.rs +++ b/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| { diff --git a/asyncband/src/event/mod.rs b/asyncband/src/event/mod.rs index 312a3a17..aaab5e83 100644 --- a/asyncband/src/event/mod.rs +++ b/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 diff --git a/asyncband/src/internal/arena.rs b/asyncband/src/internal/arena.rs index ab63e635..290a32cd 100644 --- a/asyncband/src/internal/arena.rs +++ b/asyncband/src/internal/arena.rs @@ -59,26 +59,6 @@ pub struct Arena { 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 { - first: Option, - rest: Vec, -} - -impl IntoIterator for ArenaValues { - type Item = T; - type IntoIter = std::iter::Chain, std::vec::IntoIter>; - - fn into_iter(self) -> Self::IntoIter { - self.first.into_iter().chain(self.rest) - } -} - #[derive(Debug)] enum Slot { Occupied(T), @@ -181,40 +161,32 @@ impl Arena { 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 + use { - 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 + '_ { + 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 + use { 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,7 +214,7 @@ 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); @@ -250,11 +222,24 @@ mod tests { let capacity = arena.slots.capacity(); arena.remove(second); - assert_eq!(arena.take_all().collect::>(), vec![1, 3]); + assert_eq!(arena.drain().collect::>(), 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![1, 3]); + } } diff --git a/asyncband/src/internal/countdown.rs b/asyncband/src/internal/countdown.rs index 9e331648..2f8cdce3 100644 --- a/asyncband/src/internal/countdown.rs +++ b/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); diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index a0605117..98e94586 100644 --- a/asyncband/src/internal/mod.rs +++ b/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", diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs index bb0afe56..70c33378 100644 --- a/asyncband/src/internal/semaphore.rs +++ b/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; diff --git a/asyncband/src/internal/waitset.rs b/asyncband/src/internal/waitset.rs index 13f580dd..e34a4991 100644 --- a/asyncband/src/internal/waitset.rs +++ b/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 + '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 + '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::>(); + 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(); diff --git a/asyncband/src/internal/waker_batch.rs b/asyncband/src/internal/waker_batch.rs new file mode 100644 index 00000000..0e770e35 --- /dev/null +++ b/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, + rest: Vec, +} + +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 for WakerBatch { + fn extend>(&mut self, iter: I) { + for waker in iter { + self.push(waker); + } + } +} + +impl IntoIterator 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) + } +} diff --git a/asyncband/src/watch/mod.rs b/asyncband/src/watch/mod.rs index cf68a6ba..8d938813 100644 --- a/asyncband/src/watch/mod.rs +++ b/asyncband/src/watch/mod.rs @@ -153,7 +153,7 @@ impl Drop for Sender { if state.senders != 0 { return; } - state.waiters.drain() + state.waiters.take_all() }; wake_all(wakers); }