diff --git a/asyncband/src/mpsc/unbounded/buffer.rs b/asyncband/src/mpsc/unbounded/buffer.rs index ee195a9..371bb2a 100644 --- a/asyncband/src/mpsc/unbounded/buffer.rs +++ b/asyncband/src/mpsc/unbounded/buffer.rs @@ -21,6 +21,7 @@ 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 { writable: VecDeque, @@ -28,6 +29,19 @@ pub struct Buffer { } impl Buffer { + const SEGMENT_CAPACITY: usize = if size_of::() == 0 { + usize::MAX + } else if size_of::() > SEGMENT_BYTES { + 1 + } else { + let limit = SEGMENT_BYTES / size_of::(); + // 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::>(); + pub fn new() -> Self { Self { writable: VecDeque::new(), @@ -35,41 +49,46 @@ impl Buffer { } } - fn segment_capacity() -> usize { - if size_of::() == 0 { - return usize::MAX; - } - let limit = (SEGMENT_BYTES / size_of::()).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) { + /// 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, + ) -> Option<(VecDeque, VecDeque>)> { 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::>() > 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(batch: &mut VecDeque) -> T { - if batch.len() == 1 && batch.capacity().saturating_mul(size_of::()) > 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::() > 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(buffer: &mut Buffer, batch: &mut VecDeque) -> 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() { + 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::(); + } + #[test] fn oversized_inline_values_release_the_allocation_on_the_last_receive() { let mut buffer = Buffer::new(); diff --git a/asyncband/src/mpsc/unbounded/receiver.rs b/asyncband/src/mpsc/unbounded/receiver.rs index 4b3ea76..f521cf0 100644 --- a/asyncband/src/mpsc/unbounded/receiver.rs +++ b/asyncband/src/mpsc/unbounded/receiver.rs @@ -96,9 +96,12 @@ impl UnboundedReceiver { 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 UnboundedReceiver { // 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 }