From 4ef2799dd29661b05f420e2a5ff01274f5a9ca4e Mon Sep 17 00:00:00 2001 From: Piyush Rokade <236336623+rokadepiyush49-rgb@users.noreply.github.com> Date: Sun, 8 Nov 2026 23:47:50 +0530 Subject: [PATCH] perf(semaphore): avoid allocating queue storage for the first waiter Keep the only waiter of a WaitList inline instead of allocating arena storage for it. A second concurrent waiter moves the node into the arena under the same WaiterId, and the arena is then retained for reuse. --- CHANGELOG.md | 1 + asyncband/src/internal/arena.rs | 3 + asyncband/src/internal/waitlist.rs | 156 ++++++++++++++++-- .../benches/primitives/semaphore/acquire.rs | 17 ++ tests-integration/tests/semaphore_test.rs | 131 +++++++++++++++ 5 files changed, 298 insertions(+), 10 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ca7736de..704267ba 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -45,6 +45,7 @@ All notable changes to this project will be documented in this file. * Reduce the cost of initializing an uncontended `OnceCell` by avoiding its internal wait-queue lock. * Reduce the cost of polling a pending oneshot receive again from the same task, which now keeps its registered waker instead of cloning and dropping it. * Reduce the state kept by futures waiting for a `Mutex`, `RwLock`, `Semaphore`, `Once`, `OnceCell`, or bounded broadcast capacity, making lock, permit acquisition, and bounded broadcast send futures smaller. +* Avoid allocating wait-queue storage while only one task at a time waits on a `Semaphore`, `Mutex`, `RwLock`, `Condvar`, event, or MPSC, MPMC, or SPMC channel, which makes the first wait on a new primitive cheaper. * Reduce the cost of `Barrier::wait` for participants that are not the last to arrive, which now record their arrival and register for wakeup under one lock acquisition instead of two. ## v0.7.4 (2026-10-25) diff --git a/asyncband/src/internal/arena.rs b/asyncband/src/internal/arena.rs index 2c33b5fc..c07842fd 100644 --- a/asyncband/src/internal/arena.rs +++ b/asyncband/src/internal/arena.rs @@ -29,6 +29,9 @@ use std::num::NonZeroUsize; pub struct SlotId(NonZeroUsize); impl SlotId { + /// The ID an empty arena issues for its first inserted value. + pub const FIRST: Self = Self(NonZeroUsize::MIN); + fn from_index(index: usize) -> Self { // `Slot` is non-zero-sized, so a Vec of slots cannot reach `usize::MAX` elements. let encoded = index diff --git a/asyncband/src/internal/waitlist.rs b/asyncband/src/internal/waitlist.rs index d26062ee..85e85711 100644 --- a/asyncband/src/internal/waitlist.rs +++ b/asyncband/src/internal/waitlist.rs @@ -23,7 +23,45 @@ use crate::internal::arena::SlotId; pub struct WaitList { head: Option, tail: Option, - nodes: Arena>, + nodes: Nodes, +} + +/// Node storage that stays inline until two nodes must exist at once. +/// +/// The inline node occupies the ID an empty arena issues first, so it moves into the arena under +/// the ID its owner already holds. The arena is then retained for reuse. Sharing storage between +/// the variants keeps an idle list close to the size of the arena alone. +#[derive(Debug)] +enum Nodes { + Inline(Option>), + Arena(Arena>), +} + +impl Nodes { + #[inline] + fn insert(&mut self, node: Node) -> SlotId { + match self { + Self::Inline(slot @ None) => { + *slot = Some(node); + SlotId::FIRST + } + Self::Inline(slot) => { + let first = slot.take().expect("occupied inline node must exist"); + self.move_to_arena(first, node) + } + Self::Arena(nodes) => nodes.insert(node), + } + } + + #[cold] + fn move_to_arena(&mut self, first: Node, node: Node) -> SlotId { + let mut nodes = Arena::new(); + let first_id = nodes.insert(first); + debug_assert_eq!(first_id, SlotId::FIRST); + let id = nodes.insert(node); + *self = Self::Arena(nodes); + id + } } #[derive(Clone, Copy, Debug, Eq, PartialEq)] @@ -57,7 +95,7 @@ impl WaitList { Self { head: None, tail: None, - nodes: Arena::new(), + nodes: Nodes::Inline(None), } } @@ -155,19 +193,37 @@ impl WaitList { self.node(id).links.is_none(), "waiter must be unlinked before removal" ); - self.nodes.remove(id.slot()).value + match &mut self.nodes { + Nodes::Inline(node) => node + .take() + .expect("waiter id must refer to an occupied node"), + Nodes::Arena(nodes) => nodes.remove(id.slot()), + } + .value } + #[inline] fn node(&self, id: WaiterId) -> &Node { - self.nodes - .get(id.slot()) - .expect("waiter id must refer to an occupied node") + match &self.nodes { + Nodes::Inline(node) => { + debug_assert_eq!(id.slot(), SlotId::FIRST); + node.as_ref() + } + Nodes::Arena(nodes) => nodes.get(id.slot()), + } + .expect("waiter id must refer to an occupied node") } + #[inline] fn node_mut(&mut self, id: WaiterId) -> &mut Node { - self.nodes - .get_mut(id.slot()) - .expect("waiter id must refer to an occupied node") + match &mut self.nodes { + Nodes::Inline(node) => { + debug_assert_eq!(id.slot(), SlotId::FIRST); + node.as_mut() + } + Nodes::Arena(nodes) => nodes.get_mut(id.slot()), + } + .expect("waiter id must refer to an occupied node") } fn linked_node_mut(&mut self, id: WaiterId) -> &mut Links { @@ -179,7 +235,10 @@ impl WaitList { #[cfg(test)] pub fn occupied_len(&self) -> usize { - self.nodes.len() + match &self.nodes { + Nodes::Inline(node) => usize::from(node.is_some()), + Nodes::Arena(nodes) => nodes.len(), + } } } @@ -253,4 +312,81 @@ mod tests { } assert!(waiters.is_empty()); } + + #[test] + fn a_lone_waiter_stays_inline_across_reuse() { + let mut waiters = WaitList::new(); + + for value in 0..3 { + let id = waiters.push_back(value); + assert!(matches!(waiters.nodes, Nodes::Inline(Some(_)))); + assert_eq!(waiters.unlink_first_waiter(|_| true).unwrap().0, id); + assert_eq!(waiters.remove_unlinked_waiter(id), value); + } + assert!(matches!(waiters.nodes, Nodes::Inline(None))); + } + + #[test] + fn a_second_waiter_moves_the_first_to_the_arena_under_its_id() { + let mut waiters = WaitList::new(); + let first = waiters.push_back(1); + let second = waiters.push_back(2); + assert!(matches!(waiters.nodes, Nodes::Arena(_))); + + assert_eq!(*waiters.waiter_mut(first), 1); + assert_eq!(*waiters.waiter_mut(second), 2); + for expected in [(first, 1), (second, 2)] { + let (id, value) = waiters.unlink_first_waiter(|_| true).unwrap(); + assert_eq!((id, *value), expected); + waiters.remove_unlinked_waiter(id); + } + assert!(waiters.is_empty()); + } + + #[test] + fn push_front_orders_a_new_head_before_the_inline_waiter() { + let mut waiters = WaitList::new(); + let inline = waiters.push_back(2); + let head = waiters.push_front(1); + + for expected in [(head, 1), (inline, 2)] { + let (id, value) = waiters.unlink_first_waiter(|_| true).unwrap(); + assert_eq!((id, *value), expected); + waiters.remove_unlinked_waiter(id); + } + assert!(waiters.is_empty()); + } + + #[test] + fn a_detached_inline_waiter_stays_detached_when_moved() { + let mut waiters = WaitList::new(); + let first = waiters.push_back(1); + assert_eq!(waiters.unlink_first_waiter(|_| true).unwrap().0, first); + assert!(waiters.is_empty()); + + let second = waiters.push_back(2); + assert_eq!(waiters.occupied_len(), 2); + assert_eq!(*waiters.waiter_mut(first), 1); + assert_eq!(waiters.remove_unlinked_waiter(first), 1); + + assert_eq!(waiters.unlink_first_waiter(|_| true).unwrap().0, second); + assert!(waiters.is_empty()); + assert_eq!(waiters.remove_unlinked_waiter(second), 2); + assert_eq!(waiters.occupied_len(), 0); + } + + #[test] + fn the_arena_is_retained_after_the_queue_drains() { + let mut waiters = WaitList::new(); + waiters.push_back(1); + waiters.push_back(2); + while let Some((id, _)) = waiters.unlink_first_waiter(|_| true) { + waiters.remove_unlinked_waiter(id); + } + assert!(matches!(waiters.nodes, Nodes::Arena(_))); + + let id = waiters.push_back(3); + assert!(matches!(waiters.nodes, Nodes::Arena(_))); + assert_eq!(*waiters.waiter_mut(id), 3); + } } diff --git a/benchmarks/benches/primitives/semaphore/acquire.rs b/benchmarks/benches/primitives/semaphore/acquire.rs index 9e1b1b49..33ead94c 100644 --- a/benchmarks/benches/primitives/semaphore/acquire.rs +++ b/benchmarks/benches/primitives/semaphore/acquire.rs @@ -46,6 +46,22 @@ fn handoff_permit(bencher: Bencher) { }); } +fn reused_handoff_permit(bencher: Bencher) { + let mut context = bench_context(); + let semaphore = Semaphore::new(0); + + bencher.bench_local(|| { + let mut acquire = pin!(semaphore.acquire(black_box(1))); + poll_pending(acquire.as_mut(), &mut context); + + semaphore.release(black_box(1)); + let permit = poll_pinned_ready(acquire.as_mut(), &mut context); + permit.forget(); + + black_box(semaphore.available_permits()) + }); +} + fn try_acquire_release(bencher: Bencher) { let semaphore = Semaphore::new(1); @@ -142,6 +158,7 @@ fn cancel_pending_owned_batch(bencher: Bencher, queue_depth: usize) { pub fn register(runner: &mut Runner) { primitive!(runner, handoff_permit, 3_000_000); + primitive!(runner, reused_handoff_permit, 10_000_000); primitive!(runner, try_acquire_release, 30_000_000); primitive!(runner, try_acquire_release_with_spare_permits, 60_000_000); primitive!(runner, owned_try_acquire_release, 48_000_000); diff --git a/tests-integration/tests/semaphore_test.rs b/tests-integration/tests/semaphore_test.rs index 8f6937dc..c23906a2 100644 --- a/tests-integration/tests/semaphore_test.rs +++ b/tests-integration/tests/semaphore_test.rs @@ -31,6 +31,7 @@ use std::thread; use asyncband::blocking::FutureExt; use asyncband::semaphore::Semaphore; use tests_integration::WakeCounter; +use tests_integration::assert_completes_without_deadlock; use tests_integration::expect_ready; use tests_integration::poll_with; @@ -206,6 +207,136 @@ fn cancellation_restores_permits_before_dropping_waker() { assert_eq!(Arc::strong_count(&semaphore), 1); } +#[test] +fn queued_acquires_keep_their_order_as_the_queue_grows() { + let semaphore = Semaphore::new(0); + let mut first = pin!(semaphore.acquire(2)); + let mut second = pin!(semaphore.acquire(1)); + let (first_waker, first_wakes) = WakeCounter::new(); + let (second_waker, second_wakes) = WakeCounter::new(); + + assert!(poll_with(first.as_mut(), &first_waker).is_pending()); + assert!(poll_with(second.as_mut(), &second_waker).is_pending()); + + // The head of the queue takes a partial grant, so the later single-permit waiter stays queued. + semaphore.release(1); + assert_eq!((first_wakes.count(), second_wakes.count()), (0, 0)); + + semaphore.release(1); + assert_eq!((first_wakes.count(), second_wakes.count()), (1, 0)); + expect_ready(poll_with(first.as_mut(), &first_waker)).forget(); + assert!(poll_with(second.as_mut(), &second_waker).is_pending()); + + semaphore.release(1); + assert_eq!(second_wakes.count(), 1); + expect_ready(poll_with(second.as_mut(), &second_waker)).forget(); + assert_eq!(semaphore.available_permits(), 0); +} + +#[test] +fn cancelling_a_partly_granted_first_waiter_passes_its_permits_on() { + let semaphore = Semaphore::new(0); + let mut first = Box::pin(semaphore.acquire(2)); + let mut second = pin!(semaphore.acquire(1)); + let (waker, wakes) = WakeCounter::new(); + + assert!(poll_with(first.as_mut(), Waker::noop()).is_pending()); + assert!(poll_with(second.as_mut(), &waker).is_pending()); + semaphore.release(1); + assert_eq!(wakes.count(), 0); + + drop(first); + assert_eq!(wakes.count(), 1); + expect_ready(poll_with(second.as_mut(), &waker)).forget(); + assert_eq!(semaphore.available_permits(), 0); +} + +#[test] +fn cancelling_a_fully_granted_lone_waiter_returns_its_permits() { + let semaphore = Semaphore::new(0); + let mut first = Box::pin(semaphore.acquire(2)); + let (waker, wakes) = WakeCounter::new(); + + assert!(poll_with(first.as_mut(), &waker).is_pending()); + semaphore.release(2); + assert_eq!(wakes.take(), 1); + assert_eq!(semaphore.available_permits(), 0); + + drop(first); + assert_eq!(semaphore.available_permits(), 2); + + // The cancelled waiter's slot is free for the next one. + let mut next = pin!(semaphore.acquire(3)); + assert!(poll_with(next.as_mut(), &waker).is_pending()); + assert_eq!(semaphore.available_permits(), 0); + semaphore.release(1); + assert_eq!(wakes.count(), 1); + expect_ready(poll_with(next.as_mut(), &waker)).forget(); +} + +#[test] +fn a_granted_first_waiter_keeps_its_permit_when_the_queue_grows() { + let semaphore = Semaphore::new(0); + let mut first = pin!(semaphore.acquire(1)); + let mut second = pin!(semaphore.acquire(1)); + let (first_waker, first_wakes) = WakeCounter::new(); + let (second_waker, second_wakes) = WakeCounter::new(); + + assert!(poll_with(first.as_mut(), &first_waker).is_pending()); + semaphore.release(1); + assert_eq!(first_wakes.count(), 1); + + // The second waiter arrives while the granted first node has not been collected yet. + assert!(poll_with(second.as_mut(), &second_waker).is_pending()); + let permit = expect_ready(poll_with(first.as_mut(), &first_waker)); + assert_eq!(second_wakes.count(), 0); + + drop(permit); + assert_eq!(second_wakes.count(), 1); + expect_ready(poll_with(second.as_mut(), &second_waker)).forget(); + assert_eq!(semaphore.available_permits(), 0); +} + +#[test] +fn a_first_waiter_repolled_after_the_queue_grows_uses_its_latest_waker() { + let semaphore = Semaphore::new(0); + let mut first = pin!(semaphore.acquire(1)); + let mut second = pin!(semaphore.acquire(1)); + let (old_waker, old_wakes) = WakeCounter::new(); + let (new_waker, new_wakes) = WakeCounter::new(); + + assert!(poll_with(first.as_mut(), &old_waker).is_pending()); + assert!(poll_with(second.as_mut(), Waker::noop()).is_pending()); + assert!(poll_with(first.as_mut(), &new_waker).is_pending()); + + semaphore.release(1); + assert_eq!((old_wakes.count(), new_wakes.count()), (0, 1)); + expect_ready(poll_with(first.as_mut(), &new_waker)).forget(); +} + +#[test] +fn concurrent_first_waiters_on_a_fresh_semaphore_do_not_miss_a_wakeup() { + const ROUNDS: usize = if cfg!(miri) { 8 } else { 500 }; + + assert_completes_without_deadlock(|| { + for _ in 0..ROUNDS { + let semaphore = Semaphore::new(0); + let start = Barrier::new(3); + thread::scope(|scope| { + for _ in 0..2 { + scope.spawn(|| { + start.wait(); + semaphore.acquire(1).block_on().forget(); + }); + } + start.wait(); + semaphore.release(2); + }); + assert_eq!(semaphore.available_permits(), 0); + } + }); +} + #[test] fn reduce_permits_takes_priority_over_pending_acquires() { let s = Semaphore::new(0);