perf(mpsc): streamline unbounded buffer reclamation

asyncband/src/mpsc/unbounded/buffer.rs

@@ -21,55 +21,74 @@ use std::mem;
 // Bound the inline storage retained by a partial batch. Boxed payloads belong to individual
 // messages, not these backing allocations. Empty buffers are reused without retaining peak size.
 pub const SEGMENT_BYTES: usize = 32 * 1024;
+const RETAINED_DIRECTORY_BYTES: usize = 1024;
 
 pub struct Buffer<T> {
     writable: VecDeque<T>,
     sealed: VecDeque<VecDeque<T>>,
 }
 
 impl<T> Buffer<T> {
+    const SEGMENT_CAPACITY: usize = if size_of::<T>() == 0 {
+        usize::MAX
+    } else if size_of::<T>() > SEGMENT_BYTES {
+        1
+    } else {
+        let limit = SEGMENT_BYTES / size_of::<T>();
+        // Power-of-two limits keep ordinary VecDeque growth within the segment byte budget.
+        1 << (usize::BITS - 1 - limit.leading_zeros())
+    };
+
+    // The empty directory retains segment headers, not message storage.
+    const RETAINED_DIRECTORY_SLOTS: usize = RETAINED_DIRECTORY_BYTES / size_of::<VecDeque<T>>();
+
     pub fn new() -> Self {
         Self {
             writable: VecDeque::new(),
             sealed: VecDeque::new(),
         }
     }
 
-    fn segment_capacity() -> usize {
-        if size_of::<T>() == 0 {
-            return usize::MAX;
-        }
-        let limit = (SEGMENT_BYTES / size_of::<T>()).max(1);
-        // Power-of-two limits let VecDeque grow naturally without exceeding the segment budget.
-        1 << (usize::BITS - 1 - limit.leading_zeros())
-    }
-
     pub fn push(&mut self, value: T) {
-        if self.writable.len() == Self::segment_capacity() {
-            let next = VecDeque::with_capacity(Self::segment_capacity());
+        if self.writable.len() == Self::SEGMENT_CAPACITY {
+            let next = VecDeque::with_capacity(Self::SEGMENT_CAPACITY);
             let sealed = mem::replace(&mut self.writable, next);
             self.sealed.push_back(sealed);
         }
         self.writable.push_back(value);
     }
 
-    pub fn refill(&mut self, batch: &mut VecDeque<T>) {
+    /// Returns retired allocations so the receiver can release them after unlocking.
+    #[must_use = "retired allocations must be dropped after releasing the shared lock"]
+    pub fn refill(
+        &mut self,
+        batch: &mut VecDeque<T>,
+    ) -> Option<(VecDeque<T>, VecDeque<VecDeque<T>>)> {
         debug_assert!(batch.is_empty());
         if let Some(sealed) = self.sealed.pop_front() {
-            *batch = sealed;
-            if self.sealed.is_empty() && self.sealed.capacity() * size_of::<VecDeque<T>>() > 1024 {
-                self.sealed = VecDeque::new();
-            }
-        } else if !self.writable.is_empty() {
+            let retired_batch = mem::replace(batch, sealed);
+            let retired_sealed = if self.sealed.is_empty()
+                && self.sealed.capacity() > Self::RETAINED_DIRECTORY_SLOTS
+            {
+                mem::take(&mut self.sealed)
+            } else {
+                VecDeque::new()
+            };
+            return Some((retired_batch, retired_sealed));
+        }
+        if !self.writable.is_empty() {
             mem::swap(batch, &mut self.writable);
         }
+        None
     }
 }
 
 pub fn pop_batch<T>(batch: &mut VecDeque<T>) -> T {
-    if batch.len() == 1 && batch.capacity().saturating_mul(size_of::<T>()) > SEGMENT_BYTES {
-        // Retire the allocation on the last value, outside the shared lock. Keep this as a tail
-        // expression to avoid intermediate storage for large inline values.
+    if size_of::<T>() > SEGMENT_BYTES {
+        // Ordinary segments stay within the byte budget and reuse their empty allocation.
+        // Oversized values occupy one slot per segment, so consuming one retires its allocation.
+        debug_assert_eq!(batch.len(), 1);
+        // Keep this as a tail expression to avoid intermediate storage for large inline values.
         mem::take(batch).pop_front()
     } else {
         batch.pop_front()
@@ -94,7 +113,7 @@ mod tests {
 
     fn receive<T>(buffer: &mut Buffer<T>, batch: &mut VecDeque<T>) -> T {
         if batch.is_empty() {
-            buffer.refill(batch);
+            drop(buffer.refill(batch));
         }
         pop_batch(batch)
     }
@@ -117,10 +136,39 @@ mod tests {
             assert_eq!(receive(&mut buffer, &mut batch), [value; 128]);
         }
         assert!(allocated_bytes(&buffer, &batch) <= 2 * SEGMENT_BYTES);
-        buffer.refill(&mut batch);
+        drop(buffer.refill(&mut batch));
         assert!(batch.is_empty());
     }
 
+    #[test]
+    fn ordinary_segments_stay_within_the_byte_budget_across_refills() {
+        fn check<const SIZE: usize>() {
+            let mut buffer = Buffer::new();
+            let mut batch = VecDeque::new();
+            let messages = SEGMENT_BYTES / SIZE + 2;
+            for _ in 0..2 {
+                for value in 0..messages {
+                    buffer.push([value as u8; SIZE]);
+                    for segment in std::iter::once(&buffer.writable).chain(&buffer.sealed) {
+                        assert!(segment.capacity() * SIZE <= SEGMENT_BYTES);
+                    }
+                }
+                for value in 0..messages {
+                    assert_eq!(receive(&mut buffer, &mut batch), [value as u8; SIZE]);
+                    assert!(batch.capacity() * SIZE <= SEGMENT_BYTES);
+                }
+            }
+        }
+
+        // Exercise VecDeque's initial growth and segment rounding around payload-size boundaries.
+        check::<1023>();
+        check::<1024>();
+        check::<1025>();
+        check::<{ SEGMENT_BYTES / 2 }>();
+        check::<{ SEGMENT_BYTES / 2 + 1 }>();
+        check::<SEGMENT_BYTES>();
+    }
+
     #[test]
     fn oversized_inline_values_release_the_allocation_on_the_last_receive() {
         let mut buffer = Buffer::new();

asyncband/src/mpsc/unbounded/receiver.rs

@@ -96,9 +96,12 @@ impl<T> UnboundedReceiver<T> {
         let batch = self.batch.get_mut();
         if batch.is_empty() {
             let mut state = self.shared.lock();
-            state.buffer.refill(batch);
+            let retired = state.buffer.refill(batch);
+            let disconnected = state.senders == 0;
+            drop(state);
+            drop(retired);
             if batch.is_empty() {
-                return Err(if state.senders == 0 {
+                return Err(if disconnected {
                     TryRecvError::Disconnected
                 } else {
                     TryRecvError::Empty
@@ -149,19 +152,22 @@ impl<T> UnboundedReceiver<T> {
         // Waker clone/drop callbacks can send into this channel, so run them outside the lock.
         let waker = cx.waker().clone();
         let mut state = self.shared.lock();
-        state.buffer.refill(batch);
+        let retired = state.buffer.refill(batch);
         if !batch.is_empty() {
             drop(state);
+            drop(retired);
             return Poll::Ready(Ok(pop_batch(batch)));
         }
         if state.senders == 0 {
             let old = state.recv_waker.take();
             drop(state);
+            drop(retired);
             drop(old);
             return Poll::Ready(Err(RecvError::Disconnected));
         }
         let old = state.recv_waker.replace(waker);
         drop(state);
+        drop(retired);
         drop(old);
         Poll::Pending
     }