refactor: clarify waiter arena contracts (#191)

Author: tisonkun

asyncband/src/barrier/mod.rs

@@ -77,8 +77,8 @@ use std::task::Context;
 use std::task::Poll;
 
 use crate::internal::mutex::Mutex;
-use crate::internal::waitset::WaitRegistration;
 use crate::internal::waitset::WaitSet;
+use crate::internal::waitset::WakerToken;
 
 /// A synchronization primitive for multiple tasks that need to wait for each other.
 ///
@@ -255,7 +255,7 @@ impl Barrier {
         };
 
         let fut = BarrierWait {
-            registration: None,
+            token: None,
             generation,
             barrier: self,
         };
@@ -269,7 +269,7 @@ impl Barrier {
 /// This future will complete when all tasks have reached the barrier point.
 #[must_use = "futures do nothing unless you `.await` or poll them"]
 struct BarrierWait<'a> {
-    registration: Option<WaitRegistration>,
+    token: Option<WakerToken>,
     generation: usize,
     barrier: &'a Barrier,
 }
@@ -287,7 +287,7 @@ impl Future for BarrierWait<'_> {
 
     fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
         let Self {
-            registration,
+            token,
             generation,
             barrier,
         } = self.get_mut();
@@ -296,10 +296,10 @@ impl Future for BarrierWait<'_> {
             let mut state = barrier.state.lock();
             if *generation < state.generation {
                 // Advancing the generation drains its registrations under this same lock.
-                *registration = None;
+                *token = None;
                 return Poll::Ready(());
             }
-            state.waiters.register_waker(registration, cx)
+            state.waiters.register_waker(token, cx)
         };
         drop(replaced_waker);
         Poll::Pending
@@ -308,10 +308,10 @@ impl Future for BarrierWait<'_> {
 
 impl Drop for BarrierWait<'_> {
     fn drop(&mut self) {
-        if self.registration.is_some() {
+        if self.token.is_some() {
             let removed_waker = {
                 let mut state = self.barrier.state.lock();
-                state.waiters.unregister_waker(&mut self.registration)
+                state.waiters.unregister_waker(&mut self.token)
             };
             drop(removed_waker);
         }

asyncband/src/internal/arena.rs

@@ -18,24 +18,39 @@
 use std::mem;
 use std::num::NonZeroUsize;
 
-/// A stable index into an [`Arena`] for as long as its slot remains occupied.
-#[derive(Clone, Copy, Debug, Eq, PartialEq)]
-pub struct ArenaKey(usize);
-
-impl ArenaKey {
-    /// Encodes this key so it can provide a niche when stored in an `Option`-wrapped structure.
-    pub fn encode(self) -> NonZeroUsize {
+/// Identifies a reusable slot while that slot is occupied.
+///
+/// The non-zero representation lets wrappers such as `WaiterId` retain a niche when stored in an
+/// `Option`. Slot IDs deliberately carry no generation: each consumer supplies the cheaper
+/// lifecycle rule that matches its waiter storage.
+#[derive(Clone, Copy, Eq, PartialEq)]
+pub struct SlotId(NonZeroUsize);
+
+impl SlotId {
+    fn from_index(index: usize) -> Self {
         // `Slot<T>` is non-zero-sized, so a Vec of slots cannot reach `usize::MAX` elements.
-        unsafe { NonZeroUsize::new_unchecked(self.0 + 1) }
+        let encoded = index
+            .checked_add(1)
+            .expect("arena index must fit in a non-zero usize");
+        Self(NonZeroUsize::new(encoded).expect("encoded arena index must be non-zero"))
+    }
+
+    fn index(self) -> usize {
+        self.0.get() - 1
     }
+}
 
-    /// Decodes a key produced by [`Self::encode`].
-    pub fn decode(encoded: NonZeroUsize) -> Self {
-        Self(encoded.get() - 1)
+impl std::fmt::Debug for SlotId {
+    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
+        f.debug_tuple("SlotId").field(&self.index()).finish()
     }
 }
 
 /// Minimal reusable storage for internal waiter state.
+///
+/// The occupied length equals the number of `Occupied` slots. Every `Vacant` slot appears exactly
+/// once in the singly linked vacant list, which starts at `next_vacant` and terminates at
+/// `slots.len()`. Removing a value makes its slot ID available for immediate reuse.
 #[derive(Debug)]
 pub struct Arena<T> {
     slots: Vec<Slot<T>>,
@@ -45,8 +60,12 @@ pub struct Arena<T> {
 }
 
 /// 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)]
-pub struct ArenaValues<T> {
+struct ArenaValues<T> {
     first: Option<T>,
     rest: Vec<T>,
 }
@@ -63,7 +82,7 @@ impl<T> IntoIterator for ArenaValues<T> {
 #[derive(Debug)]
 enum Slot<T> {
     Occupied(T),
-    Vacant(usize),
+    Vacant { next: usize },
 }
 
 impl<T> Arena<T> {
@@ -83,61 +102,76 @@ impl<T> Arena<T> {
         }
     }
 
-    pub fn insert(&mut self, value: T) -> ArenaKey {
-        let key = self.next_vacant;
+    pub fn insert(&mut self, value: T) -> SlotId {
+        let index = self.next_vacant;
         self.len += 1;
 
-        if key == self.slots.len() {
+        if index == self.slots.len() {
             self.slots.push(Slot::Occupied(value));
-            self.next_vacant = key + 1;
+            self.next_vacant = index + 1;
         } else {
-            self.next_vacant = match self.slots.get(key) {
-                Some(Slot::Vacant(next)) => *next,
+            self.next_vacant = match self.slots.get(index) {
+                Some(Slot::Vacant { next }) => *next,
                 Some(Slot::Occupied(_)) | None => {
                     unreachable!("arena free list must point to a vacant slot")
                 }
             };
-            self.slots[key] = Slot::Occupied(value);
+            self.slots[index] = Slot::Occupied(value);
         }
 
-        ArenaKey(key)
+        SlotId::from_index(index)
     }
 
-    pub fn get(&self, key: ArenaKey) -> Option<&T> {
-        match self.slots.get(key.0) {
+    pub fn get(&self, id: SlotId) -> Option<&T> {
+        match self.slots.get(id.index()) {
             Some(Slot::Occupied(value)) => Some(value),
-            Some(Slot::Vacant(_)) | None => None,
+            Some(Slot::Vacant { .. }) | None => None,
         }
     }
 
-    pub fn get_mut(&mut self, key: ArenaKey) -> Option<&mut T> {
-        match self.slots.get_mut(key.0) {
+    pub fn get_mut(&mut self, id: SlotId) -> Option<&mut T> {
+        match self.slots.get_mut(id.index()) {
             Some(Slot::Occupied(value)) => Some(value),
-            Some(Slot::Vacant(_)) | None => None,
+            Some(Slot::Vacant { .. }) | None => None,
         }
     }
 
-    pub fn remove(&mut self, key: ArenaKey) -> T {
-        let index = key.0;
+    /// Removes the value stored at `id`.
+    ///
+    /// # Panics
+    ///
+    /// Panics if the slot ID is out of bounds or its slot is already vacant. Either case is an
+    /// internal waiter-lifecycle violation rather than a recoverable lookup failure.
+    #[track_caller]
+    pub fn remove(&mut self, id: SlotId) -> T {
+        let index = id.index();
         let slot = self
             .slots
             .get_mut(index)
-            .expect("arena key must be in bounds");
-        let value = match mem::replace(slot, Slot::Vacant(self.next_vacant)) {
+            .expect("arena slot ID must be in bounds");
+        let value = match mem::replace(
+            slot,
+            Slot::Vacant {
+                next: self.next_vacant,
+            },
+        ) {
             Slot::Occupied(value) => value,
-            vacant @ Slot::Vacant(_) => {
+            vacant @ Slot::Vacant { .. } => {
                 *slot = vacant;
-                panic!("arena key must be occupied");
+                panic!("arena slot ID must be occupied");
             }
         };
         self.len -= 1;
         self.next_vacant = index;
         value
     }
 
-    /// Takes every occupied value while retaining the allocation for reuse.
+    /// Takes every occupied value in slot order while retaining the allocation for reuse.
+    ///
+    /// 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) -> ArenaValues<T> {
+    pub fn take_all(&mut self) -> impl Iterator<Item = T> + use<T> {
         let len = self.len;
         let mut values = ArenaValues {
             first: None,
@@ -158,7 +192,7 @@ impl<T> Arena<T> {
 
         self.next_vacant = 0;
         self.len = 0;
-        values
+        values.into_iter()
     }
 
     #[cfg(test)]
@@ -171,6 +205,11 @@ impl<T> Arena<T> {
 mod tests {
     use super::*;
 
+    #[test]
+    fn slot_id_preserves_the_option_niche() {
+        assert_eq!(size_of::<SlotId>(), size_of::<Option<SlotId>>());
+    }
+
     #[test]
     fn removed_slots_are_reused() {
         let mut arena = Arena::new();
@@ -186,19 +225,19 @@ mod tests {
     }
 
     #[test]
-    fn take_all_restarts_key_allocation() {
+    fn take_all_restarts_slot_id_allocation() {
         let mut arena = Arena::with_capacity(3);
         let first = arena.insert(1);
         let second = arena.insert(2);
         let third = arena.insert(3);
         let capacity = arena.slots.capacity();
         arena.remove(second);
 
-        assert_eq!(arena.take_all().into_iter().collect::<Vec<_>>(), vec![1, 3]);
+        assert_eq!(arena.take_all().collect::<Vec<_>>(), vec![1, 3]);
         assert_eq!(arena.len(), 0);
         assert_eq!(arena.slots.capacity(), capacity);
 
-        let keys = [arena.insert(4), arena.insert(5), arena.insert(6)];
-        assert_eq!(keys, [first, second, third]);
+        let slot_ids = [arena.insert(4), arena.insert(5), arena.insert(6)];
+        assert_eq!(slot_ids, [first, second, third]);
     }
 }

asyncband/src/internal/countdown.rs

@@ -21,8 +21,8 @@ use std::task::Context;
 use std::task::Poll;
 
 use crate::internal::mutex::Mutex;
-use crate::internal::waitset::WaitRegistration;
 use crate::internal::waitset::WaitSet;
+use crate::internal::waitset::WakerToken;
 
 #[derive(Debug)]
 pub struct CountdownState {
@@ -72,37 +72,33 @@ impl CountdownState {
     }
 
     /// Polls for zero, registering the current waker if the countdown is still active.
-    pub fn poll_wait(
-        &self,
-        registration: &mut Option<WaitRegistration>,
-        cx: &mut Context<'_>,
-    ) -> Poll<()> {
+    pub fn poll_wait(&self, token: &mut Option<WakerToken>, cx: &mut Context<'_>) -> Poll<()> {
         if self.spin_wait(16).is_ok() {
             // The zero transition owns draining this wake epoch. Avoid taking the waiter lock
             // again when the completed future is dropped.
-            *registration = None;
+            *token = None;
             return Poll::Ready(());
         }
 
         let replaced_waker = {
             let mut waiters = self.waiters.lock();
             if self.state() == 0 {
                 // A concurrent zero transition will drain after this lock is released.
-                *registration = None;
+                *token = None;
                 return Poll::Ready(());
             }
-            waiters.register_waker(registration, cx)
+            waiters.register_waker(token, cx)
         };
         drop(replaced_waker);
         Poll::Pending
     }
 
     #[inline]
-    pub fn unregister_waker(&self, registration: &mut Option<WaitRegistration>) {
-        if registration.is_some() {
+    pub fn unregister_waker(&self, token: &mut Option<WakerToken>) {
+        if token.is_some() {
             let removed_waker = {
                 let mut waiters = self.waiters.lock();
-                waiters.unregister_waker(registration)
+                waiters.unregister_waker(token)
             };
             drop(removed_waker);
         }

asyncband/src/internal/mod.rs

@@ -30,7 +30,7 @@ pub(crate) mod atomic_waker;
 // `WaitList` and `WaitSet` use different `Arena` operations. A single-primitive build therefore
 // leaves part of this shared API unused, while the all-feature build uses it.
 #[allow(dead_code)]
-pub(crate) mod arena;
+mod arena;
 
 #[cfg(any(feature = "latch", feature = "once", feature = "waitgroup"))]
 // `waitgroup` increments and decrements the countdown, while `latch` and `once` only decrement it.