perf(semaphore): register the first waiter without allocating

CHANGELOG.md

@@ -42,6 +42,7 @@ All notable changes to this project will be documented in this file.
 
 ### Improvements
 
+* Reduce the cost of registering the first waiter on a `Semaphore`, `Condvar`, `ManualResetEvent`, `AutoResetEvent`, or bounded MPMC channel: the waiter arena keeps its first slot inline and only allocates from the second concurrent waiter, removing one allocation from every fresh-queue registration.
 * Reduce intermediate payload copies in bounded MPSC receives, especially for large inline messages.
 * 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.

asyncband/src/internal/arena.rs

@@ -53,8 +53,15 @@ impl fmt::Debug for SlotId {
 /// The occupied length equals the number of `Occupied` slots. Every `Vacant` slot appears exactly
 /// once in the singly linked vacant list starting at `vacant_head`. Removing a value makes its slot
 /// ID available for immediate reuse.
+///
+/// The first occupant lives in the inline `first` slot, so registering a lone waiter never
+/// allocates. The second registration spills `first` into the slot vector, which then follows the
+/// rules above for the arena's remaining lifetime. Every arena therefore carries one inline value
+/// of `T` in exchange for an allocation-free first registration.
 #[derive(Debug)]
 pub struct Arena<T> {
+    /// The occupant of slot 0 until the vector first becomes non-empty; `None` afterwards.
+    first: Option<T>,
     slots: Vec<Slot<T>>,
     /// The head of the vacant list, or `None` when every slot is occupied.
     vacant_head: Option<SlotId>,
@@ -70,6 +77,7 @@ enum Slot<T> {
 impl<T> Arena<T> {
     pub const fn new() -> Self {
         Self {
+            first: None,
             slots: vec![],
             vacant_head: None,
             len: 0,
@@ -78,6 +86,7 @@ impl<T> Arena<T> {
 
     pub fn with_capacity(capacity: usize) -> Self {
         Self {
+            first: None,
             slots: Vec::with_capacity(capacity),
             vacant_head: None,
             len: 0,
@@ -98,6 +107,21 @@ impl<T> Arena<T> {
             };
             *slot = Slot::Occupied(value);
             id
+        } else if self.slots.is_empty() {
+            match self.first.take() {
+                None => {
+                    self.first = Some(value);
+                    SlotId::from_index(0)
+                }
+                Some(first) => {
+                    // Spill the inline slot into the vector so slot 0 keeps its ID, then append
+                    // the new value as slot 1.
+                    self.slots.push(Slot::Occupied(first));
+                    let id = SlotId::from_index(self.slots.len());
+                    self.slots.push(Slot::Occupied(value));
+                    id
+                }
+            }
         } else {
             let id = SlotId::from_index(self.slots.len());
             self.slots.push(Slot::Occupied(value));
@@ -108,25 +132,48 @@ impl<T> Arena<T> {
         id
     }
 
-    pub fn get(&self, id: SlotId) -> Option<&T> {
-        match self.slots.get(id.index()) {
+    /// Returns the value at `index`, looking through the inline slot while the vector is empty.
+    fn slot_ref(&self, index: usize) -> Option<&T> {
+        if index == 0 && self.slots.is_empty() {
+            return self.first.as_ref();
+        }
+        match self.slots.get(index) {
             Some(Slot::Occupied(value)) => Some(value),
             Some(Slot::Vacant { .. }) | None => None,
         }
     }
 
-    pub fn get_mut(&mut self, id: SlotId) -> Option<&mut T> {
-        match self.slots.get_mut(id.index()) {
+    /// Returns the value at `index` mutably, looking through the inline slot while the vector is
+    /// empty.
+    fn slot_mut(&mut self, index: usize) -> Option<&mut T> {
+        if index == 0 && self.slots.is_empty() {
+            return self.first.as_mut();
+        }
+        match self.slots.get_mut(index) {
             Some(Slot::Occupied(value)) => Some(value),
             Some(Slot::Vacant { .. }) | None => None,
         }
     }
 
+    pub fn get(&self, id: SlotId) -> Option<&T> {
+        self.slot_ref(id.index())
+    }
+
+    pub fn get_mut(&mut self, id: SlotId) -> Option<&mut T> {
+        self.slot_mut(id.index())
+    }
+
     #[cfg(test)]
     pub fn len(&self) -> usize {
         self.len
     }
 
+    /// Returns the number of heap-allocated slots; used by tests to pin the inline fast path.
+    #[cfg(test)]
+    pub fn vector_len(&self) -> usize {
+        self.slots.len()
+    }
+
     pub fn is_empty(&self) -> bool {
         self.len == 0
     }
@@ -140,6 +187,13 @@ impl<T> Arena<T> {
     #[track_caller]
     pub fn remove(&mut self, id: SlotId) -> T {
         let index = id.index();
+        if index == 0 && self.slots.is_empty() {
+            // The inline slot never joins the vacant list: while the vector is empty, slot 0 is
+            // the only ID that ever existed, so reuse goes through `first` directly.
+            let value = self.first.take().expect("arena slot ID must be occupied");
+            self.len -= 1;
+            return value;
+        }
         let slot = self
             .slots
             .get_mut(index)
@@ -168,12 +222,13 @@ impl<T> Arena<T> {
     pub fn take_all<C: FromIterator<T>>(&mut self) -> C {
         self.vacant_head = None;
         self.len = 0;
-        self.slots
-            .drain(..)
-            .filter_map(|slot| match slot {
+        self.first
+            .take()
+            .into_iter()
+            .chain(self.slots.drain(..).filter_map(|slot| match slot {
                 Slot::Occupied(value) => Some(value),
                 Slot::Vacant { .. } => None,
-            })
+            }))
             .collect()
     }
 
@@ -182,10 +237,11 @@ impl<T> Arena<T> {
     /// The returned iterator owns the backing allocation and releases it when dropped.
     #[inline]
     pub fn into_iter(self) -> impl Iterator<Item = T> {
-        self.slots.into_iter().filter_map(|slot| match slot {
+        let slots = self.slots.into_iter().filter_map(|slot| match slot {
             Slot::Occupied(value) => Some(value),
             Slot::Vacant { .. } => None,
-        })
+        });
+        self.first.into_iter().chain(slots)
     }
 }
 
@@ -212,6 +268,64 @@ mod tests {
         assert_eq!(arena.get(second), Some(&"second"));
     }
 
+    #[test]
+    fn lone_waiter_uses_the_inline_slot() {
+        let mut arena = Arena::new();
+        let id = arena.insert("only");
+
+        // The vector-only implementation this test pins against would allocate here already.
+        assert_eq!(arena.vector_len(), 0);
+
+        assert_eq!(id, SlotId::from_index(0));
+        assert_eq!(arena.get(id), Some(&"only"));
+        assert_eq!(arena.len(), 1);
+
+        assert_eq!(arena.remove(id), "only");
+        assert!(arena.is_empty());
+        assert_eq!(arena.get(id), None);
+
+        // Slot 0 is handed out again while the vector is still empty.
+        let again = arena.insert("again");
+        assert_eq!(again, id);
+        assert_eq!(arena.vector_len(), 0);
+    }
+
+    #[test]
+    fn spill_preserves_ids_and_values() {
+        let mut arena = Arena::new();
+        let first = arena.insert(1);
+        let second = arena.insert(2);
+        let third = arena.insert(3);
+
+        assert_eq!(first, SlotId::from_index(0));
+        assert_eq!(
+            (arena.get(first), arena.get(second), arena.get(third)),
+            (Some(&1), Some(&2), Some(&3))
+        );
+
+        // Removing the spilled slot 0 returns it to the vacant list, not to the inline slot.
+        assert_eq!(arena.remove(first), 1);
+        let replacement = arena.insert(10);
+        assert_eq!(replacement, first);
+        assert_eq!(arena.get(replacement), Some(&10));
+        assert_eq!(arena.get(third), Some(&3));
+    }
+
+    #[test]
+    fn take_all_starts_from_the_inline_slot_again() {
+        let mut arena = Arena::new();
+        arena.insert(1);
+        arena.insert(2);
+
+        assert_eq!(arena.take_all::<Vec<_>>(), vec![1, 2]);
+        assert!(arena.is_empty());
+
+        // The arena is back to pre-spill mode: the next insert goes to the inline slot.
+        let id = arena.insert(3);
+        assert_eq!(id, SlotId::from_index(0));
+        assert_eq!(arena.get(id), Some(&3));
+    }
+
     #[test]
     fn take_all_retains_capacity_and_restarts_slot_ids() {
         let mut arena = Arena::with_capacity(3);

asyncband/src/internal/semaphore.rs

@@ -583,4 +583,87 @@ mod tests {
         }
         assert_eq!(semaphore.waiters.lock().occupied_len(), 0);
     }
+
+    /// Counts this thread's allocations other than the 64-byte pthread-mutex boxing.
+    ///
+    /// The pthread-backed `std::sync::Mutex` heap-boxes itself on its first lock; on futex-backed
+    /// targets the same path allocates nothing. Filtering that one size makes the counter stable
+    /// across platforms while still catching any allocation introduced by waiter registration.
+    mod allocation_counting {
+        use std::alloc::GlobalAlloc;
+        use std::alloc::Layout;
+        use std::alloc::System;
+        use std::cell::Cell;
+
+        thread_local! {
+            static NON_BOXING_ALLOCATIONS: Cell<u64> = const { Cell::new(0) };
+        }
+
+        pub(super) fn count_since_reset() -> u64 {
+            NON_BOXING_ALLOCATIONS.with(Cell::get)
+        }
+
+        pub(super) fn reset() {
+            NON_BOXING_ALLOCATIONS.with(|count| count.set(0));
+        }
+
+        pub struct Counting;
+
+        unsafe impl GlobalAlloc for Counting {
+            unsafe fn alloc(&self, layout: Layout) -> *mut u8 {
+                let size = layout.size();
+                NON_BOXING_ALLOCATIONS.with(|count| {
+                    if size != 64 {
+                        count.set(count.get() + 1);
+                    }
+                });
+                unsafe { System.alloc(layout) }
+            }
+
+            unsafe fn alloc_zeroed(&self, layout: Layout) -> *mut u8 {
+                let size = layout.size();
+                NON_BOXING_ALLOCATIONS.with(|count| {
+                    if size != 64 {
+                        count.set(count.get() + 1);
+                    }
+                });
+                unsafe { System.alloc_zeroed(layout) }
+            }
+
+            unsafe fn realloc(&self, ptr: *mut u8, layout: Layout, new_size: usize) -> *mut u8 {
+                NON_BOXING_ALLOCATIONS.with(|count| {
+                    if new_size != 64 {
+                        count.set(count.get() + 1);
+                    }
+                });
+                unsafe { System.realloc(ptr, layout, new_size) }
+            }
+
+            unsafe fn dealloc(&self, ptr: *mut u8, layout: Layout) {
+                unsafe { System.dealloc(ptr, layout) }
+            }
+        }
+
+        #[global_allocator]
+        static GLOBAL_COUNTING: Counting = Counting;
+    }
+
+    #[test]
+    fn first_registration_allocates_nothing_beyond_the_mutex_boxing() {
+        let semaphore = Semaphore::new(0);
+
+        allocation_counting::reset();
+        {
+            let mut acquire = semaphore.acquire(1);
+            assert!(acquire.poll_once(Waker::noop()).is_pending());
+            // Dropping the future takes the cancellation path: unlink, remove, and reuse the
+            // waiter slot must not allocate either.
+        }
+        let count = allocation_counting::count_since_reset();
+
+        assert_eq!(
+            count, 0,
+            "registering, cancelling, and removing the first waiter must not allocate"
+        );
+    }
 }