From badf60647ba062d9cad0db5899b25a24cf367095 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 13:06:49 +0800 Subject: [PATCH 01/34] feat(mpsc): add bounded permits with sequenced publication --- CHANGELOG.md | 4 + .../mpsc/bounded/{waiters.rs => capacity.rs} | 132 +++++++++- asyncband/src/mpsc/bounded/mod.rs | 184 ++++++++++--- asyncband/src/mpsc/bounded/ring.rs | 247 +++++++----------- asyncband/src/mpsc/bounded/ring_tests.rs | 212 ++++----------- asyncband/src/mpsc/mod.rs | 1 + tests-integration/tests/mpsc_test/main.rs | 3 +- .../tests/mpsc_test/reservation.rs | 210 +++++++++++++++ 8 files changed, 623 insertions(+), 370 deletions(-) rename asyncband/src/mpsc/bounded/{waiters.rs => capacity.rs} (51%) create mode 100644 tests-integration/tests/mpsc_test/reservation.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 43ffd8e0..c7a38815 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file. ## Unreleased +### New features + +* Add bounded MPSC `reserve` and `try_reserve` methods returning a `Permit`, allowing callers to wait for capacity before constructing a message; unused permits release capacity without claiming message order. + ### Bug fixes * Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. diff --git a/asyncband/src/mpsc/bounded/waiters.rs b/asyncband/src/mpsc/bounded/capacity.rs similarity index 51% rename from asyncband/src/mpsc/bounded/waiters.rs rename to asyncband/src/mpsc/bounded/capacity.rs index 3a34e4f5..1e53dfdc 100644 --- a/asyncband/src/mpsc/bounded/waiters.rs +++ b/asyncband/src/mpsc/bounded/capacity.rs @@ -16,40 +16,116 @@ // under the License. use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Waker; +use super::SEQUENCE_STEP; +use crate::internal::cache_padded::CachePadded; use crate::internal::mutex::Mutex; use crate::internal::waitlist::WaitList; use crate::internal::waitlist::WaiterId; use crate::internal::wake_all; use crate::internal::waker_batch::WakerBatch; +use crate::mpsc::TrySendError; + +const CLOSED: usize = 1; // A wake grants a retry, not a capacity permit. Keeping notified nodes until their future // consumes the notification lets cancellation pass an unused retry to the next sender. -pub struct SendWaiters { +pub struct Capacity { + claims: CachePadded, + consumed: CachePadded, + cancelled: CachePadded, + capacity: usize, waiting: AtomicBool, queue: Mutex>>, } -impl SendWaiters { - pub fn new() -> Self { +struct Claims { + next: AtomicUsize, + returned: AtomicUsize, +} + +impl Capacity { + pub fn new(capacity: usize) -> Self { Self { + claims: CachePadded::new(Claims { + next: AtomicUsize::new(0), + returned: AtomicUsize::new(0), + }), + consumed: CachePadded::new(AtomicUsize::new(0)), + cancelled: CachePadded::new(AtomicUsize::new(0)), + capacity: capacity * SEQUENCE_STEP, waiting: AtomicBool::new(false), queue: Mutex::new(WaitList::new()), } } - pub fn waiter(&self) -> SendWaiter<'_> { - SendWaiter { + pub fn try_acquire(&self) -> Result<(), TrySendError<()>> { + let mut claimed = self.claims.next.load(Ordering::Relaxed); + let mut returned = self.claims.returned.load(Ordering::Acquire); + loop { + if claimed & CLOSED != 0 { + return Err(TrySendError::Disconnected(())); + } + if claimed.wrapping_sub(returned) >= self.capacity { + returned = self + .consumed + .load(Ordering::SeqCst) + .wrapping_add(self.cancelled.load(Ordering::SeqCst)); + if claimed.wrapping_sub(returned) >= self.capacity { + // A newer return can overtake our claim snapshot. Refresh the snapshot + // before reporting Full, including a concurrent close. + let current = self.claims.next.load(Ordering::Relaxed); + if current != claimed { + claimed = current; + continue; + } + return Err(TrySendError::Full(())); + } + // Carry the acquired consumption edge with the cached progress. Producers only + // read the consumer's changing cache line when this capacity window runs out. + self.claims.returned.store(returned, Ordering::Release); + } + match self.claims.next.compare_exchange_weak( + claimed, + claimed.wrapping_add(SEQUENCE_STEP), + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => return Ok(()), + Err(actual) => claimed = actual, + } + } + } + + pub fn consume(&self, head: usize) { + // Only the consumer advances this cursor. Producers never modify its cache line. + self.consumed.store(head, Ordering::SeqCst); + self.notify_one(); + } + + pub fn cancel(&self) { + self.cancelled.fetch_add(SEQUENCE_STEP, Ordering::SeqCst); + self.notify_one(); + } + + pub fn close(&self) { + self.claims.next.fetch_or(CLOSED, Ordering::SeqCst); + self.notify_all(); + } + + pub fn waiter(&self) -> ReserveWaiter<'_> { + ReserveWaiter { waiters: self, index: None, } } - pub fn notify_one(&self) { - // The receiver publishes its head with SeqCst before checking this flag. Registration - // publishes the flag with SeqCst before rechecking capacity (which uses a SeqCst fence). + fn notify_one(&self) { + // Releasing capacity precedes this SeqCst flag check. Registration publishes the flag + // with SeqCst before rechecking the SeqCst consumed and cancelled cursors. // Either the receiver sees the registration or the sender sees the released capacity. if !self.waiting.load(Ordering::SeqCst) { return; @@ -67,7 +143,7 @@ impl SendWaiters { } } - pub fn notify_all(&self) { + fn notify_all(&self) { let mut wakers = WakerBatch::new(); { let mut queue = self.queue.lock(); @@ -82,12 +158,12 @@ impl SendWaiters { } } -pub struct SendWaiter<'a> { - waiters: &'a SendWaiters, +pub struct ReserveWaiter<'a> { + waiters: &'a Capacity, index: Option, } -impl SendWaiter<'_> { +impl ReserveWaiter<'_> { // The caller must retry sending after registration, before returning Pending. pub fn register(&mut self, waker: &Waker) { let mut new_waker = None; @@ -145,7 +221,7 @@ impl SendWaiter<'_> { } } -impl Drop for SendWaiter<'_> { +impl Drop for ReserveWaiter<'_> { fn drop(&mut self) { if let Some(index) = self.index.take() { let waker = self.remove(index); @@ -156,3 +232,33 @@ impl Drop for SendWaiter<'_> { } } } + +#[cfg(test)] +mod tests { + use super::Capacity; + use super::Ordering; + use super::SEQUENCE_STEP; + use super::TrySendError; + + #[test] + fn consumption_and_cancellation_restore_capacity_across_counter_overflow() { + let capacity = Capacity::new(3); + // Simulate prior cancellations on the final lap without allocating billions of permits. + let position = usize::MAX - 5; + capacity.claims.next.store(position, Ordering::Relaxed); + capacity.cancelled.store(position, Ordering::Relaxed); + let mut consumed = 0; + for _ in 0..3 { + for _ in 0..3 { + capacity.try_acquire().unwrap(); + } + assert_eq!(capacity.try_acquire(), Err(TrySendError::Full(()))); + consumed += SEQUENCE_STEP; + capacity.consume(consumed); + capacity.cancel(); + capacity.cancel(); + } + capacity.close(); + assert_eq!(capacity.try_acquire(), Err(TrySendError::Disconnected(()))); + } +} diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 455e7095..f91e5354 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -20,6 +20,7 @@ use std::fmt; use std::future::poll_fn; +use std::mem; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -27,17 +28,20 @@ use std::task::Context; use std::task::Poll; use std::task::ready; +use self::capacity::Capacity; use self::ring::Ring; -use self::waiters::SendWaiters; use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; -// Ring owns capacity, publication, and waiting for the head slot. SendWaiters only schedules -// retries after receiving frees capacity; a notification does not reserve a slot. +// Capacity accounts for permits and queued messages. Ring owns FIFO publication; the receiver +// alone advances its read cursor. A public reservation does not claim a position in the ring. +mod capacity; mod ring; -mod waiters; + +// The low bit marks closure; reservation and publication cursors advance in matching units. +const SEQUENCE_STEP: usize = 2; /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// @@ -53,19 +57,19 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { let state = Arc::new(Shared { buffer: Ring::new(buffer), senders: AtomicUsize::new(1), - send_waiters: SendWaiters::new(), + capacity: Capacity::new(buffer), }); let sender = BoundedSender { state: state.clone(), }; - let receiver = BoundedReceiver { state }; + let receiver = BoundedReceiver { state, head: 0 }; (sender, receiver) } struct Shared { buffer: Ring, senders: AtomicUsize, - send_waiters: SendWaiters, + capacity: Capacity, } /// The sending endpoint of a bounded mpsc channel. @@ -107,47 +111,89 @@ impl BoundedSender { /// /// Dropping a pending `send` loses its place waiting for capacity and drops `value`; a call /// that has returned `Pending` has not sent the message. Use [`Self::try_send`] when the - /// caller must retain ownership if capacity is unavailable. + /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for + /// capacity before constructing the message. pub async fn send(&self, value: T) -> Result<(), SendError> { - let value = match self.try_send(value) { - Ok(()) => return Ok(()), - Err(TrySendError::Disconnected(value)) => return Err(SendError::new(value)), - Err(TrySendError::Full(value)) => value, - }; - let mut waiter = self.state.send_waiters.waiter(); - let mut value = Some(value); + match self.reserve().await { + Ok(permit) => permit.send(value), + Err(_) => Err(SendError::new(value)), + } + } + + /// Reserves capacity for one message before constructing it. + /// + /// A successful reservation returns a [`Permit`]. Dropping the permit releases capacity; + /// [`Permit::send`] publishes a value without waiting for space. Reservations do not establish + /// message order: other producers may send while a permit is held. + /// + /// Returns `SendError(())` if the receiver has been dropped. A permit obtained earlier does + /// not keep the receiver alive; sending with it can still return the unsent value on + /// disconnect. + /// + /// # Cancel safety + /// + /// Dropping a pending reservation removes its wait registration without consuming capacity. + /// Notifications grant a retry, so a new sender may acquire capacity before a woken waiter. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, mut rx) = asyncband::mpsc::bounded(1); + /// let permit = tx.reserve().await.unwrap(); + /// let message = String::from("constructed after capacity became available"); + /// permit.send(message).unwrap(); + /// assert_eq!( + /// rx.recv().await.unwrap(), + /// "constructed after capacity became available" + /// ); + /// # } + /// ``` + pub async fn reserve(&self) -> Result, SendError<()>> { + match self.try_reserve() { + Ok(permit) => return Ok(permit), + Err(TrySendError::Disconnected(())) => return Err(SendError::new(())), + Err(TrySendError::Full(())) => {} + } + let mut waiter = self.state.capacity.waiter(); poll_fn(|cx| { - let message = value.take().expect("send polled after completion"); - let message = match self.try_send(message) { - Ok(()) => { + match self.try_reserve() { + Ok(permit) => { waiter.finish(); - return Poll::Ready(Ok(())); + return Poll::Ready(Ok(permit)); } - Err(TrySendError::Disconnected(message)) => { + Err(TrySendError::Disconnected(())) => { waiter.finish(); - return Poll::Ready(Err(SendError::new(message))); + return Poll::Ready(Err(SendError::new(()))); } - Err(TrySendError::Full(message)) => message, - }; + Err(TrySendError::Full(())) => {} + } waiter.register(cx.waker()); - match self.try_send(message) { - Ok(()) => { + match self.try_reserve() { + Ok(permit) => { waiter.finish(); - Poll::Ready(Ok(())) + Poll::Ready(Ok(permit)) } - Err(TrySendError::Disconnected(message)) => { + Err(TrySendError::Disconnected(())) => { waiter.finish(); - Poll::Ready(Err(SendError::new(message))) - } - Err(TrySendError::Full(message)) => { - value = Some(message); - Poll::Pending + Poll::Ready(Err(SendError::new(()))) } + Err(TrySendError::Full(())) => Poll::Pending, } }) .await } + /// Reserves capacity for one message without waiting. + /// + /// Returns [`TrySendError::Full`] if queued messages and outstanding permits occupy the + /// buffer, or [`TrySendError::Disconnected`] if the receiver has been dropped. + pub fn try_reserve(&self) -> Result, TrySendError<()>> { + self.state.capacity.try_acquire()?; + Ok(Permit { sender: self }) + } + /// Attempts to send a message without waiting for capacity. /// /// A full buffer returns [`TrySendError::Full`], while a dropped receiver returns @@ -169,7 +215,54 @@ impl BoundedSender { /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - self.state.buffer.try_push(value) + match self.try_reserve() { + Ok(permit) => permit + .send(value) + .map_err(|error| TrySendError::Disconnected(error.into_inner())), + Err(TrySendError::Full(())) => Err(TrySendError::Full(value)), + Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), + } + } +} + +/// Capacity reserved for one message on a bounded channel. +/// +/// Created by [`BoundedSender::reserve`] or [`BoundedSender::try_reserve`]. Holding a permit +/// reduces available capacity but does not prevent other messages from being received. Dropping +/// it without sending releases capacity and notifies a waiting sender. +#[must_use = "dropping the permit releases its reserved capacity"] +pub struct Permit<'a, T> { + sender: &'a BoundedSender, +} + +impl fmt::Debug for Permit<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Permit").finish_non_exhaustive() + } +} + +impl Permit<'_, T> { + /// Publishes a message using this reservation, without waiting for capacity. + /// + /// If the receiver has been dropped, the returned error contains the unsent value. + pub fn send(self, value: T) -> Result<(), SendError> { + // SAFETY: This permit owns one unit of capacity. No user code runs between claiming the + // position and publishing its value; the consumer returns the capacity after reading it. + let claim = match unsafe { self.sender.state.buffer.claim() } { + Ok(claim) => claim, + Err(()) => return Err(SendError::new(value)), + }; + // Publication can wake user code that panics. Transfer capacity ownership first so + // unwinding cannot return a permit for a message that is already in the ring. + mem::forget(self); + claim.publish(value); + Ok(()) + } +} + +impl Drop for Permit<'_, T> { + fn drop(&mut self) { + self.sender.state.capacity.cancel(); } } @@ -178,6 +271,7 @@ impl BoundedSender { /// Instances are created by the [`bounded`] function. pub struct BoundedReceiver { state: Arc>, + head: usize, } impl fmt::Debug for BoundedReceiver { @@ -188,21 +282,29 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - struct DrainOnDrop<'a, T>(&'a Ring); + struct DrainOnDrop<'a, T> { + ring: &'a Ring, + head: &'a mut usize, + tail: usize, + } impl Drop for DrainOnDrop<'_, T> { fn drop(&mut self) { // SAFETY: This guard lives only within the exclusive receiver's drop, after close. - unsafe { self.0.drain() }; + unsafe { self.ring.drain(self.head, self.tail) }; } } - self.state.buffer.close(); - let drain = DrainOnDrop(&self.state.buffer); + let tail = self.state.buffer.close(); + let drain = DrainOnDrop { + ring: &self.state.buffer, + head: &mut self.head, + tail, + }; // A registered waker may own a sender; release it to break that ownership cycle. let receiver_waker = self.state.buffer.take_receiver_waker(); // Complete notifications before dropping messages. Either kind of callback may panic; // the drain guard still releases buffered values if a wake or waker drop unwinds. - self.state.send_waiters.notify_all(); + self.state.capacity.close(); drop(receiver_waker); drop(drain); } @@ -245,20 +347,20 @@ impl BoundedReceiver { fn try_recv_once(&mut self) -> Poll> { // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows. - let value = if let Some(value) = ready!(unsafe { self.state.buffer.pop() }) { + let value = if let Some(value) = ready!(unsafe { self.state.buffer.pop(&mut self.head) }) { value } else if self.state.senders.load(Ordering::Acquire) == 0 { // The final sender can enqueue between the first empty observation and decrementing // the sender count, so check the queue again before reporting disconnection. // SAFETY: The exclusive receiver borrow still guarantees a single consumer. - let Some(value) = ready!(unsafe { self.state.buffer.pop() }) else { + let Some(value) = ready!(unsafe { self.state.buffer.pop(&mut self.head) }) else { return Poll::Ready(Err(TryRecvError::Disconnected)); }; value } else { return Poll::Ready(Err(TryRecvError::Empty)); }; - self.state.send_waiters.notify_one(); + self.state.capacity.consume(self.head); Poll::Ready(Ok(value)) } diff --git a/asyncband/src/mpsc/bounded/ring.rs b/asyncband/src/mpsc/bounded/ring.rs index ace5aa76..350a2eda 100644 --- a/asyncband/src/mpsc/bounded/ring.rs +++ b/asyncband/src/mpsc/bounded/ring.rs @@ -25,21 +25,19 @@ use std::sync::atomic::fence; use std::task::Poll; use std::task::Waker; +use super::SEQUENCE_STEP; use crate::internal::cache_padded::CachePadded; use crate::internal::mutex::Mutex; -use crate::mpsc::TrySendError; + +const CLOSED: usize = 1; pub struct Ring { slots: Box<[Slot]>, - head: CachePadded, tail: CachePadded, - // This flag is usually stable while the consumer advances head. Sharing head - // for notifications would make every producer track a constantly invalidated cache line. + // Publication and receiver registration synchronize independently of capacity release. receiver_waiting: CachePadded, receiver: Mutex>, - capacity: usize, - one_lap: usize, - mark_bit: usize, + mask: usize, } struct Slot { @@ -47,9 +45,9 @@ struct Slot { value: UnsafeCell>, } -// SAFETY: A successful tail CAS gives one producer exclusive access to a slot. That producer +// SAFETY: A successful tail increment gives one producer exclusive access to a slot. That producer // initializes the value before publishing the next stamp with Release ordering. The single -// consumer reads only after acquiring that stamp and publishes the following lap before reuse. +// consumer reads only after acquiring that stamp and returns capacity before reuse. unsafe impl Sync for Slot {} // The ownership transition finishes before user code can unwind, and no stored-value reference is @@ -60,122 +58,65 @@ impl std::panic::RefUnwindSafe for Slot {} impl Ring { pub fn new(capacity: usize) -> Self { assert!(capacity <= usize::MAX / 4, "mpsc capacity is too large"); - let mark_bit = (capacity + 1).next_power_of_two(); - let one_lap = mark_bit * 2; - let slots = (0..capacity) - .map(|index| Slot { - stamp: AtomicUsize::new(index), + // Physical storage is rounded up, while Capacity enforces the exact requested limit. + // A power-of-two ring keeps indexing cheap and continuous across sequence overflow. + let storage = capacity.next_power_of_two(); + let slots = (0..storage) + .map(|_| Slot { + stamp: AtomicUsize::new(CLOSED), value: UnsafeCell::new(MaybeUninit::uninit()), }) .collect(); Self { slots, - head: CachePadded::new(AtomicUsize::new(0)), tail: CachePadded::new(AtomicUsize::new(0)), receiver_waiting: CachePadded::new(AtomicBool::new(false)), receiver: Mutex::new(None), - capacity, - one_lap, - mark_bit, + mask: storage - 1, } } - pub fn try_push(&self, value: T) -> Result<(), TrySendError> { - let mut tail = self.tail.load(Ordering::Relaxed); - let mut backoff = 0; - loop { - if tail & self.mark_bit != 0 { - return Err(TrySendError::Disconnected(value)); - } - - let index = tail & (self.mark_bit - 1); - let slot = &self.slots[index]; - let stamp = slot.stamp.load(Ordering::Acquire); - if stamp == tail { - let next_tail = self.advance(tail); - match self.tail.compare_exchange_weak( - tail, - next_tail, - Ordering::SeqCst, - Ordering::Relaxed, - ) { - Ok(_) => { - // SAFETY: The successful CAS reserved this slot exclusively, and its - // matching stamp proves the consumer completed its previous lap. - unsafe { (*slot.value.get()).write(value) }; - slot.stamp.store(tail.wrapping_add(1), Ordering::Release); - // Publication precedes the wait check; registration pairs this fence - // with a second pop before the receiver is allowed to return Pending. - fence(Ordering::SeqCst); - if self.receiver_waiting.load(Ordering::Relaxed) - && self.receiver_waiting.swap(false, Ordering::Relaxed) - { - // Claim the notification before locking so concurrent publishers - // do not all queue behind the same receiver registration. - self.wake_receiver(); - } - return Ok(()); - } - Err(actual) => tail = actual, - } - } else if stamp.wrapping_add(self.one_lap) == tail.wrapping_add(1) { - fence(Ordering::SeqCst); - if self.head.load(Ordering::Relaxed).wrapping_add(self.one_lap) == tail { - return Err(TrySendError::Full(value)); - } - tail = self.tail.load(Ordering::Relaxed); - } else { - let actual = self.tail.load(Ordering::Relaxed); - if actual == tail { - // Reserved but unpublished messages also occupy capacity. In particular, a - // capacity-one queue must report Full without waiting for its producer to - // publish the slot's stamp. - fence(Ordering::SeqCst); - if self.head.load(Ordering::Relaxed).wrapping_add(self.one_lap) == tail { - return Err(TrySendError::Full(value)); - } - } - tail = actual; - } - Self::spin(&mut backoff); + /// Claims the next FIFO position. No payload is touched until `Claim::publish`. + /// + /// # Safety + /// + /// The caller must already own one capacity permit for this ring, and transfer that permit + /// to the consumer on publication. The claim must be published without invoking user code. + pub unsafe fn claim(&self) -> Result, ()> { + // A permit may predate the previous use of this slot. Carry earlier claimants' + // capacity acquires through the sequence so even an old permit observes that read. + let position = self.tail.fetch_add(SEQUENCE_STEP, Ordering::AcqRel); + if position & CLOSED != 0 { + Err(()) + } else { + Ok(Claim { + ring: self, + position, + }) } } - /// Pending means the head slot is reserved but not published. It is distinct from an empty - /// queue: later producers may already have completed their sends. + /// Pending means the head slot is claimed but not yet published. /// /// # Safety /// - /// The caller must serialize all calls to `pop` and `drain` for this queue. - pub unsafe fn pop(&self) -> Poll> { - let mut head = self.head.load(Ordering::Relaxed); - let mut backoff = 0; - loop { - let index = head & (self.mark_bit - 1); - let slot = &self.slots[index]; - let stamp = slot.stamp.load(Ordering::Acquire); - if stamp == head.wrapping_add(1) { - let next_head = self.advance(head); - // SAFETY: Acquiring the matching stamp observes initialization by the producer. - // There is one consumer, so the value is read exactly once. - let value = unsafe { (*slot.value.get()).assume_init_read() }; - slot.stamp - .store(head.wrapping_add(self.one_lap), Ordering::Release); - self.head.store(next_head, Ordering::SeqCst); - return Poll::Ready(Some(value)); - } - - if stamp == head { - fence(Ordering::SeqCst); - if self.tail.load(Ordering::Relaxed) & !self.mark_bit == head { - return Poll::Ready(None); - } - } - if backoff == 8 { - return Poll::Pending; - } - Self::spin(&mut backoff); - head = self.head.load(Ordering::Relaxed); + /// Only the exclusive consumer may call `pop` or `drain`, with its persistent head cursor. + /// Return one capacity permit after each successful pop, after the value has been read. + pub unsafe fn pop(&self, head: &mut usize) -> Poll> { + let index = (*head / SEQUENCE_STEP) & self.mask; + let slot = &self.slots[index]; + if slot.stamp.load(Ordering::Acquire) == *head { + // SAFETY: Acquiring the published stamp observes initialization. The consumer owns + // this cursor exclusively, and capacity is not returned until after reading the value. + let value = unsafe { (*slot.value.get()).assume_init_read() }; + *head = head.wrapping_add(SEQUENCE_STEP); + return Poll::Ready(Some(value)); + } + fence(Ordering::SeqCst); + if self.tail.load(Ordering::Relaxed) & !CLOSED == *head { + Poll::Ready(None) + } else { + Poll::Pending } } @@ -214,65 +155,53 @@ impl Ring { } /// Prevents subsequent sends from reserving slots. Already reserved slots still publish. - pub fn close(&self) { - self.tail.fetch_or(self.mark_bit, Ordering::SeqCst); + pub fn close(&self) -> usize { + // Failed claims may still advance tail after close. Freeze the drain boundary at the + // close operation itself; it includes every successful claim and no rejected claims. + self.tail.fetch_or(CLOSED, Ordering::SeqCst) } - /// Drops all values after closing, including values whose publication is still in progress. + /// Drops all values after closing, including short-lived claims still being published. /// /// # Safety /// - /// The queue must be closed. The caller must serialize all calls to `pop` and `drain`. - pub unsafe fn drain(&self) { + /// `tail` must be the value returned by the first close, and `head` the consumer's cursor. + pub unsafe fn drain(&self, head: &mut usize, tail: usize) { struct DrainRemaining<'a, T> { ring: &'a Ring, + head: &'a mut usize, tail: usize, } impl Drop for DrainRemaining<'_, T> { fn drop(&mut self) { - // SAFETY: The guard is scoped to the exclusive consumer's drain of a closed ring. - unsafe { self.ring.discard_until(self.tail) }; + // SAFETY: The guard owns the consumer cursor until this closed ring is drained. + unsafe { self.ring.discard_until(self.head, self.tail) }; } } - let tail = self.tail.load(Ordering::Relaxed); - debug_assert_ne!(tail & self.mark_bit, 0); + debug_assert_eq!(tail & CLOSED, 0); let remaining = DrainRemaining { ring: self, - tail: tail & !self.mark_bit, + head, + tail, }; - // SAFETY: The caller guarantees exclusive consumer access. The guard finishes draining if - // a value's destructor panics, so messages that own senders cannot retain the closed ring. - unsafe { self.discard_until(remaining.tail) }; - } - - fn advance(&self, position: usize) -> usize { - let index = position & (self.mark_bit - 1); - if index + 1 < self.capacity { - position + 1 - } else { - let lap = position & !(self.one_lap - 1); - lap.wrapping_add(self.one_lap) - } + // SAFETY: The caller guarantees exclusive consumer access. The guard finishes draining + // if a destructor panics, including messages that own senders and would retain the ring. + unsafe { self.discard_until(remaining.head, remaining.tail) }; } - // The caller must close the queue and have exclusive consumer access before discarding values. - unsafe fn discard_until(&self, tail: usize) { - let mut head = self.head.load(Ordering::Relaxed); + // The caller owns the cursor and has prevented new claims by closing the ring. + unsafe fn discard_until(&self, head: &mut usize, tail: usize) { let mut backoff = 0; - while head != tail { - let index = head & (self.mark_bit - 1); + while *head != tail { + let index = (*head / SEQUENCE_STEP) & self.mask; let slot = &self.slots[index]; - if slot.stamp.load(Ordering::Acquire) == head.wrapping_add(1) { - let next_head = self.advance(head); - // Move the head before dropping the value so unwinding cannot drop it twice. - slot.stamp - .store(head.wrapping_add(self.one_lap), Ordering::Release); - self.head.store(next_head, Ordering::SeqCst); - // SAFETY: The acquired matching stamp proves the slot contains an initialized - // value, and advancing the single-consumer head claims it exactly once. + if slot.stamp.load(Ordering::Acquire) == *head { + // Advance before running the destructor so unwinding cannot drop a value twice. + *head = head.wrapping_add(SEQUENCE_STEP); + // SAFETY: The acquired stamp proves initialization; the cursor claims the value + // exactly once, and close prevents any producer from reusing its slot. unsafe { (*slot.value.get()).assume_init_drop() }; - head = next_head; backoff = 0; } else { Self::spin(&mut backoff); @@ -288,11 +217,29 @@ impl Ring { } } -impl Drop for Ring { - fn drop(&mut self) { - self.close(); - // SAFETY: The queue is closed and its exclusive borrow rules out concurrent access. - unsafe { self.drain() }; +// Claims never escape a synchronous send. A public Permit owns capacity without claiming a +// position, so holding or forgetting one cannot leave an unpublished hole in the ring. +pub struct Claim<'a, T> { + ring: &'a Ring, + position: usize, +} + +impl Claim<'_, T> { + pub fn publish(self, value: T) { + let slot = &self.ring.slots[(self.position / SEQUENCE_STEP) & self.ring.mask]; + // SAFETY: Claiming required a capacity permit. Its acquire observes the consumer's + // completed read on the previous lap; the tail increment grants this producer exclusive + // access. + unsafe { (*slot.value.get()).write(value) }; + slot.stamp.store(self.position, Ordering::Release); + // Paired with receiver registration: either the publisher sees the wait flag or the + // receiver's second pop observes this publication before it can return Pending. + fence(Ordering::SeqCst); + if self.ring.receiver_waiting.load(Ordering::Relaxed) + && self.ring.receiver_waiting.swap(false, Ordering::Relaxed) + { + self.ring.wake_receiver(); + } } } diff --git a/asyncband/src/mpsc/bounded/ring_tests.rs b/asyncband/src/mpsc/bounded/ring_tests.rs index d6d50433..fb68e69e 100644 --- a/asyncband/src/mpsc/bounded/ring_tests.rs +++ b/asyncband/src/mpsc/bounded/ring_tests.rs @@ -15,186 +15,68 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Poll; -use std::thread; use super::Ring; -use super::TrySendError; #[test] -fn bounded_queue_preserves_capacity_and_fifo_order() { - let queue = Ring::new(3); - for value in 0..3 { - assert!(queue.try_push(value).is_ok()); - } - assert!(matches!(queue.try_push(3), Err(TrySendError::Full(3)))); - for value in 0..3 { - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); - } - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(None)); - - for value in 3..12 { - assert!(queue.try_push(value).is_ok()); - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); - } -} - -#[test] -fn bounded_queue_does_not_report_empty_behind_an_unpublished_head() { +fn receive_waits_for_the_first_claim_even_if_a_later_claim_is_published() { let queue = Ring::new(2); - - // Pause a synthetic producer after reserving and initializing slot 0, before publishing - // its stamp. Another producer can finish sending into slot 1 in the meantime. - queue.tail.store(1, Ordering::SeqCst); - let slot = &queue.slots[0]; - // SAFETY: advancing the tail reserved this initially empty slot for the synthetic producer. - unsafe { (*slot.value.get()).write(1) }; - let later_send = queue.try_push(2); - // SAFETY: This thread is the only consumer, even while a producer is unpublished. - let receive = unsafe { queue.pop() }; - - // Finish publication before asserting so even a failed assertion can safely drop the queue. - slot.stamp.store(1, Ordering::Release); - assert!(later_send.is_ok()); - assert_eq!(receive, Poll::Pending); - // SAFETY: This thread is the only consumer. + let mut head = 0; + // SAFETY: The two claims fit in the initially empty ring. Both publish before teardown. + let (first, second) = unsafe { (queue.claim().unwrap(), queue.claim().unwrap()) }; + second.publish(2); + // SAFETY: This test owns the only consumer cursor. + let pending = unsafe { queue.pop(&mut head) }; + first.publish(1); + assert_eq!(pending, Poll::Pending); + // SAFETY: No other consumer exists, and no slot will be reused by another producer. unsafe { - assert_eq!(queue.pop(), Poll::Ready(Some(1))); - assert_eq!(queue.pop(), Poll::Ready(Some(2))); - assert_eq!(queue.pop(), Poll::Ready(None)); + assert_eq!(queue.pop(&mut head), Poll::Ready(Some(1))); + assert_eq!(queue.pop(&mut head), Poll::Ready(Some(2))); + assert_eq!(queue.pop(&mut head), Poll::Ready(None)); } } #[test] -fn unpublished_reservations_count_toward_capacity() { - let queue = Arc::new(Ring::new(1)); - queue.tail.store(queue.one_lap, Ordering::SeqCst); - let (done, completed) = std::sync::mpsc::channel(); - let producer = { - let queue = queue.clone(); - thread::spawn(move || done.send(queue.try_push(2)).unwrap()) - }; - #[cfg(not(miri))] - let result = completed.recv_timeout(std::time::Duration::from_secs(10)); - // Miri reports a deadlock directly instead of relying on an interpretation-time deadline. - #[cfg(miri)] - let result = completed.recv(); - // Finish the synthetic reservation even if the other producer stalled. This lets the - // worker and the ring's destructor finish before the failure is reported. - let slot = &queue.slots[0]; - // SAFETY: Advancing the tail above exclusively reserved the initially empty slot. - unsafe { (*slot.value.get()).write(1) }; - slot.stamp.store(1, Ordering::Release); - producer.join().unwrap(); - assert!(matches!(result, Ok(Err(TrySendError::Full(2))))); - // SAFETY: Both producers have finished and this thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(1))); -} - -#[test] -fn bounded_queue_coordinates_multiple_producers() { - let queue = Arc::new(Ring::new(4)); - let producers: Vec<_> = (0..2) - .map(|producer| { - let queue = queue.clone(); - thread::spawn(move || { - for offset in 0..32 { - let mut value = producer * 32 + offset; - loop { - match queue.try_push(value) { - Ok(()) => break, - Err(TrySendError::Full(returned)) => { - value = returned; - thread::yield_now(); - } - Err(TrySendError::Disconnected(_)) => panic!("queue disconnected"), - } - } - } - }) - }) - .collect(); - - let mut values = Vec::new(); - while values.len() < 64 { - // SAFETY: Worker threads only push; this thread is the only consumer. - if let Poll::Ready(Some(value)) = unsafe { queue.pop() } { - values.push(value); - } else { - thread::yield_now(); - } - } - for producer in producers { - producer.join().unwrap(); - } - values.sort_unstable(); - assert_eq!(values, (0..64).collect::>()); +fn closed_ring_finishes_claimed_publications_and_rejects_new_claims() { + let queue = Ring::new(2); + let mut head = 0; + // SAFETY: The initially empty ring has two available slots. + let first = unsafe { queue.claim().unwrap() }; + let tail = queue.close(); + // SAFETY: The second capacity unit has not been claimed, even though close rejects it. + let rejected = unsafe { queue.claim().is_err() }; + first.publish(String::from("claimed before close")); + assert!(rejected); + // SAFETY: This test owns the only consumer cursor, and the queue is closed. + unsafe { queue.drain(&mut head, tail) }; + // SAFETY: Repeated draining must not revisit a consumed value. + unsafe { queue.drain(&mut head, tail) }; } #[test] -fn bounded_queue_discards_wrapped_values_once_after_receiver_disconnect() { - // This has no owning fields, so a buggy second drop remains observable as count == 2 - // instead of invalidating the tracker first. - struct DropSpy<'a>(&'a AtomicUsize); - - impl<'a> Drop for DropSpy<'a> { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::Relaxed); +fn publication_stamps_survive_cursor_overflow_and_non_power_of_two_capacity() { + for capacity in [1, 3, 4] { + let queue = Ring::new(capacity); + // Start on the final lap before usize overflow; the index and close bit are both zero. + let mut head = usize::MAX - (2 * queue.slots.len() - 1); + queue.tail.store(head, Ordering::Relaxed); + for lap in 0..3 { + for index in 0..capacity { + // SAFETY: The previous lap was completely drained, so these claims fit. + unsafe { queue.claim().unwrap() }.publish((lap, index)); + } + for index in 0..capacity { + // SAFETY: The only consumer owns this cursor; all reads finish before reuse. + assert_eq!( + unsafe { queue.pop(&mut head) }, + Poll::Ready(Some((lap, index))) + ); + } + // SAFETY: The only consumer owns this cursor. + assert_eq!(unsafe { queue.pop(&mut head) }, Poll::Ready(None)); } } - - // Declare this before `queue` so the counters outlive values held by the queue. - let drops = [ - AtomicUsize::new(0), - AtomicUsize::new(0), - AtomicUsize::new(0), - AtomicUsize::new(0), - ]; - let queue = Ring::new(3); - - // Positions: 0, 1, 2 (then tail wraps to 8). - for counter in &drops[..3] { - assert!(queue.try_push(DropSpy(counter)).is_ok()); - } - - // Free slot 0, then reuse it on the next lap at position 8. - // SAFETY: This thread is the only consumer. - let popped = unsafe { queue.pop() }; - assert!(matches!(popped, Poll::Ready(Some(_)))); - drop(popped); - assert_eq!(drops[0].load(Ordering::Relaxed), 1); - assert!(queue.try_push(DropSpy(&drops[3])).is_ok()); - - // The pending range is positions 1 -> 2 -> 8 -> 9, not a contiguous integer range. - assert_eq!(queue.head.load(Ordering::Relaxed), 1); - assert_eq!(queue.tail.load(Ordering::Relaxed), queue.one_lap + 1); - - queue.close(); - // SAFETY: The queue is closed and this thread is the only consumer. - unsafe { queue.drain() }; - - // `discard_until` must dispose every value exactly once, including position 8. - for (value, counter) in drops.iter().enumerate() { - assert_eq!( - counter.load(Ordering::Relaxed), - 1, - "value {value} was dropped an unexpected number of times" - ); - } - - // Queue Drop calls discard_until again; it must see head == tail and not redrop. - drop(queue); - for (value, counter) in drops.iter().enumerate() { - assert_eq!( - counter.load(Ordering::Relaxed), - 1, - "value {value} was dropped more than once" - ); - } } diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs index 040c98b2..bad23f06 100644 --- a/asyncband/src/mpsc/mod.rs +++ b/asyncband/src/mpsc/mod.rs @@ -27,6 +27,7 @@ mod unbounded; pub use self::bounded::BoundedReceiver; pub use self::bounded::BoundedSender; +pub use self::bounded::Permit; pub use self::bounded::bounded; pub use self::error::RecvError; pub use self::error::SendError; diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index 24900214..6b1268e3 100644 --- a/tests-integration/tests/mpsc_test/main.rs +++ b/tests-integration/tests/mpsc_test/main.rs @@ -31,6 +31,7 @@ use self::support::poll_with; mod backpressure; mod callbacks; mod concurrency; +mod reservation; mod support; #[test] @@ -51,7 +52,7 @@ fn unbounded_try_recv_preserves_order_and_reports_state() { #[test] fn bounded_try_send_respects_capacity_and_order() { - for capacity in [1, 4, 16] { + for capacity in [1, 3, 4, 16] { let (tx, mut rx) = mpsc::bounded(capacity); for i in 0..capacity { diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs new file mode 100644 index 00000000..df414652 --- /dev/null +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -0,0 +1,210 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cell::Cell; +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::task::Wake; +use std::task::Waker; + +use asyncband::mpsc; +use asyncband::mpsc::TryRecvError; +use asyncband::mpsc::TrySendError; +use tests_integration::poll_once; + +use super::support::WakeCounter; +use super::support::expect_ready; +use super::support::poll_with; + +#[test] +fn held_permits_consume_capacity_without_claiming_message_order() { + for capacity in [1, 3, 64] { + let (tx, mut rx) = mpsc::bounded(capacity); + let permit = tx.try_reserve().unwrap(); + for value in 1..capacity { + tx.try_send(value).unwrap(); + } + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + assert_eq!(tx.try_send(0), Err(TrySendError::Full(0))); + for value in 1..capacity { + assert_eq!(rx.try_recv(), Ok(value)); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + permit.send(0).unwrap(); + assert_eq!(rx.try_recv(), Ok(0)); + // Repeated reservation and cancellation must restore the exact original capacity. + for _ in 0..3 { + let permits: Vec<_> = (0..capacity).map(|_| tx.try_reserve().unwrap()).collect(); + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + drop(permits); + } + } +} + +#[test] +fn dropping_a_permit_wakes_a_pending_reservation() { + let (tx, mut rx) = mpsc::bounded(1); + let held = tx.try_reserve().unwrap(); + let mut waiting = Box::pin(tx.reserve()); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(waiting.as_mut(), &waker).is_pending()); + drop(held); + assert_eq!(wakes.count(), 1); + let permit = expect_ready(poll_with(waiting.as_mut(), &waker)).unwrap(); + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + permit.send(7).unwrap(); + assert_eq!(rx.try_recv(), Ok(7)); +} + +#[test] +fn receiver_drop_does_not_wait_for_held_or_forgotten_permits() { + let (tx, mut rx) = mpsc::bounded(3); + let held = tx.try_reserve().unwrap(); + std::mem::forget(tx.try_reserve().unwrap()); + tx.try_send(String::from("ready")).unwrap(); + assert_eq!(rx.try_recv().unwrap(), "ready"); + tx.try_send(String::from("discarded on close")).unwrap(); + let mut waiting = Box::pin(tx.reserve()); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(waiting.as_mut(), &waker).is_pending()); + drop(rx); + assert_eq!(wakes.count(), 1); + assert!(expect_ready(poll_with(waiting.as_mut(), &waker)).is_err()); + assert!(matches!( + tx.try_reserve(), + Err(TrySendError::Disconnected(())) + )); + assert_eq!( + held.send(String::from("unsent")).unwrap_err().into_inner(), + "unsent" + ); +} + +#[test] +fn a_permit_can_publish_send_only_payloads_from_another_thread() { + let (tx, mut rx) = mpsc::bounded(1); + let permit = tx.try_reserve().unwrap(); + std::thread::scope(|scope| { + scope + .spawn(move || permit.send(Cell::new(42)).unwrap()) + .join() + .unwrap(); + }); + assert_eq!(rx.try_recv().unwrap().get(), 42); +} + +#[test] +fn a_panicking_publication_wake_cannot_return_capacity_twice() { + struct PanicOnWake; + impl Wake for PanicOnWake { + fn wake(self: Arc) { + panic!("publication wake"); + } + } + let (tx, mut rx) = mpsc::bounded(1); + let permit = tx.try_reserve().unwrap(); + let waker = Waker::from(Arc::new(PanicOnWake)); + let mut receive = Box::pin(rx.recv()); + assert!(poll_with(receive.as_mut(), &waker).is_pending()); + assert!(catch_unwind(AssertUnwindSafe(|| permit.send(1))).is_err()); + assert_eq!(tx.try_send(2), Err(TrySendError::Full(2))); + assert_eq!(expect_ready(poll_once(receive.as_mut())), Ok(1)); + drop(receive); + tx.try_send(2).unwrap(); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn an_old_permit_observes_consumption_before_reusing_a_slot() { + let (tx, mut rx) = mpsc::bounded(2); + let old = tx.try_reserve().unwrap(); + let recycled = AtomicBool::new(false); + std::thread::scope(|scope| { + let recycled = &recycled; + let producer = scope.spawn(move || { + // Coordinate the schedule without supplying the happens-before edge that the + // channel itself must provide between the previous read and this slot's reuse. + while !recycled.load(Ordering::Relaxed) { + std::thread::yield_now(); + } + old.send(String::from("reused")).unwrap(); + }); + for value in ["first", "second"] { + tx.try_send(String::from(value)).unwrap(); + assert_eq!(rx.try_recv().unwrap(), value); + } + recycled.store(true, Ordering::Relaxed); + producer.join().unwrap(); + }); + assert_eq!(rx.try_recv().unwrap(), "reused"); +} + +#[test] +fn concurrent_cancellation_preserves_capacity_and_message_order() { + const PRODUCERS: usize = 3; + const MESSAGES: usize = if cfg!(miri) { 8 } else { 256 }; + let (tx, mut rx) = mpsc::bounded(3); + let mut out_of_order = 0; + std::thread::scope(|scope| { + let mut workers = Vec::new(); + for producer in 0..PRODUCERS { + let tx = tx.clone(); + workers.push(scope.spawn(move || { + for sequence in 0..MESSAGES { + for cancel in [true, false] { + let permit = loop { + match tx.try_reserve() { + Ok(permit) => break permit, + Err(TrySendError::Full(())) => std::thread::yield_now(), + Err(TrySendError::Disconnected(())) => panic!("receiver is alive"), + } + }; + if cancel { + drop(permit); + } else { + permit.send((producer, sequence)).unwrap(); + } + } + } + })); + } + let mut next = [0; PRODUCERS]; + let mut count = 0; + while count < PRODUCERS * MESSAGES { + match rx.try_recv() { + Ok((producer, sequence)) => { + out_of_order += usize::from(next[producer] != sequence); + next[producer] += 1; + count += 1; + } + Err(TryRecvError::Empty) => std::thread::yield_now(), + Err(TryRecvError::Disconnected) => panic!("senders are alive"), + } + } + for worker in workers { + worker.join().unwrap(); + } + }); + // Drain and join before asserting so a regression cannot strand producers on a full ring. + assert_eq!(out_of_order, 0); + let permits: Vec<_> = (0..3).map(|_| tx.try_reserve().unwrap()).collect(); + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + drop(permits); +} From 22f6c6e380f85f4dcbc5faffcfb285c6bd69531c Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 13:07:00 +0800 Subject: [PATCH 02/34] test(mpsc): benchmark reservations and inline bounded messages --- benchmarks/ecosystem/mpsc/adapters.rs | 106 ++++++++++---------- benchmarks/ecosystem/mpsc/bounded.rs | 102 +++++++++++++++++++ benchmarks/ecosystem/mpsc/mod.rs | 1 + benchmarks/ecosystem/mpsc/reservation.rs | 121 +++++++++++++++++++++++ 4 files changed, 277 insertions(+), 53 deletions(-) create mode 100644 benchmarks/ecosystem/mpsc/reservation.rs diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs index 506f1678..f9fe81b6 100644 --- a/benchmarks/ecosystem/mpsc/adapters.rs +++ b/benchmarks/ecosystem/mpsc/adapters.rs @@ -26,19 +26,19 @@ pub struct Tokio; pub struct AsyncChannel; pub struct Flume; -pub trait BoundedMpsc: Send + Sync + 'static { +pub trait BoundedMpsc: Send + Sync + 'static { type Sender: Clone + Send + Sync + 'static; type Receiver: Send + 'static; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver); - fn try_send(sender: &Self::Sender, value: usize); - fn try_recv(receiver: &mut Self::Receiver) -> usize; - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>); - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize; - fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; - fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; - fn send_blocking(sender: &Self::Sender, value: usize); - fn recv_blocking(receiver: &mut Self::Receiver) -> usize; + fn try_send(sender: &Self::Sender, value: T); + fn try_recv(receiver: &mut Self::Receiver) -> T; + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>); + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T; + fn send_async(sender: &Self::Sender, value: T) -> impl Future + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; + fn send_blocking(sender: &Self::Sender, value: T); + fn recv_blocking(receiver: &mut Self::Receiver) -> T; } pub trait UnboundedMpsc: Send + Sync + 'static { @@ -53,166 +53,166 @@ pub trait UnboundedMpsc: Send + Sync + 'static { fn recv_blocking(receiver: &mut Self::Receiver) -> T; } -impl BoundedMpsc for Asyncband { - type Receiver = asyncband::mpsc::BoundedReceiver; - type Sender = asyncband::mpsc::BoundedSender; +impl BoundedMpsc for Asyncband { + type Receiver = asyncband::mpsc::BoundedReceiver; + type Sender = asyncband::mpsc::BoundedSender; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { asyncband::mpsc::bounded(capacity) } - fn try_send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { poll_ready(sender.send(value), context).unwrap(); } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - async fn send_async(sender: &Self::Sender, value: usize) { + async fn send_async(sender: &Self::Sender, value: T) { sender.send(value).await.unwrap(); } - async fn recv_async(receiver: &mut Self::Receiver) -> usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv().await.unwrap() } - fn send_blocking(sender: &Self::Sender, value: usize) { + fn send_blocking(sender: &Self::Sender, value: T) { pollster::block_on(sender.send(value)).unwrap(); } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl BoundedMpsc for Tokio { - type Receiver = tokio::sync::mpsc::Receiver; - type Sender = tokio::sync::mpsc::Sender; +impl BoundedMpsc for Tokio { + type Receiver = tokio::sync::mpsc::Receiver; + type Sender = tokio::sync::mpsc::Sender; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { tokio::sync::mpsc::channel(capacity) } - fn try_send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { poll_ready(sender.send(value), context).unwrap(); } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - async fn send_async(sender: &Self::Sender, value: usize) { + async fn send_async(sender: &Self::Sender, value: T) { sender.send(value).await.unwrap(); } - async fn recv_async(receiver: &mut Self::Receiver) -> usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv().await.unwrap() } - fn send_blocking(sender: &Self::Sender, value: usize) { + fn send_blocking(sender: &Self::Sender, value: T) { pollster::block_on(sender.send(value)).unwrap(); } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl BoundedMpsc for AsyncChannel { - type Receiver = async_channel::Receiver; - type Sender = async_channel::Sender; +impl BoundedMpsc for AsyncChannel { + type Receiver = async_channel::Receiver; + type Sender = async_channel::Sender; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { async_channel::bounded(capacity) } - fn try_send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { poll_ready(sender.send(value), context).unwrap(); } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - async fn send_async(sender: &Self::Sender, value: usize) { + async fn send_async(sender: &Self::Sender, value: T) { sender.send(value).await.unwrap(); } - async fn recv_async(receiver: &mut Self::Receiver) -> usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv().await.unwrap() } - fn send_blocking(sender: &Self::Sender, value: usize) { + fn send_blocking(sender: &Self::Sender, value: T) { pollster::block_on(sender.send(value)).unwrap(); } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl BoundedMpsc for Flume { - type Receiver = flume::Receiver; - type Sender = flume::Sender; +impl BoundedMpsc for Flume { + type Receiver = flume::Receiver; + type Sender = flume::Sender; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { flume::bounded(capacity) } - fn try_send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { poll_ready(sender.send_async(value), context).unwrap(); } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv_async(), context).unwrap() } - async fn send_async(sender: &Self::Sender, value: usize) { + async fn send_async(sender: &Self::Sender, value: T) { sender.send_async(value).await.unwrap(); } - async fn recv_async(receiver: &mut Self::Receiver) -> usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv_async().await.unwrap() } - fn send_blocking(sender: &Self::Sender, value: usize) { + fn send_blocking(sender: &Self::Sender, value: T) { pollster::block_on(sender.send_async(value)).unwrap(); } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv_async()).unwrap() } } diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index d388cc67..42d875fc 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -15,6 +15,13 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; +use std::future::poll_fn; +use std::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + use divan::Bencher; use divan::black_box; use divan::counter::ItemsCount; @@ -119,3 +126,98 @@ fn scheduled( batch.run(); bencher.bench_local(|| batch.run()); } + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + consts = [64, 4096], + args = [1, 8], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn scheduled_inline, const CAPACITY: usize>( + bencher: Bencher, + producers: usize, +) { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .build() + .unwrap(); + let (sender, mut receiver) = C::channel(CAPACITY); + let start: Vec<_> = (0..producers) + .map(|_| Arc::new(tokio::sync::Notify::new())) + .collect(); + let stop = Arc::new(AtomicBool::new(false)); + let workers: Vec<_> = start + .iter() + .enumerate() + .map(|(producer, start)| { + let sender = sender.clone(); + let start = start.clone(); + let stop = stop.clone(); + runtime.spawn(async move { + let mut sequence = 0u64; + loop { + start.notified().await; + if stop.load(Ordering::Acquire) { + break; + } + for _ in 0..BATCH_MESSAGES / producers { + let mut value = [1; 1024]; + value[..8].copy_from_slice(&(producer as u64).to_le_bytes()); + value[8..16].copy_from_slice(&sequence.to_le_bytes()); + C::send_async(&sender, black_box(value)).await; + sequence += 1; + } + } + }) + }) + .collect(); + drop(sender); + let mut expected = vec![0u64; producers]; + let mut run = || { + runtime.block_on(async { + let first = { + let mut receive = pin!(tokio::task::unconstrained(C::recv_async(&mut receiver))); + let mut released = false; + poll_fn(|cx| { + let result = receive.as_mut().poll(cx); + if !released { + assert!(result.is_pending(), "each sample starts with an empty wait"); + released = true; + for producer in &start { + producer.notify_one(); + } + } + result + }) + .await + }; + let mut value = first; + for received in 0..BATCH_MESSAGES { + let producer = u64::from_le_bytes(value[..8].try_into().unwrap()) as usize; + let sequence = u64::from_le_bytes(value[8..16].try_into().unwrap()); + assert_eq!(sequence, expected[producer]); + expected[producer] += 1; + assert_eq!(black_box(value)[1023], 1); + if received + 1 < BATCH_MESSAGES { + value = C::recv_async(&mut receiver).await; + } + } + assert!(expected.iter().all(|count| *count == expected[0])); + }) + }; + // Reuse tasks and the channel. Include the initial empty wait, backpressure, and payload + // movement; verify per-producer order and payload integrity on every measured sample. + run(); + bencher.bench_local(run); + stop.store(true, Ordering::Release); + for producer in &start { + producer.notify_one(); + } + runtime.block_on(async { + for worker in workers { + worker.await.unwrap(); + } + }); +} diff --git a/benchmarks/ecosystem/mpsc/mod.rs b/benchmarks/ecosystem/mpsc/mod.rs index dd09282b..ed522cc3 100644 --- a/benchmarks/ecosystem/mpsc/mod.rs +++ b/benchmarks/ecosystem/mpsc/mod.rs @@ -17,5 +17,6 @@ mod adapters; mod bounded; +mod reservation; mod support; mod unbounded; diff --git a/benchmarks/ecosystem/mpsc/reservation.rs b/benchmarks/ecosystem/mpsc/reservation.rs new file mode 100644 index 00000000..b8f3a275 --- /dev/null +++ b/benchmarks/ecosystem/mpsc/reservation.rs @@ -0,0 +1,121 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::future::Future; +use std::marker::PhantomData; + +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use super::adapters::Asyncband; +use super::adapters::BoundedMpsc; +use super::adapters::Tokio; +use super::support::BATCH_MESSAGES; +use super::support::ConcurrentMpsc; +use super::support::RepeatedTasks; +use crate::support::bench_context; +use crate::support::poll_ready; + +trait Reservable: BoundedMpsc { + type Permit<'a>: Send; + fn try_reserve(sender: &Self::Sender) -> Self::Permit<'_>; + fn reserve(sender: &Self::Sender) -> impl Future> + Send; + fn publish(permit: Self::Permit<'_>, value: usize); +} + +impl Reservable for Asyncband { + type Permit<'a> = asyncband::mpsc::Permit<'a, usize>; + fn try_reserve(sender: &Self::Sender) -> Self::Permit<'_> { + sender.try_reserve().unwrap() + } + async fn reserve(sender: &Self::Sender) -> Self::Permit<'_> { + sender.reserve().await.unwrap() + } + fn publish(permit: Self::Permit<'_>, value: usize) { + permit.send(value).unwrap(); + } +} + +impl Reservable for Tokio { + type Permit<'a> = tokio::sync::mpsc::Permit<'a, usize>; + fn try_reserve(sender: &Self::Sender) -> Self::Permit<'_> { + sender.try_reserve().unwrap() + } + async fn reserve(sender: &Self::Sender) -> Self::Permit<'_> { + sender.reserve().await.unwrap() + } + fn publish(permit: Self::Permit<'_>, value: usize) { + permit.send(value); + } +} + +#[divan::bench(types = [Asyncband, Tokio])] +fn reserve_publish_receive(bencher: Bencher) { + let (sender, mut receiver) = C::channel(64); + let mut context = bench_context(); + bencher.bench_local(|| { + let permit = poll_ready(C::reserve(&sender), &mut context); + C::publish(permit, black_box(usize::MAX)); + black_box(C::try_recv(&mut receiver)) + }); +} + +#[divan::bench(types = [Asyncband, Tokio])] +fn cancel_reserved_capacity(bencher: Bencher) { + let (sender, _receiver) = C::channel(64); + bencher.bench_local(|| drop(black_box(C::try_reserve(&sender)))); +} + +struct Reserved(PhantomData); + +impl ConcurrentMpsc for Reserved { + type Sender = C::Sender; + type Receiver = C::Receiver; + fn channel() -> (Self::Sender, Self::Receiver) { + C::channel(CAPACITY) + } + fn send(sender: &Self::Sender, value: usize) { + C::publish(pollster::block_on(C::reserve(sender)), value); + } + fn recv(receiver: &mut Self::Receiver) -> usize { + C::recv_blocking(receiver) + } + async fn send_async(sender: &Self::Sender, value: usize) { + C::publish(C::reserve(sender).await, value); + } + async fn recv_async(receiver: &mut Self::Receiver) -> usize { + C::recv_async(receiver).await + } +} + +#[divan::bench( + types = [Asyncband, Tokio], + consts = [64, 4096], + args = [(1, 0), (8, 4)], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn scheduled( + bencher: Bencher, + (producers, workers): (usize, usize), +) { + let mut batch = RepeatedTasks::>::new(producers, workers); + batch.run(); + bencher.bench_local(|| batch.run()); +} From 82d6992028b2b25fd42b50ead4e41e08b7440491 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 13:07:00 +0800 Subject: [PATCH 03/34] chore: remove the unused internal atomic waker --- LICENSE | 12 - asyncband/src/internal/atomic_waker.rs | 464 ------------------------- asyncband/src/internal/mod.rs | 5 - 3 files changed, 481 deletions(-) delete mode 100644 asyncband/src/internal/atomic_waker.rs diff --git a/LICENSE b/LICENSE index 942ddcee..3bd8900a 100644 --- a/LICENSE +++ b/LICENSE @@ -377,18 +377,6 @@ the Apache-2.0 option for the incorporated portions. Asyncband does not provide the upstream crate's synchronized receive operations and simplifies the incorporated implementation accordingly. -Portions of asyncband/src/internal/atomic_waker.rs are derived from futures-rs -0.3.34 at the following exact revision and source path: - - https://github.com/rust-lang/futures-rs/blob/705e6b5c0f06535b1aac1cb1989a172b3d45be8c/futures-core/src/task/__internal/atomic_waker.rs - -futures-rs is licensed under Apache-2.0 or MIT. Apache Asyncband uses the -Apache-2.0 option for the incorporated portions. The upstream source carries -the following copyright notices: - - Copyright (c) 2016 Alex Crichton - Copyright (c) 2017 The Tokio Authors - The polling loop in asyncband/src/blocking/executor.rs is adapted from Pollster 1.0.1 at the following exact revision and source path: diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs deleted file mode 100644 index a6dcd20f..00000000 --- a/asyncband/src/internal/atomic_waker.rs +++ /dev/null @@ -1,464 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -// This file contains a state machine derived from futures-rs 0.3.34 and panic-recovery behavior -// informed by Tokio 1.53.1. -// Asyncband uses the Apache-2.0 license option for code incorporated from futures-rs. -// The incorporated code has been modified for use in Apache Asyncband. -// Upstream sources: -// https://github.com/rust-lang/futures-rs/blob/705e6b5c0f06535b1aac1cb1989a172b3d45be8c/futures-core/src/task/__internal/atomic_waker.rs -// https://github.com/tokio-rs/tokio/blob/75fef53d0a8590c2d1dbb63672aa7b7d1ef51155/tokio/src/sync/task/atomic_waker.rs - -use std::cell::UnsafeCell; -use std::panic::AssertUnwindSafe; -use std::panic::RefUnwindSafe; -use std::panic::UnwindSafe; -use std::panic::catch_unwind; -use std::panic::resume_unwind; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Waker; - -const WAITING: usize = 0; -const REGISTERING: usize = 0b01; -const WAKING: usize = 0b10; - -/// A single-registerer, multi-notifier cell for task wake-up. -/// -/// The atomic state both grants exclusive access to `waker` and records one coalesced wake request. -/// The operation that moves the state out of `WAITING` remains the only slot owner until it returns -/// the state to `WAITING`. -/// -/// * `WAITING`: the slot is unlocked and may contain a registered waker. -/// * `REGISTERING`: `register` exclusively owns the slot and no concurrent wake is pending. -/// * `WAKING`: `wake` exclusively owns the slot. A racing `register` self-wakes without touching -/// the slot. -/// * `REGISTERING | WAKING`: `register` still owns the slot and must complete a concurrent wake -/// before returning to `WAITING`. -/// -/// Valid state transitions are: -/// -/// ```text -/// register: WAITING ----------------Acquire CAS---------------> REGISTERING -/// REGISTERING ------------AcqRel CAS----------------> WAITING -/// -/// wake: WAITING ----------------AcqRel fetch_or-----------> WAKING -/// WAKING -----------------Release swap--------------> WAITING -/// -/// race: REGISTERING ------------AcqRel fetch_or-----------> REGISTERING | WAKING -/// REGISTERING | WAKING ---AcqRel swap---------------> WAITING -/// ``` -/// -/// Additional calls to `wake` while `WAKING` is set are coalesced. A wake completed before a -/// registration starts is not remembered, so callers must register before rechecking the condition -/// that determines whether to return `Pending`. -/// -/// Every transition that acquires slot ownership has an Acquire operation paired with the previous -/// owner's Release transition to `WAITING`. The Release half of `wake` also publishes the caller's -/// preceding condition update; a racing `register` acquires that publication before it returns. -pub struct AtomicWaker { - state: AtomicUsize, - waker: UnsafeCell>, -} - -// SAFETY: `state` grants exclusive access to `waker`, and losing concurrent registrations do not -// touch the slot. `Waker` itself is `Send + Sync`. -unsafe impl Sync for AtomicWaker {} - -// `Waker` callbacks may unwind, but no panic leaves a state bit owned by the unwinding operation. A -// failed clone leaves the old slot intact and completes any raced wake, while wake and drop -// callbacks run after that operation's critical section has been released. -impl RefUnwindSafe for AtomicWaker {} -impl UnwindSafe for AtomicWaker {} - -impl AtomicWaker { - #[inline] - pub const fn new() -> Self { - Self { - state: AtomicUsize::new(WAITING), - waker: UnsafeCell::new(None), - } - } - - /// Registers `waker`, replacing a previously registered task if it differs. - /// - /// Calls to this method must not overlap. It may run concurrently with any number of calls to - /// [`wake`](Self::wake). - #[inline] - pub fn register(&self, waker: &Waker) { - // ORDERING: On success, Acquire pairs with the Release operation that last returned the - // state to WAITING and transfers exclusive ownership of the waker slot to this thread. On - // failure, Acquire matters when this reads WAKING from a notifier's AcqRel fetch_or: it - // receives the condition update that preceded that wake before this method returns. - match self - .state - .compare_exchange(WAITING, REGISTERING, Ordering::Acquire, Ordering::Acquire) - .unwrap_or_else(|state| state) - { - WAITING => { - // SAFETY: changing WAITING to REGISTERING grants this thread exclusive access to - // the waker slot until the state is returned to WAITING. - unsafe { self.register_locked(waker) } - } - WAKING => { - // A concurrent wake owns the slot. Self-waking ensures that this registration is - // not lost even though it cannot replace the slot right now. - waker.wake_by_ref(); - } - state => { - // Concurrent registration violates this type's contract. Ignoring the losing - // registration preserves memory safety and lets the winner provide notification. - debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); - } - } - } - - /// Registers a waker after this thread has acquired the REGISTERING state. - /// - /// # Safety - /// - /// The caller must have changed `state` from WAITING to REGISTERING and must be the only - /// thread accessing `waker`. - #[inline] - unsafe fn register_locked(&self, waker: &Waker) { - // Avoid both cloning and dropping the common case where an executor polls the receiver - // repeatedly with the same task waker. - let needs_replacement = match unsafe { &*self.waker.get() } { - Some(current) => !current.will_wake(waker), - None => true, - }; - - let mut clone_panic = None; - let old_waker = if needs_replacement { - match catch_unwind(AssertUnwindSafe(|| waker.clone())) { - Ok(new_waker) => unsafe { (*self.waker.get()).replace(new_waker) }, - Err(payload) => { - clone_panic = Some(payload); - None - } - } - } else { - None - }; - - // ORDERING: Release publishes a newly registered waker when the CAS succeeds. If it fails, - // Acquire receives the concurrent notifier's Release publication before the wake is - // completed below. AcqRel is the weakest success ordering that permits an Acquire failure - // ordering, although its Acquire half is not otherwise relied upon on the success path. - let concurrent_wake = match self.state.compare_exchange( - REGISTERING, - WAITING, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => None, - Err(state) => { - debug_assert_eq!(state, REGISTERING | WAKING); - - // SAFETY: REGISTERING remains set, so this thread still owns the waker slot. - let registered = unsafe { (*self.waker.get()).take() }; - - // ORDERING: Acquire receives all coalesced wake publications. Release publishes - // the empty slot and makes it available to the next register or wake operation. - self.state.swap(WAITING, Ordering::AcqRel); - registered - } - }; - - if let Some(payload) = clone_panic { - // Preserve the original clone panic while still completing a wake that raced with it. - if let Some(waker) = concurrent_wake { - let _ = catch_unwind(AssertUnwindSafe(|| waker.wake())); - } - resume_unwind(payload); - } - - // User waker code runs only after the state machine is back in WAITING, so a panic cannot - // leave the cell locked. If the wake raced with a replacement, notify both tasks: the - // concurrent call may have targeted the old registration, while future progress relies on - // the new one. A panic from the superseded waker must not prevent the new task from waking. - if let Some(waker) = concurrent_wake { - if let Some(old_waker) = old_waker { - let _ = catch_unwind(AssertUnwindSafe(|| old_waker.wake())); - } - waker.wake(); - } else { - // Drop a replaced waker only after releasing the state lock. - drop(old_waker); - } - } - - /// Wakes and removes the most recently registered waker, if any. - #[inline] - pub fn wake(&self) { - if let Some(waker) = self.take() { - waker.wake(); - } - } - - /// Removes the registered waker if this call acquires the slot. A concurrent registration or - /// wake may instead take responsibility for notifying it. - #[inline] - pub fn take(&self) -> Option { - // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the - // previous owner. Release publishes the condition update that the caller performed before - // calling wake, including when a registering thread already owns the slot. - match self.state.fetch_or(WAKING, Ordering::AcqRel) { - WAITING => { - // SAFETY: changing WAITING to WAKING grants this thread exclusive access to the - // waker slot until the state is returned to WAITING. - let waker = unsafe { (*self.waker.get()).take() }; - - // ORDERING: Release publishes the emptied slot before another operation acquires - // it. The fetch_or above already performed the required Acquire operation. - let old_state = self.state.swap(WAITING, Ordering::Release); - debug_assert_eq!(old_state, WAKING); - waker - } - state => { - // The thread registering a waker observes WAKING and completes this notification, - // or another waking thread has already taken responsibility for it. - debug_assert!( - state == REGISTERING || state == REGISTERING | WAKING || state == WAKING - ); - None - } - } - } -} - -#[cfg(test)] -mod tests { - use std::ptr; - use std::sync::Arc; - use std::sync::atomic::AtomicBool; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - use std::task::RawWaker; - use std::task::RawWakerVTable; - use std::task::Wake; - - use super::*; - - struct WakeCounter(AtomicUsize); - - impl Wake for WakeCounter { - fn wake(self: Arc) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } - - #[cfg(panic = "unwind")] - fn clone_panicking_waker() -> Waker { - static VTABLE: RawWakerVTable = RawWakerVTable::new( - |_| panic!("clone failed"), - |_| unreachable!(), - |_| unreachable!(), - |_| {}, - ); - - unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &VTABLE)) } - } - - #[test] - fn wake_notifies_once() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.register(&waker); - atomic_waker.wake(); - atomic_waker.wake(); - - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn reregistering_same_task_does_not_clone_waker() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.register(&waker); - let registered_refs = Arc::strong_count(&counter); - atomic_waker.register(&waker); - - assert_eq!(Arc::strong_count(&counter), registered_refs); - } - - #[test] - fn wake_before_register_is_not_remembered() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.wake(); - atomic_waker.register(&waker); - - assert_eq!(counter.0.load(Ordering::Relaxed), 0); - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn wake_during_replacement_notifies_old_and_new_tasks() { - let old_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let old_waker = Waker::from(old_counter.clone()); - let new_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let new_waker = Waker::from(new_counter.clone()); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&old_waker); - - assert_eq!( - atomic_waker.state.compare_exchange( - WAITING, - REGISTERING, - Ordering::AcqRel, - Ordering::Acquire, - ), - Ok(WAITING) - ); - std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); - - // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching - // the slot. Calling the helper completes the interrupted registration. - unsafe { atomic_waker.register_locked(&new_waker) }; - - assert_eq!(old_counter.0.load(Ordering::Relaxed), 1); - assert_eq!(new_counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn failed_wake_synchronizes_with_next_registration() { - for _ in 0..1_000 { - let did_publish = AtomicBool::new(false); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(Waker::noop()); - - std::thread::scope(|scope| { - let wake = scope.spawn(|| { - did_publish.store(true, Ordering::Relaxed); - atomic_waker.take() - }); - - let local_waker = atomic_waker.take(); - atomic_waker.register(Waker::noop()); - - let publication_is_visible = did_publish.load(Ordering::Relaxed); - let concurrent_thread_took_waker = wake.join().unwrap().is_some(); - assert!(publication_is_visible || concurrent_thread_took_waker); - drop(local_waker); - }); - } - } - - #[cfg(panic = "unwind")] - #[test] - fn clone_panic_does_not_poison_state() { - let atomic_waker = AtomicWaker::new(); - - assert!( - catch_unwind(|| { - atomic_waker.register(&clone_panicking_waker()); - }) - .is_err() - ); - - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - atomic_waker.register(&Waker::from(counter.clone())); - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[cfg(panic = "unwind")] - #[test] - fn clone_panic_completes_concurrent_wake() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&Waker::from(counter.clone())); - - assert_eq!( - atomic_waker.state.compare_exchange( - WAITING, - REGISTERING, - Ordering::Acquire, - Ordering::Acquire, - ), - Ok(WAITING) - ); - std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); - - // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching - // the state. Calling the helper completes the interrupted registration. - assert!( - catch_unwind(|| unsafe { - atomic_waker.register_locked(&clone_panicking_waker()); - }) - .is_err() - ); - - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - - let next_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - atomic_waker.register(&Waker::from(next_counter.clone())); - atomic_waker.wake(); - assert_eq!(next_counter.0.load(Ordering::Relaxed), 1); - } - - #[cfg(panic = "unwind")] - #[test] - fn drop_panic_does_not_poison_state() { - unsafe fn clone_drop_panicker(data: *const ()) -> RawWaker { - RawWaker::new(data, &DROP_PANICKING_VTABLE) - } - - unsafe fn wake_drop_panicker(_: *const ()) {} - - unsafe fn drop_drop_panicker(data: *const ()) { - // SAFETY: the test keeps the pointed-to AtomicBool alive until every derived waker has - // been dropped. - let should_panic = unsafe { &*data.cast::() }; - if should_panic.swap(false, Ordering::Relaxed) { - panic!("drop failed"); - } - } - - static DROP_PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( - clone_drop_panicker, - wake_drop_panicker, - wake_drop_panicker, - drop_drop_panicker, - ); - - let should_panic = AtomicBool::new(true); - let old_waker = unsafe { - Waker::from_raw(RawWaker::new( - ptr::from_ref(&should_panic).cast(), - &DROP_PANICKING_VTABLE, - )) - }; - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let new_waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&old_waker); - - assert!(catch_unwind(AssertUnwindSafe(|| atomic_waker.register(&new_waker))).is_err()); - - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } -} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 045a1c2b..2bf8ac68 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -49,11 +49,6 @@ pub(crate) fn wake_all(mut wakers: impl Iterator) { } } -// MPSC owns its receiver wait protocol; the general-purpose waker currently has no production -// users. -#[cfg(test)] -pub(crate) mod atomic_waker; - #[cfg(any( feature = "barrier", feature = "broadcast", From 151e8a86e570a919c951cea7c207d7514c481149 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 21:14:12 +0800 Subject: [PATCH 04/34] refactor(mpsc): simplify bounded state under one mutex --- CHANGELOG.md | 3 +- Cargo.lock | 11 + Cargo.toml | 1 + asyncband/src/internal/cache_padded.rs | 55 --- asyncband/src/internal/mod.rs | 3 - asyncband/src/mpsc/bounded/capacity.rs | 264 ----------- asyncband/src/mpsc/bounded/mod.rs | 421 ++++++++++++------ asyncband/src/mpsc/bounded/ring.rs | 248 ----------- asyncband/src/mpsc/bounded/ring_tests.rs | 82 ---- benchmarks/Cargo.toml | 1 + benchmarks/ecosystem/mpsc/adapters.rs | 42 ++ benchmarks/ecosystem/mpsc/bounded.rs | 17 +- .../tests/mpsc_test/backpressure.rs | 4 + .../tests/mpsc_test/callbacks.rs | 78 +++- .../tests/mpsc_test/reservation.rs | 90 ++-- tests-integration/tests/mpsc_test/support.rs | 18 + 16 files changed, 479 insertions(+), 859 deletions(-) delete mode 100644 asyncband/src/internal/cache_padded.rs delete mode 100644 asyncband/src/mpsc/bounded/capacity.rs delete mode 100644 asyncband/src/mpsc/bounded/ring.rs delete mode 100644 asyncband/src/mpsc/bounded/ring_tests.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index c7a38815..7113dc6b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ All notable changes to this project will be documented in this file. ### New features -* Add bounded MPSC `reserve` and `try_reserve` methods returning a `Permit`, allowing callers to wait for capacity before constructing a message; unused permits release capacity without claiming message order. +* Add bounded MPSC `reserve` and `try_reserve` methods returning a `Permit`, allowing callers to wait for capacity before constructing a message; pending sends and reservations receive capacity in wait-queue order, and unused permits release capacity without claiming message order. ### Bug fixes @@ -17,7 +17,6 @@ All notable changes to this project will be documented in this file. ### Improvements * Finish releasing buffered bounded MPSC messages even if one message destructor panics. -* Reduce bounded MPSC contention when senders or the receiver are not waiting, improving throughput without changing capacity or cancellation semantics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. ## v0.7.2 diff --git a/Cargo.lock b/Cargo.lock index c6baa36c..78f58915 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -105,6 +105,7 @@ dependencies = [ "asyncband", "divan", "flume", + "kanal", "pollster", "tokio", "waitgroup", @@ -564,6 +565,16 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" +[[package]] +name = "kanal" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e3953adf0cd667798b396c2fa13552d6d9b3269d7dd1154c4c416442d1ff574" +dependencies = [ + "futures-core", + "lock_api", +] + [[package]] name = "libc" version = "0.2.189" diff --git a/Cargo.toml b/Cargo.toml index 05e963ce..dbcbeda2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,6 +41,7 @@ cargo_metadata = { version = "0.23.1" } clap = { version = "4.6.5" } divan = { version = "0.1.21" } flume = { version = "0.12.0", default-features = false } +kanal = { version = "0.1.1" } pollster = { version = "1.0.1" } semver = { version = "1.0.28" } serde = { version = "1.0.229" } diff --git a/asyncband/src/internal/cache_padded.rs b/asyncband/src/internal/cache_padded.rs deleted file mode 100644 index 7268030d..00000000 --- a/asyncband/src/internal/cache_padded.rs +++ /dev/null @@ -1,55 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -// Use conservative architecture estimates, not a guarantee about every CPU's cache line. -// Keep 128 bytes for large ARM/PowerPC lines and adjacent-line prefetching on x86-64, -// 256 bytes for s390x, and at least 64 bytes elsewhere. -#[cfg_attr(target_arch = "s390x", repr(align(256)))] -#[cfg_attr( - any( - target_arch = "aarch64", - target_arch = "arm64ec", - target_arch = "powerpc64", - target_arch = "x86_64", - ), - repr(align(128)) -)] -#[cfg_attr( - not(any( - target_arch = "s390x", - target_arch = "aarch64", - target_arch = "arm64ec", - target_arch = "powerpc64", - target_arch = "x86_64", - )), - repr(align(64)) -)] -pub struct CachePadded(T); - -impl CachePadded { - pub const fn new(value: T) -> Self { - Self(value) - } -} - -impl std::ops::Deref for CachePadded { - type Target = T; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 2bf8ac68..0252aa83 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -67,9 +67,6 @@ pub(crate) fn wake_all(mut wakers: impl Iterator) { #[allow(dead_code)] pub(crate) mod arena; -#[cfg(feature = "mpsc")] -pub(crate) mod cache_padded; - #[cfg(any(feature = "latch", feature = "once"))] pub(crate) mod countdown; diff --git a/asyncband/src/mpsc/bounded/capacity.rs b/asyncband/src/mpsc/bounded/capacity.rs deleted file mode 100644 index 1e53dfdc..00000000 --- a/asyncband/src/mpsc/bounded/capacity.rs +++ /dev/null @@ -1,264 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Waker; - -use super::SEQUENCE_STEP; -use crate::internal::cache_padded::CachePadded; -use crate::internal::mutex::Mutex; -use crate::internal::waitlist::WaitList; -use crate::internal::waitlist::WaiterId; -use crate::internal::wake_all; -use crate::internal::waker_batch::WakerBatch; -use crate::mpsc::TrySendError; - -const CLOSED: usize = 1; - -// A wake grants a retry, not a capacity permit. Keeping notified nodes until their future -// consumes the notification lets cancellation pass an unused retry to the next sender. -pub struct Capacity { - claims: CachePadded, - consumed: CachePadded, - cancelled: CachePadded, - capacity: usize, - waiting: AtomicBool, - queue: Mutex>>, -} - -struct Claims { - next: AtomicUsize, - returned: AtomicUsize, -} - -impl Capacity { - pub fn new(capacity: usize) -> Self { - Self { - claims: CachePadded::new(Claims { - next: AtomicUsize::new(0), - returned: AtomicUsize::new(0), - }), - consumed: CachePadded::new(AtomicUsize::new(0)), - cancelled: CachePadded::new(AtomicUsize::new(0)), - capacity: capacity * SEQUENCE_STEP, - waiting: AtomicBool::new(false), - queue: Mutex::new(WaitList::new()), - } - } - - pub fn try_acquire(&self) -> Result<(), TrySendError<()>> { - let mut claimed = self.claims.next.load(Ordering::Relaxed); - let mut returned = self.claims.returned.load(Ordering::Acquire); - loop { - if claimed & CLOSED != 0 { - return Err(TrySendError::Disconnected(())); - } - if claimed.wrapping_sub(returned) >= self.capacity { - returned = self - .consumed - .load(Ordering::SeqCst) - .wrapping_add(self.cancelled.load(Ordering::SeqCst)); - if claimed.wrapping_sub(returned) >= self.capacity { - // A newer return can overtake our claim snapshot. Refresh the snapshot - // before reporting Full, including a concurrent close. - let current = self.claims.next.load(Ordering::Relaxed); - if current != claimed { - claimed = current; - continue; - } - return Err(TrySendError::Full(())); - } - // Carry the acquired consumption edge with the cached progress. Producers only - // read the consumer's changing cache line when this capacity window runs out. - self.claims.returned.store(returned, Ordering::Release); - } - match self.claims.next.compare_exchange_weak( - claimed, - claimed.wrapping_add(SEQUENCE_STEP), - Ordering::AcqRel, - Ordering::Relaxed, - ) { - Ok(_) => return Ok(()), - Err(actual) => claimed = actual, - } - } - } - - pub fn consume(&self, head: usize) { - // Only the consumer advances this cursor. Producers never modify its cache line. - self.consumed.store(head, Ordering::SeqCst); - self.notify_one(); - } - - pub fn cancel(&self) { - self.cancelled.fetch_add(SEQUENCE_STEP, Ordering::SeqCst); - self.notify_one(); - } - - pub fn close(&self) { - self.claims.next.fetch_or(CLOSED, Ordering::SeqCst); - self.notify_all(); - } - - pub fn waiter(&self) -> ReserveWaiter<'_> { - ReserveWaiter { - waiters: self, - index: None, - } - } - - fn notify_one(&self) { - // Releasing capacity precedes this SeqCst flag check. Registration publishes the flag - // with SeqCst before rechecking the SeqCst consumed and cancelled cursors. - // Either the receiver sees the registration or the sender sees the released capacity. - if !self.waiting.load(Ordering::SeqCst) { - return; - } - let waker = { - let mut queue = self.queue.lock(); - let waker = queue - .unlink_first_waiter(|_| true) - .and_then(|(_, waker)| waker.take()); - self.waiting.store(!queue.is_empty(), Ordering::SeqCst); - waker - }; - if let Some(waker) = waker { - waker.wake(); - } - } - - fn notify_all(&self) { - let mut wakers = WakerBatch::new(); - { - let mut queue = self.queue.lock(); - while let Some((_, waker)) = queue.unlink_first_waiter(|_| true) { - if let Some(waker) = waker.take() { - wakers.push(waker); - } - } - self.waiting.store(false, Ordering::SeqCst); - } - wake_all(wakers.into_iter()); - } -} - -pub struct ReserveWaiter<'a> { - waiters: &'a Capacity, - index: Option, -} - -impl ReserveWaiter<'_> { - // The caller must retry sending after registration, before returning Pending. - pub fn register(&mut self, waker: &Waker) { - let mut new_waker = None; - loop { - let mut queue = self.waiters.queue.lock(); - if let Some(index) = self.index { - if queue - .waiter_mut(index) - .as_ref() - .is_some_and(|current| current.will_wake(waker)) - { - return; - } - } - let Some(waker) = new_waker.take() else { - // Waker callbacks may reenter the channel, including clone and drop callbacks. - drop(queue); - new_waker = Some(waker.clone()); - continue; - }; - let old_waker = if let Some(index) = self.index { - let node = queue.waiter_mut(index); - if node.is_some() { - node.replace(waker) - } else { - queue.remove_unlinked_waiter(index); - self.index = Some(queue.push_back(Some(waker))); - None - } - } else { - self.index = Some(queue.push_back(Some(waker))); - None - }; - self.waiters.waiting.store(true, Ordering::SeqCst); - drop(queue); - drop(old_waker); - return; - } - } - - pub fn finish(&mut self) { - if let Some(index) = self.index.take() { - drop(self.remove(index)); - } - } - - fn remove(&self, index: WaiterId) -> Option { - let mut queue = self.waiters.queue.lock(); - queue.unlink_waiter(index, |_| true); - let waker = queue.remove_unlinked_waiter(index); - self.waiters - .waiting - .store(!queue.is_empty(), Ordering::SeqCst); - waker - } -} - -impl Drop for ReserveWaiter<'_> { - fn drop(&mut self) { - if let Some(index) = self.index.take() { - let waker = self.remove(index); - if waker.is_none() { - self.waiters.notify_one(); - } - drop(waker); - } - } -} - -#[cfg(test)] -mod tests { - use super::Capacity; - use super::Ordering; - use super::SEQUENCE_STEP; - use super::TrySendError; - - #[test] - fn consumption_and_cancellation_restore_capacity_across_counter_overflow() { - let capacity = Capacity::new(3); - // Simulate prior cancellations on the final lap without allocating billions of permits. - let position = usize::MAX - 5; - capacity.claims.next.store(position, Ordering::Relaxed); - capacity.cancelled.store(position, Ordering::Relaxed); - let mut consumed = 0; - for _ in 0..3 { - for _ in 0..3 { - capacity.try_acquire().unwrap(); - } - assert_eq!(capacity.try_acquire(), Err(TrySendError::Full(()))); - consumed += SEQUENCE_STEP; - capacity.consume(consumed); - capacity.cancel(); - capacity.cancel(); - } - capacity.close(); - assert_eq!(capacity.try_acquire(), Err(TrySendError::Disconnected(()))); - } -} diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index f91e5354..dab1f7e5 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -18,70 +18,198 @@ //! A bounded multi-producer, single-consumer queue for sending values between asynchronous //! tasks with backpressure control. +use std::collections::VecDeque; use std::fmt; use std::future::poll_fn; use std::mem; use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; -use std::task::ready; +use std::task::Waker; -use self::capacity::Capacity; -use self::ring::Ring; use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; - -// Capacity accounts for permits and queued messages. Ring owns FIFO publication; the receiver -// alone advances its read cursor. A public reservation does not claim a position in the ring. -mod capacity; -mod ring; - -// The low bit marks closure; reservation and publication cursors advance in matching units. -const SEQUENCE_STEP: usize = 2; +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaitList; +use crate::internal::waitlist::WaiterId; +use crate::internal::wake_all; +use crate::internal::waker_batch::WakerBatch; /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// /// [`BoundedSender::send`] waits for capacity when the buffer is full. Receiving a message releases -/// one slot for a waiting sender. +/// one slot for a waiting sender. Capacity is granted in the order that pending sends and +/// reservations enter the wait queue; new senders cannot take an already granted slot. /// /// # Panics /// -/// Panics if `buffer` is zero. +/// Panics if `buffer` is zero or the preallocated message buffer exceeds the allocation size +/// limit. There is no additional channel-specific capacity limit. #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); - let state = Arc::new(Shared { - buffer: Ring::new(buffer), - senders: AtomicUsize::new(1), - capacity: Capacity::new(buffer), - }); + let state = Arc::new(Mutex::new(State { + queue: VecDeque::with_capacity(buffer), + available: buffer, + receiver_open: true, + senders: 1, + receiver_waker: None, + waiters: WaitList::new(), + })); let sender = BoundedSender { state: state.clone(), }; - let receiver = BoundedReceiver { state, head: 0 }; + let receiver = BoundedReceiver { state }; (sender, receiver) } -struct Shared { - buffer: Ring, - senders: AtomicUsize, - capacity: Capacity, +// All transitions happen under one lock. Capacity belongs to exactly one of: `available`, a +// queued message, a live Permit, or a granted waiter. Waker callbacks and payload destruction +// run after unlocking; neither a pending send nor a public Permit owns a queue position. +struct State { + queue: VecDeque, + available: usize, + receiver_open: bool, + senders: usize, + receiver_waker: Option, + waiters: WaitList, +} + +impl State { + fn acquire(&mut self) -> Result<(), TrySendError<()>> { + if !self.receiver_open { + Err(TrySendError::Disconnected(())) + } else if self.available == 0 { + Err(TrySendError::Full(())) + } else { + self.available -= 1; + Ok(()) + } + } + + fn release(&mut self) -> Option { + if let Some((_, waiter)) = self.waiters.unlink_first_waiter(|_| true) { + // Keep the detached node until its future claims or cancels this grant. + waiter.granted = true; + return waiter.waker.take(); + } + self.available += 1; + None + } + + fn pop(&mut self) -> Result<(T, Option), TryRecvError> { + if let Some(value) = self.queue.pop_front() { + Ok((value, self.release())) + } else if self.senders == 0 { + Err(TryRecvError::Disconnected) + } else { + Err(TryRecvError::Empty) + } + } +} + +struct Waiter { + granted: bool, + waker: Option, +} + +struct Reservation<'a, T> { + sender: &'a BoundedSender, + index: Option, +} + +impl<'a, T> Reservation<'a, T> { + fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { + let mut cloned_waker = None; + loop { + let mut state = self.sender.state.lock(); + if !state.receiver_open { + // Drop removes any remaining registration, including an unused grant. + return Poll::Ready(Err(SendError::new(()))); + } + if let Some(index) = self.index { + let waiter = state.waiters.waiter_mut(index); + if waiter.granted { + let waiter = state.waiters.remove_unlinked_waiter(index); + self.index = None; + let permit = Permit { + sender: Some(self.sender), + }; + drop(state); + drop(waiter); + // A clone callback may have freed capacity. Establish ownership before + // dropping the unused clone, whose destructor can also run user code. + drop(cloned_waker); + return Poll::Ready(Ok(permit)); + } + if waiter + .waker + .as_ref() + .is_some_and(|w| w.will_wake(cx.waker())) + { + return Poll::Pending; + } + if let Some(waker) = cloned_waker.take() { + let old = waiter.waker.replace(waker); + drop(state); + drop(old); + return Poll::Pending; + } + } else if state.available != 0 { + state.available -= 1; + let permit = Permit { + sender: Some(self.sender), + }; + drop(state); + drop(cloned_waker); + return Poll::Ready(Ok(permit)); + } else if let Some(waker) = cloned_waker.take() { + self.index = Some(state.waiters.push_back(Waiter { + granted: false, + waker: Some(waker), + })); + return Poll::Pending; + } + drop(state); + // Clone outside the lock, then recheck capacity and closure before registering. + cloned_waker = Some(cx.waker().clone()); + } + } +} + +impl Drop for Reservation<'_, T> { + fn drop(&mut self) { + let Some(index) = self.index else { return }; + let (waiter, wake) = { + let mut state = self.sender.state.lock(); + state.waiters.unlink_waiter(index, |_| true); + let waiter = state.waiters.remove_unlinked_waiter(index); + let wake = if waiter.granted { + state.release() + } else { + None + }; + (waiter, wake) + }; + if let Some(waker) = wake { + waker.wake(); + } + drop(waiter); + } } /// The sending endpoint of a bounded mpsc channel. /// /// Instances are created by the [`bounded`] function. pub struct BoundedSender { - state: Arc>, + state: Arc>>, } impl Clone for BoundedSender { fn clone(&self) -> Self { - self.state.senders.fetch_add(1, Ordering::Release); + self.state.lock().senders += 1; BoundedSender { state: self.state.clone(), } @@ -96,8 +224,17 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { - if self.state.senders.fetch_sub(1, Ordering::AcqRel) == 1 { - self.state.buffer.wake_receiver(); + let wake = { + let mut state = self.state.lock(); + state.senders -= 1; + if state.senders == 0 { + state.receiver_waker.take() + } else { + None + } + }; + if let Some(waker) = wake { + waker.wake(); } } } @@ -114,6 +251,24 @@ impl BoundedSender { /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for /// capacity before constructing the message. pub async fn send(&self, value: T) -> Result<(), SendError> { + // Publish directly so a ready payload does not travel through try_send's large error + // return value. Capacity and publication still share one critical section. + { + let mut state = self.state.lock(); + match state.acquire() { + Ok(()) => { + state.queue.push_back(value); + let wake = state.receiver_waker.take(); + drop(state); + if let Some(waker) = wake { + waker.wake(); + } + return Ok(()); + } + Err(TrySendError::Disconnected(())) => return Err(SendError::new(value)), + Err(TrySendError::Full(())) => {} + } + } match self.reserve().await { Ok(permit) => permit.send(value), Err(_) => Err(SendError::new(value)), @@ -132,8 +287,8 @@ impl BoundedSender { /// /// # Cancel safety /// - /// Dropping a pending reservation removes its wait registration without consuming capacity. - /// Notifications grant a retry, so a new sender may acquire capacity before a woken waiter. + /// Dropping a pending reservation loses its place in the wait queue. If capacity has already + /// been granted, it is released to the next waiter or made available to a new sender. /// /// # Examples /// @@ -151,38 +306,11 @@ impl BoundedSender { /// # } /// ``` pub async fn reserve(&self) -> Result, SendError<()>> { - match self.try_reserve() { - Ok(permit) => return Ok(permit), - Err(TrySendError::Disconnected(())) => return Err(SendError::new(())), - Err(TrySendError::Full(())) => {} - } - let mut waiter = self.state.capacity.waiter(); - poll_fn(|cx| { - match self.try_reserve() { - Ok(permit) => { - waiter.finish(); - return Poll::Ready(Ok(permit)); - } - Err(TrySendError::Disconnected(())) => { - waiter.finish(); - return Poll::Ready(Err(SendError::new(()))); - } - Err(TrySendError::Full(())) => {} - } - waiter.register(cx.waker()); - match self.try_reserve() { - Ok(permit) => { - waiter.finish(); - Poll::Ready(Ok(permit)) - } - Err(TrySendError::Disconnected(())) => { - waiter.finish(); - Poll::Ready(Err(SendError::new(()))) - } - Err(TrySendError::Full(())) => Poll::Pending, - } - }) - .await + let mut reservation = Reservation { + sender: self, + index: None, + }; + poll_fn(|cx| reservation.poll(cx)).await } /// Reserves capacity for one message without waiting. @@ -190,8 +318,8 @@ impl BoundedSender { /// Returns [`TrySendError::Full`] if queued messages and outstanding permits occupy the /// buffer, or [`TrySendError::Disconnected`] if the receiver has been dropped. pub fn try_reserve(&self) -> Result, TrySendError<()>> { - self.state.capacity.try_acquire()?; - Ok(Permit { sender: self }) + self.state.lock().acquire()?; + Ok(Permit { sender: Some(self) }) } /// Attempts to send a message without waiting for capacity. @@ -215,10 +343,17 @@ impl BoundedSender { /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - match self.try_reserve() { - Ok(permit) => permit - .send(value) - .map_err(|error| TrySendError::Disconnected(error.into_inner())), + let mut state = self.state.lock(); + match state.acquire() { + Ok(()) => { + state.queue.push_back(value); + let wake = state.receiver_waker.take(); + drop(state); + if let Some(waker) = wake { + waker.wake(); + } + Ok(()) + } Err(TrySendError::Full(())) => Err(TrySendError::Full(value)), Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), } @@ -232,7 +367,7 @@ impl BoundedSender { /// it without sending releases capacity and notifies a waiting sender. #[must_use = "dropping the permit releases its reserved capacity"] pub struct Permit<'a, T> { - sender: &'a BoundedSender, + sender: Option<&'a BoundedSender>, } impl fmt::Debug for Permit<'_, T> { @@ -245,24 +380,31 @@ impl Permit<'_, T> { /// Publishes a message using this reservation, without waiting for capacity. /// /// If the receiver has been dropped, the returned error contains the unsent value. - pub fn send(self, value: T) -> Result<(), SendError> { - // SAFETY: This permit owns one unit of capacity. No user code runs between claiming the - // position and publishing its value; the consumer returns the capacity after reading it. - let claim = match unsafe { self.sender.state.buffer.claim() } { - Ok(claim) => claim, - Err(()) => return Err(SendError::new(value)), - }; - // Publication can wake user code that panics. Transfer capacity ownership first so - // unwinding cannot return a permit for a message that is already in the ring. - mem::forget(self); - claim.publish(value); + pub fn send(mut self, value: T) -> Result<(), SendError> { + let mut state = self.sender.unwrap().state.lock(); + if !state.receiver_open { + return Err(SendError::new(value)); + } + state.queue.push_back(value); + // The queued message owns the capacity before any wake callback can panic. + self.sender = None; + let wake = state.receiver_waker.take(); + drop(state); + if let Some(waker) = wake { + waker.wake(); + } Ok(()) } } impl Drop for Permit<'_, T> { fn drop(&mut self) { - self.sender.state.capacity.cancel(); + if let Some(sender) = self.sender { + let wake = sender.state.lock().release(); + if let Some(waker) = wake { + waker.wake(); + } + } } } @@ -270,8 +412,7 @@ impl Drop for Permit<'_, T> { /// /// Instances are created by the [`bounded`] function. pub struct BoundedReceiver { - state: Arc>, - head: usize, + state: Arc>>, } impl fmt::Debug for BoundedReceiver { @@ -282,40 +423,30 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - struct DrainOnDrop<'a, T> { - ring: &'a Ring, - head: &'a mut usize, - tail: usize, - } - impl Drop for DrainOnDrop<'_, T> { - fn drop(&mut self) { - // SAFETY: This guard lives only within the exclusive receiver's drop, after close. - unsafe { self.ring.drain(self.head, self.tail) }; + let (queue, receiver_waker, wakers) = { + let mut state = self.state.lock(); + state.receiver_open = false; + let queue = mem::take(&mut state.queue); + state.available += queue.len(); + let receiver_waker = state.receiver_waker.take(); + let mut wakers = WakerBatch::new(); + while let Some((_, waiter)) = state.waiters.unlink_first_waiter(|_| true) { + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } } - } - - let tail = self.state.buffer.close(); - let drain = DrainOnDrop { - ring: &self.state.buffer, - head: &mut self.head, - tail, + (queue, receiver_waker, wakers) }; - // A registered waker may own a sender; release it to break that ownership cycle. - let receiver_waker = self.state.buffer.take_receiver_waker(); - // Complete notifications before dropping messages. Either kind of callback may panic; - // the drain guard still releases buffered values if a wake or waker drop unwinds. - self.state.capacity.close(); + // Local ownership also drains the queue if a wake or waker destructor unwinds. + wake_all(wakers.into_iter()); drop(receiver_waker); - drop(drain); + drop(queue); } } impl BoundedReceiver { /// Attempts to receive the next queued value without waiting for a new message. /// - /// A producer already publishing a queued message may delay this call until publication - /// finishes. Use [`Self::recv`] to yield asynchronously while publication is in progress. - /// /// Receiving a value frees one buffer slot. An empty channel returns [`TryRecvError::Empty`] /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has /// been dropped and all queued values have been consumed. @@ -337,31 +468,11 @@ impl BoundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - loop { - if let Poll::Ready(result) = self.try_recv_once() { - return result; - } - std::thread::yield_now(); + let (value, wake) = self.state.lock().pop()?; + if let Some(waker) = wake { + waker.wake(); } - } - - fn try_recv_once(&mut self) -> Poll> { - // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows. - let value = if let Some(value) = ready!(unsafe { self.state.buffer.pop(&mut self.head) }) { - value - } else if self.state.senders.load(Ordering::Acquire) == 0 { - // The final sender can enqueue between the first empty observation and decrementing - // the sender count, so check the queue again before reporting disconnection. - // SAFETY: The exclusive receiver borrow still guarantees a single consumer. - let Some(value) = ready!(unsafe { self.state.buffer.pop(&mut self.head) }) else { - return Poll::Ready(Err(TryRecvError::Disconnected)); - }; - value - } else { - return Poll::Ready(Err(TryRecvError::Empty)); - }; - self.state.capacity.consume(self.head); - Poll::Ready(Ok(value)) + Ok(value) } /// Waits for and receives the next value, freeing one buffer slot. @@ -398,22 +509,40 @@ impl BoundedReceiver { } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - match self.try_recv_once() { - Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)), - Poll::Ready(Err(TryRecvError::Disconnected)) => { - Poll::Ready(Err(RecvError::Disconnected)) - } - Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => { - self.state.buffer.register_receiver(cx.waker()); - - match self.try_recv_once() { - Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)), - Poll::Ready(Err(TryRecvError::Disconnected)) => { - Poll::Ready(Err(RecvError::Disconnected)) + let mut cloned_waker = None; + loop { + let mut state = self.state.lock(); + match state.pop() { + Ok((value, wake)) => { + drop(state); + if let Some(waker) = wake { + waker.wake(); } - Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => Poll::Pending, + return Poll::Ready(Ok(value)); + } + Err(TryRecvError::Disconnected) => { + let old = state.receiver_waker.take(); + drop(state); + drop(old); + return Poll::Ready(Err(RecvError::Disconnected)); } + Err(TryRecvError::Empty) => {} + } + if state + .receiver_waker + .as_ref() + .is_some_and(|w| w.will_wake(cx.waker())) + { + return Poll::Pending; + } + if let Some(waker) = cloned_waker.take() { + let old = state.receiver_waker.replace(waker); + drop(state); + drop(old); + return Poll::Pending; } + drop(state); + cloned_waker = Some(cx.waker().clone()); } } } diff --git a/asyncband/src/mpsc/bounded/ring.rs b/asyncband/src/mpsc/bounded/ring.rs deleted file mode 100644 index 350a2eda..00000000 --- a/asyncband/src/mpsc/bounded/ring.rs +++ /dev/null @@ -1,248 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::cell::UnsafeCell; -use std::hint::spin_loop; -use std::mem::MaybeUninit; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::sync::atomic::fence; -use std::task::Poll; -use std::task::Waker; - -use super::SEQUENCE_STEP; -use crate::internal::cache_padded::CachePadded; -use crate::internal::mutex::Mutex; - -const CLOSED: usize = 1; - -pub struct Ring { - slots: Box<[Slot]>, - tail: CachePadded, - // Publication and receiver registration synchronize independently of capacity release. - receiver_waiting: CachePadded, - receiver: Mutex>, - mask: usize, -} - -struct Slot { - stamp: AtomicUsize, - value: UnsafeCell>, -} - -// SAFETY: A successful tail increment gives one producer exclusive access to a slot. That producer -// initializes the value before publishing the next stamp with Release ordering. The single -// consumer reads only after acquiring that stamp and returns capacity before reuse. -unsafe impl Sync for Slot {} - -// The ownership transition finishes before user code can unwind, and no stored-value reference is -// exposed. -impl std::panic::UnwindSafe for Slot {} -impl std::panic::RefUnwindSafe for Slot {} - -impl Ring { - pub fn new(capacity: usize) -> Self { - assert!(capacity <= usize::MAX / 4, "mpsc capacity is too large"); - // Physical storage is rounded up, while Capacity enforces the exact requested limit. - // A power-of-two ring keeps indexing cheap and continuous across sequence overflow. - let storage = capacity.next_power_of_two(); - let slots = (0..storage) - .map(|_| Slot { - stamp: AtomicUsize::new(CLOSED), - value: UnsafeCell::new(MaybeUninit::uninit()), - }) - .collect(); - Self { - slots, - tail: CachePadded::new(AtomicUsize::new(0)), - receiver_waiting: CachePadded::new(AtomicBool::new(false)), - receiver: Mutex::new(None), - mask: storage - 1, - } - } - - /// Claims the next FIFO position. No payload is touched until `Claim::publish`. - /// - /// # Safety - /// - /// The caller must already own one capacity permit for this ring, and transfer that permit - /// to the consumer on publication. The claim must be published without invoking user code. - pub unsafe fn claim(&self) -> Result, ()> { - // A permit may predate the previous use of this slot. Carry earlier claimants' - // capacity acquires through the sequence so even an old permit observes that read. - let position = self.tail.fetch_add(SEQUENCE_STEP, Ordering::AcqRel); - if position & CLOSED != 0 { - Err(()) - } else { - Ok(Claim { - ring: self, - position, - }) - } - } - - /// Pending means the head slot is claimed but not yet published. - /// - /// # Safety - /// - /// Only the exclusive consumer may call `pop` or `drain`, with its persistent head cursor. - /// Return one capacity permit after each successful pop, after the value has been read. - pub unsafe fn pop(&self, head: &mut usize) -> Poll> { - let index = (*head / SEQUENCE_STEP) & self.mask; - let slot = &self.slots[index]; - if slot.stamp.load(Ordering::Acquire) == *head { - // SAFETY: Acquiring the published stamp observes initialization. The consumer owns - // this cursor exclusively, and capacity is not returned until after reading the value. - let value = unsafe { (*slot.value.get()).assume_init_read() }; - *head = head.wrapping_add(SEQUENCE_STEP); - return Poll::Ready(Some(value)); - } - fence(Ordering::SeqCst); - if self.tail.load(Ordering::Relaxed) & !CLOSED == *head { - Poll::Ready(None) - } else { - Poll::Pending - } - } - - /// Registers the exclusive receiver, which must retry `pop` before returning Pending. - pub fn register_receiver(&self, waker: &Waker) { - let mut receiver = self.receiver.lock(); - let old_waker = if receiver.as_ref().is_some_and(|old| old.will_wake(waker)) { - None - } else { - // Only the receiver registers. Producers can take the old waker while we clone, - // but cannot install a replacement. Clone/drop callbacks may send into this channel. - drop(receiver); - let waker = waker.clone(); - receiver = self.receiver.lock(); - receiver.replace(waker) - }; - self.receiver_waiting.store(true, Ordering::Relaxed); - // Paired with the publisher's fence: either it sees this flag, or the receiver's - // subsequent pop sees its stamp. Taking a waker clears the flag under the same lock, - // so it cannot erase a newer registration without also taking responsibility for it. - fence(Ordering::SeqCst); - drop(receiver); - drop(old_waker); - } - - pub fn take_receiver_waker(&self) -> Option { - let mut receiver = self.receiver.lock(); - self.receiver_waiting.store(false, Ordering::Relaxed); - receiver.take() - } - - pub fn wake_receiver(&self) { - if let Some(waker) = self.take_receiver_waker() { - waker.wake(); - } - } - - /// Prevents subsequent sends from reserving slots. Already reserved slots still publish. - pub fn close(&self) -> usize { - // Failed claims may still advance tail after close. Freeze the drain boundary at the - // close operation itself; it includes every successful claim and no rejected claims. - self.tail.fetch_or(CLOSED, Ordering::SeqCst) - } - - /// Drops all values after closing, including short-lived claims still being published. - /// - /// # Safety - /// - /// `tail` must be the value returned by the first close, and `head` the consumer's cursor. - pub unsafe fn drain(&self, head: &mut usize, tail: usize) { - struct DrainRemaining<'a, T> { - ring: &'a Ring, - head: &'a mut usize, - tail: usize, - } - impl Drop for DrainRemaining<'_, T> { - fn drop(&mut self) { - // SAFETY: The guard owns the consumer cursor until this closed ring is drained. - unsafe { self.ring.discard_until(self.head, self.tail) }; - } - } - - debug_assert_eq!(tail & CLOSED, 0); - let remaining = DrainRemaining { - ring: self, - head, - tail, - }; - // SAFETY: The caller guarantees exclusive consumer access. The guard finishes draining - // if a destructor panics, including messages that own senders and would retain the ring. - unsafe { self.discard_until(remaining.head, remaining.tail) }; - } - - // The caller owns the cursor and has prevented new claims by closing the ring. - unsafe fn discard_until(&self, head: &mut usize, tail: usize) { - let mut backoff = 0; - while *head != tail { - let index = (*head / SEQUENCE_STEP) & self.mask; - let slot = &self.slots[index]; - if slot.stamp.load(Ordering::Acquire) == *head { - // Advance before running the destructor so unwinding cannot drop a value twice. - *head = head.wrapping_add(SEQUENCE_STEP); - // SAFETY: The acquired stamp proves initialization; the cursor claims the value - // exactly once, and close prevents any producer from reusing its slot. - unsafe { (*slot.value.get()).assume_init_drop() }; - backoff = 0; - } else { - Self::spin(&mut backoff); - } - } - } - - fn spin(step: &mut u32) { - for _ in 0..(*step).min(6).pow(2) { - spin_loop(); - } - *step = (*step).saturating_add(1); - } -} - -// Claims never escape a synchronous send. A public Permit owns capacity without claiming a -// position, so holding or forgetting one cannot leave an unpublished hole in the ring. -pub struct Claim<'a, T> { - ring: &'a Ring, - position: usize, -} - -impl Claim<'_, T> { - pub fn publish(self, value: T) { - let slot = &self.ring.slots[(self.position / SEQUENCE_STEP) & self.ring.mask]; - // SAFETY: Claiming required a capacity permit. Its acquire observes the consumer's - // completed read on the previous lap; the tail increment grants this producer exclusive - // access. - unsafe { (*slot.value.get()).write(value) }; - slot.stamp.store(self.position, Ordering::Release); - // Paired with receiver registration: either the publisher sees the wait flag or the - // receiver's second pop observes this publication before it can return Pending. - fence(Ordering::SeqCst); - if self.ring.receiver_waiting.load(Ordering::Relaxed) - && self.ring.receiver_waiting.swap(false, Ordering::Relaxed) - { - self.ring.wake_receiver(); - } - } -} - -#[cfg(test)] -#[path = "ring_tests.rs"] -mod tests; diff --git a/asyncband/src/mpsc/bounded/ring_tests.rs b/asyncband/src/mpsc/bounded/ring_tests.rs deleted file mode 100644 index fb68e69e..00000000 --- a/asyncband/src/mpsc/bounded/ring_tests.rs +++ /dev/null @@ -1,82 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::atomic::Ordering; -use std::task::Poll; - -use super::Ring; - -#[test] -fn receive_waits_for_the_first_claim_even_if_a_later_claim_is_published() { - let queue = Ring::new(2); - let mut head = 0; - // SAFETY: The two claims fit in the initially empty ring. Both publish before teardown. - let (first, second) = unsafe { (queue.claim().unwrap(), queue.claim().unwrap()) }; - second.publish(2); - // SAFETY: This test owns the only consumer cursor. - let pending = unsafe { queue.pop(&mut head) }; - first.publish(1); - assert_eq!(pending, Poll::Pending); - // SAFETY: No other consumer exists, and no slot will be reused by another producer. - unsafe { - assert_eq!(queue.pop(&mut head), Poll::Ready(Some(1))); - assert_eq!(queue.pop(&mut head), Poll::Ready(Some(2))); - assert_eq!(queue.pop(&mut head), Poll::Ready(None)); - } -} - -#[test] -fn closed_ring_finishes_claimed_publications_and_rejects_new_claims() { - let queue = Ring::new(2); - let mut head = 0; - // SAFETY: The initially empty ring has two available slots. - let first = unsafe { queue.claim().unwrap() }; - let tail = queue.close(); - // SAFETY: The second capacity unit has not been claimed, even though close rejects it. - let rejected = unsafe { queue.claim().is_err() }; - first.publish(String::from("claimed before close")); - assert!(rejected); - // SAFETY: This test owns the only consumer cursor, and the queue is closed. - unsafe { queue.drain(&mut head, tail) }; - // SAFETY: Repeated draining must not revisit a consumed value. - unsafe { queue.drain(&mut head, tail) }; -} - -#[test] -fn publication_stamps_survive_cursor_overflow_and_non_power_of_two_capacity() { - for capacity in [1, 3, 4] { - let queue = Ring::new(capacity); - // Start on the final lap before usize overflow; the index and close bit are both zero. - let mut head = usize::MAX - (2 * queue.slots.len() - 1); - queue.tail.store(head, Ordering::Relaxed); - for lap in 0..3 { - for index in 0..capacity { - // SAFETY: The previous lap was completely drained, so these claims fit. - unsafe { queue.claim().unwrap() }.publish((lap, index)); - } - for index in 0..capacity { - // SAFETY: The only consumer owns this cursor; all reads finish before reuse. - assert_eq!( - unsafe { queue.pop(&mut head) }, - Poll::Ready(Some((lap, index))) - ); - } - // SAFETY: The only consumer owns this cursor. - assert_eq!(unsafe { queue.pop(&mut head) }, Poll::Ready(None)); - } - } -} diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index 19e5ac33..d7e6e467 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -49,6 +49,7 @@ asyncband = { workspace = true, features = [ ] } divan = { workspace = true } flume = { workspace = true, features = ["async"] } +kanal = { workspace = true } pollster = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "sync"] } waitgroup = { workspace = true } diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs index f9fe81b6..8dd61b7f 100644 --- a/benchmarks/ecosystem/mpsc/adapters.rs +++ b/benchmarks/ecosystem/mpsc/adapters.rs @@ -25,6 +25,7 @@ pub struct Asyncband; pub struct Tokio; pub struct AsyncChannel; pub struct Flume; +pub struct Kanal; pub trait BoundedMpsc: Send + Sync + 'static { type Sender: Clone + Send + Sync + 'static; @@ -217,6 +218,47 @@ impl BoundedMpsc for Flume { } } +impl BoundedMpsc for Kanal { + type Receiver = kanal::AsyncReceiver; + type Sender = kanal::AsyncSender; + + fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { + kanal::bounded_async(capacity) + } + + fn try_send(sender: &Self::Sender, value: T) { + assert!(sender.try_send(value).unwrap()); + } + + fn try_recv(receiver: &mut Self::Receiver) -> T { + receiver.try_recv().unwrap().unwrap() + } + + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { + poll_ready(sender.send(value), context).unwrap(); + } + + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { + poll_ready(receiver.recv(), context).unwrap() + } + + async fn send_async(sender: &Self::Sender, value: T) { + sender.send(value).await.unwrap(); + } + + async fn recv_async(receiver: &mut Self::Receiver) -> T { + receiver.recv().await.unwrap() + } + + fn send_blocking(sender: &Self::Sender, value: T) { + pollster::block_on(sender.send(value)).unwrap(); + } + + fn recv_blocking(receiver: &mut Self::Receiver) -> T { + pollster::block_on(receiver.recv()).unwrap() + } +} + impl UnboundedMpsc for Asyncband { type Receiver = asyncband::mpsc::UnboundedReceiver; type Sender = asyncband::mpsc::UnboundedSender; diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index 42d875fc..dc595a89 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -30,6 +30,7 @@ use super::adapters::AsyncChannel; use super::adapters::Asyncband; use super::adapters::BoundedMpsc; use super::adapters::Flume; +use super::adapters::Kanal; use super::adapters::Tokio; use super::support::BATCH_MESSAGES; use super::support::BOUNDED_CAPACITY; @@ -40,7 +41,7 @@ use super::support::RepeatedBatch; use super::support::RepeatedTasks; use crate::support::bench_context; -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal])] fn try_round_trip(bencher: Bencher) { let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); @@ -50,7 +51,7 @@ fn try_round_trip(bencher: Bencher) { }); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal])] fn ready_round_trip(bencher: Bencher) { let mut context = bench_context(); let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); @@ -62,7 +63,7 @@ fn ready_round_trip(bencher: Bencher) { } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], args = PRODUCER_COUNTS, sample_count = 20, sample_size = 1, @@ -75,7 +76,7 @@ fn concurrent(bencher: Bencher, producer_count: usize) { } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], args = PRODUCER_COUNTS, sample_count = 50, sample_size = 1, @@ -87,14 +88,14 @@ fn sustained(bencher: Bencher, producer_count: usize) { bencher.bench_local(|| batch.run()); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal])] fn clone_drop_sender(bencher: Bencher) { let (sender, _receiver) = C::channel(BOUNDED_CAPACITY); bencher.bench_local(|| drop(black_box(sender.clone()))); } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], consts = [1, 4096], args = PRODUCER_COUNTS, sample_count = 50, @@ -111,7 +112,7 @@ fn sustained_capacity( } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], consts = [1, 64, 4096], args = [(1, 0), (4, 0), (1, 4), (4, 4), (8, 4)], sample_count = 50, @@ -128,7 +129,7 @@ fn scheduled( } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume], + types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], consts = [64, 4096], args = [1, 8], sample_count = 50, diff --git a/tests-integration/tests/mpsc_test/backpressure.rs b/tests-integration/tests/mpsc_test/backpressure.rs index ba21cfc0..2db99193 100644 --- a/tests-integration/tests/mpsc_test/backpressure.rs +++ b/tests-integration/tests/mpsc_test/backpressure.rs @@ -67,6 +67,10 @@ fn cancelling_a_sender_preserves_capacity_and_notifies_the_next_waiter() { assert_eq!(second_wakes.count(), 0); } drop(first); + // Even a granted send must not publish its value until it is polled to completion. + if cancel_after_notification { + assert_eq!(rx.try_recv(), Err(mpsc::TryRecvError::Empty)); + } if !cancel_after_notification { assert_eq!(first_wakes.count(), 0); assert_eq!(second_wakes.count(), 0); diff --git a/tests-integration/tests/mpsc_test/callbacks.rs b/tests-integration/tests/mpsc_test/callbacks.rs index 2f654fcd..3f1bed5e 100644 --- a/tests-integration/tests/mpsc_test/callbacks.rs +++ b/tests-integration/tests/mpsc_test/callbacks.rs @@ -35,6 +35,7 @@ use super::support::assert_completes_without_deadlock; use super::support::expect_ready; use super::support::poll_with; use super::support::waker_on_clone; +use super::support::waker_on_drop; struct HoldSender { _sender: S, @@ -155,35 +156,68 @@ fn wake_callbacks_can_send_into_the_same_channel() { } #[test] -fn unbounded_replaced_and_disconnected_wakers_can_send() { - struct SendOnDrop { - sender: mpsc::UnboundedSender, - disconnected: bool, - drops: Arc, - } - - // The final waker drop must run a callback, even though waking itself does nothing. - #[allow(clippy::manual_noop_waker)] - impl Wake for SendOnDrop { - fn wake(self: Arc) {} - } +fn bounded_waiter_waker_replacement_and_cancellation_can_reenter() { + assert_completes_without_deadlock(|| { + for replace in [false, true] { + let (tx, mut rx) = mpsc::bounded(1); + tx.try_send(0).unwrap(); + let drops = Arc::new(AtomicUsize::new(0)); + let waker = waker_on_drop({ + let tx = tx.clone(); + let drops = drops.clone(); + move || { + assert_eq!(tx.try_send(9), Err(mpsc::TrySendError::Full(9))); + drops.fetch_add(1, Ordering::Relaxed); + } + }); + let mut send = Box::pin(tx.send(1)); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + drop(waker); + if replace { + assert!(poll_once(send.as_mut()).is_pending()); + assert_eq!(drops.load(Ordering::Relaxed), 1); + } + drop(send); + assert_eq!(drops.load(Ordering::Relaxed), 1); + assert_eq!(rx.try_recv(), Ok(0)); + tx.try_send(2).unwrap(); + assert_eq!(rx.try_recv(), Ok(2)); + } + }); +} - impl Drop for SendOnDrop { - fn drop(&mut self) { - assert_eq!(self.sender.send(7).is_err(), self.disconnected); - self.drops.fetch_add(1, Ordering::Relaxed); +#[test] +fn bounded_receiver_waker_replacement_can_send() { + assert_completes_without_deadlock(|| { + let (tx, mut rx) = mpsc::bounded(1); + let waker = waker_on_drop(move || tx.try_send(7).unwrap()); + assert!(poll_with(Box::pin(rx.recv()).as_mut(), &waker).is_pending()); + drop(waker); + let (waker, wakes) = WakeCounter::new(); + let poll = poll_with(Box::pin(rx.recv()).as_mut(), &waker); + if poll.is_pending() { + assert!(wakes.count() > 0); + assert_eq!(rx.try_recv(), Ok(7)); + } else { + assert_eq!(poll, Poll::Ready(Ok(7))); } - } + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); + }); +} +#[test] +fn unbounded_replaced_and_disconnected_wakers_can_send() { assert_completes_without_deadlock(|| { for disconnected in [false, true] { let (tx, mut rx) = mpsc::unbounded(); let drops = Arc::new(AtomicUsize::new(0)); - let waker = Waker::from(Arc::new(SendOnDrop { - sender: tx, - disconnected, - drops: drops.clone(), - })); + let waker = waker_on_drop({ + let drops = drops.clone(); + move || { + assert_eq!(tx.send(7).is_err(), disconnected); + drops.fetch_add(1, Ordering::Relaxed); + } + }); assert!( Box::pin(rx.recv()) .as_mut() diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index df414652..de3a5eb3 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -19,8 +19,6 @@ use std::cell::Cell; use std::panic::AssertUnwindSafe; use std::panic::catch_unwind; use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; use std::task::Wake; use std::task::Waker; @@ -59,7 +57,27 @@ fn held_permits_consume_capacity_without_claiming_message_order() { } #[test] -fn dropping_a_permit_wakes_a_pending_reservation() { +fn zero_sized_messages_support_the_full_capacity_range() { + for capacity in [usize::MAX / 4 + 1, usize::MAX / 2 + 1, usize::MAX] { + let (tx, mut rx) = mpsc::bounded::<()>(capacity); + let permit = tx.try_reserve().unwrap(); + tx.try_send(()).unwrap(); + assert_eq!(rx.try_recv(), Ok(())); + drop(permit); + tx.try_reserve().unwrap().send(()).unwrap(); + // Closing restores buffered capacity before outstanding permits are dropped. + let held = tx.try_reserve().unwrap(); + drop(rx); + drop(held); + assert!(matches!( + tx.try_reserve(), + Err(TrySendError::Disconnected(())) + )); + } +} + +#[test] +fn released_capacity_is_granted_to_the_oldest_waiter() { let (tx, mut rx) = mpsc::bounded(1); let held = tx.try_reserve().unwrap(); let mut waiting = Box::pin(tx.reserve()); @@ -67,12 +85,51 @@ fn dropping_a_permit_wakes_a_pending_reservation() { assert!(poll_with(waiting.as_mut(), &waker).is_pending()); drop(held); assert_eq!(wakes.count(), 1); + // The waiting future owns the released slot even before the executor polls it again. + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + assert_eq!(tx.try_send(9), Err(TrySendError::Full(9))); let permit = expect_ready(poll_with(waiting.as_mut(), &waker)).unwrap(); assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); permit.send(7).unwrap(); assert_eq!(rx.try_recv(), Ok(7)); } +#[test] +fn cancelling_a_granted_reservation_passes_capacity_to_a_waiting_send() { + let (tx, mut rx) = mpsc::bounded(1); + let held = tx.try_reserve().unwrap(); + let mut reservation = Box::pin(tx.reserve()); + let mut send = Box::pin(tx.send(7)); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_once(reservation.as_mut()).is_pending()); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + drop(held); + assert_eq!(wakes.count(), 0); + drop(reservation); + assert_eq!(wakes.count(), 1); + assert_eq!(tx.try_send(9), Err(TrySendError::Full(9))); + assert_eq!(expect_ready(poll_once(send.as_mut())), Ok(())); + assert_eq!(rx.try_recv(), Ok(7)); + let held = tx.try_reserve().unwrap(); + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + drop(held); +} + +#[test] +fn closing_after_a_grant_returns_the_unsent_message() { + let (tx, mut rx) = mpsc::bounded(1); + tx.try_send(String::from("queued")).unwrap(); + let mut send = Box::pin(tx.send(String::from("unsent"))); + let mut reservation = Box::pin(tx.reserve()); + assert!(poll_once(send.as_mut()).is_pending()); + assert!(poll_once(reservation.as_mut()).is_pending()); + assert_eq!(rx.try_recv().unwrap(), "queued"); + drop(rx); + let error = expect_ready(poll_once(send.as_mut())).unwrap_err(); + assert_eq!(error.into_inner(), "unsent"); + assert!(expect_ready(poll_once(reservation.as_mut())).is_err()); +} + #[test] fn receiver_drop_does_not_wait_for_held_or_forgotten_permits() { let (tx, mut rx) = mpsc::bounded(3); @@ -131,31 +188,6 @@ fn a_panicking_publication_wake_cannot_return_capacity_twice() { assert_eq!(rx.try_recv(), Ok(2)); } -#[test] -fn an_old_permit_observes_consumption_before_reusing_a_slot() { - let (tx, mut rx) = mpsc::bounded(2); - let old = tx.try_reserve().unwrap(); - let recycled = AtomicBool::new(false); - std::thread::scope(|scope| { - let recycled = &recycled; - let producer = scope.spawn(move || { - // Coordinate the schedule without supplying the happens-before edge that the - // channel itself must provide between the previous read and this slot's reuse. - while !recycled.load(Ordering::Relaxed) { - std::thread::yield_now(); - } - old.send(String::from("reused")).unwrap(); - }); - for value in ["first", "second"] { - tx.try_send(String::from(value)).unwrap(); - assert_eq!(rx.try_recv().unwrap(), value); - } - recycled.store(true, Ordering::Relaxed); - producer.join().unwrap(); - }); - assert_eq!(rx.try_recv().unwrap(), "reused"); -} - #[test] fn concurrent_cancellation_preserves_capacity_and_message_order() { const PRODUCERS: usize = 3; @@ -202,7 +234,7 @@ fn concurrent_cancellation_preserves_capacity_and_message_order() { worker.join().unwrap(); } }); - // Drain and join before asserting so a regression cannot strand producers on a full ring. + // Drain and join before asserting so a regression cannot strand producers on a full channel. assert_eq!(out_of_order, 0); let permits: Vec<_> = (0..3).map(|_| tx.try_reserve().unwrap()).collect(); assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); diff --git a/tests-integration/tests/mpsc_test/support.rs b/tests-integration/tests/mpsc_test/support.rs index 360cb0ab..fdeb8367 100644 --- a/tests-integration/tests/mpsc_test/support.rs +++ b/tests-integration/tests/mpsc_test/support.rs @@ -63,6 +63,24 @@ pub fn poll_with(future: Pin<&mut F>, waker: &Waker) -> Poll Waker { + struct OnDrop(Box); + + // Only destruction runs the callback; waking consumes the reference as usual. + #[allow(clippy::manual_noop_waker)] + impl Wake for OnDrop { + fn wake(self: Arc) {} + } + + impl Drop for OnDrop { + fn drop(&mut self) { + (self.0)(); + } + } + + Waker::from(Arc::new(OnDrop(Box::new(callback)))) +} + // RawWaker is needed only to exercise clone callbacks, which the safe Wake trait cannot override. pub fn waker_on_clone(callback: impl Fn() + Send + Sync + 'static) -> Waker { struct OnClone(Box); From d085334c32114f43b13ca37959ff112c87a340aa Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 08:57:50 +0800 Subject: [PATCH 05/34] refactor(mpsc): separate bounded capacity and ring publication Reserve capacity independently from the FIFO ticket. Claim a ticket only inside synchronous send, then initialize and publish the slot without holding the capacity wait queue lock. Keep publication and close ownership in a per-slot state. Cover cursor wrap, paused publishers, held permits, and payload cleanup with native and Miri tests. Restore the unmodified AtomicWaker with its attribution. --- LICENSE | 12 + asyncband/src/internal/atomic_waker.rs | 464 ++++++++++++++++++ asyncband/src/internal/cache_padded.rs | 55 +++ asyncband/src/internal/mod.rs | 6 + asyncband/src/mpsc/bounded/buffer.rs | 263 ++++++++++ asyncband/src/mpsc/bounded/buffer_tests.rs | 221 +++++++++ asyncband/src/mpsc/bounded/mod.rs | 364 +++++++------- .../tests/mpsc_test/reservation.rs | 25 + 8 files changed, 1241 insertions(+), 169 deletions(-) create mode 100644 asyncband/src/internal/atomic_waker.rs create mode 100644 asyncband/src/internal/cache_padded.rs create mode 100644 asyncband/src/mpsc/bounded/buffer.rs create mode 100644 asyncband/src/mpsc/bounded/buffer_tests.rs diff --git a/LICENSE b/LICENSE index 3bd8900a..942ddcee 100644 --- a/LICENSE +++ b/LICENSE @@ -377,6 +377,18 @@ the Apache-2.0 option for the incorporated portions. Asyncband does not provide the upstream crate's synchronized receive operations and simplifies the incorporated implementation accordingly. +Portions of asyncband/src/internal/atomic_waker.rs are derived from futures-rs +0.3.34 at the following exact revision and source path: + + https://github.com/rust-lang/futures-rs/blob/705e6b5c0f06535b1aac1cb1989a172b3d45be8c/futures-core/src/task/__internal/atomic_waker.rs + +futures-rs is licensed under Apache-2.0 or MIT. Apache Asyncband uses the +Apache-2.0 option for the incorporated portions. The upstream source carries +the following copyright notices: + + Copyright (c) 2016 Alex Crichton + Copyright (c) 2017 The Tokio Authors + The polling loop in asyncband/src/blocking/executor.rs is adapted from Pollster 1.0.1 at the following exact revision and source path: diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs new file mode 100644 index 00000000..a6dcd20f --- /dev/null +++ b/asyncband/src/internal/atomic_waker.rs @@ -0,0 +1,464 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// This file contains a state machine derived from futures-rs 0.3.34 and panic-recovery behavior +// informed by Tokio 1.53.1. +// Asyncband uses the Apache-2.0 license option for code incorporated from futures-rs. +// The incorporated code has been modified for use in Apache Asyncband. +// Upstream sources: +// https://github.com/rust-lang/futures-rs/blob/705e6b5c0f06535b1aac1cb1989a172b3d45be8c/futures-core/src/task/__internal/atomic_waker.rs +// https://github.com/tokio-rs/tokio/blob/75fef53d0a8590c2d1dbb63672aa7b7d1ef51155/tokio/src/sync/task/atomic_waker.rs + +use std::cell::UnsafeCell; +use std::panic::AssertUnwindSafe; +use std::panic::RefUnwindSafe; +use std::panic::UnwindSafe; +use std::panic::catch_unwind; +use std::panic::resume_unwind; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Waker; + +const WAITING: usize = 0; +const REGISTERING: usize = 0b01; +const WAKING: usize = 0b10; + +/// A single-registerer, multi-notifier cell for task wake-up. +/// +/// The atomic state both grants exclusive access to `waker` and records one coalesced wake request. +/// The operation that moves the state out of `WAITING` remains the only slot owner until it returns +/// the state to `WAITING`. +/// +/// * `WAITING`: the slot is unlocked and may contain a registered waker. +/// * `REGISTERING`: `register` exclusively owns the slot and no concurrent wake is pending. +/// * `WAKING`: `wake` exclusively owns the slot. A racing `register` self-wakes without touching +/// the slot. +/// * `REGISTERING | WAKING`: `register` still owns the slot and must complete a concurrent wake +/// before returning to `WAITING`. +/// +/// Valid state transitions are: +/// +/// ```text +/// register: WAITING ----------------Acquire CAS---------------> REGISTERING +/// REGISTERING ------------AcqRel CAS----------------> WAITING +/// +/// wake: WAITING ----------------AcqRel fetch_or-----------> WAKING +/// WAKING -----------------Release swap--------------> WAITING +/// +/// race: REGISTERING ------------AcqRel fetch_or-----------> REGISTERING | WAKING +/// REGISTERING | WAKING ---AcqRel swap---------------> WAITING +/// ``` +/// +/// Additional calls to `wake` while `WAKING` is set are coalesced. A wake completed before a +/// registration starts is not remembered, so callers must register before rechecking the condition +/// that determines whether to return `Pending`. +/// +/// Every transition that acquires slot ownership has an Acquire operation paired with the previous +/// owner's Release transition to `WAITING`. The Release half of `wake` also publishes the caller's +/// preceding condition update; a racing `register` acquires that publication before it returns. +pub struct AtomicWaker { + state: AtomicUsize, + waker: UnsafeCell>, +} + +// SAFETY: `state` grants exclusive access to `waker`, and losing concurrent registrations do not +// touch the slot. `Waker` itself is `Send + Sync`. +unsafe impl Sync for AtomicWaker {} + +// `Waker` callbacks may unwind, but no panic leaves a state bit owned by the unwinding operation. A +// failed clone leaves the old slot intact and completes any raced wake, while wake and drop +// callbacks run after that operation's critical section has been released. +impl RefUnwindSafe for AtomicWaker {} +impl UnwindSafe for AtomicWaker {} + +impl AtomicWaker { + #[inline] + pub const fn new() -> Self { + Self { + state: AtomicUsize::new(WAITING), + waker: UnsafeCell::new(None), + } + } + + /// Registers `waker`, replacing a previously registered task if it differs. + /// + /// Calls to this method must not overlap. It may run concurrently with any number of calls to + /// [`wake`](Self::wake). + #[inline] + pub fn register(&self, waker: &Waker) { + // ORDERING: On success, Acquire pairs with the Release operation that last returned the + // state to WAITING and transfers exclusive ownership of the waker slot to this thread. On + // failure, Acquire matters when this reads WAKING from a notifier's AcqRel fetch_or: it + // receives the condition update that preceded that wake before this method returns. + match self + .state + .compare_exchange(WAITING, REGISTERING, Ordering::Acquire, Ordering::Acquire) + .unwrap_or_else(|state| state) + { + WAITING => { + // SAFETY: changing WAITING to REGISTERING grants this thread exclusive access to + // the waker slot until the state is returned to WAITING. + unsafe { self.register_locked(waker) } + } + WAKING => { + // A concurrent wake owns the slot. Self-waking ensures that this registration is + // not lost even though it cannot replace the slot right now. + waker.wake_by_ref(); + } + state => { + // Concurrent registration violates this type's contract. Ignoring the losing + // registration preserves memory safety and lets the winner provide notification. + debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); + } + } + } + + /// Registers a waker after this thread has acquired the REGISTERING state. + /// + /// # Safety + /// + /// The caller must have changed `state` from WAITING to REGISTERING and must be the only + /// thread accessing `waker`. + #[inline] + unsafe fn register_locked(&self, waker: &Waker) { + // Avoid both cloning and dropping the common case where an executor polls the receiver + // repeatedly with the same task waker. + let needs_replacement = match unsafe { &*self.waker.get() } { + Some(current) => !current.will_wake(waker), + None => true, + }; + + let mut clone_panic = None; + let old_waker = if needs_replacement { + match catch_unwind(AssertUnwindSafe(|| waker.clone())) { + Ok(new_waker) => unsafe { (*self.waker.get()).replace(new_waker) }, + Err(payload) => { + clone_panic = Some(payload); + None + } + } + } else { + None + }; + + // ORDERING: Release publishes a newly registered waker when the CAS succeeds. If it fails, + // Acquire receives the concurrent notifier's Release publication before the wake is + // completed below. AcqRel is the weakest success ordering that permits an Acquire failure + // ordering, although its Acquire half is not otherwise relied upon on the success path. + let concurrent_wake = match self.state.compare_exchange( + REGISTERING, + WAITING, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => None, + Err(state) => { + debug_assert_eq!(state, REGISTERING | WAKING); + + // SAFETY: REGISTERING remains set, so this thread still owns the waker slot. + let registered = unsafe { (*self.waker.get()).take() }; + + // ORDERING: Acquire receives all coalesced wake publications. Release publishes + // the empty slot and makes it available to the next register or wake operation. + self.state.swap(WAITING, Ordering::AcqRel); + registered + } + }; + + if let Some(payload) = clone_panic { + // Preserve the original clone panic while still completing a wake that raced with it. + if let Some(waker) = concurrent_wake { + let _ = catch_unwind(AssertUnwindSafe(|| waker.wake())); + } + resume_unwind(payload); + } + + // User waker code runs only after the state machine is back in WAITING, so a panic cannot + // leave the cell locked. If the wake raced with a replacement, notify both tasks: the + // concurrent call may have targeted the old registration, while future progress relies on + // the new one. A panic from the superseded waker must not prevent the new task from waking. + if let Some(waker) = concurrent_wake { + if let Some(old_waker) = old_waker { + let _ = catch_unwind(AssertUnwindSafe(|| old_waker.wake())); + } + waker.wake(); + } else { + // Drop a replaced waker only after releasing the state lock. + drop(old_waker); + } + } + + /// Wakes and removes the most recently registered waker, if any. + #[inline] + pub fn wake(&self) { + if let Some(waker) = self.take() { + waker.wake(); + } + } + + /// Removes the registered waker if this call acquires the slot. A concurrent registration or + /// wake may instead take responsibility for notifying it. + #[inline] + pub fn take(&self) -> Option { + // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the + // previous owner. Release publishes the condition update that the caller performed before + // calling wake, including when a registering thread already owns the slot. + match self.state.fetch_or(WAKING, Ordering::AcqRel) { + WAITING => { + // SAFETY: changing WAITING to WAKING grants this thread exclusive access to the + // waker slot until the state is returned to WAITING. + let waker = unsafe { (*self.waker.get()).take() }; + + // ORDERING: Release publishes the emptied slot before another operation acquires + // it. The fetch_or above already performed the required Acquire operation. + let old_state = self.state.swap(WAITING, Ordering::Release); + debug_assert_eq!(old_state, WAKING); + waker + } + state => { + // The thread registering a waker observes WAKING and completes this notification, + // or another waking thread has already taken responsibility for it. + debug_assert!( + state == REGISTERING || state == REGISTERING | WAKING || state == WAKING + ); + None + } + } + } +} + +#[cfg(test)] +mod tests { + use std::ptr; + use std::sync::Arc; + use std::sync::atomic::AtomicBool; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::task::RawWaker; + use std::task::RawWakerVTable; + use std::task::Wake; + + use super::*; + + struct WakeCounter(AtomicUsize); + + impl Wake for WakeCounter { + fn wake(self: Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + #[cfg(panic = "unwind")] + fn clone_panicking_waker() -> Waker { + static VTABLE: RawWakerVTable = RawWakerVTable::new( + |_| panic!("clone failed"), + |_| unreachable!(), + |_| unreachable!(), + |_| {}, + ); + + unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &VTABLE)) } + } + + #[test] + fn wake_notifies_once() { + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let atomic_waker = AtomicWaker::new(); + + atomic_waker.register(&waker); + atomic_waker.wake(); + atomic_waker.wake(); + + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + } + + #[test] + fn reregistering_same_task_does_not_clone_waker() { + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let atomic_waker = AtomicWaker::new(); + + atomic_waker.register(&waker); + let registered_refs = Arc::strong_count(&counter); + atomic_waker.register(&waker); + + assert_eq!(Arc::strong_count(&counter), registered_refs); + } + + #[test] + fn wake_before_register_is_not_remembered() { + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + let atomic_waker = AtomicWaker::new(); + + atomic_waker.wake(); + atomic_waker.register(&waker); + + assert_eq!(counter.0.load(Ordering::Relaxed), 0); + atomic_waker.wake(); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + } + + #[test] + fn wake_during_replacement_notifies_old_and_new_tasks() { + let old_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let old_waker = Waker::from(old_counter.clone()); + let new_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let new_waker = Waker::from(new_counter.clone()); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(&old_waker); + + assert_eq!( + atomic_waker.state.compare_exchange( + WAITING, + REGISTERING, + Ordering::AcqRel, + Ordering::Acquire, + ), + Ok(WAITING) + ); + std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); + + // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching + // the slot. Calling the helper completes the interrupted registration. + unsafe { atomic_waker.register_locked(&new_waker) }; + + assert_eq!(old_counter.0.load(Ordering::Relaxed), 1); + assert_eq!(new_counter.0.load(Ordering::Relaxed), 1); + } + + #[test] + fn failed_wake_synchronizes_with_next_registration() { + for _ in 0..1_000 { + let did_publish = AtomicBool::new(false); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(Waker::noop()); + + std::thread::scope(|scope| { + let wake = scope.spawn(|| { + did_publish.store(true, Ordering::Relaxed); + atomic_waker.take() + }); + + let local_waker = atomic_waker.take(); + atomic_waker.register(Waker::noop()); + + let publication_is_visible = did_publish.load(Ordering::Relaxed); + let concurrent_thread_took_waker = wake.join().unwrap().is_some(); + assert!(publication_is_visible || concurrent_thread_took_waker); + drop(local_waker); + }); + } + } + + #[cfg(panic = "unwind")] + #[test] + fn clone_panic_does_not_poison_state() { + let atomic_waker = AtomicWaker::new(); + + assert!( + catch_unwind(|| { + atomic_waker.register(&clone_panicking_waker()); + }) + .is_err() + ); + + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + atomic_waker.register(&Waker::from(counter.clone())); + atomic_waker.wake(); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + } + + #[cfg(panic = "unwind")] + #[test] + fn clone_panic_completes_concurrent_wake() { + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(&Waker::from(counter.clone())); + + assert_eq!( + atomic_waker.state.compare_exchange( + WAITING, + REGISTERING, + Ordering::Acquire, + Ordering::Acquire, + ), + Ok(WAITING) + ); + std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); + + // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching + // the state. Calling the helper completes the interrupted registration. + assert!( + catch_unwind(|| unsafe { + atomic_waker.register_locked(&clone_panicking_waker()); + }) + .is_err() + ); + + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + + let next_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + atomic_waker.register(&Waker::from(next_counter.clone())); + atomic_waker.wake(); + assert_eq!(next_counter.0.load(Ordering::Relaxed), 1); + } + + #[cfg(panic = "unwind")] + #[test] + fn drop_panic_does_not_poison_state() { + unsafe fn clone_drop_panicker(data: *const ()) -> RawWaker { + RawWaker::new(data, &DROP_PANICKING_VTABLE) + } + + unsafe fn wake_drop_panicker(_: *const ()) {} + + unsafe fn drop_drop_panicker(data: *const ()) { + // SAFETY: the test keeps the pointed-to AtomicBool alive until every derived waker has + // been dropped. + let should_panic = unsafe { &*data.cast::() }; + if should_panic.swap(false, Ordering::Relaxed) { + panic!("drop failed"); + } + } + + static DROP_PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( + clone_drop_panicker, + wake_drop_panicker, + wake_drop_panicker, + drop_drop_panicker, + ); + + let should_panic = AtomicBool::new(true); + let old_waker = unsafe { + Waker::from_raw(RawWaker::new( + ptr::from_ref(&should_panic).cast(), + &DROP_PANICKING_VTABLE, + )) + }; + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let new_waker = Waker::from(counter.clone()); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(&old_waker); + + assert!(catch_unwind(AssertUnwindSafe(|| atomic_waker.register(&new_waker))).is_err()); + + atomic_waker.wake(); + assert_eq!(counter.0.load(Ordering::Relaxed), 1); + } +} diff --git a/asyncband/src/internal/cache_padded.rs b/asyncband/src/internal/cache_padded.rs new file mode 100644 index 00000000..7268030d --- /dev/null +++ b/asyncband/src/internal/cache_padded.rs @@ -0,0 +1,55 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Use conservative architecture estimates, not a guarantee about every CPU's cache line. +// Keep 128 bytes for large ARM/PowerPC lines and adjacent-line prefetching on x86-64, +// 256 bytes for s390x, and at least 64 bytes elsewhere. +#[cfg_attr(target_arch = "s390x", repr(align(256)))] +#[cfg_attr( + any( + target_arch = "aarch64", + target_arch = "arm64ec", + target_arch = "powerpc64", + target_arch = "x86_64", + ), + repr(align(128)) +)] +#[cfg_attr( + not(any( + target_arch = "s390x", + target_arch = "aarch64", + target_arch = "arm64ec", + target_arch = "powerpc64", + target_arch = "x86_64", + )), + repr(align(64)) +)] +pub struct CachePadded(T); + +impl CachePadded { + pub const fn new(value: T) -> Self { + Self(value) + } +} + +impl std::ops::Deref for CachePadded { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 0252aa83..6271cade 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -49,6 +49,12 @@ pub(crate) fn wake_all(mut wakers: impl Iterator) { } } +#[cfg(feature = "mpsc")] +pub(crate) mod atomic_waker; + +#[cfg(feature = "mpsc")] +pub(crate) mod cache_padded; + #[cfg(any( feature = "barrier", feature = "broadcast", diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs new file mode 100644 index 00000000..87bb8bcc --- /dev/null +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -0,0 +1,263 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Capacity, position, and publication are separate ownership transitions: +//! +//! - A permit owns capacity, but holds no position until synchronous `push` claims a ticket. +//! - The ticket gives one producer a slot. `READY` publishes its initialized value to the receiver. +//! - The receiver finishes reading before returning capacity. AcqRel ticket increments carry that +//! reuse ordering even to a producer that acquired its permit on an earlier lap. +//! - Close competes with publication on the slot state. The drain owns `READY` values; a producer +//! that encounters `CLOSED` owns its unpublished value. Neither waits for the other to resume. +//! +//! Only the non-cloneable receiver advances the read cursor. All endpoints retain the shared +//! allocation, so a publisher's slot stays alive even when receiver drop closes it concurrently. + +use std::cell::UnsafeCell; +use std::mem; +use std::mem::MaybeUninit; +use std::ptr::NonNull; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU8; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Poll; + +use crate::internal::cache_padded::CachePadded; +use crate::internal::mutex::Mutex; + +const EMPTY: u8 = 0; +const READY: u8 = 1; +const CLOSED: u8 = 2; + +pub struct Buffer { + slots: Box<[Slot]>, + tail: CachePadded, + closed: AtomicBool, + // ZSTs need no positions or per-slot flags. Counting them separately also allows every + // nonzero usize capacity without allocating publication metadata for nonexistent bytes. + zero_sized: Mutex, +} + +struct Slot { + state: AtomicU8, + value: UnsafeCell>, +} + +// SAFETY: Capacity and the tail ticket give a producer exclusive ownership of an empty slot. +// Release publication transfers its value to the exclusive consumer. Closing an unpublished +// slot leaves its value with the producer; closing a READY slot transfers it to the drain. +unsafe impl Sync for Slot {} + +// No reference to a stored value escapes. Every value is removed from the slot's ownership +// before running a callback or destructor that might panic. +impl std::panic::UnwindSafe for Slot {} +impl std::panic::RefUnwindSafe for Slot {} + +impl Buffer { + pub fn new(capacity: usize) -> Self { + let slots = if mem::size_of::() == 0 { + Box::default() + } else { + (0..capacity.next_power_of_two()) + .map(|_| Slot { + state: AtomicU8::new(EMPTY), + value: UnsafeCell::new(MaybeUninit::uninit()), + }) + .collect() + }; + Self { + slots, + tail: CachePadded::new(AtomicUsize::new(0)), + closed: AtomicBool::new(false), + zero_sized: Mutex::new(0), + } + } + + fn slot(&self, position: usize) -> &Slot { + // Power-of-two storage preserves indexing when the full-width ticket wraps. The + // semaphore still enforces the exact requested capacity, including non-powers of two. + &self.slots[position & (self.slots.len() - 1)] + } + + fn claim(&self) -> Result { + if self.closed.load(Ordering::Acquire) { + return Err(()); + } + // Closing may race after this check. It marks every physical slot CLOSED, so even a + // delayed claimant will recover its own value instead of publishing into a dead queue. + Ok(self.tail.fetch_add(1, Ordering::AcqRel)) + } + + /// Writes and publishes one message. Closing may instead return the unsent value. + /// + /// # Safety + /// + /// Own one capacity permit before calling; release it only after a failed push or after + /// the consumer reads the published value. No user code runs between claim and publication. + pub unsafe fn push(&self, value: T) -> Result<(), T> { + if mem::size_of::() == 0 { + let mut queued = self.zero_sized.lock(); + if self.closed.load(Ordering::Acquire) { + return Err(value); + } + *queued += 1; + mem::forget(value); + return Ok(()); + } + let Ok(position) = self.claim() else { + return Err(value); + }; + // SAFETY: The caller owns capacity and the atomic increment assigned this position. + unsafe { self.publish(position, value) } + } + + unsafe fn publish(&self, position: usize, value: T) -> Result<(), T> { + let slot = self.slot(position); + // SAFETY: Capacity prevents wrapping over unread slots. AcqRel tail increments carry prior + // claimants' capacity-acquire edges even when this producer held its permit for a long + // time. The previous consumer has therefore finished reading before this write. + unsafe { (*slot.value.get()).write(value) }; + match slot + .state + .compare_exchange(EMPTY, READY, Ordering::Release, Ordering::Acquire) + { + Ok(_) => Ok(()), + Err(state) => { + debug_assert_eq!(state, CLOSED); + // SAFETY: Close saw an unpublished slot and did not read it. Failed publication + // leaves exclusive ownership with this producer, including during receiver drop. + Err(unsafe { (*slot.value.get()).assume_init_read() }) + } + } + } + + /// Pending means a producer claimed the head but has not published it yet. + /// + /// # Safety + /// + /// Only the exclusive consumer may call this, using its persistent cursor. Release one + /// capacity permit after each successful pop, after the value has been read completely. + pub unsafe fn pop(&self, head: &mut usize) -> Poll> { + if mem::size_of::() == 0 { + let mut queued = self.zero_sized.lock(); + return if *queued == 0 { + Poll::Ready(None) + } else { + *queued -= 1; + // SAFETY: A queued value proves that this ZST is inhabited and owns one value. + Poll::Ready(Some(unsafe { Self::read_zero_sized() })) + }; + } + let slot = self.slot(*head); + if slot.state.load(Ordering::Acquire) == READY { + // SAFETY: Publication initialized the value, and only this consumer can read it. + // Capacity is still held until this method has returned the value to its caller. + let value = unsafe { (*slot.value.get()).assume_init_read() }; + slot.state.store(EMPTY, Ordering::Release); + *head = head.wrapping_add(1); + Poll::Ready(Some(value)) + } else if self.tail.load(Ordering::Acquire) == *head { + Poll::Ready(None) + } else { + Poll::Pending + } + } + + /// Stops new claims and returns ownership of published values to a drain guard. + /// + /// # Safety + /// + /// Only the exclusive consumer may close the buffer, once, using its current cursor. + pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { + self.closed.store(true, Ordering::Release); + let remaining = if mem::size_of::() == 0 { + mem::take(&mut *self.zero_sized.lock()) + } else { + // Cover every physical slot: a producer may have passed the open check but not + // obtained its ticket yet. Such a late claim must also find a CLOSED slot. + self.slots.len() + }; + Drain { + buffer: self, + position: head, + remaining, + } + } + + unsafe fn read_zero_sized() -> T { + // SAFETY: The caller owns an initialized, inhabited ZST. Reading it accesses no bytes; + // dangling supplies a non-null, correctly aligned pointer, as in a ZST Vec. + unsafe { NonNull::::dangling().as_ptr().read() } + } +} + +pub struct Drain<'a, T> { + buffer: &'a Buffer, + position: usize, + remaining: usize, +} + +impl Iterator for Drain<'_, T> { + type Item = T; + + fn next(&mut self) -> Option { + while self.remaining != 0 { + let position = self.position; + self.remaining -= 1; + self.position = self.position.wrapping_add(1); + if mem::size_of::() == 0 { + // SAFETY: Closing transferred this many initialized ZST values to the drain. + return Some(unsafe { Buffer::::read_zero_sized() }); + } + let slot = self.buffer.slot(position); + if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { + // SAFETY: The drain won ownership of a published value. The cursor and state + // already advanced, so a panicking destructor cannot cause a second read. + return Some(unsafe { (*slot.value.get()).assume_init_read() }); + } + // An unpublished slot stays owned by its producer, which will observe CLOSED and + // recover its value. The shared Arc keeps this allocation alive until that finishes. + } + None + } +} + +impl Drop for Drain<'_, T> { + fn drop(&mut self) { + struct Remaining<'a, 'b, T>(&'a mut Drain<'b, T>); + + impl Drop for Remaining<'_, '_, T> { + fn drop(&mut self) { + for value in self.0.by_ref() { + drop(value); + } + } + } + + // A guard inside Drop is necessary: Drop itself is not called again if a payload's + // destructor panics while this normal drain is running. + let remaining = Remaining(self); + for value in remaining.0.by_ref() { + drop(value); + } + } +} + +#[cfg(test)] +#[path = "buffer_tests.rs"] +mod tests; diff --git a/asyncband/src/mpsc/bounded/buffer_tests.rs b/asyncband/src/mpsc/bounded/buffer_tests.rs new file mode 100644 index 00000000..43690029 --- /dev/null +++ b/asyncband/src/mpsc/bounded/buffer_tests.rs @@ -0,0 +1,221 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; +use std::sync::Barrier; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; +use std::thread; + +use crate::mpsc::BoundedSender; +use crate::mpsc::Permit; +use crate::mpsc::TryRecvError; +use crate::mpsc::bounded; + +// Exercise the scheduling window inside synchronous send, while retaining a real capacity +// permit. Ordinary callers cannot split a claim from its publication. +fn publish_claimed(mut permit: Permit<'_, T>, position: usize, value: T) -> Result<(), T> { + let shared = &permit.sender.unwrap().shared; + // SAFETY: The test claimed this position while holding the same capacity permit. + unsafe { shared.buffer.publish(position, value) }?; + permit.sender = None; + shared.rx_waker.wake(); + Ok(()) +} + +#[test] +fn a_claimed_head_waits_for_publication_across_laps() { + for capacity in [1, 3, 7] { + for initial in [0, usize::MAX - 1] { + let (tx, mut rx) = bounded(capacity); + // Start an empty ring near ticket overflow instead of running usize::MAX sends. + tx.shared.buffer.tail.store(initial, Ordering::Relaxed); + rx.head = initial; + let mut cx = Context::from_waker(Waker::noop()); + for lap in 0..8 { + let permit = tx.try_reserve().unwrap(); + let position = tx.shared.buffer.claim().unwrap(); + for offset in 1..capacity { + tx.try_send(lap * capacity + offset).unwrap(); + } + // A full ring must differ from an empty one even if no head value is ready yet. + assert!(rx.poll_recv(&mut cx).is_pending()); + publish_claimed(permit, position, lap * capacity).unwrap(); + for offset in 0..capacity { + assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } + } + } +} + +#[test] +fn a_claim_delayed_past_close_returns_its_value() { + let (tx, rx) = bounded(3); + let permit = tx.try_reserve().unwrap(); + let drops = Arc::new(AtomicUsize::new(0)); + let value = Payload { + bytes: [7; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let allocation = Arc::downgrade(&tx.shared); + // Pause after claim's open check, then resume its atomic ticket allocation after close. + assert!(!tx.shared.buffer.closed.load(Ordering::Acquire)); + drop(rx); + let position = tx.shared.buffer.tail.fetch_add(1, Ordering::AcqRel); + let unsent = publish_claimed(permit, position, value).unwrap_err(); + assert_eq!(unsent.bytes, [7; 1024]); + drop(unsent); + assert_eq!(drops.load(Ordering::Relaxed), 1); + drop(tx); + assert!(allocation.upgrade().is_none()); +} + +#[derive(Debug)] +#[repr(align(128))] +struct Payload { + bytes: [u8; 1024], + drops: Arc, + // Queued messages must not keep the shared allocation alive through a sender cycle. + _sender: BoundedSender, +} + +impl Drop for Payload { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::Relaxed); + } +} + +#[test] +fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { + let (tx, rx) = bounded(2); + let allocation = Arc::downgrade(&tx.shared); + let drops = Arc::new(AtomicUsize::new(0)); + let paused = Barrier::new(2); + let (resume_tx, resume_rx) = std::sync::mpsc::channel(); + let (closed_tx, closed_rx) = std::sync::mpsc::channel(); + + thread::scope(|scope| { + let sender = &tx; + let drops = &drops; + let paused = &paused; + let publisher = scope.spawn(move || { + let permit = sender.try_reserve().unwrap(); + let position = sender.shared.buffer.claim().unwrap(); + let value = Payload { + bytes: [1; 1024], + drops: drops.clone(), + _sender: sender.clone(), + }; + paused.wait(); + resume_rx.recv().unwrap(); + let unsent = publish_claimed(permit, position, value).unwrap_err(); + assert_eq!(unsent.bytes, [1; 1024]); + drop(unsent); + }); + paused.wait(); + tx.try_send(Payload { + bytes: [2; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }) + .unwrap(); + let closer = scope.spawn(move || { + drop(rx); + closed_tx.send(()).unwrap(); + }); + #[cfg(not(miri))] + let closed = closed_rx.recv_timeout(std::time::Duration::from_secs(10)); + #[cfg(miri)] + let closed = closed_rx.recv(); + let dropped_before_resume = drops.load(Ordering::Relaxed); + // Unblock the publisher before asserting so a failed close cannot strand the scope. + resume_tx.send(()).unwrap(); + publisher.join().unwrap(); + closer.join().unwrap(); + assert!(closed.is_ok(), "close waited for the paused publisher"); + assert_eq!(dropped_before_resume, 1); + }); + + assert_eq!(drops.load(Ordering::Relaxed), 2); + drop(tx); + assert!(allocation.upgrade().is_none()); +} + +#[test] +fn publication_racing_with_close_drops_every_payload_once() { + for _ in 0..if cfg!(miri) { 8 } else { 128 } { + let (tx, rx) = bounded(3); + let allocation = Arc::downgrade(&tx.shared); + let drops = Arc::new(AtomicUsize::new(0)); + let start = Barrier::new(4); + thread::scope(|scope| { + for byte in 0..3 { + let permit = tx.try_reserve().unwrap(); + let value = Payload { + bytes: [byte; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let start = &start; + scope.spawn(move || { + start.wait(); + if let Err(error) = permit.send(value) { + let value = error.into_inner(); + assert_eq!(value.bytes, [byte; 1024]); + drop(value); + } + }); + } + start.wait(); + drop(rx); + }); + assert_eq!(drops.load(Ordering::Relaxed), 3); + drop(tx); + assert!(allocation.upgrade().is_none()); + } +} + +#[test] +fn an_old_permit_can_publish_after_other_producers_wrap_the_ring() { + let (tx, mut rx) = bounded(3); + let old = tx.try_reserve().unwrap(); + for lap in 0..16 { + for offset in 0..2 { + tx.try_send([lap * 2 + offset; 1024]).unwrap(); + } + for offset in 0..2 { + assert_eq!(rx.try_recv(), Ok([lap * 2 + offset; 1024])); + } + } + thread::scope(|scope| { + scope + .spawn(move || old.send([42; 1024]).unwrap()) + .join() + .unwrap(); + }); + assert_eq!( + rx.poll_recv(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok([42; 1024])) + ); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); +} diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index dab1f7e5..dfad2896 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -18,31 +18,40 @@ //! A bounded multi-producer, single-consumer queue for sending values between asynchronous //! tasks with backpressure control. -use std::collections::VecDeque; use std::fmt; use std::future::poll_fn; -use std::mem; use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; use std::task::Waker; +use self::buffer::Buffer; use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; +use crate::internal::atomic_waker::AtomicWaker; +use crate::internal::cache_padded::CachePadded; use crate::internal::mutex::Mutex; use crate::internal::waitlist::WaitList; use crate::internal::waitlist::WaiterId; use crate::internal::wake_all; use crate::internal::waker_batch::WakerBatch; +mod buffer; + /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// /// [`BoundedSender::send`] waits for capacity when the buffer is full. Receiving a message releases /// one slot for a waiting sender. Capacity is granted in the order that pending sends and /// reservations enter the wait queue; new senders cannot take an already granted slot. /// +/// Storage for nonzero-sized messages is preallocated and rounded up to a power of two; the +/// channel's capacity remains exactly `buffer`. Zero-sized messages need no per-slot storage. +/// /// # Panics /// /// Panics if `buffer` is zero or the preallocated message buffer exceeds the allocation size @@ -50,63 +59,96 @@ use crate::internal::waker_batch::WakerBatch; #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); - let state = Arc::new(Mutex::new(State { - queue: VecDeque::with_capacity(buffer), - available: buffer, - receiver_open: true, - senders: 1, - receiver_waker: None, - waiters: WaitList::new(), - })); + let shared = Arc::new(Shared { + senders: AtomicUsize::new(1), + tx_permits: CachePadded::new(Semaphore::new(buffer)), + rx_waker: CachePadded::new(AtomicWaker::new()), + buffer: Buffer::new(buffer), + }); let sender = BoundedSender { - state: state.clone(), + shared: shared.clone(), }; - let receiver = BoundedReceiver { state }; + let receiver = BoundedReceiver { shared, head: 0 }; (sender, receiver) } -// All transitions happen under one lock. Capacity belongs to exactly one of: `available`, a -// queued message, a live Permit, or a granted waiter. Waker callbacks and payload destruction -// run after unlocking; neither a pending send nor a public Permit owns a queue position. -struct State { - queue: VecDeque, - available: usize, - receiver_open: bool, - senders: usize, - receiver_waker: Option, - waiters: WaitList, +struct Shared { + senders: AtomicUsize, + tx_permits: CachePadded, + rx_waker: CachePadded, + buffer: Buffer, } -impl State { - fn acquire(&mut self) -> Result<(), TrySendError<()>> { - if !self.receiver_open { - Err(TrySendError::Disconnected(())) - } else if self.available == 0 { - Err(TrySendError::Full(())) - } else { - self.available -= 1; - Ok(()) +// This channel-local semaphore grants one permit at a time and can close its wait queue. +// The general-purpose semaphore has neither a close operation nor acquisition errors. +struct Semaphore { + available: AtomicUsize, + closed: AtomicBool, + waiters: Mutex>, +} + +impl Semaphore { + fn new(available: usize) -> Self { + Self { + available: AtomicUsize::new(available), + closed: AtomicBool::new(false), + waiters: Mutex::new(WaitList::new()), } } - fn release(&mut self) -> Option { - if let Some((_, waiter)) = self.waiters.unlink_first_waiter(|_| true) { - // Keep the detached node until its future claims or cancels this grant. + fn try_acquire(&self) -> Result<(), TrySendError<()>> { + if self.closed.load(Ordering::Acquire) { + return Err(TrySendError::Disconnected(())); + } + let mut available = self.available.load(Ordering::Relaxed); + loop { + if available == 0 { + return Err(TrySendError::Full(())); + } + match self.available.compare_exchange_weak( + available, + available - 1, + Ordering::Acquire, + Ordering::Relaxed, + ) { + Ok(_) => return Ok(()), + Err(actual) => available = actual, + } + } + } + + fn release(&self) { + let wake = self.release_locked(&mut self.waiters.lock()); + if let Some(waker) = wake { + waker.wake(); + } + } + + fn release_locked(&self, waiters: &mut WaitList) -> Option { + if self.closed.load(Ordering::Relaxed) { + return None; + } + if let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { + // Grant ownership before waking; new arrivals cannot steal this capacity. waiter.granted = true; return waiter.waker.take(); } - self.available += 1; + // Only releases add permits, and all releases hold the wait queue lock. A linked + // waiter therefore always sees zero available permits until it receives its own grant. + self.available.fetch_add(1, Ordering::Release); None } - fn pop(&mut self) -> Result<(T, Option), TryRecvError> { - if let Some(value) = self.queue.pop_front() { - Ok((value, self.release())) - } else if self.senders == 0 { - Err(TryRecvError::Disconnected) - } else { - Err(TryRecvError::Empty) + fn close(&self) -> WakerBatch { + let mut waiters = self.waiters.lock(); + self.closed.store(true, Ordering::Release); + let mut wakers = WakerBatch::new(); + while let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } } + wakers } } @@ -122,25 +164,40 @@ struct Reservation<'a, T> { impl<'a, T> Reservation<'a, T> { fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { + let semaphore = &self.sender.shared.tx_permits; let mut cloned_waker = None; loop { - let mut state = self.sender.state.lock(); - if !state.receiver_open { + if self.index.is_none() { + match semaphore.try_acquire() { + Ok(()) => { + let permit = Permit { + sender: Some(self.sender), + }; + // The permit owns capacity before an unused cloned waker can panic. + drop(cloned_waker); + return Poll::Ready(Ok(permit)); + } + Err(TrySendError::Disconnected(())) => { + return Poll::Ready(Err(SendError::new(()))); + } + Err(TrySendError::Full(())) => {} + } + } + let mut waiters = semaphore.waiters.lock(); + if semaphore.closed.load(Ordering::Relaxed) { // Drop removes any remaining registration, including an unused grant. return Poll::Ready(Err(SendError::new(()))); } if let Some(index) = self.index { - let waiter = state.waiters.waiter_mut(index); + let waiter = waiters.waiter_mut(index); if waiter.granted { - let waiter = state.waiters.remove_unlinked_waiter(index); + let waiter = waiters.remove_unlinked_waiter(index); self.index = None; let permit = Permit { sender: Some(self.sender), }; - drop(state); + drop(waiters); drop(waiter); - // A clone callback may have freed capacity. Establish ownership before - // dropping the unused clone, whose destructor can also run user code. drop(cloned_waker); return Poll::Ready(Ok(permit)); } @@ -153,26 +210,27 @@ impl<'a, T> Reservation<'a, T> { } if let Some(waker) = cloned_waker.take() { let old = waiter.waker.replace(waker); - drop(state); + drop(waiters); drop(old); return Poll::Pending; } - } else if state.available != 0 { - state.available -= 1; + } else if semaphore.try_acquire().is_ok() { + // A release may have raced with the fast path; recheck under the queue lock + // before committing to wait so no permit can be stranded without a wake. let permit = Permit { sender: Some(self.sender), }; - drop(state); + drop(waiters); drop(cloned_waker); return Poll::Ready(Ok(permit)); } else if let Some(waker) = cloned_waker.take() { - self.index = Some(state.waiters.push_back(Waiter { + self.index = Some(waiters.push_back(Waiter { granted: false, waker: Some(waker), })); return Poll::Pending; } - drop(state); + drop(waiters); // Clone outside the lock, then recheck capacity and closure before registering. cloned_waker = Some(cx.waker().clone()); } @@ -182,12 +240,13 @@ impl<'a, T> Reservation<'a, T> { impl Drop for Reservation<'_, T> { fn drop(&mut self) { let Some(index) = self.index else { return }; + let semaphore = &self.sender.shared.tx_permits; let (waiter, wake) = { - let mut state = self.sender.state.lock(); - state.waiters.unlink_waiter(index, |_| true); - let waiter = state.waiters.remove_unlinked_waiter(index); + let mut waiters = semaphore.waiters.lock(); + waiters.unlink_waiter(index, |_| true); + let waiter = waiters.remove_unlinked_waiter(index); let wake = if waiter.granted { - state.release() + semaphore.release_locked(&mut waiters) } else { None }; @@ -204,14 +263,14 @@ impl Drop for Reservation<'_, T> { /// /// Instances are created by the [`bounded`] function. pub struct BoundedSender { - state: Arc>>, + shared: Arc>, } impl Clone for BoundedSender { fn clone(&self) -> Self { - self.state.lock().senders += 1; + self.shared.senders.fetch_add(1, Ordering::Relaxed); BoundedSender { - state: self.state.clone(), + shared: self.shared.clone(), } } } @@ -224,17 +283,8 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { - let wake = { - let mut state = self.state.lock(); - state.senders -= 1; - if state.senders == 0 { - state.receiver_waker.take() - } else { - None - } - }; - if let Some(waker) = wake { - waker.wake(); + if self.shared.senders.fetch_sub(1, Ordering::AcqRel) == 1 { + self.shared.rx_waker.wake(); } } } @@ -251,24 +301,6 @@ impl BoundedSender { /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for /// capacity before constructing the message. pub async fn send(&self, value: T) -> Result<(), SendError> { - // Publish directly so a ready payload does not travel through try_send's large error - // return value. Capacity and publication still share one critical section. - { - let mut state = self.state.lock(); - match state.acquire() { - Ok(()) => { - state.queue.push_back(value); - let wake = state.receiver_waker.take(); - drop(state); - if let Some(waker) = wake { - waker.wake(); - } - return Ok(()); - } - Err(TrySendError::Disconnected(())) => return Err(SendError::new(value)), - Err(TrySendError::Full(())) => {} - } - } match self.reserve().await { Ok(permit) => permit.send(value), Err(_) => Err(SendError::new(value)), @@ -318,7 +350,7 @@ impl BoundedSender { /// Returns [`TrySendError::Full`] if queued messages and outstanding permits occupy the /// buffer, or [`TrySendError::Disconnected`] if the receiver has been dropped. pub fn try_reserve(&self) -> Result, TrySendError<()>> { - self.state.lock().acquire()?; + self.shared.tx_permits.try_acquire()?; Ok(Permit { sender: Some(self) }) } @@ -343,17 +375,10 @@ impl BoundedSender { /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - let mut state = self.state.lock(); - match state.acquire() { - Ok(()) => { - state.queue.push_back(value); - let wake = state.receiver_waker.take(); - drop(state); - if let Some(waker) = wake { - waker.wake(); - } - Ok(()) - } + match self.try_reserve() { + Ok(permit) => permit + .send(value) + .map_err(|error| TrySendError::Disconnected(error.into_inner())), Err(TrySendError::Full(())) => Err(TrySendError::Full(value)), Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), } @@ -381,18 +406,13 @@ impl Permit<'_, T> { /// /// If the receiver has been dropped, the returned error contains the unsent value. pub fn send(mut self, value: T) -> Result<(), SendError> { - let mut state = self.sender.unwrap().state.lock(); - if !state.receiver_open { - return Err(SendError::new(value)); - } - state.queue.push_back(value); - // The queued message owns the capacity before any wake callback can panic. + let shared = &self.sender.unwrap().shared; + // SAFETY: This permit owns one capacity unit. Claiming a slot and writing it is a + // synchronous operation with no user callbacks or await points between the two. + unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; + // Publication owns the capacity before a wake callback can panic. self.sender = None; - let wake = state.receiver_waker.take(); - drop(state); - if let Some(waker) = wake { - waker.wake(); - } + shared.rx_waker.wake(); Ok(()) } } @@ -400,10 +420,7 @@ impl Permit<'_, T> { impl Drop for Permit<'_, T> { fn drop(&mut self) { if let Some(sender) = self.sender { - let wake = sender.state.lock().release(); - if let Some(waker) = wake { - waker.wake(); - } + sender.shared.tx_permits.release(); } } } @@ -411,8 +428,11 @@ impl Drop for Permit<'_, T> { /// The receiving endpoint of a bounded mpsc channel. /// /// Instances are created by the [`bounded`] function. +/// Dropping the receiver discards queued values. The backing allocation remains alive until +/// all endpoints are dropped, so a concurrent sender can safely finish returning an unsent value. pub struct BoundedReceiver { - state: Arc>>, + shared: Arc>, + head: usize, } impl fmt::Debug for BoundedReceiver { @@ -423,24 +443,14 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - let (queue, receiver_waker, wakers) = { - let mut state = self.state.lock(); - state.receiver_open = false; - let queue = mem::take(&mut state.queue); - state.available += queue.len(); - let receiver_waker = state.receiver_waker.take(); - let mut wakers = WakerBatch::new(); - while let Some((_, waiter)) = state.waiters.unlink_first_waiter(|_| true) { - if let Some(waker) = waiter.waker.take() { - wakers.push(waker); - } - } - (queue, receiver_waker, wakers) - }; - // Local ownership also drains the queue if a wake or waker destructor unwinds. + // SAFETY: Receiver ownership provides exclusive access to the consumption cursor. + // The drain first prevents new claims. Its destructor completes cleanup on unwinding. + let drain = unsafe { self.shared.buffer.close(self.head) }; + let wakers = self.shared.tx_permits.close(); + let receiver_waker = self.shared.rx_waker.take(); wake_all(wakers.into_iter()); drop(receiver_waker); - drop(queue); + drop(drain); } } @@ -451,6 +461,9 @@ impl BoundedReceiver { /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has /// been dropped and all queued values have been consumed. /// + /// If a producer is still completing a synchronous publication at the queue head, this + /// method waits for that publication. Use [`Self::recv`] to wait asynchronously instead. + /// /// # Examples /// /// ``` @@ -468,11 +481,45 @@ impl BoundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - let (value, wake) = self.state.lock().pop()?; - if let Some(waker) = wake { - waker.wake(); + let mut spins = 0; + loop { + match self.try_pop() { + Poll::Ready(result) => return result, + Poll::Pending => { + // A synchronous publisher already owns the head. Reporting Empty here + // could hide a later send that has completed. Async recv parks instead. + if spins < 32 { + std::hint::spin_loop(); + spins += 1; + } else { + std::thread::yield_now(); + } + } + } + } + } + + fn try_pop(&mut self) -> Poll> { + let mut disconnected = false; + loop { + // SAFETY: Only this receiver owns head. Capacity is released after the buffer + // finishes reading and advances the cursor, so no producer can overwrite the value. + match unsafe { self.shared.buffer.pop(&mut self.head) } { + Poll::Ready(Some(value)) => { + self.shared.tx_permits.release(); + return Poll::Ready(Ok(value)); + } + Poll::Ready(None) if disconnected => { + return Poll::Ready(Err(TryRecvError::Disconnected)); + } + Poll::Ready(None) if self.shared.senders.load(Ordering::Acquire) == 0 => { + // Acquire the last sender's completed publications before checking again. + disconnected = true; + } + Poll::Ready(None) => return Poll::Ready(Err(TryRecvError::Empty)), + Poll::Pending => return Poll::Pending, + } } - Ok(value) } /// Waits for and receives the next value, freeing one buffer slot. @@ -509,40 +556,19 @@ impl BoundedReceiver { } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - let mut cloned_waker = None; - loop { - let mut state = self.state.lock(); - match state.pop() { - Ok((value, wake)) => { - drop(state); - if let Some(waker) = wake { - waker.wake(); - } - return Poll::Ready(Ok(value)); - } - Err(TryRecvError::Disconnected) => { - let old = state.receiver_waker.take(); - drop(state); - drop(old); + for registered in [false, true] { + match self.try_pop() { + Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), + Poll::Ready(Err(TryRecvError::Disconnected)) => { + drop(self.shared.rx_waker.take()); return Poll::Ready(Err(RecvError::Disconnected)); } - Err(TryRecvError::Empty) => {} + Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {} } - if state - .receiver_waker - .as_ref() - .is_some_and(|w| w.will_wake(cx.waker())) - { - return Poll::Pending; - } - if let Some(waker) = cloned_waker.take() { - let old = state.receiver_waker.replace(waker); - drop(state); - drop(old); - return Poll::Pending; + if !registered { + self.shared.rx_waker.register(cx.waker()); } - drop(state); - cloned_waker = Some(cx.waker().clone()); } + Poll::Pending } } diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index de3a5eb3..e26c1e81 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -76,6 +76,31 @@ fn zero_sized_messages_support_the_full_capacity_range() { } } +#[test] +fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + static DROPS: AtomicUsize = AtomicUsize::new(0); + struct Message; + impl Drop for Message { + fn drop(&mut self) { + DROPS.fetch_add(1, Ordering::Relaxed); + } + } + + let (tx, mut rx) = mpsc::bounded(usize::MAX); + for _ in 0..3 { + assert!(tx.try_send(Message).is_ok()); + } + drop(rx.try_recv().unwrap()); + assert_eq!(DROPS.load(Ordering::Relaxed), 1); + drop(rx); + assert_eq!(DROPS.load(Ordering::Relaxed), 3); + drop(tx.try_send(Message).err().unwrap().into_inner()); + assert_eq!(DROPS.load(Ordering::Relaxed), 4); +} + #[test] fn released_capacity_is_granted_to_the_oldest_waiter() { let (tx, mut rx) = mpsc::bounded(1); From 7ad35ade673d9fa43d7da033a88aacf6596fc90c Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 08:58:36 +0800 Subject: [PATCH 06/34] test(mpsc): distinguish task and external receiver benchmarks Run producer and receiver tasks on the same executor by default, and measure a receiver on the block_on caller thread as a separate workload. Use the same reusable-task harness and start protocol for usize and inline 1 KiB messages, including capacity-one and reservation cases. --- benchmarks/ecosystem/mpsc/bounded.rs | 129 ++++++------------ benchmarks/ecosystem/mpsc/reservation.rs | 5 +- benchmarks/ecosystem/mpsc/support.rs | 158 ++++++++++++++++++----- 3 files changed, 170 insertions(+), 122 deletions(-) diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index dc595a89..95635419 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -15,13 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::future::Future; -use std::future::poll_fn; -use std::pin::pin; -use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; - use divan::Bencher; use divan::black_box; use divan::counter::ItemsCount; @@ -114,7 +107,7 @@ fn sustained_capacity( #[divan::bench( types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], consts = [1, 64, 4096], - args = [(1, 0), (4, 0), (1, 4), (4, 4), (8, 4)], + args = [(1, 0), (8, 0), (1, 4), (8, 4)], sample_count = 50, sample_size = 1, counter = ItemsCount::new(BATCH_MESSAGES), @@ -130,95 +123,49 @@ fn scheduled( #[divan::bench( types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], - consts = [64, 4096], - args = [1, 8], + consts = [1, 64, 4096], + args = [(1, 0), (8, 0), (1, 4), (8, 4)], sample_count = 50, sample_size = 1, counter = ItemsCount::new(BATCH_MESSAGES), )] fn scheduled_inline, const CAPACITY: usize>( + bencher: Bencher, + (producers, workers): (usize, usize), +) { + let mut batch = RepeatedTasks::>::new(producers, workers); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + consts = [1, 64, 4096], + args = [1, 8], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn external_receiver(bencher: Bencher, producers: usize) { + let mut batch = RepeatedTasks::>::external_receiver(producers, 4); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + consts = [1, 64, 4096], + args = [1, 8], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn external_receiver_inline, const CAPACITY: usize>( bencher: Bencher, producers: usize, ) { - let runtime = tokio::runtime::Builder::new_multi_thread() - .worker_threads(4) - .build() - .unwrap(); - let (sender, mut receiver) = C::channel(CAPACITY); - let start: Vec<_> = (0..producers) - .map(|_| Arc::new(tokio::sync::Notify::new())) - .collect(); - let stop = Arc::new(AtomicBool::new(false)); - let workers: Vec<_> = start - .iter() - .enumerate() - .map(|(producer, start)| { - let sender = sender.clone(); - let start = start.clone(); - let stop = stop.clone(); - runtime.spawn(async move { - let mut sequence = 0u64; - loop { - start.notified().await; - if stop.load(Ordering::Acquire) { - break; - } - for _ in 0..BATCH_MESSAGES / producers { - let mut value = [1; 1024]; - value[..8].copy_from_slice(&(producer as u64).to_le_bytes()); - value[8..16].copy_from_slice(&sequence.to_le_bytes()); - C::send_async(&sender, black_box(value)).await; - sequence += 1; - } - } - }) - }) - .collect(); - drop(sender); - let mut expected = vec![0u64; producers]; - let mut run = || { - runtime.block_on(async { - let first = { - let mut receive = pin!(tokio::task::unconstrained(C::recv_async(&mut receiver))); - let mut released = false; - poll_fn(|cx| { - let result = receive.as_mut().poll(cx); - if !released { - assert!(result.is_pending(), "each sample starts with an empty wait"); - released = true; - for producer in &start { - producer.notify_one(); - } - } - result - }) - .await - }; - let mut value = first; - for received in 0..BATCH_MESSAGES { - let producer = u64::from_le_bytes(value[..8].try_into().unwrap()) as usize; - let sequence = u64::from_le_bytes(value[8..16].try_into().unwrap()); - assert_eq!(sequence, expected[producer]); - expected[producer] += 1; - assert_eq!(black_box(value)[1023], 1); - if received + 1 < BATCH_MESSAGES { - value = C::recv_async(&mut receiver).await; - } - } - assert!(expected.iter().all(|count| *count == expected[0])); - }) - }; - // Reuse tasks and the channel. Include the initial empty wait, backpressure, and payload - // movement; verify per-producer order and payload integrity on every measured sample. - run(); - bencher.bench_local(run); - stop.store(true, Ordering::Release); - for producer in &start { - producer.notify_one(); - } - runtime.block_on(async { - for worker in workers { - worker.await.unwrap(); - } - }); + let mut batch = + RepeatedTasks::>::external_receiver(producers, 4); + batch.run(); + bencher.bench_local(|| batch.run()); } diff --git a/benchmarks/ecosystem/mpsc/reservation.rs b/benchmarks/ecosystem/mpsc/reservation.rs index b8f3a275..5952d2a6 100644 --- a/benchmarks/ecosystem/mpsc/reservation.rs +++ b/benchmarks/ecosystem/mpsc/reservation.rs @@ -84,6 +84,7 @@ fn cancel_reserved_capacity(bencher: Bencher) { struct Reserved(PhantomData); impl ConcurrentMpsc for Reserved { + type Message = usize; type Sender = C::Sender; type Receiver = C::Receiver; fn channel() -> (Self::Sender, Self::Receiver) { @@ -105,8 +106,8 @@ impl ConcurrentMpsc for Reserved Self; + fn sequence(self) -> usize; +} + +impl Message for usize { + fn new(sequence: usize) -> Self { + sequence + } + + fn sequence(self) -> usize { + self + } +} + +impl Message for [u8; 1024] { + fn new(sequence: usize) -> Self { + let mut value = [1; 1024]; + value[..size_of::()].copy_from_slice(&sequence.to_le_bytes()); + value + } + + fn sequence(self) -> usize { + let value = black_box(self); + assert_eq!(value[1023], 1); + usize::from_le_bytes(value[..size_of::()].try_into().unwrap()) + } +} + pub trait ConcurrentMpsc: Send + Sync + 'static { + type Message: Message; type Sender: Clone + Send + Sync + 'static; type Receiver: Send + 'static; fn channel() -> (Self::Sender, Self::Receiver); - fn send(sender: &Self::Sender, value: usize); - fn recv(receiver: &mut Self::Receiver) -> usize; - fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; - fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; + fn send(sender: &Self::Sender, value: Self::Message); + fn recv(receiver: &mut Self::Receiver) -> Self::Message; + fn send_async(sender: &Self::Sender, value: Self::Message) -> impl Future + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; } -pub struct Bounded(PhantomData); +pub struct Bounded( + PhantomData (C, T)>, +); -impl ConcurrentMpsc for Bounded { +impl, const CAPACITY: usize, T: Message> ConcurrentMpsc + for Bounded +{ + type Message = T; type Receiver = C::Receiver; type Sender = C::Sender; @@ -54,19 +89,19 @@ impl ConcurrentMpsc for Bounded usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { C::recv_async(receiver).await } - fn recv(receiver: &mut Self::Receiver) -> usize { + fn recv(receiver: &mut Self::Receiver) -> T { C::recv_blocking(receiver) } } @@ -74,6 +109,7 @@ impl ConcurrentMpsc for Bounded(PhantomData); impl ConcurrentMpsc for Unbounded { + type Message = usize; type Receiver = C::Receiver; type Sender = C::Sender; @@ -98,13 +134,13 @@ impl ConcurrentMpsc for Unbounded { } } -pub struct ConcurrentBatch { +pub struct ConcurrentBatch> { receiver: C::Receiver, start: Arc, workers: Vec>, } -impl ConcurrentBatch { +impl> ConcurrentBatch { pub fn new(producer_count: usize) -> Self { assert_eq!(BATCH_MESSAGES % producer_count, 0); @@ -143,7 +179,7 @@ impl ConcurrentBatch { } } -impl Drop for ConcurrentBatch { +impl> Drop for ConcurrentBatch { fn drop(&mut self) { let panicking = thread::panicking(); for worker in self.workers.drain(..) { @@ -156,14 +192,14 @@ impl Drop for ConcurrentBatch { } // Reuse worker threads and channel storage so steady-state samples exclude thread creation. -pub struct RepeatedBatch { +pub struct RepeatedBatch> { receiver: C::Receiver, start: Arc, stop: Arc, workers: Vec>, } -impl RepeatedBatch { +impl> RepeatedBatch { pub fn new(producer_count: usize) -> Self { assert_eq!(BATCH_MESSAGES % producer_count, 0); let (sender, receiver) = C::channel(); @@ -208,7 +244,7 @@ impl RepeatedBatch { } } -impl Drop for RepeatedBatch { +impl> Drop for RepeatedBatch { fn drop(&mut self) { self.stop.store(true, Ordering::Release); self.start.wait(); @@ -218,11 +254,21 @@ impl Drop for RepeatedBatch { } } -// Exercise executor wakeups as well as channel traffic. Reuse tasks, threads, and channel storage -// across samples; a current-thread runtime also exposes polling that monopolizes the executor. +// A spawned receiver shares the executor's scheduling with producers. Keeping the receiver in +// block_on instead measures worker-to-caller thread handoffs, which is a separate workload. +enum Receiver { + Task { + start: Arc, + completed: tokio::sync::mpsc::UnboundedReceiver, + }, + External(C::Receiver), +} + +// Reuse every task and the channel. The small control exchange happens once per 16,384-message +// batch; it never forwards measured messages. Both payload sizes use this same start protocol. pub struct RepeatedTasks { runtime: tokio::runtime::Runtime, - receiver: C::Receiver, + receiver: Receiver, start: Arc, stop: Arc, workers: Vec>, @@ -230,6 +276,18 @@ pub struct RepeatedTasks { impl RepeatedTasks { pub fn new(producer_count: usize, worker_threads: usize) -> Self { + Self::with_receiver(producer_count, worker_threads, false) + } + + pub fn external_receiver(producer_count: usize, worker_threads: usize) -> Self { + Self::with_receiver(producer_count, worker_threads, true) + } + + fn with_receiver( + producer_count: usize, + worker_threads: usize, + external_receiver: bool, + ) -> Self { assert_eq!(BATCH_MESSAGES % producer_count, 0); let runtime = if worker_threads == 0 { tokio::runtime::Builder::new_current_thread() @@ -241,11 +299,11 @@ impl RepeatedTasks { .build() .unwrap() }; - let (sender, receiver) = C::channel(); + let (sender, mut receiver) = C::channel(); let start = Arc::new(tokio::sync::Barrier::new(producer_count + 1)); let stop = Arc::new(AtomicBool::new(false)); let messages_per_producer = BATCH_MESSAGES / producer_count; - let workers = (0..producer_count) + let mut workers: Vec<_> = (0..producer_count) .map(|producer| { let sender = sender.clone(); let start = start.clone(); @@ -258,13 +316,38 @@ impl RepeatedTasks { } let first = producer * messages_per_producer; for offset in 0..messages_per_producer { - C::send_async(&sender, black_box(first + offset)).await; + C::send_async(&sender, black_box(C::Message::new(first + offset))) + .await; } } }) }) .collect(); drop(sender); + let receiver = if external_receiver { + Receiver::External(receiver) + } else { + let request = Arc::new(tokio::sync::Notify::new()); + let (completed_tx, completed) = tokio::sync::mpsc::unbounded_channel(); + let request_rx = request.clone(); + let start = start.clone(); + let stop = stop.clone(); + workers.push(runtime.spawn(async move { + loop { + request_rx.notified().await; + start.wait().await; + if stop.load(Ordering::Acquire) { + break; + } + let checksum = receive_batch::(&mut receiver).await; + completed_tx.send(checksum).unwrap(); + } + })); + Receiver::Task { + start: request, + completed, + } + }; Self { runtime, receiver, @@ -276,13 +359,16 @@ impl RepeatedTasks { pub fn run(&mut self) -> usize { self.runtime.block_on(async { - self.start.wait().await; - let mut checksum = 0usize; - for _ in 0..BATCH_MESSAGES { - checksum = checksum.wrapping_add(C::recv_async(&mut self.receiver).await); + match &mut self.receiver { + Receiver::Task { start, completed } => { + start.notify_one(); + completed.recv().await.expect("benchmark receiver panicked") + } + Receiver::External(receiver) => { + self.start.wait().await; + receive_batch::(receiver).await + } } - assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); - black_box(checksum) }) } } @@ -291,10 +377,24 @@ impl Drop for RepeatedTasks { fn drop(&mut self) { self.stop.store(true, Ordering::Release); self.runtime.block_on(async { - self.start.wait().await; + match &self.receiver { + Receiver::Task { start, .. } => start.notify_one(), + Receiver::External(_) => { + self.start.wait().await; + } + } for worker in self.workers.drain(..) { worker.await.expect("benchmark producer panicked"); } }); } } + +async fn receive_batch(receiver: &mut C::Receiver) -> usize { + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(C::recv_async(receiver).await.sequence()); + } + assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); + black_box(checksum) +} From 31c64cce7fe471850e7d18f3ad49c4e0f42c1bba Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 09:07:59 +0800 Subject: [PATCH 07/34] fixup Signed-off-by: tison --- asyncband/src/mpsc/bounded/buffer.rs | 10 +++++----- asyncband/src/mpsc/unbounded/buffer.rs | 10 ++++------ 2 files changed, 9 insertions(+), 11 deletions(-) diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index 87bb8bcc..bd9b98ea 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -70,7 +70,7 @@ impl std::panic::RefUnwindSafe for Slot {} impl Buffer { pub fn new(capacity: usize) -> Self { - let slots = if mem::size_of::() == 0 { + let slots = if size_of::() == 0 { Box::default() } else { (0..capacity.next_power_of_two()) @@ -110,7 +110,7 @@ impl Buffer { /// Own one capacity permit before calling; release it only after a failed push or after /// the consumer reads the published value. No user code runs between claim and publication. pub unsafe fn push(&self, value: T) -> Result<(), T> { - if mem::size_of::() == 0 { + if size_of::() == 0 { let mut queued = self.zero_sized.lock(); if self.closed.load(Ordering::Acquire) { return Err(value); @@ -153,7 +153,7 @@ impl Buffer { /// Only the exclusive consumer may call this, using its persistent cursor. Release one /// capacity permit after each successful pop, after the value has been read completely. pub unsafe fn pop(&self, head: &mut usize) -> Poll> { - if mem::size_of::() == 0 { + if size_of::() == 0 { let mut queued = self.zero_sized.lock(); return if *queued == 0 { Poll::Ready(None) @@ -185,7 +185,7 @@ impl Buffer { /// Only the exclusive consumer may close the buffer, once, using its current cursor. pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { self.closed.store(true, Ordering::Release); - let remaining = if mem::size_of::() == 0 { + let remaining = if size_of::() == 0 { mem::take(&mut *self.zero_sized.lock()) } else { // Cover every physical slot: a producer may have passed the open check but not @@ -220,7 +220,7 @@ impl Iterator for Drain<'_, T> { let position = self.position; self.remaining -= 1; self.position = self.position.wrapping_add(1); - if mem::size_of::() == 0 { + if size_of::() == 0 { // SAFETY: Closing transferred this many initialized ZST values to the drain. return Some(unsafe { Buffer::::read_zero_sized() }); } diff --git a/asyncband/src/mpsc/unbounded/buffer.rs b/asyncband/src/mpsc/unbounded/buffer.rs index 7c114495..8df2640f 100644 --- a/asyncband/src/mpsc/unbounded/buffer.rs +++ b/asyncband/src/mpsc/unbounded/buffer.rs @@ -38,10 +38,10 @@ impl Buffer { } fn segment_capacity() -> usize { - if mem::size_of::() == 0 { + if size_of::() == 0 { return usize::MAX; } - let limit = (SEGMENT_BYTES / mem::size_of::()).max(1); + 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()) } @@ -65,9 +65,7 @@ impl Buffer { // Keep one empty segment for the next producer rollover. Every other consumed // segment is released, so retained payload storage does not track peak occupancy. self.spare = mem::replace(batch, sealed); - if self.sealed.is_empty() - && self.sealed.capacity() * mem::size_of::>() > 1024 - { + if self.sealed.is_empty() && self.sealed.capacity() * size_of::>() > 1024 { self.sealed = VecDeque::new(); } } else if !self.writable.is_empty() { @@ -78,7 +76,7 @@ impl Buffer { } pub fn pop_batch(batch: &mut VecDeque) -> T { - if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::()) > SEGMENT_BYTES { + if batch.len() == 1 && batch.capacity().saturating_mul(size_of::()) > SEGMENT_BYTES { // Retire the allocation on the last value, outside the inbox lock. Keep this as a tail // expression to avoid intermediate storage for large inline values. mem::take(batch).pop_front() From df2b96a96781fa4e52c0ee5f7c0daf3e5e434202 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 09:33:15 +0800 Subject: [PATCH 08/34] test(mpsc): remove Kanal from ecosystem benchmarks Remove the Kanal adapter and dependency because its pending-operation cancellation semantics differ from the bounded MPSC contract. --- Cargo.lock | 11 ------- Cargo.toml | 1 - benchmarks/Cargo.toml | 1 - benchmarks/ecosystem/mpsc/adapters.rs | 42 --------------------------- benchmarks/ecosystem/mpsc/bounded.rs | 21 +++++++------- 5 files changed, 10 insertions(+), 66 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 78f58915..c6baa36c 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -105,7 +105,6 @@ dependencies = [ "asyncband", "divan", "flume", - "kanal", "pollster", "tokio", "waitgroup", @@ -565,16 +564,6 @@ version = "1.0.18" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" -[[package]] -name = "kanal" -version = "0.1.1" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9e3953adf0cd667798b396c2fa13552d6d9b3269d7dd1154c4c416442d1ff574" -dependencies = [ - "futures-core", - "lock_api", -] - [[package]] name = "libc" version = "0.2.189" diff --git a/Cargo.toml b/Cargo.toml index dbcbeda2..05e963ce 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -41,7 +41,6 @@ cargo_metadata = { version = "0.23.1" } clap = { version = "4.6.5" } divan = { version = "0.1.21" } flume = { version = "0.12.0", default-features = false } -kanal = { version = "0.1.1" } pollster = { version = "1.0.1" } semver = { version = "1.0.28" } serde = { version = "1.0.229" } diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index d7e6e467..19e5ac33 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -49,7 +49,6 @@ asyncband = { workspace = true, features = [ ] } divan = { workspace = true } flume = { workspace = true, features = ["async"] } -kanal = { workspace = true } pollster = { workspace = true } tokio = { workspace = true, features = ["rt-multi-thread", "sync"] } waitgroup = { workspace = true } diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs index 8dd61b7f..f9fe81b6 100644 --- a/benchmarks/ecosystem/mpsc/adapters.rs +++ b/benchmarks/ecosystem/mpsc/adapters.rs @@ -25,7 +25,6 @@ pub struct Asyncband; pub struct Tokio; pub struct AsyncChannel; pub struct Flume; -pub struct Kanal; pub trait BoundedMpsc: Send + Sync + 'static { type Sender: Clone + Send + Sync + 'static; @@ -218,47 +217,6 @@ impl BoundedMpsc for Flume { } } -impl BoundedMpsc for Kanal { - type Receiver = kanal::AsyncReceiver; - type Sender = kanal::AsyncSender; - - fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { - kanal::bounded_async(capacity) - } - - fn try_send(sender: &Self::Sender, value: T) { - assert!(sender.try_send(value).unwrap()); - } - - fn try_recv(receiver: &mut Self::Receiver) -> T { - receiver.try_recv().unwrap().unwrap() - } - - fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { - poll_ready(sender.send(value), context).unwrap(); - } - - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { - poll_ready(receiver.recv(), context).unwrap() - } - - async fn send_async(sender: &Self::Sender, value: T) { - sender.send(value).await.unwrap(); - } - - async fn recv_async(receiver: &mut Self::Receiver) -> T { - receiver.recv().await.unwrap() - } - - fn send_blocking(sender: &Self::Sender, value: T) { - pollster::block_on(sender.send(value)).unwrap(); - } - - fn recv_blocking(receiver: &mut Self::Receiver) -> T { - pollster::block_on(receiver.recv()).unwrap() - } -} - impl UnboundedMpsc for Asyncband { type Receiver = asyncband::mpsc::UnboundedReceiver; type Sender = asyncband::mpsc::UnboundedSender; diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index 95635419..c483e216 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -23,7 +23,6 @@ use super::adapters::AsyncChannel; use super::adapters::Asyncband; use super::adapters::BoundedMpsc; use super::adapters::Flume; -use super::adapters::Kanal; use super::adapters::Tokio; use super::support::BATCH_MESSAGES; use super::support::BOUNDED_CAPACITY; @@ -34,7 +33,7 @@ use super::support::RepeatedBatch; use super::support::RepeatedTasks; use crate::support::bench_context; -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] fn try_round_trip(bencher: Bencher) { let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); @@ -44,7 +43,7 @@ fn try_round_trip(bencher: Bencher) { }); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] fn ready_round_trip(bencher: Bencher) { let mut context = bench_context(); let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); @@ -56,7 +55,7 @@ fn ready_round_trip(bencher: Bencher) { } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + types = [Asyncband, Tokio, AsyncChannel, Flume], args = PRODUCER_COUNTS, sample_count = 20, sample_size = 1, @@ -69,7 +68,7 @@ fn concurrent(bencher: Bencher, producer_count: usize) { } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + types = [Asyncband, Tokio, AsyncChannel, Flume], args = PRODUCER_COUNTS, sample_count = 50, sample_size = 1, @@ -81,14 +80,14 @@ fn sustained(bencher: Bencher, producer_count: usize) { bencher.bench_local(|| batch.run()); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] fn clone_drop_sender(bencher: Bencher) { let (sender, _receiver) = C::channel(BOUNDED_CAPACITY); bencher.bench_local(|| drop(black_box(sender.clone()))); } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + types = [Asyncband, Tokio, AsyncChannel, Flume], consts = [1, 4096], args = PRODUCER_COUNTS, sample_count = 50, @@ -105,7 +104,7 @@ fn sustained_capacity( } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + types = [Asyncband, Tokio, AsyncChannel, Flume], consts = [1, 64, 4096], args = [(1, 0), (8, 0), (1, 4), (8, 4)], sample_count = 50, @@ -122,7 +121,7 @@ fn scheduled( } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + types = [Asyncband, Tokio, AsyncChannel, Flume], consts = [1, 64, 4096], args = [(1, 0), (8, 0), (1, 4), (8, 4)], sample_count = 50, @@ -139,7 +138,7 @@ fn scheduled_inline, const CAPACITY: usize>( } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + types = [Asyncband, Tokio, AsyncChannel, Flume], consts = [1, 64, 4096], args = [1, 8], sample_count = 50, @@ -153,7 +152,7 @@ fn external_receiver(bencher: Bencher, pr } #[divan::bench( - types = [Asyncband, Tokio, AsyncChannel, Flume, Kanal], + types = [Asyncband, Tokio, AsyncChannel, Flume], consts = [1, 64, 4096], args = [1, 8], sample_count = 50, From 04d2f0859456e1a76ef512b488159695758242a3 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 09:55:30 +0800 Subject: [PATCH 09/34] refactor(mpsc): back bounded storage with std sync_channel --- asyncband/src/mpsc/bounded/buffer.rs | 263 ------------------ asyncband/src/mpsc/bounded/buffer_tests.rs | 221 --------------- asyncband/src/mpsc/bounded/mod.rs | 92 +++--- asyncband/src/mpsc/bounded/storage.rs | 118 ++++++++ asyncband/src/mpsc/bounded/tests.rs | 103 +++++++ asyncband/src/mpsc/bounded/zero_sized.rs | 60 ++++ benchmarks/ecosystem/mpsc/bounded.rs | 10 + .../tests/mpsc_test/reservation.rs | 33 +++ 8 files changed, 357 insertions(+), 543 deletions(-) delete mode 100644 asyncband/src/mpsc/bounded/buffer.rs delete mode 100644 asyncband/src/mpsc/bounded/buffer_tests.rs create mode 100644 asyncband/src/mpsc/bounded/storage.rs create mode 100644 asyncband/src/mpsc/bounded/tests.rs create mode 100644 asyncband/src/mpsc/bounded/zero_sized.rs diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs deleted file mode 100644 index bd9b98ea..00000000 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ /dev/null @@ -1,263 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Capacity, position, and publication are separate ownership transitions: -//! -//! - A permit owns capacity, but holds no position until synchronous `push` claims a ticket. -//! - The ticket gives one producer a slot. `READY` publishes its initialized value to the receiver. -//! - The receiver finishes reading before returning capacity. AcqRel ticket increments carry that -//! reuse ordering even to a producer that acquired its permit on an earlier lap. -//! - Close competes with publication on the slot state. The drain owns `READY` values; a producer -//! that encounters `CLOSED` owns its unpublished value. Neither waits for the other to resume. -//! -//! Only the non-cloneable receiver advances the read cursor. All endpoints retain the shared -//! allocation, so a publisher's slot stays alive even when receiver drop closes it concurrently. - -use std::cell::UnsafeCell; -use std::mem; -use std::mem::MaybeUninit; -use std::ptr::NonNull; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicU8; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Poll; - -use crate::internal::cache_padded::CachePadded; -use crate::internal::mutex::Mutex; - -const EMPTY: u8 = 0; -const READY: u8 = 1; -const CLOSED: u8 = 2; - -pub struct Buffer { - slots: Box<[Slot]>, - tail: CachePadded, - closed: AtomicBool, - // ZSTs need no positions or per-slot flags. Counting them separately also allows every - // nonzero usize capacity without allocating publication metadata for nonexistent bytes. - zero_sized: Mutex, -} - -struct Slot { - state: AtomicU8, - value: UnsafeCell>, -} - -// SAFETY: Capacity and the tail ticket give a producer exclusive ownership of an empty slot. -// Release publication transfers its value to the exclusive consumer. Closing an unpublished -// slot leaves its value with the producer; closing a READY slot transfers it to the drain. -unsafe impl Sync for Slot {} - -// No reference to a stored value escapes. Every value is removed from the slot's ownership -// before running a callback or destructor that might panic. -impl std::panic::UnwindSafe for Slot {} -impl std::panic::RefUnwindSafe for Slot {} - -impl Buffer { - pub fn new(capacity: usize) -> Self { - let slots = if size_of::() == 0 { - Box::default() - } else { - (0..capacity.next_power_of_two()) - .map(|_| Slot { - state: AtomicU8::new(EMPTY), - value: UnsafeCell::new(MaybeUninit::uninit()), - }) - .collect() - }; - Self { - slots, - tail: CachePadded::new(AtomicUsize::new(0)), - closed: AtomicBool::new(false), - zero_sized: Mutex::new(0), - } - } - - fn slot(&self, position: usize) -> &Slot { - // Power-of-two storage preserves indexing when the full-width ticket wraps. The - // semaphore still enforces the exact requested capacity, including non-powers of two. - &self.slots[position & (self.slots.len() - 1)] - } - - fn claim(&self) -> Result { - if self.closed.load(Ordering::Acquire) { - return Err(()); - } - // Closing may race after this check. It marks every physical slot CLOSED, so even a - // delayed claimant will recover its own value instead of publishing into a dead queue. - Ok(self.tail.fetch_add(1, Ordering::AcqRel)) - } - - /// Writes and publishes one message. Closing may instead return the unsent value. - /// - /// # Safety - /// - /// Own one capacity permit before calling; release it only after a failed push or after - /// the consumer reads the published value. No user code runs between claim and publication. - pub unsafe fn push(&self, value: T) -> Result<(), T> { - if size_of::() == 0 { - let mut queued = self.zero_sized.lock(); - if self.closed.load(Ordering::Acquire) { - return Err(value); - } - *queued += 1; - mem::forget(value); - return Ok(()); - } - let Ok(position) = self.claim() else { - return Err(value); - }; - // SAFETY: The caller owns capacity and the atomic increment assigned this position. - unsafe { self.publish(position, value) } - } - - unsafe fn publish(&self, position: usize, value: T) -> Result<(), T> { - let slot = self.slot(position); - // SAFETY: Capacity prevents wrapping over unread slots. AcqRel tail increments carry prior - // claimants' capacity-acquire edges even when this producer held its permit for a long - // time. The previous consumer has therefore finished reading before this write. - unsafe { (*slot.value.get()).write(value) }; - match slot - .state - .compare_exchange(EMPTY, READY, Ordering::Release, Ordering::Acquire) - { - Ok(_) => Ok(()), - Err(state) => { - debug_assert_eq!(state, CLOSED); - // SAFETY: Close saw an unpublished slot and did not read it. Failed publication - // leaves exclusive ownership with this producer, including during receiver drop. - Err(unsafe { (*slot.value.get()).assume_init_read() }) - } - } - } - - /// Pending means a producer claimed the head but has not published it yet. - /// - /// # Safety - /// - /// Only the exclusive consumer may call this, using its persistent cursor. Release one - /// capacity permit after each successful pop, after the value has been read completely. - pub unsafe fn pop(&self, head: &mut usize) -> Poll> { - if size_of::() == 0 { - let mut queued = self.zero_sized.lock(); - return if *queued == 0 { - Poll::Ready(None) - } else { - *queued -= 1; - // SAFETY: A queued value proves that this ZST is inhabited and owns one value. - Poll::Ready(Some(unsafe { Self::read_zero_sized() })) - }; - } - let slot = self.slot(*head); - if slot.state.load(Ordering::Acquire) == READY { - // SAFETY: Publication initialized the value, and only this consumer can read it. - // Capacity is still held until this method has returned the value to its caller. - let value = unsafe { (*slot.value.get()).assume_init_read() }; - slot.state.store(EMPTY, Ordering::Release); - *head = head.wrapping_add(1); - Poll::Ready(Some(value)) - } else if self.tail.load(Ordering::Acquire) == *head { - Poll::Ready(None) - } else { - Poll::Pending - } - } - - /// Stops new claims and returns ownership of published values to a drain guard. - /// - /// # Safety - /// - /// Only the exclusive consumer may close the buffer, once, using its current cursor. - pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { - self.closed.store(true, Ordering::Release); - let remaining = if size_of::() == 0 { - mem::take(&mut *self.zero_sized.lock()) - } else { - // Cover every physical slot: a producer may have passed the open check but not - // obtained its ticket yet. Such a late claim must also find a CLOSED slot. - self.slots.len() - }; - Drain { - buffer: self, - position: head, - remaining, - } - } - - unsafe fn read_zero_sized() -> T { - // SAFETY: The caller owns an initialized, inhabited ZST. Reading it accesses no bytes; - // dangling supplies a non-null, correctly aligned pointer, as in a ZST Vec. - unsafe { NonNull::::dangling().as_ptr().read() } - } -} - -pub struct Drain<'a, T> { - buffer: &'a Buffer, - position: usize, - remaining: usize, -} - -impl Iterator for Drain<'_, T> { - type Item = T; - - fn next(&mut self) -> Option { - while self.remaining != 0 { - let position = self.position; - self.remaining -= 1; - self.position = self.position.wrapping_add(1); - if size_of::() == 0 { - // SAFETY: Closing transferred this many initialized ZST values to the drain. - return Some(unsafe { Buffer::::read_zero_sized() }); - } - let slot = self.buffer.slot(position); - if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { - // SAFETY: The drain won ownership of a published value. The cursor and state - // already advanced, so a panicking destructor cannot cause a second read. - return Some(unsafe { (*slot.value.get()).assume_init_read() }); - } - // An unpublished slot stays owned by its producer, which will observe CLOSED and - // recover its value. The shared Arc keeps this allocation alive until that finishes. - } - None - } -} - -impl Drop for Drain<'_, T> { - fn drop(&mut self) { - struct Remaining<'a, 'b, T>(&'a mut Drain<'b, T>); - - impl Drop for Remaining<'_, '_, T> { - fn drop(&mut self) { - for value in self.0.by_ref() { - drop(value); - } - } - } - - // A guard inside Drop is necessary: Drop itself is not called again if a payload's - // destructor panics while this normal drain is running. - let remaining = Remaining(self); - for value in remaining.0.by_ref() { - drop(value); - } - } -} - -#[cfg(test)] -#[path = "buffer_tests.rs"] -mod tests; diff --git a/asyncband/src/mpsc/bounded/buffer_tests.rs b/asyncband/src/mpsc/bounded/buffer_tests.rs deleted file mode 100644 index 43690029..00000000 --- a/asyncband/src/mpsc/bounded/buffer_tests.rs +++ /dev/null @@ -1,221 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; -use std::sync::Barrier; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use std::task::Waker; -use std::thread; - -use crate::mpsc::BoundedSender; -use crate::mpsc::Permit; -use crate::mpsc::TryRecvError; -use crate::mpsc::bounded; - -// Exercise the scheduling window inside synchronous send, while retaining a real capacity -// permit. Ordinary callers cannot split a claim from its publication. -fn publish_claimed(mut permit: Permit<'_, T>, position: usize, value: T) -> Result<(), T> { - let shared = &permit.sender.unwrap().shared; - // SAFETY: The test claimed this position while holding the same capacity permit. - unsafe { shared.buffer.publish(position, value) }?; - permit.sender = None; - shared.rx_waker.wake(); - Ok(()) -} - -#[test] -fn a_claimed_head_waits_for_publication_across_laps() { - for capacity in [1, 3, 7] { - for initial in [0, usize::MAX - 1] { - let (tx, mut rx) = bounded(capacity); - // Start an empty ring near ticket overflow instead of running usize::MAX sends. - tx.shared.buffer.tail.store(initial, Ordering::Relaxed); - rx.head = initial; - let mut cx = Context::from_waker(Waker::noop()); - for lap in 0..8 { - let permit = tx.try_reserve().unwrap(); - let position = tx.shared.buffer.claim().unwrap(); - for offset in 1..capacity { - tx.try_send(lap * capacity + offset).unwrap(); - } - // A full ring must differ from an empty one even if no head value is ready yet. - assert!(rx.poll_recv(&mut cx).is_pending()); - publish_claimed(permit, position, lap * capacity).unwrap(); - for offset in 0..capacity { - assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); - } - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - } - } - } -} - -#[test] -fn a_claim_delayed_past_close_returns_its_value() { - let (tx, rx) = bounded(3); - let permit = tx.try_reserve().unwrap(); - let drops = Arc::new(AtomicUsize::new(0)); - let value = Payload { - bytes: [7; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }; - let allocation = Arc::downgrade(&tx.shared); - // Pause after claim's open check, then resume its atomic ticket allocation after close. - assert!(!tx.shared.buffer.closed.load(Ordering::Acquire)); - drop(rx); - let position = tx.shared.buffer.tail.fetch_add(1, Ordering::AcqRel); - let unsent = publish_claimed(permit, position, value).unwrap_err(); - assert_eq!(unsent.bytes, [7; 1024]); - drop(unsent); - assert_eq!(drops.load(Ordering::Relaxed), 1); - drop(tx); - assert!(allocation.upgrade().is_none()); -} - -#[derive(Debug)] -#[repr(align(128))] -struct Payload { - bytes: [u8; 1024], - drops: Arc, - // Queued messages must not keep the shared allocation alive through a sender cycle. - _sender: BoundedSender, -} - -impl Drop for Payload { - fn drop(&mut self) { - self.drops.fetch_add(1, Ordering::Relaxed); - } -} - -#[test] -fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { - let (tx, rx) = bounded(2); - let allocation = Arc::downgrade(&tx.shared); - let drops = Arc::new(AtomicUsize::new(0)); - let paused = Barrier::new(2); - let (resume_tx, resume_rx) = std::sync::mpsc::channel(); - let (closed_tx, closed_rx) = std::sync::mpsc::channel(); - - thread::scope(|scope| { - let sender = &tx; - let drops = &drops; - let paused = &paused; - let publisher = scope.spawn(move || { - let permit = sender.try_reserve().unwrap(); - let position = sender.shared.buffer.claim().unwrap(); - let value = Payload { - bytes: [1; 1024], - drops: drops.clone(), - _sender: sender.clone(), - }; - paused.wait(); - resume_rx.recv().unwrap(); - let unsent = publish_claimed(permit, position, value).unwrap_err(); - assert_eq!(unsent.bytes, [1; 1024]); - drop(unsent); - }); - paused.wait(); - tx.try_send(Payload { - bytes: [2; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }) - .unwrap(); - let closer = scope.spawn(move || { - drop(rx); - closed_tx.send(()).unwrap(); - }); - #[cfg(not(miri))] - let closed = closed_rx.recv_timeout(std::time::Duration::from_secs(10)); - #[cfg(miri)] - let closed = closed_rx.recv(); - let dropped_before_resume = drops.load(Ordering::Relaxed); - // Unblock the publisher before asserting so a failed close cannot strand the scope. - resume_tx.send(()).unwrap(); - publisher.join().unwrap(); - closer.join().unwrap(); - assert!(closed.is_ok(), "close waited for the paused publisher"); - assert_eq!(dropped_before_resume, 1); - }); - - assert_eq!(drops.load(Ordering::Relaxed), 2); - drop(tx); - assert!(allocation.upgrade().is_none()); -} - -#[test] -fn publication_racing_with_close_drops_every_payload_once() { - for _ in 0..if cfg!(miri) { 8 } else { 128 } { - let (tx, rx) = bounded(3); - let allocation = Arc::downgrade(&tx.shared); - let drops = Arc::new(AtomicUsize::new(0)); - let start = Barrier::new(4); - thread::scope(|scope| { - for byte in 0..3 { - let permit = tx.try_reserve().unwrap(); - let value = Payload { - bytes: [byte; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }; - let start = &start; - scope.spawn(move || { - start.wait(); - if let Err(error) = permit.send(value) { - let value = error.into_inner(); - assert_eq!(value.bytes, [byte; 1024]); - drop(value); - } - }); - } - start.wait(); - drop(rx); - }); - assert_eq!(drops.load(Ordering::Relaxed), 3); - drop(tx); - assert!(allocation.upgrade().is_none()); - } -} - -#[test] -fn an_old_permit_can_publish_after_other_producers_wrap_the_ring() { - let (tx, mut rx) = bounded(3); - let old = tx.try_reserve().unwrap(); - for lap in 0..16 { - for offset in 0..2 { - tx.try_send([lap * 2 + offset; 1024]).unwrap(); - } - for offset in 0..2 { - assert_eq!(rx.try_recv(), Ok([lap * 2 + offset; 1024])); - } - } - thread::scope(|scope| { - scope - .spawn(move || old.send([42; 1024]).unwrap()) - .join() - .unwrap(); - }); - assert_eq!( - rx.poll_recv(&mut Context::from_waker(Waker::noop())), - Poll::Ready(Ok([42; 1024])) - ); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); -} diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index dfad2896..04522338 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -28,7 +28,6 @@ use std::task::Context; use std::task::Poll; use std::task::Waker; -use self::buffer::Buffer; use super::RecvError; use super::SendError; use super::TryRecvError; @@ -41,7 +40,11 @@ use crate::internal::waitlist::WaiterId; use crate::internal::wake_all; use crate::internal::waker_batch::WakerBatch; -mod buffer; +mod storage; +mod zero_sized; + +#[cfg(test)] +mod tests; /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// @@ -49,8 +52,8 @@ mod buffer; /// one slot for a waiting sender. Capacity is granted in the order that pending sends and /// reservations enter the wait queue; new senders cannot take an already granted slot. /// -/// Storage for nonzero-sized messages is preallocated and rounded up to a power of two; the -/// channel's capacity remains exactly `buffer`. Zero-sized messages need no per-slot storage. +/// Storage for nonzero-sized messages is preallocated. The channel's capacity is exactly +/// `buffer`. Zero-sized messages need no per-slot storage. /// /// # Panics /// @@ -59,16 +62,17 @@ mod buffer; #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); + let (sender, receiver) = storage::channel(buffer); let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), rx_waker: CachePadded::new(AtomicWaker::new()), - buffer: Buffer::new(buffer), + sender, }); let sender = BoundedSender { shared: shared.clone(), }; - let receiver = BoundedReceiver { shared, head: 0 }; + let receiver = BoundedReceiver { shared, receiver }; (sender, receiver) } @@ -76,7 +80,7 @@ struct Shared { senders: AtomicUsize, tx_permits: CachePadded, rx_waker: CachePadded, - buffer: Buffer, + sender: storage::Sender, } // This channel-local semaphore grants one permit at a time and can close its wait queue. @@ -407,9 +411,10 @@ impl Permit<'_, T> { /// If the receiver has been dropped, the returned error contains the unsent value. pub fn send(mut self, value: T) -> Result<(), SendError> { let shared = &self.sender.unwrap().shared; - // SAFETY: This permit owns one capacity unit. Claiming a slot and writing it is a - // synchronous operation with no user callbacks or await points between the two. - unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; + if shared.tx_permits.closed.load(Ordering::Acquire) { + return Err(SendError::new(value)); + } + shared.sender.send(value).map_err(SendError::new)?; // Publication owns the capacity before a wake callback can panic. self.sender = None; shared.rx_waker.wake(); @@ -428,11 +433,10 @@ impl Drop for Permit<'_, T> { /// The receiving endpoint of a bounded mpsc channel. /// /// Instances are created by the [`bounded`] function. -/// Dropping the receiver discards queued values. The backing allocation remains alive until -/// all endpoints are dropped, so a concurrent sender can safely finish returning an unsent value. +/// Dropping the receiver discards queued values and wakes senders waiting for capacity. pub struct BoundedReceiver { shared: Arc>, - head: usize, + receiver: storage::Receiver, } impl fmt::Debug for BoundedReceiver { @@ -443,14 +447,11 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - // SAFETY: Receiver ownership provides exclusive access to the consumption cursor. - // The drain first prevents new claims. Its destructor completes cleanup on unwinding. - let drain = unsafe { self.shared.buffer.close(self.head) }; let wakers = self.shared.tx_permits.close(); + self.shared.sender.close(); let receiver_waker = self.shared.rx_waker.take(); wake_all(wakers.into_iter()); drop(receiver_waker); - drop(drain); } } @@ -461,9 +462,6 @@ impl BoundedReceiver { /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has /// been dropped and all queued values have been consumed. /// - /// If a producer is still completing a synchronous publication at the queue head, this - /// method waits for that publication. Use [`Self::recv`] to wait asynchronously instead. - /// /// # Examples /// /// ``` @@ -481,44 +479,20 @@ impl BoundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - let mut spins = 0; - loop { - match self.try_pop() { - Poll::Ready(result) => return result, - Poll::Pending => { - // A synchronous publisher already owns the head. Reporting Empty here - // could hide a later send that has completed. Async recv parks instead. - if spins < 32 { - std::hint::spin_loop(); - spins += 1; - } else { - std::thread::yield_now(); - } - } - } - } - } - - fn try_pop(&mut self) -> Poll> { let mut disconnected = false; loop { - // SAFETY: Only this receiver owns head. Capacity is released after the buffer - // finishes reading and advances the cursor, so no producer can overwrite the value. - match unsafe { self.shared.buffer.pop(&mut self.head) } { - Poll::Ready(Some(value)) => { - self.shared.tx_permits.release(); - return Poll::Ready(Ok(value)); - } - Poll::Ready(None) if disconnected => { - return Poll::Ready(Err(TryRecvError::Disconnected)); - } - Poll::Ready(None) if self.shared.senders.load(Ordering::Acquire) == 0 => { - // Acquire the last sender's completed publications before checking again. - disconnected = true; - } - Poll::Ready(None) => return Poll::Ready(Err(TryRecvError::Empty)), - Poll::Pending => return Poll::Pending, + if let Some(value) = self.receiver.recv() { + self.shared.tx_permits.release(); + return Ok(value); + } + if disconnected { + return Err(TryRecvError::Disconnected); + } + if self.shared.senders.load(Ordering::Acquire) != 0 { + return Err(TryRecvError::Empty); } + // Acquire the last sender's completed publications before checking again. + disconnected = true; } } @@ -557,13 +531,13 @@ impl BoundedReceiver { fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { for registered in [false, true] { - match self.try_pop() { - Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), - Poll::Ready(Err(TryRecvError::Disconnected)) => { + match self.try_recv() { + Ok(value) => return Poll::Ready(Ok(value)), + Err(TryRecvError::Disconnected) => { drop(self.shared.rx_waker.take()); return Poll::Ready(Err(RecvError::Disconnected)); } - Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {} + Err(TryRecvError::Empty) => {} } if !registered { self.shared.rx_waker.register(cx.waker()); diff --git a/asyncband/src/mpsc/bounded/storage.rs b/asyncband/src/mpsc/bounded/storage.rs new file mode 100644 index 00000000..64f02e58 --- /dev/null +++ b/asyncband/src/mpsc/bounded/storage.rs @@ -0,0 +1,118 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; +use std::sync::PoisonError; +use std::sync::RwLock; +use std::sync::mpsc; + +use super::zero_sized; +use crate::internal::mutex::Mutex; + +pub fn channel(capacity: usize) -> (Sender, Receiver) { + if size_of::() == 0 { + let channel = Arc::new(zero_sized::Channel::new()); + ( + Sender::ZeroSized(channel.clone()), + Receiver::ZeroSized(channel), + ) + } else { + let (sender, receiver) = mpsc::sync_channel(capacity); + ( + Sender::Standard(RwLock::new(Some(sender))), + Receiver::Standard(Mutex::new(receiver)), + ) + } +} + +pub enum Sender { + Standard(RwLock>>), + ZeroSized(Arc>), +} + +impl Sender { + // The caller owns capacity until the receiver removes this message. + pub fn send(&self, value: T) -> Result<(), T> { + match self { + Self::Standard(sender) => { + let sender = sender.read().unwrap_or_else(PoisonError::into_inner); + let Some(sender) = sender.as_ref() else { + return Err(value); + }; + match sender.try_send(value) { + Ok(()) => Ok(()), + Err(mpsc::TrySendError::Disconnected(value)) => Err(value), + Err(mpsc::TrySendError::Full(_)) => { + unreachable!("a reserved capacity unit must fit in the backing channel") + } + } + } + Self::ZeroSized(channel) => channel.send(value), + } + } + + pub fn close(&self) { + match self { + Self::Standard(sender) => { + // Taking the sole STD sender waits for synchronous publications to finish and + // prevents new ones. The receiver can then drain without racing a late message. + let sender = sender + .write() + .unwrap_or_else(PoisonError::into_inner) + .take(); + drop(sender); + } + Self::ZeroSized(channel) => channel.close(), + } + } +} + +pub enum Receiver { + // The mutex preserves Sync for Send-only payloads. Receiving uses exclusive get_mut access. + Standard(Mutex>), + ZeroSized(Arc>), +} + +impl Receiver { + pub fn recv(&mut self) -> Option { + match self { + // The sender lives in shared state; endpoint disconnection is tracked by its owner. + Self::Standard(receiver) => receiver.get_mut().try_recv().ok(), + Self::ZeroSized(channel) => channel.recv(), + } + } +} + +impl Drop for Receiver { + fn drop(&mut self) { + // The owner closes the sender first. STD's receiver destructor may leak remaining + // messages if a payload destructor panics, so remove each value before dropping it. + // This guard finishes the drain during unwinding from the first destructor panic. + struct Drain<'a, T>(&'a mut Receiver); + impl Drop for Drain<'_, T> { + fn drop(&mut self) { + while let Some(value) = self.0.recv() { + drop(value); + } + } + } + let remaining = Drain(self); + while let Some(value) = remaining.0.recv() { + drop(value); + } + } +} diff --git a/asyncband/src/mpsc/bounded/tests.rs b/asyncband/src/mpsc/bounded/tests.rs new file mode 100644 index 00000000..96e2241e --- /dev/null +++ b/asyncband/src/mpsc/bounded/tests.rs @@ -0,0 +1,103 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; +use std::sync::Barrier; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; +use std::thread; + +use crate::mpsc::BoundedSender; +use crate::mpsc::TryRecvError; +use crate::mpsc::bounded; + +#[derive(Debug)] +#[repr(align(128))] +struct Payload { + bytes: [u8; 1024], + drops: Arc, + // Queued messages must not keep the shared allocation alive through a sender cycle. + _sender: BoundedSender, +} + +impl Drop for Payload { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::Relaxed); + } +} + +#[test] +fn publication_racing_with_close_drops_every_payload_once() { + for _ in 0..if cfg!(miri) { 8 } else { 128 } { + let (tx, rx) = bounded(3); + let allocation = Arc::downgrade(&tx.shared); + let drops = Arc::new(AtomicUsize::new(0)); + let start = Barrier::new(4); + thread::scope(|scope| { + for byte in 0..3 { + let permit = tx.try_reserve().unwrap(); + let value = Payload { + bytes: [byte; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let start = &start; + scope.spawn(move || { + start.wait(); + if let Err(error) = permit.send(value) { + let value = error.into_inner(); + assert_eq!(value.bytes, [byte; 1024]); + drop(value); + } + }); + } + start.wait(); + drop(rx); + }); + assert_eq!(drops.load(Ordering::Relaxed), 3); + drop(tx); + assert!(allocation.upgrade().is_none()); + } +} + +#[test] +fn a_held_permit_remains_usable_after_other_senders_make_progress() { + let (tx, mut rx) = bounded(3); + let old = tx.try_reserve().unwrap(); + for lap in 0..16 { + for offset in 0..2 { + tx.try_send([lap * 2 + offset; 1024]).unwrap(); + } + for offset in 0..2 { + assert_eq!(rx.try_recv(), Ok([lap * 2 + offset; 1024])); + } + } + thread::scope(|scope| { + scope + .spawn(move || old.send([42; 1024]).unwrap()) + .join() + .unwrap(); + }); + assert_eq!( + rx.poll_recv(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok([42; 1024])) + ); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); +} diff --git a/asyncband/src/mpsc/bounded/zero_sized.rs b/asyncband/src/mpsc/bounded/zero_sized.rs new file mode 100644 index 00000000..a55d312d --- /dev/null +++ b/asyncband/src/mpsc/bounded/zero_sized.rs @@ -0,0 +1,60 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::VecDeque; + +use crate::internal::mutex::Mutex; + +// VecDeque stores only a count for ZSTs, while handling alignment and destructors safely. +// Keeping this separate avoids allocating a standard-channel stamp for every zero-sized value. +pub struct Channel { + state: Mutex>, +} + +struct State { + values: VecDeque, + closed: bool, +} + +impl Channel { + pub fn new() -> Self { + debug_assert_eq!(size_of::(), 0); + Self { + state: Mutex::new(State { + values: VecDeque::new(), + closed: false, + }), + } + } + + pub fn send(&self, value: T) -> Result<(), T> { + let mut state = self.state.lock(); + if state.closed { + return Err(value); + } + state.values.push_back(value); + Ok(()) + } + + pub fn recv(&self) -> Option { + self.state.lock().values.pop_front() + } + + pub fn close(&self) { + self.state.lock().closed = true; + } +} diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index c483e216..8c21947b 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -43,6 +43,16 @@ fn try_round_trip(bencher: Bencher) { }); } +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +fn try_round_trip_zero_sized>(bencher: Bencher) { + let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); + + bencher.bench_local(|| { + C::try_send(black_box(&sender), ()); + C::try_recv(black_box(&mut receiver)); + }); +} + #[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] fn ready_round_trip(bencher: Bencher) { let mut context = bench_context(); diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index e26c1e81..0b49c0bc 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -82,6 +82,7 @@ fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { use std::sync::atomic::Ordering; static DROPS: AtomicUsize = AtomicUsize::new(0); + #[repr(align(128))] struct Message; impl Drop for Message { fn drop(&mut self) { @@ -101,6 +102,38 @@ fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { assert_eq!(DROPS.load(Ordering::Relaxed), 4); } +#[cfg(panic = "unwind")] +#[test] +fn zero_sized_messages_preserve_backpressure_and_cleanup_after_a_drop_panic() { + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + + static DROPS: AtomicUsize = AtomicUsize::new(0); + #[repr(align(128))] + struct Message; + impl Drop for Message { + fn drop(&mut self) { + if DROPS.fetch_add(1, Ordering::Relaxed) == 0 { + panic!("first zero-sized message destructor"); + } + } + } + + let (tx, rx) = mpsc::bounded(3); + for _ in 0..3 { + assert!(tx.try_send(Message).is_ok()); + } + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + let mut waiting = Box::pin(tx.reserve()); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(waiting.as_mut(), &waker).is_pending()); + + assert!(catch_unwind(|| drop(rx)).is_err()); + assert_eq!(DROPS.load(Ordering::Relaxed), 3); + assert_eq!(wakes.count(), 1); + assert!(expect_ready(poll_with(waiting.as_mut(), &waker)).is_err()); +} + #[test] fn released_capacity_is_granted_to_the_oldest_waiter() { let (tx, mut rx) = mpsc::bounded(1); From 9cb5d1edf8e028b2f2f1f1acda80a23ae82129b5 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 10:09:07 +0800 Subject: [PATCH 10/34] Revert "refactor(mpsc): back bounded storage with std sync_channel" This reverts commit 04d2f0859456e1a76ef512b488159695758242a3. --- asyncband/src/mpsc/bounded/buffer.rs | 263 ++++++++++++++++++ asyncband/src/mpsc/bounded/buffer_tests.rs | 221 +++++++++++++++ asyncband/src/mpsc/bounded/mod.rs | 92 +++--- asyncband/src/mpsc/bounded/storage.rs | 118 -------- asyncband/src/mpsc/bounded/tests.rs | 103 ------- asyncband/src/mpsc/bounded/zero_sized.rs | 60 ---- benchmarks/ecosystem/mpsc/bounded.rs | 10 - .../tests/mpsc_test/reservation.rs | 33 --- 8 files changed, 543 insertions(+), 357 deletions(-) create mode 100644 asyncband/src/mpsc/bounded/buffer.rs create mode 100644 asyncband/src/mpsc/bounded/buffer_tests.rs delete mode 100644 asyncband/src/mpsc/bounded/storage.rs delete mode 100644 asyncband/src/mpsc/bounded/tests.rs delete mode 100644 asyncband/src/mpsc/bounded/zero_sized.rs diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs new file mode 100644 index 00000000..bd9b98ea --- /dev/null +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -0,0 +1,263 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Capacity, position, and publication are separate ownership transitions: +//! +//! - A permit owns capacity, but holds no position until synchronous `push` claims a ticket. +//! - The ticket gives one producer a slot. `READY` publishes its initialized value to the receiver. +//! - The receiver finishes reading before returning capacity. AcqRel ticket increments carry that +//! reuse ordering even to a producer that acquired its permit on an earlier lap. +//! - Close competes with publication on the slot state. The drain owns `READY` values; a producer +//! that encounters `CLOSED` owns its unpublished value. Neither waits for the other to resume. +//! +//! Only the non-cloneable receiver advances the read cursor. All endpoints retain the shared +//! allocation, so a publisher's slot stays alive even when receiver drop closes it concurrently. + +use std::cell::UnsafeCell; +use std::mem; +use std::mem::MaybeUninit; +use std::ptr::NonNull; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU8; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Poll; + +use crate::internal::cache_padded::CachePadded; +use crate::internal::mutex::Mutex; + +const EMPTY: u8 = 0; +const READY: u8 = 1; +const CLOSED: u8 = 2; + +pub struct Buffer { + slots: Box<[Slot]>, + tail: CachePadded, + closed: AtomicBool, + // ZSTs need no positions or per-slot flags. Counting them separately also allows every + // nonzero usize capacity without allocating publication metadata for nonexistent bytes. + zero_sized: Mutex, +} + +struct Slot { + state: AtomicU8, + value: UnsafeCell>, +} + +// SAFETY: Capacity and the tail ticket give a producer exclusive ownership of an empty slot. +// Release publication transfers its value to the exclusive consumer. Closing an unpublished +// slot leaves its value with the producer; closing a READY slot transfers it to the drain. +unsafe impl Sync for Slot {} + +// No reference to a stored value escapes. Every value is removed from the slot's ownership +// before running a callback or destructor that might panic. +impl std::panic::UnwindSafe for Slot {} +impl std::panic::RefUnwindSafe for Slot {} + +impl Buffer { + pub fn new(capacity: usize) -> Self { + let slots = if size_of::() == 0 { + Box::default() + } else { + (0..capacity.next_power_of_two()) + .map(|_| Slot { + state: AtomicU8::new(EMPTY), + value: UnsafeCell::new(MaybeUninit::uninit()), + }) + .collect() + }; + Self { + slots, + tail: CachePadded::new(AtomicUsize::new(0)), + closed: AtomicBool::new(false), + zero_sized: Mutex::new(0), + } + } + + fn slot(&self, position: usize) -> &Slot { + // Power-of-two storage preserves indexing when the full-width ticket wraps. The + // semaphore still enforces the exact requested capacity, including non-powers of two. + &self.slots[position & (self.slots.len() - 1)] + } + + fn claim(&self) -> Result { + if self.closed.load(Ordering::Acquire) { + return Err(()); + } + // Closing may race after this check. It marks every physical slot CLOSED, so even a + // delayed claimant will recover its own value instead of publishing into a dead queue. + Ok(self.tail.fetch_add(1, Ordering::AcqRel)) + } + + /// Writes and publishes one message. Closing may instead return the unsent value. + /// + /// # Safety + /// + /// Own one capacity permit before calling; release it only after a failed push or after + /// the consumer reads the published value. No user code runs between claim and publication. + pub unsafe fn push(&self, value: T) -> Result<(), T> { + if size_of::() == 0 { + let mut queued = self.zero_sized.lock(); + if self.closed.load(Ordering::Acquire) { + return Err(value); + } + *queued += 1; + mem::forget(value); + return Ok(()); + } + let Ok(position) = self.claim() else { + return Err(value); + }; + // SAFETY: The caller owns capacity and the atomic increment assigned this position. + unsafe { self.publish(position, value) } + } + + unsafe fn publish(&self, position: usize, value: T) -> Result<(), T> { + let slot = self.slot(position); + // SAFETY: Capacity prevents wrapping over unread slots. AcqRel tail increments carry prior + // claimants' capacity-acquire edges even when this producer held its permit for a long + // time. The previous consumer has therefore finished reading before this write. + unsafe { (*slot.value.get()).write(value) }; + match slot + .state + .compare_exchange(EMPTY, READY, Ordering::Release, Ordering::Acquire) + { + Ok(_) => Ok(()), + Err(state) => { + debug_assert_eq!(state, CLOSED); + // SAFETY: Close saw an unpublished slot and did not read it. Failed publication + // leaves exclusive ownership with this producer, including during receiver drop. + Err(unsafe { (*slot.value.get()).assume_init_read() }) + } + } + } + + /// Pending means a producer claimed the head but has not published it yet. + /// + /// # Safety + /// + /// Only the exclusive consumer may call this, using its persistent cursor. Release one + /// capacity permit after each successful pop, after the value has been read completely. + pub unsafe fn pop(&self, head: &mut usize) -> Poll> { + if size_of::() == 0 { + let mut queued = self.zero_sized.lock(); + return if *queued == 0 { + Poll::Ready(None) + } else { + *queued -= 1; + // SAFETY: A queued value proves that this ZST is inhabited and owns one value. + Poll::Ready(Some(unsafe { Self::read_zero_sized() })) + }; + } + let slot = self.slot(*head); + if slot.state.load(Ordering::Acquire) == READY { + // SAFETY: Publication initialized the value, and only this consumer can read it. + // Capacity is still held until this method has returned the value to its caller. + let value = unsafe { (*slot.value.get()).assume_init_read() }; + slot.state.store(EMPTY, Ordering::Release); + *head = head.wrapping_add(1); + Poll::Ready(Some(value)) + } else if self.tail.load(Ordering::Acquire) == *head { + Poll::Ready(None) + } else { + Poll::Pending + } + } + + /// Stops new claims and returns ownership of published values to a drain guard. + /// + /// # Safety + /// + /// Only the exclusive consumer may close the buffer, once, using its current cursor. + pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { + self.closed.store(true, Ordering::Release); + let remaining = if size_of::() == 0 { + mem::take(&mut *self.zero_sized.lock()) + } else { + // Cover every physical slot: a producer may have passed the open check but not + // obtained its ticket yet. Such a late claim must also find a CLOSED slot. + self.slots.len() + }; + Drain { + buffer: self, + position: head, + remaining, + } + } + + unsafe fn read_zero_sized() -> T { + // SAFETY: The caller owns an initialized, inhabited ZST. Reading it accesses no bytes; + // dangling supplies a non-null, correctly aligned pointer, as in a ZST Vec. + unsafe { NonNull::::dangling().as_ptr().read() } + } +} + +pub struct Drain<'a, T> { + buffer: &'a Buffer, + position: usize, + remaining: usize, +} + +impl Iterator for Drain<'_, T> { + type Item = T; + + fn next(&mut self) -> Option { + while self.remaining != 0 { + let position = self.position; + self.remaining -= 1; + self.position = self.position.wrapping_add(1); + if size_of::() == 0 { + // SAFETY: Closing transferred this many initialized ZST values to the drain. + return Some(unsafe { Buffer::::read_zero_sized() }); + } + let slot = self.buffer.slot(position); + if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { + // SAFETY: The drain won ownership of a published value. The cursor and state + // already advanced, so a panicking destructor cannot cause a second read. + return Some(unsafe { (*slot.value.get()).assume_init_read() }); + } + // An unpublished slot stays owned by its producer, which will observe CLOSED and + // recover its value. The shared Arc keeps this allocation alive until that finishes. + } + None + } +} + +impl Drop for Drain<'_, T> { + fn drop(&mut self) { + struct Remaining<'a, 'b, T>(&'a mut Drain<'b, T>); + + impl Drop for Remaining<'_, '_, T> { + fn drop(&mut self) { + for value in self.0.by_ref() { + drop(value); + } + } + } + + // A guard inside Drop is necessary: Drop itself is not called again if a payload's + // destructor panics while this normal drain is running. + let remaining = Remaining(self); + for value in remaining.0.by_ref() { + drop(value); + } + } +} + +#[cfg(test)] +#[path = "buffer_tests.rs"] +mod tests; diff --git a/asyncband/src/mpsc/bounded/buffer_tests.rs b/asyncband/src/mpsc/bounded/buffer_tests.rs new file mode 100644 index 00000000..43690029 --- /dev/null +++ b/asyncband/src/mpsc/bounded/buffer_tests.rs @@ -0,0 +1,221 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; +use std::sync::Barrier; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; +use std::thread; + +use crate::mpsc::BoundedSender; +use crate::mpsc::Permit; +use crate::mpsc::TryRecvError; +use crate::mpsc::bounded; + +// Exercise the scheduling window inside synchronous send, while retaining a real capacity +// permit. Ordinary callers cannot split a claim from its publication. +fn publish_claimed(mut permit: Permit<'_, T>, position: usize, value: T) -> Result<(), T> { + let shared = &permit.sender.unwrap().shared; + // SAFETY: The test claimed this position while holding the same capacity permit. + unsafe { shared.buffer.publish(position, value) }?; + permit.sender = None; + shared.rx_waker.wake(); + Ok(()) +} + +#[test] +fn a_claimed_head_waits_for_publication_across_laps() { + for capacity in [1, 3, 7] { + for initial in [0, usize::MAX - 1] { + let (tx, mut rx) = bounded(capacity); + // Start an empty ring near ticket overflow instead of running usize::MAX sends. + tx.shared.buffer.tail.store(initial, Ordering::Relaxed); + rx.head = initial; + let mut cx = Context::from_waker(Waker::noop()); + for lap in 0..8 { + let permit = tx.try_reserve().unwrap(); + let position = tx.shared.buffer.claim().unwrap(); + for offset in 1..capacity { + tx.try_send(lap * capacity + offset).unwrap(); + } + // A full ring must differ from an empty one even if no head value is ready yet. + assert!(rx.poll_recv(&mut cx).is_pending()); + publish_claimed(permit, position, lap * capacity).unwrap(); + for offset in 0..capacity { + assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } + } + } +} + +#[test] +fn a_claim_delayed_past_close_returns_its_value() { + let (tx, rx) = bounded(3); + let permit = tx.try_reserve().unwrap(); + let drops = Arc::new(AtomicUsize::new(0)); + let value = Payload { + bytes: [7; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let allocation = Arc::downgrade(&tx.shared); + // Pause after claim's open check, then resume its atomic ticket allocation after close. + assert!(!tx.shared.buffer.closed.load(Ordering::Acquire)); + drop(rx); + let position = tx.shared.buffer.tail.fetch_add(1, Ordering::AcqRel); + let unsent = publish_claimed(permit, position, value).unwrap_err(); + assert_eq!(unsent.bytes, [7; 1024]); + drop(unsent); + assert_eq!(drops.load(Ordering::Relaxed), 1); + drop(tx); + assert!(allocation.upgrade().is_none()); +} + +#[derive(Debug)] +#[repr(align(128))] +struct Payload { + bytes: [u8; 1024], + drops: Arc, + // Queued messages must not keep the shared allocation alive through a sender cycle. + _sender: BoundedSender, +} + +impl Drop for Payload { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::Relaxed); + } +} + +#[test] +fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { + let (tx, rx) = bounded(2); + let allocation = Arc::downgrade(&tx.shared); + let drops = Arc::new(AtomicUsize::new(0)); + let paused = Barrier::new(2); + let (resume_tx, resume_rx) = std::sync::mpsc::channel(); + let (closed_tx, closed_rx) = std::sync::mpsc::channel(); + + thread::scope(|scope| { + let sender = &tx; + let drops = &drops; + let paused = &paused; + let publisher = scope.spawn(move || { + let permit = sender.try_reserve().unwrap(); + let position = sender.shared.buffer.claim().unwrap(); + let value = Payload { + bytes: [1; 1024], + drops: drops.clone(), + _sender: sender.clone(), + }; + paused.wait(); + resume_rx.recv().unwrap(); + let unsent = publish_claimed(permit, position, value).unwrap_err(); + assert_eq!(unsent.bytes, [1; 1024]); + drop(unsent); + }); + paused.wait(); + tx.try_send(Payload { + bytes: [2; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }) + .unwrap(); + let closer = scope.spawn(move || { + drop(rx); + closed_tx.send(()).unwrap(); + }); + #[cfg(not(miri))] + let closed = closed_rx.recv_timeout(std::time::Duration::from_secs(10)); + #[cfg(miri)] + let closed = closed_rx.recv(); + let dropped_before_resume = drops.load(Ordering::Relaxed); + // Unblock the publisher before asserting so a failed close cannot strand the scope. + resume_tx.send(()).unwrap(); + publisher.join().unwrap(); + closer.join().unwrap(); + assert!(closed.is_ok(), "close waited for the paused publisher"); + assert_eq!(dropped_before_resume, 1); + }); + + assert_eq!(drops.load(Ordering::Relaxed), 2); + drop(tx); + assert!(allocation.upgrade().is_none()); +} + +#[test] +fn publication_racing_with_close_drops_every_payload_once() { + for _ in 0..if cfg!(miri) { 8 } else { 128 } { + let (tx, rx) = bounded(3); + let allocation = Arc::downgrade(&tx.shared); + let drops = Arc::new(AtomicUsize::new(0)); + let start = Barrier::new(4); + thread::scope(|scope| { + for byte in 0..3 { + let permit = tx.try_reserve().unwrap(); + let value = Payload { + bytes: [byte; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let start = &start; + scope.spawn(move || { + start.wait(); + if let Err(error) = permit.send(value) { + let value = error.into_inner(); + assert_eq!(value.bytes, [byte; 1024]); + drop(value); + } + }); + } + start.wait(); + drop(rx); + }); + assert_eq!(drops.load(Ordering::Relaxed), 3); + drop(tx); + assert!(allocation.upgrade().is_none()); + } +} + +#[test] +fn an_old_permit_can_publish_after_other_producers_wrap_the_ring() { + let (tx, mut rx) = bounded(3); + let old = tx.try_reserve().unwrap(); + for lap in 0..16 { + for offset in 0..2 { + tx.try_send([lap * 2 + offset; 1024]).unwrap(); + } + for offset in 0..2 { + assert_eq!(rx.try_recv(), Ok([lap * 2 + offset; 1024])); + } + } + thread::scope(|scope| { + scope + .spawn(move || old.send([42; 1024]).unwrap()) + .join() + .unwrap(); + }); + assert_eq!( + rx.poll_recv(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok([42; 1024])) + ); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); +} diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 04522338..dfad2896 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -28,6 +28,7 @@ use std::task::Context; use std::task::Poll; use std::task::Waker; +use self::buffer::Buffer; use super::RecvError; use super::SendError; use super::TryRecvError; @@ -40,11 +41,7 @@ use crate::internal::waitlist::WaiterId; use crate::internal::wake_all; use crate::internal::waker_batch::WakerBatch; -mod storage; -mod zero_sized; - -#[cfg(test)] -mod tests; +mod buffer; /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// @@ -52,8 +49,8 @@ mod tests; /// one slot for a waiting sender. Capacity is granted in the order that pending sends and /// reservations enter the wait queue; new senders cannot take an already granted slot. /// -/// Storage for nonzero-sized messages is preallocated. The channel's capacity is exactly -/// `buffer`. Zero-sized messages need no per-slot storage. +/// Storage for nonzero-sized messages is preallocated and rounded up to a power of two; the +/// channel's capacity remains exactly `buffer`. Zero-sized messages need no per-slot storage. /// /// # Panics /// @@ -62,17 +59,16 @@ mod tests; #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); - let (sender, receiver) = storage::channel(buffer); let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), rx_waker: CachePadded::new(AtomicWaker::new()), - sender, + buffer: Buffer::new(buffer), }); let sender = BoundedSender { shared: shared.clone(), }; - let receiver = BoundedReceiver { shared, receiver }; + let receiver = BoundedReceiver { shared, head: 0 }; (sender, receiver) } @@ -80,7 +76,7 @@ struct Shared { senders: AtomicUsize, tx_permits: CachePadded, rx_waker: CachePadded, - sender: storage::Sender, + buffer: Buffer, } // This channel-local semaphore grants one permit at a time and can close its wait queue. @@ -411,10 +407,9 @@ impl Permit<'_, T> { /// If the receiver has been dropped, the returned error contains the unsent value. pub fn send(mut self, value: T) -> Result<(), SendError> { let shared = &self.sender.unwrap().shared; - if shared.tx_permits.closed.load(Ordering::Acquire) { - return Err(SendError::new(value)); - } - shared.sender.send(value).map_err(SendError::new)?; + // SAFETY: This permit owns one capacity unit. Claiming a slot and writing it is a + // synchronous operation with no user callbacks or await points between the two. + unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; // Publication owns the capacity before a wake callback can panic. self.sender = None; shared.rx_waker.wake(); @@ -433,10 +428,11 @@ impl Drop for Permit<'_, T> { /// The receiving endpoint of a bounded mpsc channel. /// /// Instances are created by the [`bounded`] function. -/// Dropping the receiver discards queued values and wakes senders waiting for capacity. +/// Dropping the receiver discards queued values. The backing allocation remains alive until +/// all endpoints are dropped, so a concurrent sender can safely finish returning an unsent value. pub struct BoundedReceiver { shared: Arc>, - receiver: storage::Receiver, + head: usize, } impl fmt::Debug for BoundedReceiver { @@ -447,11 +443,14 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { + // SAFETY: Receiver ownership provides exclusive access to the consumption cursor. + // The drain first prevents new claims. Its destructor completes cleanup on unwinding. + let drain = unsafe { self.shared.buffer.close(self.head) }; let wakers = self.shared.tx_permits.close(); - self.shared.sender.close(); let receiver_waker = self.shared.rx_waker.take(); wake_all(wakers.into_iter()); drop(receiver_waker); + drop(drain); } } @@ -462,6 +461,9 @@ impl BoundedReceiver { /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has /// been dropped and all queued values have been consumed. /// + /// If a producer is still completing a synchronous publication at the queue head, this + /// method waits for that publication. Use [`Self::recv`] to wait asynchronously instead. + /// /// # Examples /// /// ``` @@ -479,20 +481,44 @@ impl BoundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - let mut disconnected = false; + let mut spins = 0; loop { - if let Some(value) = self.receiver.recv() { - self.shared.tx_permits.release(); - return Ok(value); - } - if disconnected { - return Err(TryRecvError::Disconnected); + match self.try_pop() { + Poll::Ready(result) => return result, + Poll::Pending => { + // A synchronous publisher already owns the head. Reporting Empty here + // could hide a later send that has completed. Async recv parks instead. + if spins < 32 { + std::hint::spin_loop(); + spins += 1; + } else { + std::thread::yield_now(); + } + } } - if self.shared.senders.load(Ordering::Acquire) != 0 { - return Err(TryRecvError::Empty); + } + } + + fn try_pop(&mut self) -> Poll> { + let mut disconnected = false; + loop { + // SAFETY: Only this receiver owns head. Capacity is released after the buffer + // finishes reading and advances the cursor, so no producer can overwrite the value. + match unsafe { self.shared.buffer.pop(&mut self.head) } { + Poll::Ready(Some(value)) => { + self.shared.tx_permits.release(); + return Poll::Ready(Ok(value)); + } + Poll::Ready(None) if disconnected => { + return Poll::Ready(Err(TryRecvError::Disconnected)); + } + Poll::Ready(None) if self.shared.senders.load(Ordering::Acquire) == 0 => { + // Acquire the last sender's completed publications before checking again. + disconnected = true; + } + Poll::Ready(None) => return Poll::Ready(Err(TryRecvError::Empty)), + Poll::Pending => return Poll::Pending, } - // Acquire the last sender's completed publications before checking again. - disconnected = true; } } @@ -531,13 +557,13 @@ impl BoundedReceiver { fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { for registered in [false, true] { - match self.try_recv() { - Ok(value) => return Poll::Ready(Ok(value)), - Err(TryRecvError::Disconnected) => { + match self.try_pop() { + Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), + Poll::Ready(Err(TryRecvError::Disconnected)) => { drop(self.shared.rx_waker.take()); return Poll::Ready(Err(RecvError::Disconnected)); } - Err(TryRecvError::Empty) => {} + Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {} } if !registered { self.shared.rx_waker.register(cx.waker()); diff --git a/asyncband/src/mpsc/bounded/storage.rs b/asyncband/src/mpsc/bounded/storage.rs deleted file mode 100644 index 64f02e58..00000000 --- a/asyncband/src/mpsc/bounded/storage.rs +++ /dev/null @@ -1,118 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; -use std::sync::PoisonError; -use std::sync::RwLock; -use std::sync::mpsc; - -use super::zero_sized; -use crate::internal::mutex::Mutex; - -pub fn channel(capacity: usize) -> (Sender, Receiver) { - if size_of::() == 0 { - let channel = Arc::new(zero_sized::Channel::new()); - ( - Sender::ZeroSized(channel.clone()), - Receiver::ZeroSized(channel), - ) - } else { - let (sender, receiver) = mpsc::sync_channel(capacity); - ( - Sender::Standard(RwLock::new(Some(sender))), - Receiver::Standard(Mutex::new(receiver)), - ) - } -} - -pub enum Sender { - Standard(RwLock>>), - ZeroSized(Arc>), -} - -impl Sender { - // The caller owns capacity until the receiver removes this message. - pub fn send(&self, value: T) -> Result<(), T> { - match self { - Self::Standard(sender) => { - let sender = sender.read().unwrap_or_else(PoisonError::into_inner); - let Some(sender) = sender.as_ref() else { - return Err(value); - }; - match sender.try_send(value) { - Ok(()) => Ok(()), - Err(mpsc::TrySendError::Disconnected(value)) => Err(value), - Err(mpsc::TrySendError::Full(_)) => { - unreachable!("a reserved capacity unit must fit in the backing channel") - } - } - } - Self::ZeroSized(channel) => channel.send(value), - } - } - - pub fn close(&self) { - match self { - Self::Standard(sender) => { - // Taking the sole STD sender waits for synchronous publications to finish and - // prevents new ones. The receiver can then drain without racing a late message. - let sender = sender - .write() - .unwrap_or_else(PoisonError::into_inner) - .take(); - drop(sender); - } - Self::ZeroSized(channel) => channel.close(), - } - } -} - -pub enum Receiver { - // The mutex preserves Sync for Send-only payloads. Receiving uses exclusive get_mut access. - Standard(Mutex>), - ZeroSized(Arc>), -} - -impl Receiver { - pub fn recv(&mut self) -> Option { - match self { - // The sender lives in shared state; endpoint disconnection is tracked by its owner. - Self::Standard(receiver) => receiver.get_mut().try_recv().ok(), - Self::ZeroSized(channel) => channel.recv(), - } - } -} - -impl Drop for Receiver { - fn drop(&mut self) { - // The owner closes the sender first. STD's receiver destructor may leak remaining - // messages if a payload destructor panics, so remove each value before dropping it. - // This guard finishes the drain during unwinding from the first destructor panic. - struct Drain<'a, T>(&'a mut Receiver); - impl Drop for Drain<'_, T> { - fn drop(&mut self) { - while let Some(value) = self.0.recv() { - drop(value); - } - } - } - let remaining = Drain(self); - while let Some(value) = remaining.0.recv() { - drop(value); - } - } -} diff --git a/asyncband/src/mpsc/bounded/tests.rs b/asyncband/src/mpsc/bounded/tests.rs deleted file mode 100644 index 96e2241e..00000000 --- a/asyncband/src/mpsc/bounded/tests.rs +++ /dev/null @@ -1,103 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; -use std::sync::Barrier; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use std::task::Waker; -use std::thread; - -use crate::mpsc::BoundedSender; -use crate::mpsc::TryRecvError; -use crate::mpsc::bounded; - -#[derive(Debug)] -#[repr(align(128))] -struct Payload { - bytes: [u8; 1024], - drops: Arc, - // Queued messages must not keep the shared allocation alive through a sender cycle. - _sender: BoundedSender, -} - -impl Drop for Payload { - fn drop(&mut self) { - self.drops.fetch_add(1, Ordering::Relaxed); - } -} - -#[test] -fn publication_racing_with_close_drops_every_payload_once() { - for _ in 0..if cfg!(miri) { 8 } else { 128 } { - let (tx, rx) = bounded(3); - let allocation = Arc::downgrade(&tx.shared); - let drops = Arc::new(AtomicUsize::new(0)); - let start = Barrier::new(4); - thread::scope(|scope| { - for byte in 0..3 { - let permit = tx.try_reserve().unwrap(); - let value = Payload { - bytes: [byte; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }; - let start = &start; - scope.spawn(move || { - start.wait(); - if let Err(error) = permit.send(value) { - let value = error.into_inner(); - assert_eq!(value.bytes, [byte; 1024]); - drop(value); - } - }); - } - start.wait(); - drop(rx); - }); - assert_eq!(drops.load(Ordering::Relaxed), 3); - drop(tx); - assert!(allocation.upgrade().is_none()); - } -} - -#[test] -fn a_held_permit_remains_usable_after_other_senders_make_progress() { - let (tx, mut rx) = bounded(3); - let old = tx.try_reserve().unwrap(); - for lap in 0..16 { - for offset in 0..2 { - tx.try_send([lap * 2 + offset; 1024]).unwrap(); - } - for offset in 0..2 { - assert_eq!(rx.try_recv(), Ok([lap * 2 + offset; 1024])); - } - } - thread::scope(|scope| { - scope - .spawn(move || old.send([42; 1024]).unwrap()) - .join() - .unwrap(); - }); - assert_eq!( - rx.poll_recv(&mut Context::from_waker(Waker::noop())), - Poll::Ready(Ok([42; 1024])) - ); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); -} diff --git a/asyncband/src/mpsc/bounded/zero_sized.rs b/asyncband/src/mpsc/bounded/zero_sized.rs deleted file mode 100644 index a55d312d..00000000 --- a/asyncband/src/mpsc/bounded/zero_sized.rs +++ /dev/null @@ -1,60 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::collections::VecDeque; - -use crate::internal::mutex::Mutex; - -// VecDeque stores only a count for ZSTs, while handling alignment and destructors safely. -// Keeping this separate avoids allocating a standard-channel stamp for every zero-sized value. -pub struct Channel { - state: Mutex>, -} - -struct State { - values: VecDeque, - closed: bool, -} - -impl Channel { - pub fn new() -> Self { - debug_assert_eq!(size_of::(), 0); - Self { - state: Mutex::new(State { - values: VecDeque::new(), - closed: false, - }), - } - } - - pub fn send(&self, value: T) -> Result<(), T> { - let mut state = self.state.lock(); - if state.closed { - return Err(value); - } - state.values.push_back(value); - Ok(()) - } - - pub fn recv(&self) -> Option { - self.state.lock().values.pop_front() - } - - pub fn close(&self) { - self.state.lock().closed = true; - } -} diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index 8c21947b..c483e216 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -43,16 +43,6 @@ fn try_round_trip(bencher: Bencher) { }); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] -fn try_round_trip_zero_sized>(bencher: Bencher) { - let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); - - bencher.bench_local(|| { - C::try_send(black_box(&sender), ()); - C::try_recv(black_box(&mut receiver)); - }); -} - #[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] fn ready_round_trip(bencher: Bencher) { let mut context = bench_context(); diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index 0b49c0bc..e26c1e81 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -82,7 +82,6 @@ fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { use std::sync::atomic::Ordering; static DROPS: AtomicUsize = AtomicUsize::new(0); - #[repr(align(128))] struct Message; impl Drop for Message { fn drop(&mut self) { @@ -102,38 +101,6 @@ fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { assert_eq!(DROPS.load(Ordering::Relaxed), 4); } -#[cfg(panic = "unwind")] -#[test] -fn zero_sized_messages_preserve_backpressure_and_cleanup_after_a_drop_panic() { - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - - static DROPS: AtomicUsize = AtomicUsize::new(0); - #[repr(align(128))] - struct Message; - impl Drop for Message { - fn drop(&mut self) { - if DROPS.fetch_add(1, Ordering::Relaxed) == 0 { - panic!("first zero-sized message destructor"); - } - } - } - - let (tx, rx) = mpsc::bounded(3); - for _ in 0..3 { - assert!(tx.try_send(Message).is_ok()); - } - assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); - let mut waiting = Box::pin(tx.reserve()); - let (waker, wakes) = WakeCounter::new(); - assert!(poll_with(waiting.as_mut(), &waker).is_pending()); - - assert!(catch_unwind(|| drop(rx)).is_err()); - assert_eq!(DROPS.load(Ordering::Relaxed), 3); - assert_eq!(wakes.count(), 1); - assert!(expect_ready(poll_with(waiting.as_mut(), &waker)).is_err()); -} - #[test] fn released_capacity_is_granted_to_the_oldest_waiter() { let (tx, mut rx) = mpsc::bounded(1); From 7282b14734896d86ef01529ea7bca8496a96e1bc Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 10:26:44 +0800 Subject: [PATCH 11/34] fix(mpsc): reject oversized bounded capacity up front Rounded-up slot storage overflows a power of two and the shared permit counter packs channel state into two flag bits, so capacity is now bounded at usize::MAX >> 2 with an explicit panic message, replacing opaque arithmetic or allocation failures. Zero-sized messages allocate no slots and remain limited only by the permit counter. --- CHANGELOG.md | 1 + asyncband/src/mpsc/bounded/buffer.rs | 18 +++++++++++++ asyncband/src/mpsc/bounded/mod.rs | 15 +++++++++-- tests-integration/tests/mpsc_test/main.rs | 26 +++++++++++++++++++ .../tests/mpsc_test/reservation.rs | 4 +-- 5 files changed, 60 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 7113dc6b..d79caad7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,6 +16,7 @@ All notable changes to this project will be documented in this file. ### Improvements +* Reject bounded MPSC capacities above `usize::MAX >> 2` up front with an explicit panic message instead of an opaque arithmetic or allocation failure; zero-sized messages need no slot storage and remain limited only by the permit counter. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index bd9b98ea..28a8a9f5 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -69,6 +69,24 @@ impl std::panic::UnwindSafe for Slot {} impl std::panic::RefUnwindSafe for Slot {} impl Buffer { + /// Asserts that the preallocated slot storage for `capacity` fits one allocation. + /// + /// Zero-sized messages allocate no slots. Callers must already have bounded `capacity` so + /// that rounding up to a power of two cannot overflow. + pub fn check_allocation(capacity: usize) { + if size_of::() == 0 { + return; + } + let within_limit = capacity + .next_power_of_two() + .checked_mul(size_of::>()) + .is_some_and(|bytes| bytes <= isize::MAX as usize); + assert!( + within_limit, + "mpsc bounded channel capacity {capacity} exceeds the allocation limit" + ); + } + pub fn new(capacity: usize) -> Self { let slots = if size_of::() == 0 { Box::default() diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index dfad2896..9e3c1796 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -54,11 +54,16 @@ mod buffer; /// /// # Panics /// -/// Panics if `buffer` is zero or the preallocated message buffer exceeds the allocation size -/// limit. There is no additional channel-specific capacity limit. +/// Panics if `buffer` is zero or exceeds the maximum capacity of `usize::MAX >> 2`, or if the +/// rounded-up message buffer would exceed the allocation size limit. #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); + assert!( + buffer <= MAX_CAPACITY, + "mpsc bounded channel capacity {buffer} exceeds the maximum of {MAX_CAPACITY}" + ); + Buffer::::check_allocation(buffer); let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), @@ -79,6 +84,12 @@ struct Shared { buffer: Buffer, } +/// The largest capacity accepted by [`bounded`]. +/// +/// The shared permit counter packs channel state into two flag bits, which also keeps the +/// rounded-up slot storage from overflowing a power of two. +const MAX_CAPACITY: usize = usize::MAX >> 2; + // This channel-local semaphore grants one permit at a time and can close its wait queue. // The general-purpose semaphore has neither a close operation nor acquisition errors. struct Semaphore { diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index 6b1268e3..524d70b8 100644 --- a/tests-integration/tests/mpsc_test/main.rs +++ b/tests-integration/tests/mpsc_test/main.rs @@ -174,3 +174,29 @@ fn receives_wake_for_messages_and_the_last_sender_drop() { Poll::Ready(Err(RecvError::Disconnected)) ); } + +#[test] +#[should_panic(expected = "mpsc bounded channel requires buffer > 0")] +fn bounded_rejects_zero_capacity() { + let _ = mpsc::bounded::(0); +} + +#[test] +#[should_panic(expected = "exceeds the maximum")] +fn bounded_rejects_capacity_above_the_maximum() { + let _ = mpsc::bounded::((usize::MAX >> 2) + 1); +} + +#[test] +#[should_panic(expected = "exceeds the allocation limit")] +fn bounded_rejects_capacity_above_the_allocation_limit() { + // Within the maximum capacity, but the rounded-up slot storage cannot fit one allocation. + let _ = mpsc::bounded::<[u64; 4]>(usize::MAX >> 2); +} + +#[test] +fn bounded_zero_sized_messages_need_no_slot_storage() { + let (tx, mut rx) = mpsc::bounded::<()>(usize::MAX >> 2); + tx.try_send(()).unwrap(); + assert_eq!(rx.try_recv(), Ok(())); +} diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index e26c1e81..bda19072 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -58,7 +58,7 @@ fn held_permits_consume_capacity_without_claiming_message_order() { #[test] fn zero_sized_messages_support_the_full_capacity_range() { - for capacity in [usize::MAX / 4 + 1, usize::MAX / 2 + 1, usize::MAX] { + for capacity in [(usize::MAX >> 2) - 1, usize::MAX >> 2] { let (tx, mut rx) = mpsc::bounded::<()>(capacity); let permit = tx.try_reserve().unwrap(); tx.try_send(()).unwrap(); @@ -89,7 +89,7 @@ fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { } } - let (tx, mut rx) = mpsc::bounded(usize::MAX); + let (tx, mut rx) = mpsc::bounded(usize::MAX >> 2); for _ in 0..3 { assert!(tx.try_send(Message).is_ok()); } From a8cf7497e34154b1e3605a245c4fe34eaac7e900 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 10:33:28 +0800 Subject: [PATCH 12/34] perf(mpsc): pack bounded semaphore state into one atomic The permit counter, closed flag, and a may-have-waiters flag now share a single atomic state. Acquire and release are each one lock-free operation when no sender waits; releases only take the wait-queue lock to hand capacity directly to the oldest waiter, preserving registration order. A registration sets the waiting flag before its final capacity recheck so a racing release switches to the locked path and no permit is stranded without a wake. --- asyncband/src/mpsc/bounded/mod.rs | 159 ++++++++++++++++++++---------- 1 file changed, 105 insertions(+), 54 deletions(-) diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 9e3c1796..df923980 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -21,7 +21,6 @@ use std::fmt; use std::future::poll_fn; use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -90,45 +89,89 @@ struct Shared { /// rounded-up slot storage from overflowing a power of two. const MAX_CAPACITY: usize = usize::MAX >> 2; -// This channel-local semaphore grants one permit at a time and can close its wait queue. +// This channel-local semaphore packs its permit counter and channel state into one atomic. // The general-purpose semaphore has neither a close operation nor acquisition errors. +// +// `state` holds the available permits shifted left by two, plus two flag bits: +// +// * `CLOSED`: the receiver is gone. No permits are issued or returned, and waiters drain with an +// error. +// * `WAITING`: the wait queue may be non-empty. Releases then take the locked path and grant the +// permit directly to the oldest waiter instead of returning it to the counter, so capacity is +// handed out in registration order and new arrivals cannot steal an already granted slot. +// +// With neither flag set, acquire and release are single lock-free operations on `state`. +// Wait-queue mutations always hold the queue lock; a registration sets `WAITING` before its +// final capacity recheck, which switches any racing release to the locked path and strands +// no permit without a wake. struct Semaphore { - available: AtomicUsize, - closed: AtomicBool, + state: AtomicUsize, waiters: Mutex>, } +const CLOSED: usize = 0b01; +const WAITING: usize = 0b10; +const PERMIT: usize = 0b100; + impl Semaphore { fn new(available: usize) -> Self { Self { - available: AtomicUsize::new(available), - closed: AtomicBool::new(false), + state: AtomicUsize::new(available * PERMIT), waiters: Mutex::new(WaitList::new()), } } fn try_acquire(&self) -> Result<(), TrySendError<()>> { - if self.closed.load(Ordering::Acquire) { - return Err(TrySendError::Disconnected(())); - } - let mut available = self.available.load(Ordering::Relaxed); + let mut state = self.state.load(Ordering::Acquire); loop { - if available == 0 { + if state & CLOSED != 0 { + return Err(TrySendError::Disconnected(())); + } + if state < PERMIT { return Err(TrySendError::Full(())); } - match self.available.compare_exchange_weak( - available, - available - 1, + match self.state.compare_exchange_weak( + state, + state - PERMIT, + Ordering::Acquire, Ordering::Acquire, - Ordering::Relaxed, ) { Ok(_) => return Ok(()), - Err(actual) => available = actual, + Err(actual) => state = actual, } } } + fn is_closed(&self) -> bool { + self.state.load(Ordering::Acquire) & CLOSED != 0 + } + + fn set_waiting(&self) { + self.state.fetch_or(WAITING, Ordering::AcqRel); + } + + fn clear_waiting(&self) { + self.state.fetch_and(!WAITING, Ordering::Release); + } + fn release(&self) { + // Fast path: with no waiting sender and no close in sight, the permit goes straight + // back to the counter. + let mut state = self.state.load(Ordering::Relaxed); + loop { + if state & (WAITING | CLOSED) != 0 { + break; + } + match self.state.compare_exchange_weak( + state, + state + PERMIT, + Ordering::Release, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(actual) => state = actual, + } + } let wake = self.release_locked(&mut self.waiters.lock()); if let Some(waker) = wake { waker.wake(); @@ -136,23 +179,29 @@ impl Semaphore { } fn release_locked(&self, waiters: &mut WaitList) -> Option { - if self.closed.load(Ordering::Relaxed) { + if self.is_closed() { return None; } if let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { // Grant ownership before waking; new arrivals cannot steal this capacity. waiter.granted = true; - return waiter.waker.take(); + let waker = waiter.waker.take(); + if waiters.is_empty() { + self.clear_waiting(); + } + return waker; } - // Only releases add permits, and all releases hold the wait queue lock. A linked - // waiter therefore always sees zero available permits until it receives its own grant. - self.available.fetch_add(1, Ordering::Release); + // The queue is empty. Only releases add permits, and the counter grows with the + // queue locked, so a linked waiter always sees zero available permits until it + // receives its own grant. An outstanding grant already owns its capacity. + self.state.fetch_add(PERMIT, Ordering::Release); + self.clear_waiting(); None } fn close(&self) -> WakerBatch { let mut waiters = self.waiters.lock(); - self.closed.store(true, Ordering::Release); + self.state.fetch_or(CLOSED, Ordering::AcqRel); let mut wakers = WakerBatch::new(); while let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { if let Some(waker) = waiter.waker.take() { @@ -177,40 +226,33 @@ impl<'a, T> Reservation<'a, T> { fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { let semaphore = &self.sender.shared.tx_permits; let mut cloned_waker = None; - loop { + let result = loop { if self.index.is_none() { match semaphore.try_acquire() { Ok(()) => { - let permit = Permit { + break Ok(Permit { sender: Some(self.sender), - }; - // The permit owns capacity before an unused cloned waker can panic. - drop(cloned_waker); - return Poll::Ready(Ok(permit)); - } - Err(TrySendError::Disconnected(())) => { - return Poll::Ready(Err(SendError::new(()))); + }); } + Err(TrySendError::Disconnected(())) => break Err(SendError::new(())), Err(TrySendError::Full(())) => {} } } let mut waiters = semaphore.waiters.lock(); - if semaphore.closed.load(Ordering::Relaxed) { + if semaphore.is_closed() { // Drop removes any remaining registration, including an unused grant. - return Poll::Ready(Err(SendError::new(()))); + break Err(SendError::new(())); } if let Some(index) = self.index { let waiter = waiters.waiter_mut(index); if waiter.granted { let waiter = waiters.remove_unlinked_waiter(index); self.index = None; - let permit = Permit { - sender: Some(self.sender), - }; drop(waiters); drop(waiter); - drop(cloned_waker); - return Poll::Ready(Ok(permit)); + break Ok(Permit { + sender: Some(self.sender), + }); } if waiter .waker @@ -225,26 +267,32 @@ impl<'a, T> Reservation<'a, T> { drop(old); return Poll::Pending; } - } else if semaphore.try_acquire().is_ok() { - // A release may have raced with the fast path; recheck under the queue lock - // before committing to wait so no permit can be stranded without a wake. - let permit = Permit { - sender: Some(self.sender), - }; - drop(waiters); - drop(cloned_waker); - return Poll::Ready(Ok(permit)); - } else if let Some(waker) = cloned_waker.take() { - self.index = Some(waiters.push_back(Waiter { - granted: false, - waker: Some(waker), - })); - return Poll::Pending; + } else { + // Set WAITING before the final capacity recheck: a racing release switches + // to the locked path, so no permit can be stranded without a wake. Waiting + // senders already in the queue take priority over this recheck. + semaphore.set_waiting(); + if waiters.is_empty() && semaphore.try_acquire().is_ok() { + semaphore.clear_waiting(); + break Ok(Permit { + sender: Some(self.sender), + }); + } + if let Some(waker) = cloned_waker.take() { + self.index = Some(waiters.push_back(Waiter { + granted: false, + waker: Some(waker), + })); + return Poll::Pending; + } } drop(waiters); // Clone outside the lock, then recheck capacity and closure before registering. cloned_waker = Some(cx.waker().clone()); - } + }; + // The permit owns capacity before an unused cloned waker can panic. + drop(cloned_waker); + Poll::Ready(result) } } @@ -259,6 +307,9 @@ impl Drop for Reservation<'_, T> { let wake = if waiter.granted { semaphore.release_locked(&mut waiters) } else { + if waiters.is_empty() { + semaphore.clear_waiting(); + } None }; (waiter, wake) From b2c4ff62c7324868de9d4deb879e82b1cd58333d Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 10:38:12 +0800 Subject: [PATCH 13/34] perf(mpsc): skip the receiver wake when it is not parked Publishing a message previously ran the receiver waker's two read-modify-write operations even when no receiver was parked. A publishable parked flag now gates the wake: both flag accesses are SeqCst, so a skipped wake implies the receiver's post-store recheck observes the published message and no wake is lost. --- asyncband/src/mpsc/bounded/buffer_tests.rs | 2 +- asyncband/src/mpsc/bounded/mod.rs | 45 ++++++++++++++++++---- 2 files changed, 39 insertions(+), 8 deletions(-) diff --git a/asyncband/src/mpsc/bounded/buffer_tests.rs b/asyncband/src/mpsc/bounded/buffer_tests.rs index 43690029..ac33fa39 100644 --- a/asyncband/src/mpsc/bounded/buffer_tests.rs +++ b/asyncband/src/mpsc/bounded/buffer_tests.rs @@ -36,7 +36,7 @@ fn publish_claimed(mut permit: Permit<'_, T>, position: usize, value: T) -> R // SAFETY: The test claimed this position while holding the same capacity permit. unsafe { shared.buffer.publish(position, value) }?; permit.sender = None; - shared.rx_waker.wake(); + shared.rx_wake.wake_parked(); Ok(()) } diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index df923980..923f09ed 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -21,6 +21,7 @@ use std::fmt; use std::future::poll_fn; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -66,7 +67,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), - rx_waker: CachePadded::new(AtomicWaker::new()), + rx_wake: CachePadded::new(RxWake::new()), buffer: Buffer::new(buffer), }); let sender = BoundedSender { @@ -79,10 +80,36 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { struct Shared { senders: AtomicUsize, tx_permits: CachePadded, - rx_waker: CachePadded, + rx_wake: CachePadded, buffer: Buffer, } +/// The receiver's wake registration and its publishable intent to park. +/// +/// `parked` and each slot publication are both SeqCst operations. A producer that observes +/// `parked` unset therefore precedes the receiver's store in the total order, so its +/// publication happens-before the receiver's post-store recheck and cannot be missed. A +/// skipped wake thus always pairs with a recheck that observes the published message. +struct RxWake { + waker: AtomicWaker, + parked: AtomicBool, +} + +impl RxWake { + fn new() -> Self { + Self { + waker: AtomicWaker::new(), + parked: AtomicBool::new(false), + } + } + + fn wake_parked(&self) { + if self.parked.swap(false, Ordering::SeqCst) { + self.waker.wake(); + } + } +} + /// The largest capacity accepted by [`bounded`]. /// /// The shared permit counter packs channel state into two flag bits, which also keeps the @@ -346,7 +373,7 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { if self.shared.senders.fetch_sub(1, Ordering::AcqRel) == 1 { - self.shared.rx_waker.wake(); + self.shared.rx_wake.wake_parked(); } } } @@ -474,7 +501,7 @@ impl Permit<'_, T> { unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; // Publication owns the capacity before a wake callback can panic. self.sender = None; - shared.rx_waker.wake(); + shared.rx_wake.wake_parked(); Ok(()) } } @@ -509,7 +536,7 @@ impl Drop for BoundedReceiver { // The drain first prevents new claims. Its destructor completes cleanup on unwinding. let drain = unsafe { self.shared.buffer.close(self.head) }; let wakers = self.shared.tx_permits.close(); - let receiver_waker = self.shared.rx_waker.take(); + let receiver_waker = self.shared.rx_wake.waker.take(); wake_all(wakers.into_iter()); drop(receiver_waker); drop(drain); @@ -622,13 +649,17 @@ impl BoundedReceiver { match self.try_pop() { Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), Poll::Ready(Err(TryRecvError::Disconnected)) => { - drop(self.shared.rx_waker.take()); + drop(self.shared.rx_wake.waker.take()); return Poll::Ready(Err(RecvError::Disconnected)); } Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {} } if !registered { - self.shared.rx_waker.register(cx.waker()); + // Publish the intent to park before rechecking the queue: a producer that + // observes this flag after publishing completes the wake, and one that does + // not has its publication observed by the recheck below. + self.shared.rx_wake.parked.store(true, Ordering::SeqCst); + self.shared.rx_wake.waker.register(cx.waker()); } } Poll::Pending From 24ae4c8a99507d0eec48ac99ff8493140730754e Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 10:40:55 +0800 Subject: [PATCH 14/34] perf(mpsc): count zero-sized bounded messages atomically The zero-sized path packs its message count and a closed flag into one atomic word, replacing the last mutex inside the bounded buffer. Close and publication stay atomic: a publication ahead of the flag is counted in the drain, and every later one observes the flag and fails. --- asyncband/src/mpsc/bounded/buffer.rs | 45 ++++++++++++++++++++-------- 1 file changed, 32 insertions(+), 13 deletions(-) diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index 28a8a9f5..8a18fa02 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -38,11 +38,11 @@ use std::sync::atomic::Ordering; use std::task::Poll; use crate::internal::cache_padded::CachePadded; -use crate::internal::mutex::Mutex; const EMPTY: u8 = 0; const READY: u8 = 1; const CLOSED: u8 = 2; +const ZERO_SIZED_CLOSED: usize = 1 << (usize::BITS - 1); pub struct Buffer { slots: Box<[Slot]>, @@ -50,7 +50,8 @@ pub struct Buffer { closed: AtomicBool, // ZSTs need no positions or per-slot flags. Counting them separately also allows every // nonzero usize capacity without allocating publication metadata for nonexistent bytes. - zero_sized: Mutex, + // The count packs a closed flag into its top bit so publication and close stay atomic. + zero_sized: AtomicUsize, } struct Slot { @@ -102,7 +103,7 @@ impl Buffer { slots, tail: CachePadded::new(AtomicUsize::new(0)), closed: AtomicBool::new(false), - zero_sized: Mutex::new(0), + zero_sized: AtomicUsize::new(0), } } @@ -129,13 +130,25 @@ impl Buffer { /// the consumer reads the published value. No user code runs between claim and publication. pub unsafe fn push(&self, value: T) -> Result<(), T> { if size_of::() == 0 { - let mut queued = self.zero_sized.lock(); - if self.closed.load(Ordering::Acquire) { - return Err(value); + let mut queued = self.zero_sized.load(Ordering::Acquire); + loop { + if queued & ZERO_SIZED_CLOSED != 0 { + return Err(value); + } + // The capacity limit keeps the count far below the closed flag bit. + match self.zero_sized.compare_exchange_weak( + queued, + queued + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => { + mem::forget(value); + return Ok(()); + } + Err(actual) => queued = actual, + } } - *queued += 1; - mem::forget(value); - return Ok(()); } let Ok(position) = self.claim() else { return Err(value); @@ -172,11 +185,13 @@ impl Buffer { /// capacity permit after each successful pop, after the value has been read completely. pub unsafe fn pop(&self, head: &mut usize) -> Poll> { if size_of::() == 0 { - let mut queued = self.zero_sized.lock(); - return if *queued == 0 { + let queued = self.zero_sized.load(Ordering::Acquire) & !ZERO_SIZED_CLOSED; + return if queued == 0 { Poll::Ready(None) } else { - *queued -= 1; + // Only this consumer decrements, and producers can only add: the count + // observed above is a lower bound, so this cannot wrap. + self.zero_sized.fetch_sub(1, Ordering::AcqRel); // SAFETY: A queued value proves that this ZST is inhabited and owns one value. Poll::Ready(Some(unsafe { Self::read_zero_sized() })) }; @@ -204,7 +219,11 @@ impl Buffer { pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { self.closed.store(true, Ordering::Release); let remaining = if size_of::() == 0 { - mem::take(&mut *self.zero_sized.lock()) + // Counting stops with the closed flag: a publication that raced ahead of it is + // included in the count, and every later one observes the flag and fails. + self.zero_sized + .fetch_or(ZERO_SIZED_CLOSED, Ordering::AcqRel) + & !ZERO_SIZED_CLOSED } else { // Cover every physical slot: a producer may have passed the open check but not // obtained its ticket yet. Such a late claim must also find a CLOSED slot. From 03c2336e06e01c9ad77d09bc90f6b5d7fc002f62 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 10:51:06 +0800 Subject: [PATCH 15/34] Revert "perf(mpsc): skip the receiver wake when it is not parked" The parked flag is unsound in the C++ memory model: slot publication is only a Release operation, so the SeqCst total order on the flag does not bind it. Miri finds a store-buffering interleaving where the producer observes the flag unset (skipping the wake) while the receiver's recheck still observes the pre-publication slot, stranding a parked receiver (concurrency::publication_racing_with_receiver_registration_cannot_lose_wakeup, concurrency::last_sender_drop_racing_with_receiver_registration_cannot_lose_wakeup). Closing the hole needs a SeqCst fence per publish, which costs as much as the wake it replaces. Keep the unconditional wake; Tokio does the same. --- asyncband/src/mpsc/bounded/buffer_tests.rs | 2 +- asyncband/src/mpsc/bounded/mod.rs | 45 ++++------------------ 2 files changed, 8 insertions(+), 39 deletions(-) diff --git a/asyncband/src/mpsc/bounded/buffer_tests.rs b/asyncband/src/mpsc/bounded/buffer_tests.rs index ac33fa39..43690029 100644 --- a/asyncband/src/mpsc/bounded/buffer_tests.rs +++ b/asyncband/src/mpsc/bounded/buffer_tests.rs @@ -36,7 +36,7 @@ fn publish_claimed(mut permit: Permit<'_, T>, position: usize, value: T) -> R // SAFETY: The test claimed this position while holding the same capacity permit. unsafe { shared.buffer.publish(position, value) }?; permit.sender = None; - shared.rx_wake.wake_parked(); + shared.rx_waker.wake(); Ok(()) } diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 923f09ed..df923980 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -21,7 +21,6 @@ use std::fmt; use std::future::poll_fn; use std::sync::Arc; -use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -67,7 +66,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), - rx_wake: CachePadded::new(RxWake::new()), + rx_waker: CachePadded::new(AtomicWaker::new()), buffer: Buffer::new(buffer), }); let sender = BoundedSender { @@ -80,36 +79,10 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { struct Shared { senders: AtomicUsize, tx_permits: CachePadded, - rx_wake: CachePadded, + rx_waker: CachePadded, buffer: Buffer, } -/// The receiver's wake registration and its publishable intent to park. -/// -/// `parked` and each slot publication are both SeqCst operations. A producer that observes -/// `parked` unset therefore precedes the receiver's store in the total order, so its -/// publication happens-before the receiver's post-store recheck and cannot be missed. A -/// skipped wake thus always pairs with a recheck that observes the published message. -struct RxWake { - waker: AtomicWaker, - parked: AtomicBool, -} - -impl RxWake { - fn new() -> Self { - Self { - waker: AtomicWaker::new(), - parked: AtomicBool::new(false), - } - } - - fn wake_parked(&self) { - if self.parked.swap(false, Ordering::SeqCst) { - self.waker.wake(); - } - } -} - /// The largest capacity accepted by [`bounded`]. /// /// The shared permit counter packs channel state into two flag bits, which also keeps the @@ -373,7 +346,7 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { if self.shared.senders.fetch_sub(1, Ordering::AcqRel) == 1 { - self.shared.rx_wake.wake_parked(); + self.shared.rx_waker.wake(); } } } @@ -501,7 +474,7 @@ impl Permit<'_, T> { unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; // Publication owns the capacity before a wake callback can panic. self.sender = None; - shared.rx_wake.wake_parked(); + shared.rx_waker.wake(); Ok(()) } } @@ -536,7 +509,7 @@ impl Drop for BoundedReceiver { // The drain first prevents new claims. Its destructor completes cleanup on unwinding. let drain = unsafe { self.shared.buffer.close(self.head) }; let wakers = self.shared.tx_permits.close(); - let receiver_waker = self.shared.rx_wake.waker.take(); + let receiver_waker = self.shared.rx_waker.take(); wake_all(wakers.into_iter()); drop(receiver_waker); drop(drain); @@ -649,17 +622,13 @@ impl BoundedReceiver { match self.try_pop() { Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), Poll::Ready(Err(TryRecvError::Disconnected)) => { - drop(self.shared.rx_wake.waker.take()); + drop(self.shared.rx_waker.take()); return Poll::Ready(Err(RecvError::Disconnected)); } Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {} } if !registered { - // Publish the intent to park before rechecking the queue: a producer that - // observes this flag after publishing completes the wake, and one that does - // not has its publication observed by the recheck below. - self.shared.rx_wake.parked.store(true, Ordering::SeqCst); - self.shared.rx_wake.waker.register(cx.waker()); + self.shared.rx_waker.register(cx.waker()); } } Poll::Pending From bcd7129c25e664ff395edfad2226a0af71edfeed Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 11:00:46 +0800 Subject: [PATCH 16/34] test(mpsc): pin sample sizes for nanosecond-scale benches Divan's auto-tuning starts at one iteration per sample and stops once a sample exceeds 100x timer precision. A cold first call (lazy initialization, cache misses) can cross that threshold immediately, ending tuning at sample_size = 1. Every sample then quantizes to one timer tick (41 ns on this machine): try_round_trip reported 41.74 ns medians for Tokio and AsyncChannel, and cancel_reserved_capacity reported 40.75 ns for Asyncband, with means dominated by cold outliers. Which bench hits this varies from run to run, so pin sample_size on all nanosecond-scale benches, in line with the sizes auto-tuning converges to on healthy runs. Medians now agree with independent measurements (e.g. cancel_reserved_capacity: Asyncband 5.11 ns, Tokio 7.10 ns). --- benchmarks/ecosystem/mpsc/bounded.rs | 10 +++++++--- benchmarks/ecosystem/mpsc/reservation.rs | 5 +++-- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index c483e216..44d2cd05 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -33,7 +33,11 @@ use super::support::RepeatedBatch; use super::support::RepeatedTasks; use crate::support::bench_context; -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +// Nanosecond-scale benches pin `sample_size`: divan's auto-tuning starts at 1 +// iteration per sample and stops once a sample exceeds 100x timer precision, +// so a cold first call (lazy initialization, cache misses) can end tuning +// immediately and quantize every sample to one timer tick (41 ns on macOS). +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume], sample_size = 512)] fn try_round_trip(bencher: Bencher) { let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); @@ -43,7 +47,7 @@ fn try_round_trip(bencher: Bencher) { }); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume], sample_size = 256)] fn ready_round_trip(bencher: Bencher) { let mut context = bench_context(); let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); @@ -80,7 +84,7 @@ fn sustained(bencher: Bencher, producer_count: usize) { bencher.bench_local(|| batch.run()); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume], sample_size = 2048)] fn clone_drop_sender(bencher: Bencher) { let (sender, _receiver) = C::channel(BOUNDED_CAPACITY); bencher.bench_local(|| drop(black_box(sender.clone()))); diff --git a/benchmarks/ecosystem/mpsc/reservation.rs b/benchmarks/ecosystem/mpsc/reservation.rs index 5952d2a6..32defce4 100644 --- a/benchmarks/ecosystem/mpsc/reservation.rs +++ b/benchmarks/ecosystem/mpsc/reservation.rs @@ -64,7 +64,8 @@ impl Reservable for Tokio { } } -#[divan::bench(types = [Asyncband, Tokio])] +// `sample_size` is pinned for the reason documented in `bounded.rs`. +#[divan::bench(types = [Asyncband, Tokio], sample_size = 512)] fn reserve_publish_receive(bencher: Bencher) { let (sender, mut receiver) = C::channel(64); let mut context = bench_context(); @@ -75,7 +76,7 @@ fn reserve_publish_receive(bencher: Bencher) { }); } -#[divan::bench(types = [Asyncband, Tokio])] +#[divan::bench(types = [Asyncband, Tokio], sample_size = 1024)] fn cancel_reserved_capacity(bencher: Bencher) { let (sender, _receiver) = C::channel(64); bencher.bench_local(|| drop(black_box(C::try_reserve(&sender)))); From 385876b5e3287dc37860d025c3fa13362a33314c Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 11:01:27 +0800 Subject: [PATCH 17/34] docs: note bounded MPSC throughput improvements in the changelog --- CHANGELOG.md | 1 + 1 file changed, 1 insertion(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index d79caad7..196552ef 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,6 +19,7 @@ All notable changes to this project will be documented in this file. * Reject bounded MPSC capacities above `usize::MAX >> 2` up front with an explicit panic message instead of an opaque arithmetic or allocation failure; zero-sized messages need no slot storage and remain limited only by the permit counter. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. +* Improve bounded MPSC throughput: acquiring and releasing capacity no longer takes an internal lock while no sender is waiting, and zero-sized messages no longer take the buffer lock. ## v0.7.2 From 9febe2adcb832ca1a36f1c6a4c9d8ecf606383f3 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 11:49:37 +0800 Subject: [PATCH 18/34] refactor(mpsc): drop the bounded allocation pre-check The pre-flight check duplicates what allocation already reports: collecting the slot array panics with "capacity overflow" when the layout exceeds the allocation limit, and larger requests abort like any other Rust allocation failure. Keep only the `usize::MAX >> 2` capacity bound, which protects the packed permit counter and the power-of-two rounding. --- CHANGELOG.md | 2 +- asyncband/src/mpsc/bounded/buffer.rs | 18 ------------------ asyncband/src/mpsc/bounded/mod.rs | 4 +--- tests-integration/tests/mpsc_test/main.rs | 7 ------- 4 files changed, 2 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 196552ef..e06fc8a9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,7 @@ All notable changes to this project will be documented in this file. ### Improvements -* Reject bounded MPSC capacities above `usize::MAX >> 2` up front with an explicit panic message instead of an opaque arithmetic or allocation failure; zero-sized messages need no slot storage and remain limited only by the permit counter. +* Reject bounded MPSC capacities above `usize::MAX >> 2` up front with an explicit panic message instead of an opaque arithmetic overflow; zero-sized messages need no slot storage and remain limited only by the permit counter. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. * Improve bounded MPSC throughput: acquiring and releasing capacity no longer takes an internal lock while no sender is waiting, and zero-sized messages no longer take the buffer lock. diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index 8a18fa02..aa300b71 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -70,24 +70,6 @@ impl std::panic::UnwindSafe for Slot {} impl std::panic::RefUnwindSafe for Slot {} impl Buffer { - /// Asserts that the preallocated slot storage for `capacity` fits one allocation. - /// - /// Zero-sized messages allocate no slots. Callers must already have bounded `capacity` so - /// that rounding up to a power of two cannot overflow. - pub fn check_allocation(capacity: usize) { - if size_of::() == 0 { - return; - } - let within_limit = capacity - .next_power_of_two() - .checked_mul(size_of::>()) - .is_some_and(|bytes| bytes <= isize::MAX as usize); - assert!( - within_limit, - "mpsc bounded channel capacity {capacity} exceeds the allocation limit" - ); - } - pub fn new(capacity: usize) -> Self { let slots = if size_of::() == 0 { Box::default() diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index df923980..481abab5 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -53,8 +53,7 @@ mod buffer; /// /// # Panics /// -/// Panics if `buffer` is zero or exceeds the maximum capacity of `usize::MAX >> 2`, or if the -/// rounded-up message buffer would exceed the allocation size limit. +/// Panics if `buffer` is zero or exceeds the maximum capacity of `usize::MAX >> 2`. #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); @@ -62,7 +61,6 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { buffer <= MAX_CAPACITY, "mpsc bounded channel capacity {buffer} exceeds the maximum of {MAX_CAPACITY}" ); - Buffer::::check_allocation(buffer); let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index 524d70b8..3a6079f3 100644 --- a/tests-integration/tests/mpsc_test/main.rs +++ b/tests-integration/tests/mpsc_test/main.rs @@ -187,13 +187,6 @@ fn bounded_rejects_capacity_above_the_maximum() { let _ = mpsc::bounded::((usize::MAX >> 2) + 1); } -#[test] -#[should_panic(expected = "exceeds the allocation limit")] -fn bounded_rejects_capacity_above_the_allocation_limit() { - // Within the maximum capacity, but the rounded-up slot storage cannot fit one allocation. - let _ = mpsc::bounded::<[u64; 4]>(usize::MAX >> 2); -} - #[test] fn bounded_zero_sized_messages_need_no_slot_storage() { let (tx, mut rx) = mpsc::bounded::<()>(usize::MAX >> 2); From ee540c8eea3277b000e67ed624a1c322c0f9eefc Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 11:51:28 +0800 Subject: [PATCH 19/34] perf(mpsc): restore the receiver waker state with a plain store After a wake claims the WAKING bit, the state can only be WAKING: registration enters from WAITING and concurrent wakes keep the bit set, so the restoring swap always reads the value it replaces. A Release store publishes the emptied slot the same way and drops one read-modify-write from every message publication. --- asyncband/src/internal/atomic_waker.rs | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs index a6dcd20f..a8c0448b 100644 --- a/asyncband/src/internal/atomic_waker.rs +++ b/asyncband/src/internal/atomic_waker.rs @@ -57,7 +57,7 @@ const WAKING: usize = 0b10; /// REGISTERING ------------AcqRel CAS----------------> WAITING /// /// wake: WAITING ----------------AcqRel fetch_or-----------> WAKING -/// WAKING -----------------Release swap--------------> WAITING +/// WAKING -----------------Release store-------------> WAITING /// /// race: REGISTERING ------------AcqRel fetch_or-----------> REGISTERING | WAKING /// REGISTERING | WAKING ---AcqRel swap---------------> WAITING @@ -224,9 +224,15 @@ impl AtomicWaker { let waker = unsafe { (*self.waker.get()).take() }; // ORDERING: Release publishes the emptied slot before another operation acquires - // it. The fetch_or above already performed the required Acquire operation. - let old_state = self.state.swap(WAITING, Ordering::Release); - debug_assert_eq!(old_state, WAKING); + // it. The fetch_or above already performed the required Acquire operation. A + // plain store suffices: only this claim moves the state out of WAKING, because + // registration enters from WAITING and concurrent wakes keep the bit set. Debug + // builds pay for a swap to assert that invariant. + if cfg!(debug_assertions) { + debug_assert_eq!(self.state.swap(WAITING, Ordering::Release), WAKING); + } else { + self.state.store(WAITING, Ordering::Release); + } waker } state => { From afb7c7fb4e7fb61ff823e422c8d83109ca5977b3 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 11:54:01 +0800 Subject: [PATCH 20/34] refactor(mpsc): give zero-sized bounded storage its own type The packed queue length moves out of the slot ring into a ZeroSized storage value whose three operations own its invariants: publication fails once the closed flag is set, consumption relies on the count being a lower bound, and close transfers the remaining count to the drain. Dispatch on size_of::() == 0 is const-folded, so generated code is unchanged. --- asyncband/src/mpsc/bounded/buffer.rs | 104 +++++++++++++++++---------- 1 file changed, 68 insertions(+), 36 deletions(-) diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index aa300b71..ef4cfdbf 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -48,10 +48,7 @@ pub struct Buffer { slots: Box<[Slot]>, tail: CachePadded, closed: AtomicBool, - // ZSTs need no positions or per-slot flags. Counting them separately also allows every - // nonzero usize capacity without allocating publication metadata for nonexistent bytes. - // The count packs a closed flag into its top bit so publication and close stay atomic. - zero_sized: AtomicUsize, + zero_sized: ZeroSized, } struct Slot { @@ -69,6 +66,62 @@ unsafe impl Sync for Slot {} impl std::panic::UnwindSafe for Slot {} impl std::panic::RefUnwindSafe for Slot {} +/// Queue storage for zero-sized messages, which need no slots, positions, or per-slot flags. +/// Counting them separately also allows every nonzero usize capacity without allocating +/// publication metadata for nonexistent bytes. +/// +/// The queue is entirely its length. The count packs a closed flag into its top bit so that +/// publication and close stay atomic: a publication that raced ahead of the flag is included +/// in the drained count, and every later one observes the flag and fails. +struct ZeroSized { + queued: AtomicUsize, +} + +impl ZeroSized { + fn new() -> Self { + Self { + queued: AtomicUsize::new(0), + } + } + + /// Accounts for one published message, returning `false` once the queue is closed. + fn push(&self) -> bool { + let mut queued = self.queued.load(Ordering::Acquire); + loop { + if queued & ZERO_SIZED_CLOSED != 0 { + return false; + } + // The capacity limit keeps the count far below the closed flag bit. + match self.queued.compare_exchange_weak( + queued, + queued + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(actual) => queued = actual, + } + } + } + + /// Accounts for one consumed message, returning `false` when the queue was observed empty. + fn pop(&self) -> bool { + let queued = self.queued.load(Ordering::Acquire) & !ZERO_SIZED_CLOSED; + if queued == 0 { + return false; + } + // Only the consumer decrements, and producers can only add: the count observed above + // is a lower bound, so this cannot wrap. + self.queued.fetch_sub(1, Ordering::AcqRel); + true + } + + /// Stops publication and returns the queue length transferred to the drain. + fn close(&self) -> usize { + self.queued.fetch_or(ZERO_SIZED_CLOSED, Ordering::AcqRel) & !ZERO_SIZED_CLOSED + } +} + impl Buffer { pub fn new(capacity: usize) -> Self { let slots = if size_of::() == 0 { @@ -85,7 +138,7 @@ impl Buffer { slots, tail: CachePadded::new(AtomicUsize::new(0)), closed: AtomicBool::new(false), - zero_sized: AtomicUsize::new(0), + zero_sized: ZeroSized::new(), } } @@ -112,25 +165,12 @@ impl Buffer { /// the consumer reads the published value. No user code runs between claim and publication. pub unsafe fn push(&self, value: T) -> Result<(), T> { if size_of::() == 0 { - let mut queued = self.zero_sized.load(Ordering::Acquire); - loop { - if queued & ZERO_SIZED_CLOSED != 0 { - return Err(value); - } - // The capacity limit keeps the count far below the closed flag bit. - match self.zero_sized.compare_exchange_weak( - queued, - queued + 1, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => { - mem::forget(value); - return Ok(()); - } - Err(actual) => queued = actual, - } - } + return if self.zero_sized.push() { + mem::forget(value); + Ok(()) + } else { + Err(value) + }; } let Ok(position) = self.claim() else { return Err(value); @@ -167,15 +207,11 @@ impl Buffer { /// capacity permit after each successful pop, after the value has been read completely. pub unsafe fn pop(&self, head: &mut usize) -> Poll> { if size_of::() == 0 { - let queued = self.zero_sized.load(Ordering::Acquire) & !ZERO_SIZED_CLOSED; - return if queued == 0 { - Poll::Ready(None) - } else { - // Only this consumer decrements, and producers can only add: the count - // observed above is a lower bound, so this cannot wrap. - self.zero_sized.fetch_sub(1, Ordering::AcqRel); + return if self.zero_sized.pop() { // SAFETY: A queued value proves that this ZST is inhabited and owns one value. Poll::Ready(Some(unsafe { Self::read_zero_sized() })) + } else { + Poll::Ready(None) }; } let slot = self.slot(*head); @@ -201,11 +237,7 @@ impl Buffer { pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { self.closed.store(true, Ordering::Release); let remaining = if size_of::() == 0 { - // Counting stops with the closed flag: a publication that raced ahead of it is - // included in the count, and every later one observes the flag and fails. - self.zero_sized - .fetch_or(ZERO_SIZED_CLOSED, Ordering::AcqRel) - & !ZERO_SIZED_CLOSED + self.zero_sized.close() } else { // Cover every physical slot: a producer may have passed the open check but not // obtained its ticket yet. Such a late claim must also find a CLOSED slot. From 99373fe6ae21871e16a2f8d314b0e256dca2f349 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 12:23:41 +0800 Subject: [PATCH 21/34] refactor(mpsc): encode bounded semaphore sentinels as values, not bits WAITING and CLOSED become the two largest usize values instead of flag bits below a shifted permit counter. The usable range is no longer narrowed by the encoding, so the capacity ceiling rises from usize::MAX >> 2 to usize::MAX >> 1, now bound by the power-of-two slot rounding and the zero-sized queue's packed closed bit. Installing WAITING is a compare exchange over the exhausted counter: a permit that arrives first wins and is picked up by the registering sender's recheck, keeping the no-stranded-permit invariant. Hot-path operation counts are unchanged. --- CHANGELOG.md | 2 +- asyncband/src/mpsc/bounded/mod.rs | 80 ++++++++++++------- tests-integration/tests/mpsc_test/main.rs | 4 +- .../tests/mpsc_test/reservation.rs | 4 +- 4 files changed, 54 insertions(+), 36 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e06fc8a9..3b133146 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,7 @@ All notable changes to this project will be documented in this file. ### Improvements -* Reject bounded MPSC capacities above `usize::MAX >> 2` up front with an explicit panic message instead of an opaque arithmetic overflow; zero-sized messages need no slot storage and remain limited only by the permit counter. +* Reject bounded MPSC capacities above `usize::MAX >> 1` up front with an explicit panic message instead of an opaque arithmetic overflow; zero-sized messages need no slot storage and remain limited only by the permit counter. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. * Improve bounded MPSC throughput: acquiring and releasing capacity no longer takes an internal lock while no sender is waiting, and zero-sized messages no longer take the buffer lock. diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 481abab5..1c5b1de2 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -53,7 +53,7 @@ mod buffer; /// /// # Panics /// -/// Panics if `buffer` is zero or exceeds the maximum capacity of `usize::MAX >> 2`. +/// Panics if `buffer` is zero or exceeds the maximum capacity of `usize::MAX >> 1`. #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); @@ -83,38 +83,40 @@ struct Shared { /// The largest capacity accepted by [`bounded`]. /// -/// The shared permit counter packs channel state into two flag bits, which also keeps the +/// The shared permit counter reserves two sentinel values above the usable range, and the +/// zero-sized queue length packs a closed flag into its top bit. This bound also keeps the /// rounded-up slot storage from overflowing a power of two. -const MAX_CAPACITY: usize = usize::MAX >> 2; +const MAX_CAPACITY: usize = usize::MAX >> 1; -// This channel-local semaphore packs its permit counter and channel state into one atomic. +// This channel-local semaphore keeps its permit counter and channel state in one atomic. // The general-purpose semaphore has neither a close operation nor acquisition errors. // -// `state` holds the available permits shifted left by two, plus two flag bits: +// `state` is the available permit count, plus two sentinel values at the top of the range: // // * `CLOSED`: the receiver is gone. No permits are issued or returned, and waiters drain with an // error. // * `WAITING`: the wait queue may be non-empty. Releases then take the locked path and grant the // permit directly to the oldest waiter instead of returning it to the counter, so capacity is -// handed out in registration order and new arrivals cannot steal an already granted slot. +// handed out in registration order and new arrivals cannot steal an already granted slot. The +// counter is zero while this sentinel stands: waiters only register after observing exhaustion, +// and grants bypass the counter. // -// With neither flag set, acquire and release are single lock-free operations on `state`. -// Wait-queue mutations always hold the queue lock; a registration sets `WAITING` before its -// final capacity recheck, which switches any racing release to the locked path and strands -// no permit without a wake. +// With neither sentinel installed, acquire and release are single lock-free operations on `state`. +// Wait-queue mutations always hold the queue lock; a registration installs `WAITING` before its +// final capacity recheck, which switches any racing release to the locked path and strands no +// permit without a wake. struct Semaphore { state: AtomicUsize, waiters: Mutex>, } -const CLOSED: usize = 0b01; -const WAITING: usize = 0b10; -const PERMIT: usize = 0b100; +const CLOSED: usize = usize::MAX; +const WAITING: usize = usize::MAX - 1; impl Semaphore { fn new(available: usize) -> Self { Self { - state: AtomicUsize::new(available * PERMIT), + state: AtomicUsize::new(available), waiters: Mutex::new(WaitList::new()), } } @@ -122,15 +124,15 @@ impl Semaphore { fn try_acquire(&self) -> Result<(), TrySendError<()>> { let mut state = self.state.load(Ordering::Acquire); loop { - if state & CLOSED != 0 { + if state == CLOSED { return Err(TrySendError::Disconnected(())); } - if state < PERMIT { + if state == WAITING || state == 0 { return Err(TrySendError::Full(())); } match self.state.compare_exchange_weak( state, - state - PERMIT, + state - 1, Ordering::Acquire, Ordering::Acquire, ) { @@ -141,15 +143,22 @@ impl Semaphore { } fn is_closed(&self) -> bool { - self.state.load(Ordering::Acquire) & CLOSED != 0 + self.state.load(Ordering::Acquire) == CLOSED } + // Installs WAITING over an exhausted counter. A permit that arrived first wins the compare + // exchange, and the caller's recheck under the queue lock picks it up instead. fn set_waiting(&self) { - self.state.fetch_or(WAITING, Ordering::AcqRel); + let _ = self + .state + .compare_exchange(0, WAITING, Ordering::AcqRel, Ordering::Acquire); } + // Removes WAITING, keeping whatever count a racing grant restoration left behind. fn clear_waiting(&self) { - self.state.fetch_and(!WAITING, Ordering::Release); + let _ = self + .state + .compare_exchange(WAITING, 0, Ordering::Release, Ordering::Relaxed); } fn release(&self) { @@ -157,12 +166,12 @@ impl Semaphore { // back to the counter. let mut state = self.state.load(Ordering::Relaxed); loop { - if state & (WAITING | CLOSED) != 0 { + if state == WAITING || state == CLOSED { break; } match self.state.compare_exchange_weak( state, - state + PERMIT, + state + 1, Ordering::Release, Ordering::Relaxed, ) { @@ -189,17 +198,24 @@ impl Semaphore { } return waker; } - // The queue is empty. Only releases add permits, and the counter grows with the - // queue locked, so a linked waiter always sees zero available permits until it - // receives its own grant. An outstanding grant already owns its capacity. - self.state.fetch_add(PERMIT, Ordering::Release); - self.clear_waiting(); + // The queue is empty: return the permit to the counter. An outstanding grant already + // owns its capacity. Adding to a plain count is safe because only lock-holding + // operations install a sentinel, and this operation holds the lock; WAITING itself + // must be displaced rather than incremented, because WAITING + 1 is CLOSED. + if self.state.load(Ordering::Relaxed) == WAITING { + let _displaced = + self.state + .compare_exchange(WAITING, 1, Ordering::Release, Ordering::Relaxed); + debug_assert_eq!(_displaced, Ok(WAITING)); + } else { + self.state.fetch_add(1, Ordering::Release); + } None } fn close(&self) -> WakerBatch { let mut waiters = self.waiters.lock(); - self.state.fetch_or(CLOSED, Ordering::AcqRel); + self.state.store(CLOSED, Ordering::Release); let mut wakers = WakerBatch::new(); while let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { if let Some(waker) = waiter.waker.take() { @@ -266,9 +282,11 @@ impl<'a, T> Reservation<'a, T> { return Poll::Pending; } } else { - // Set WAITING before the final capacity recheck: a racing release switches - // to the locked path, so no permit can be stranded without a wake. Waiting - // senders already in the queue take priority over this recheck. + // Install WAITING before the final capacity recheck: if a permit arrived + // first, the installation loses the compare exchange and the recheck picks + // the permit up; otherwise a racing release switches to the locked path, so + // no permit can be stranded without a wake. Waiting senders already in the + // queue take priority over this recheck. semaphore.set_waiting(); if waiters.is_empty() && semaphore.try_acquire().is_ok() { semaphore.clear_waiting(); diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index 3a6079f3..a7e60077 100644 --- a/tests-integration/tests/mpsc_test/main.rs +++ b/tests-integration/tests/mpsc_test/main.rs @@ -184,12 +184,12 @@ fn bounded_rejects_zero_capacity() { #[test] #[should_panic(expected = "exceeds the maximum")] fn bounded_rejects_capacity_above_the_maximum() { - let _ = mpsc::bounded::((usize::MAX >> 2) + 1); + let _ = mpsc::bounded::((usize::MAX >> 1) + 1); } #[test] fn bounded_zero_sized_messages_need_no_slot_storage() { - let (tx, mut rx) = mpsc::bounded::<()>(usize::MAX >> 2); + let (tx, mut rx) = mpsc::bounded::<()>(usize::MAX >> 1); tx.try_send(()).unwrap(); assert_eq!(rx.try_recv(), Ok(())); } diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index bda19072..6f470832 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -58,7 +58,7 @@ fn held_permits_consume_capacity_without_claiming_message_order() { #[test] fn zero_sized_messages_support_the_full_capacity_range() { - for capacity in [(usize::MAX >> 2) - 1, usize::MAX >> 2] { + for capacity in [(usize::MAX >> 1) - 1, usize::MAX >> 1] { let (tx, mut rx) = mpsc::bounded::<()>(capacity); let permit = tx.try_reserve().unwrap(); tx.try_send(()).unwrap(); @@ -89,7 +89,7 @@ fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { } } - let (tx, mut rx) = mpsc::bounded(usize::MAX >> 2); + let (tx, mut rx) = mpsc::bounded(usize::MAX >> 1); for _ in 0..3 { assert!(tx.try_send(Message).is_ok()); } From 682bfabf62a83e3e4b90524f4edeccbfd5d0ee92 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 12:52:43 +0800 Subject: [PATCH 22/34] perf(mpsc): drop cache padding from the receiver waker The receiver waker is written at most once per park cycle, while the permit counter and slot ticket are read-modify-written by every producer on every message. Padding the quiet word only spreads the shared allocation over more cache lines. Interleaved runs of the 8-producer sustained-capacity bench show padding nothing loses ~12% throughput, while padding only the two contended words is at parity with padding all three (slightly ahead in every round). --- asyncband/src/mpsc/bounded/mod.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 1c5b1de2..bca78e09 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -64,7 +64,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), - rx_waker: CachePadded::new(AtomicWaker::new()), + rx_waker: AtomicWaker::new(), buffer: Buffer::new(buffer), }); let sender = BoundedSender { @@ -77,7 +77,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { struct Shared { senders: AtomicUsize, tx_permits: CachePadded, - rx_waker: CachePadded, + rx_waker: AtomicWaker, buffer: Buffer, } From aec619f5a9ba1a27b46c76dedcce887f70e8258e Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 13:16:50 +0800 Subject: [PATCH 23/34] refactor(mpsc): organize the bounded channel by endpoint MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split the 650-line bounded module along the channel's own structure: mod.rs wires up the channel, sender.rs and receiver.rs hold the two public endpoints, and the two internal mechanisms get one file each — semaphore.rs for capacity, buffer.rs for storage. The sender's wait machinery moves into the semaphore as Acquire, the in-flight counterpart of acquire(), so all sentinel and wait-queue interlocking lives in one place. Verb-named operations replace noun-shaped internals: Reservation becomes Acquire, and the receiver's single-attempt core try_pop becomes pull. --- asyncband/src/mpsc/bounded/buffer_tests.rs | 39 +- asyncband/src/mpsc/bounded/mod.rs | 590 +-------------------- asyncband/src/mpsc/bounded/receiver.rs | 187 +++++++ asyncband/src/mpsc/bounded/semaphore.rs | 286 ++++++++++ asyncband/src/mpsc/bounded/sender.rs | 198 +++++++ 5 files changed, 704 insertions(+), 596 deletions(-) create mode 100644 asyncband/src/mpsc/bounded/receiver.rs create mode 100644 asyncband/src/mpsc/bounded/semaphore.rs create mode 100644 asyncband/src/mpsc/bounded/sender.rs diff --git a/asyncband/src/mpsc/bounded/buffer_tests.rs b/asyncband/src/mpsc/bounded/buffer_tests.rs index 43690029..6ceba628 100644 --- a/asyncband/src/mpsc/bounded/buffer_tests.rs +++ b/asyncband/src/mpsc/bounded/buffer_tests.rs @@ -31,12 +31,17 @@ use crate::mpsc::bounded; // Exercise the scheduling window inside synchronous send, while retaining a real capacity // permit. Ordinary callers cannot split a claim from its publication. -fn publish_claimed(mut permit: Permit<'_, T>, position: usize, value: T) -> Result<(), T> { - let shared = &permit.sender.unwrap().shared; +fn publish_claimed( + tx: &BoundedSender, + permit: Permit<'_, T>, + position: usize, + value: T, +) -> Result<(), T> { // SAFETY: The test claimed this position while holding the same capacity permit. - unsafe { shared.buffer.publish(position, value) }?; - permit.sender = None; - shared.rx_waker.wake(); + unsafe { tx.shared().buffer.publish(position, value) }?; + // Publication owns the capacity now; forgetting skips the permit's release on drop. + std::mem::forget(permit); + tx.shared().rx_waker.wake(); Ok(()) } @@ -46,18 +51,18 @@ fn a_claimed_head_waits_for_publication_across_laps() { for initial in [0, usize::MAX - 1] { let (tx, mut rx) = bounded(capacity); // Start an empty ring near ticket overflow instead of running usize::MAX sends. - tx.shared.buffer.tail.store(initial, Ordering::Relaxed); - rx.head = initial; + tx.shared().buffer.tail.store(initial, Ordering::Relaxed); + rx.set_head(initial); let mut cx = Context::from_waker(Waker::noop()); for lap in 0..8 { let permit = tx.try_reserve().unwrap(); - let position = tx.shared.buffer.claim().unwrap(); + let position = tx.shared().buffer.claim().unwrap(); for offset in 1..capacity { tx.try_send(lap * capacity + offset).unwrap(); } // A full ring must differ from an empty one even if no head value is ready yet. assert!(rx.poll_recv(&mut cx).is_pending()); - publish_claimed(permit, position, lap * capacity).unwrap(); + publish_claimed(&tx, permit, position, lap * capacity).unwrap(); for offset in 0..capacity { assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); } @@ -77,12 +82,12 @@ fn a_claim_delayed_past_close_returns_its_value() { drops: drops.clone(), _sender: tx.clone(), }; - let allocation = Arc::downgrade(&tx.shared); + let allocation = Arc::downgrade(tx.shared()); // Pause after claim's open check, then resume its atomic ticket allocation after close. - assert!(!tx.shared.buffer.closed.load(Ordering::Acquire)); + assert!(!tx.shared().buffer.closed.load(Ordering::Acquire)); drop(rx); - let position = tx.shared.buffer.tail.fetch_add(1, Ordering::AcqRel); - let unsent = publish_claimed(permit, position, value).unwrap_err(); + let position = tx.shared().buffer.tail.fetch_add(1, Ordering::AcqRel); + let unsent = publish_claimed(&tx, permit, position, value).unwrap_err(); assert_eq!(unsent.bytes, [7; 1024]); drop(unsent); assert_eq!(drops.load(Ordering::Relaxed), 1); @@ -108,7 +113,7 @@ impl Drop for Payload { #[test] fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { let (tx, rx) = bounded(2); - let allocation = Arc::downgrade(&tx.shared); + let allocation = Arc::downgrade(tx.shared()); let drops = Arc::new(AtomicUsize::new(0)); let paused = Barrier::new(2); let (resume_tx, resume_rx) = std::sync::mpsc::channel(); @@ -120,7 +125,7 @@ fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { let paused = &paused; let publisher = scope.spawn(move || { let permit = sender.try_reserve().unwrap(); - let position = sender.shared.buffer.claim().unwrap(); + let position = sender.shared().buffer.claim().unwrap(); let value = Payload { bytes: [1; 1024], drops: drops.clone(), @@ -128,7 +133,7 @@ fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { }; paused.wait(); resume_rx.recv().unwrap(); - let unsent = publish_claimed(permit, position, value).unwrap_err(); + let unsent = publish_claimed(sender, permit, position, value).unwrap_err(); assert_eq!(unsent.bytes, [1; 1024]); drop(unsent); }); @@ -165,7 +170,7 @@ fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { fn publication_racing_with_close_drops_every_payload_once() { for _ in 0..if cfg!(miri) { 8 } else { 128 } { let (tx, rx) = bounded(3); - let allocation = Arc::downgrade(&tx.shared); + let allocation = Arc::downgrade(tx.shared()); let drops = Arc::new(AtomicUsize::new(0)); let start = Barrier::new(4); thread::scope(|scope| { diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index bca78e09..914f7560 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -18,29 +18,26 @@ //! A bounded multi-producer, single-consumer queue for sending values between asynchronous //! tasks with backpressure control. -use std::fmt; -use std::future::poll_fn; use std::sync::Arc; use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use std::task::Waker; use self::buffer::Buffer; +use self::semaphore::Semaphore; use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; use crate::internal::atomic_waker::AtomicWaker; use crate::internal::cache_padded::CachePadded; -use crate::internal::mutex::Mutex; -use crate::internal::waitlist::WaitList; -use crate::internal::waitlist::WaiterId; -use crate::internal::wake_all; -use crate::internal::waker_batch::WakerBatch; mod buffer; +mod receiver; +mod semaphore; +mod sender; + +pub use self::receiver::BoundedReceiver; +pub use self::sender::BoundedSender; +pub use self::sender::Permit; /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// @@ -67,14 +64,12 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { rx_waker: AtomicWaker::new(), buffer: Buffer::new(buffer), }); - let sender = BoundedSender { - shared: shared.clone(), - }; - let receiver = BoundedReceiver { shared, head: 0 }; + let sender = BoundedSender::new(shared.clone()); + let receiver = BoundedReceiver::new(shared); (sender, receiver) } -struct Shared { +pub struct Shared { senders: AtomicUsize, tx_permits: CachePadded, rx_waker: AtomicWaker, @@ -87,566 +82,3 @@ struct Shared { /// zero-sized queue length packs a closed flag into its top bit. This bound also keeps the /// rounded-up slot storage from overflowing a power of two. const MAX_CAPACITY: usize = usize::MAX >> 1; - -// This channel-local semaphore keeps its permit counter and channel state in one atomic. -// The general-purpose semaphore has neither a close operation nor acquisition errors. -// -// `state` is the available permit count, plus two sentinel values at the top of the range: -// -// * `CLOSED`: the receiver is gone. No permits are issued or returned, and waiters drain with an -// error. -// * `WAITING`: the wait queue may be non-empty. Releases then take the locked path and grant the -// permit directly to the oldest waiter instead of returning it to the counter, so capacity is -// handed out in registration order and new arrivals cannot steal an already granted slot. The -// counter is zero while this sentinel stands: waiters only register after observing exhaustion, -// and grants bypass the counter. -// -// With neither sentinel installed, acquire and release are single lock-free operations on `state`. -// Wait-queue mutations always hold the queue lock; a registration installs `WAITING` before its -// final capacity recheck, which switches any racing release to the locked path and strands no -// permit without a wake. -struct Semaphore { - state: AtomicUsize, - waiters: Mutex>, -} - -const CLOSED: usize = usize::MAX; -const WAITING: usize = usize::MAX - 1; - -impl Semaphore { - fn new(available: usize) -> Self { - Self { - state: AtomicUsize::new(available), - waiters: Mutex::new(WaitList::new()), - } - } - - fn try_acquire(&self) -> Result<(), TrySendError<()>> { - let mut state = self.state.load(Ordering::Acquire); - loop { - if state == CLOSED { - return Err(TrySendError::Disconnected(())); - } - if state == WAITING || state == 0 { - return Err(TrySendError::Full(())); - } - match self.state.compare_exchange_weak( - state, - state - 1, - Ordering::Acquire, - Ordering::Acquire, - ) { - Ok(_) => return Ok(()), - Err(actual) => state = actual, - } - } - } - - fn is_closed(&self) -> bool { - self.state.load(Ordering::Acquire) == CLOSED - } - - // Installs WAITING over an exhausted counter. A permit that arrived first wins the compare - // exchange, and the caller's recheck under the queue lock picks it up instead. - fn set_waiting(&self) { - let _ = self - .state - .compare_exchange(0, WAITING, Ordering::AcqRel, Ordering::Acquire); - } - - // Removes WAITING, keeping whatever count a racing grant restoration left behind. - fn clear_waiting(&self) { - let _ = self - .state - .compare_exchange(WAITING, 0, Ordering::Release, Ordering::Relaxed); - } - - fn release(&self) { - // Fast path: with no waiting sender and no close in sight, the permit goes straight - // back to the counter. - let mut state = self.state.load(Ordering::Relaxed); - loop { - if state == WAITING || state == CLOSED { - break; - } - match self.state.compare_exchange_weak( - state, - state + 1, - Ordering::Release, - Ordering::Relaxed, - ) { - Ok(_) => return, - Err(actual) => state = actual, - } - } - let wake = self.release_locked(&mut self.waiters.lock()); - if let Some(waker) = wake { - waker.wake(); - } - } - - fn release_locked(&self, waiters: &mut WaitList) -> Option { - if self.is_closed() { - return None; - } - if let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { - // Grant ownership before waking; new arrivals cannot steal this capacity. - waiter.granted = true; - let waker = waiter.waker.take(); - if waiters.is_empty() { - self.clear_waiting(); - } - return waker; - } - // The queue is empty: return the permit to the counter. An outstanding grant already - // owns its capacity. Adding to a plain count is safe because only lock-holding - // operations install a sentinel, and this operation holds the lock; WAITING itself - // must be displaced rather than incremented, because WAITING + 1 is CLOSED. - if self.state.load(Ordering::Relaxed) == WAITING { - let _displaced = - self.state - .compare_exchange(WAITING, 1, Ordering::Release, Ordering::Relaxed); - debug_assert_eq!(_displaced, Ok(WAITING)); - } else { - self.state.fetch_add(1, Ordering::Release); - } - None - } - - fn close(&self) -> WakerBatch { - let mut waiters = self.waiters.lock(); - self.state.store(CLOSED, Ordering::Release); - let mut wakers = WakerBatch::new(); - while let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { - if let Some(waker) = waiter.waker.take() { - wakers.push(waker); - } - } - wakers - } -} - -struct Waiter { - granted: bool, - waker: Option, -} - -struct Reservation<'a, T> { - sender: &'a BoundedSender, - index: Option, -} - -impl<'a, T> Reservation<'a, T> { - fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { - let semaphore = &self.sender.shared.tx_permits; - let mut cloned_waker = None; - let result = loop { - if self.index.is_none() { - match semaphore.try_acquire() { - Ok(()) => { - break Ok(Permit { - sender: Some(self.sender), - }); - } - Err(TrySendError::Disconnected(())) => break Err(SendError::new(())), - Err(TrySendError::Full(())) => {} - } - } - let mut waiters = semaphore.waiters.lock(); - if semaphore.is_closed() { - // Drop removes any remaining registration, including an unused grant. - break Err(SendError::new(())); - } - if let Some(index) = self.index { - let waiter = waiters.waiter_mut(index); - if waiter.granted { - let waiter = waiters.remove_unlinked_waiter(index); - self.index = None; - drop(waiters); - drop(waiter); - break Ok(Permit { - sender: Some(self.sender), - }); - } - if waiter - .waker - .as_ref() - .is_some_and(|w| w.will_wake(cx.waker())) - { - return Poll::Pending; - } - if let Some(waker) = cloned_waker.take() { - let old = waiter.waker.replace(waker); - drop(waiters); - drop(old); - return Poll::Pending; - } - } else { - // Install WAITING before the final capacity recheck: if a permit arrived - // first, the installation loses the compare exchange and the recheck picks - // the permit up; otherwise a racing release switches to the locked path, so - // no permit can be stranded without a wake. Waiting senders already in the - // queue take priority over this recheck. - semaphore.set_waiting(); - if waiters.is_empty() && semaphore.try_acquire().is_ok() { - semaphore.clear_waiting(); - break Ok(Permit { - sender: Some(self.sender), - }); - } - if let Some(waker) = cloned_waker.take() { - self.index = Some(waiters.push_back(Waiter { - granted: false, - waker: Some(waker), - })); - return Poll::Pending; - } - } - drop(waiters); - // Clone outside the lock, then recheck capacity and closure before registering. - cloned_waker = Some(cx.waker().clone()); - }; - // The permit owns capacity before an unused cloned waker can panic. - drop(cloned_waker); - Poll::Ready(result) - } -} - -impl Drop for Reservation<'_, T> { - fn drop(&mut self) { - let Some(index) = self.index else { return }; - let semaphore = &self.sender.shared.tx_permits; - let (waiter, wake) = { - let mut waiters = semaphore.waiters.lock(); - waiters.unlink_waiter(index, |_| true); - let waiter = waiters.remove_unlinked_waiter(index); - let wake = if waiter.granted { - semaphore.release_locked(&mut waiters) - } else { - if waiters.is_empty() { - semaphore.clear_waiting(); - } - None - }; - (waiter, wake) - }; - if let Some(waker) = wake { - waker.wake(); - } - drop(waiter); - } -} - -/// The sending endpoint of a bounded mpsc channel. -/// -/// Instances are created by the [`bounded`] function. -pub struct BoundedSender { - shared: Arc>, -} - -impl Clone for BoundedSender { - fn clone(&self) -> Self { - self.shared.senders.fetch_add(1, Ordering::Relaxed); - BoundedSender { - shared: self.shared.clone(), - } - } -} - -impl fmt::Debug for BoundedSender { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("BoundedSender").finish_non_exhaustive() - } -} - -impl Drop for BoundedSender { - fn drop(&mut self) { - if self.shared.senders.fetch_sub(1, Ordering::AcqRel) == 1 { - self.shared.rx_waker.wake(); - } - } -} - -impl BoundedSender { - /// Sends a message, waiting until the channel has capacity when necessary. - /// - /// If the receiver has been dropped, the returned error contains `value`. - /// - /// # Cancel safety - /// - /// Dropping a pending `send` loses its place waiting for capacity and drops `value`; a call - /// that has returned `Pending` has not sent the message. Use [`Self::try_send`] when the - /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for - /// capacity before constructing the message. - pub async fn send(&self, value: T) -> Result<(), SendError> { - match self.reserve().await { - Ok(permit) => permit.send(value), - Err(_) => Err(SendError::new(value)), - } - } - - /// Reserves capacity for one message before constructing it. - /// - /// A successful reservation returns a [`Permit`]. Dropping the permit releases capacity; - /// [`Permit::send`] publishes a value without waiting for space. Reservations do not establish - /// message order: other producers may send while a permit is held. - /// - /// Returns `SendError(())` if the receiver has been dropped. A permit obtained earlier does - /// not keep the receiver alive; sending with it can still return the unsent value on - /// disconnect. - /// - /// # Cancel safety - /// - /// Dropping a pending reservation loses its place in the wait queue. If capacity has already - /// been granted, it is released to the next waiter or made available to a new sender. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// let (tx, mut rx) = asyncband::mpsc::bounded(1); - /// let permit = tx.reserve().await.unwrap(); - /// let message = String::from("constructed after capacity became available"); - /// permit.send(message).unwrap(); - /// assert_eq!( - /// rx.recv().await.unwrap(), - /// "constructed after capacity became available" - /// ); - /// # } - /// ``` - pub async fn reserve(&self) -> Result, SendError<()>> { - let mut reservation = Reservation { - sender: self, - index: None, - }; - poll_fn(|cx| reservation.poll(cx)).await - } - - /// Reserves capacity for one message without waiting. - /// - /// Returns [`TrySendError::Full`] if queued messages and outstanding permits occupy the - /// buffer, or [`TrySendError::Disconnected`] if the receiver has been dropped. - pub fn try_reserve(&self) -> Result, TrySendError<()>> { - self.shared.tx_permits.try_acquire()?; - Ok(Permit { sender: Some(self) }) - } - - /// Attempts to send a message without waiting for capacity. - /// - /// A full buffer returns [`TrySendError::Full`], while a dropped receiver returns - /// [`TrySendError::Disconnected`]. Both errors return ownership of the unsent value. - /// - /// # Examples - /// - /// ``` - /// use asyncband::mpsc::TrySendError; - /// use asyncband::mpsc::bounded; - /// - /// let (tx, mut rx) = bounded(1); - /// tx.try_send(10).unwrap(); - /// assert_eq!(tx.try_send(20), Err(TrySendError::Full(20))); - /// - /// assert_eq!(rx.try_recv(), Ok(10)); - /// tx.try_send(20).unwrap(); - /// drop(rx); - /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); - /// ``` - pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - match self.try_reserve() { - Ok(permit) => permit - .send(value) - .map_err(|error| TrySendError::Disconnected(error.into_inner())), - Err(TrySendError::Full(())) => Err(TrySendError::Full(value)), - Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), - } - } -} - -/// Capacity reserved for one message on a bounded channel. -/// -/// Created by [`BoundedSender::reserve`] or [`BoundedSender::try_reserve`]. Holding a permit -/// reduces available capacity but does not prevent other messages from being received. Dropping -/// it without sending releases capacity and notifies a waiting sender. -#[must_use = "dropping the permit releases its reserved capacity"] -pub struct Permit<'a, T> { - sender: Option<&'a BoundedSender>, -} - -impl fmt::Debug for Permit<'_, T> { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("Permit").finish_non_exhaustive() - } -} - -impl Permit<'_, T> { - /// Publishes a message using this reservation, without waiting for capacity. - /// - /// If the receiver has been dropped, the returned error contains the unsent value. - pub fn send(mut self, value: T) -> Result<(), SendError> { - let shared = &self.sender.unwrap().shared; - // SAFETY: This permit owns one capacity unit. Claiming a slot and writing it is a - // synchronous operation with no user callbacks or await points between the two. - unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; - // Publication owns the capacity before a wake callback can panic. - self.sender = None; - shared.rx_waker.wake(); - Ok(()) - } -} - -impl Drop for Permit<'_, T> { - fn drop(&mut self) { - if let Some(sender) = self.sender { - sender.shared.tx_permits.release(); - } - } -} - -/// The receiving endpoint of a bounded mpsc channel. -/// -/// Instances are created by the [`bounded`] function. -/// Dropping the receiver discards queued values. The backing allocation remains alive until -/// all endpoints are dropped, so a concurrent sender can safely finish returning an unsent value. -pub struct BoundedReceiver { - shared: Arc>, - head: usize, -} - -impl fmt::Debug for BoundedReceiver { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("BoundedReceiver").finish_non_exhaustive() - } -} - -impl Drop for BoundedReceiver { - fn drop(&mut self) { - // SAFETY: Receiver ownership provides exclusive access to the consumption cursor. - // The drain first prevents new claims. Its destructor completes cleanup on unwinding. - let drain = unsafe { self.shared.buffer.close(self.head) }; - let wakers = self.shared.tx_permits.close(); - let receiver_waker = self.shared.rx_waker.take(); - wake_all(wakers.into_iter()); - drop(receiver_waker); - drop(drain); - } -} - -impl BoundedReceiver { - /// Attempts to receive the next queued value without waiting for a new message. - /// - /// Receiving a value frees one buffer slot. An empty channel returns [`TryRecvError::Empty`] - /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has - /// been dropped and all queued values have been consumed. - /// - /// If a producer is still completing a synchronous publication at the queue head, this - /// method waits for that publication. Use [`Self::recv`] to wait asynchronously instead. - /// - /// # Examples - /// - /// ``` - /// use asyncband::mpsc::TryRecvError; - /// use asyncband::mpsc::bounded; - /// - /// let (tx, mut rx) = bounded(2); - /// tx.try_send("first").unwrap(); - /// tx.try_send("second").unwrap(); - /// - /// assert_eq!(rx.try_recv(), Ok("first")); - /// assert_eq!(rx.try_recv(), Ok("second")); - /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - /// drop(tx); - /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); - /// ``` - pub fn try_recv(&mut self) -> Result { - let mut spins = 0; - loop { - match self.try_pop() { - Poll::Ready(result) => return result, - Poll::Pending => { - // A synchronous publisher already owns the head. Reporting Empty here - // could hide a later send that has completed. Async recv parks instead. - if spins < 32 { - std::hint::spin_loop(); - spins += 1; - } else { - std::thread::yield_now(); - } - } - } - } - } - - fn try_pop(&mut self) -> Poll> { - let mut disconnected = false; - loop { - // SAFETY: Only this receiver owns head. Capacity is released after the buffer - // finishes reading and advances the cursor, so no producer can overwrite the value. - match unsafe { self.shared.buffer.pop(&mut self.head) } { - Poll::Ready(Some(value)) => { - self.shared.tx_permits.release(); - return Poll::Ready(Ok(value)); - } - Poll::Ready(None) if disconnected => { - return Poll::Ready(Err(TryRecvError::Disconnected)); - } - Poll::Ready(None) if self.shared.senders.load(Ordering::Acquire) == 0 => { - // Acquire the last sender's completed publications before checking again. - disconnected = true; - } - Poll::Ready(None) => return Poll::Ready(Err(TryRecvError::Empty)), - Poll::Pending => return Poll::Pending, - } - } - } - - /// Waits for and receives the next value, freeing one buffer slot. - /// - /// If no value is queued, this method waits until a sender adds one or the last sender is - /// dropped. It returns [`RecvError::Disconnected`] only after all senders are gone and the - /// buffer has been drained. - /// - /// # Cancel safety - /// - /// Dropping a pending `recv` does not remove a message from the channel. A later receive - /// operation can still observe the next queued value, so `recv` may safely be raced with other - /// futures in a selection construct. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc; - /// let (tx, mut rx) = mpsc::bounded(2); - /// - /// tx.send("first").await.unwrap(); - /// tx.send("second").await.unwrap(); - /// drop(tx); - /// - /// assert_eq!(rx.recv().await, Ok("first")); - /// assert_eq!(rx.recv().await, Ok("second")); - /// assert_eq!(rx.recv().await, Err(mpsc::RecvError::Disconnected)); - /// # } - /// ``` - pub async fn recv(&mut self) -> Result { - poll_fn(|cx| self.poll_recv(cx)).await - } - - fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - for registered in [false, true] { - match self.try_pop() { - Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), - Poll::Ready(Err(TryRecvError::Disconnected)) => { - drop(self.shared.rx_waker.take()); - return Poll::Ready(Err(RecvError::Disconnected)); - } - Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {} - } - if !registered { - self.shared.rx_waker.register(cx.waker()); - } - } - Poll::Pending - } -} diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs new file mode 100644 index 00000000..9a3cbd73 --- /dev/null +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -0,0 +1,187 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt; +use std::future::poll_fn; +use std::sync::Arc; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; + +use super::RecvError; +use super::Shared; +use super::TryRecvError; +use crate::internal::wake_all; + +/// The receiving endpoint of a bounded mpsc channel. +/// +/// Instances are created by the [`bounded`](super::bounded) function. +/// Dropping the receiver discards queued values. The backing allocation remains alive until +/// all endpoints are dropped, so a concurrent sender can safely finish returning an unsent value. +pub struct BoundedReceiver { + shared: Arc>, + head: usize, +} + +impl fmt::Debug for BoundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedReceiver").finish_non_exhaustive() + } +} + +impl Drop for BoundedReceiver { + fn drop(&mut self) { + // SAFETY: Receiver ownership provides exclusive access to the consumption cursor. + // The drain first prevents new claims. Its destructor completes cleanup on unwinding. + let drain = unsafe { self.shared.buffer.close(self.head) }; + let wakers = self.shared.tx_permits.close(); + let receiver_waker = self.shared.rx_waker.take(); + wake_all(wakers.into_iter()); + drop(receiver_waker); + drop(drain); + } +} + +impl BoundedReceiver { + pub(crate) fn new(shared: Arc>) -> Self { + Self { shared, head: 0 } + } + + /// Attempts to receive the next queued value without waiting for a new message. + /// + /// Receiving a value frees one buffer slot. An empty channel returns [`TryRecvError::Empty`] + /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has + /// been dropped and all queued values have been consumed. + /// + /// If a producer is still completing a synchronous publication at the queue head, this + /// method waits for that publication. Use [`Self::recv`] to wait asynchronously instead. + /// + /// # Examples + /// + /// ``` + /// use asyncband::mpsc::TryRecvError; + /// use asyncband::mpsc::bounded; + /// + /// let (tx, mut rx) = bounded(2); + /// tx.try_send("first").unwrap(); + /// tx.try_send("second").unwrap(); + /// + /// assert_eq!(rx.try_recv(), Ok("first")); + /// assert_eq!(rx.try_recv(), Ok("second")); + /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + /// drop(tx); + /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); + /// ``` + pub fn try_recv(&mut self) -> Result { + let mut spins = 0; + loop { + match self.pull() { + Poll::Ready(result) => return result, + Poll::Pending => { + // A synchronous publisher already owns the head. Reporting Empty here + // could hide a later send that has completed. Async recv parks instead. + if spins < 32 { + std::hint::spin_loop(); + spins += 1; + } else { + std::thread::yield_now(); + } + } + } + } + } + + /// One attempt to take the head value: a message, an empty-or-disconnected classification, + /// or `Pending` while a claimed head waits for its publication. + fn pull(&mut self) -> Poll> { + let mut disconnected = false; + loop { + // SAFETY: Only this receiver owns head. Capacity is released after the buffer + // finishes reading and advances the cursor, so no producer can overwrite the value. + match unsafe { self.shared.buffer.pop(&mut self.head) } { + Poll::Ready(Some(value)) => { + self.shared.tx_permits.release(); + return Poll::Ready(Ok(value)); + } + Poll::Ready(None) if disconnected => { + return Poll::Ready(Err(TryRecvError::Disconnected)); + } + Poll::Ready(None) if self.shared.senders.load(Ordering::Acquire) == 0 => { + // Acquire the last sender's completed publications before checking again. + disconnected = true; + } + Poll::Ready(None) => return Poll::Ready(Err(TryRecvError::Empty)), + Poll::Pending => return Poll::Pending, + } + } + } + + /// Waits for and receives the next value, freeing one buffer slot. + /// + /// If no value is queued, this method waits until a sender adds one or the last sender is + /// dropped. It returns [`RecvError::Disconnected`] only after all senders are gone and the + /// buffer has been drained. + /// + /// # Cancel safety + /// + /// Dropping a pending `recv` does not remove a message from the channel. A later receive + /// operation can still observe the next queued value, so `recv` may safely be raced with other + /// futures in a selection construct. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::mpsc; + /// let (tx, mut rx) = mpsc::bounded(2); + /// + /// tx.send("first").await.unwrap(); + /// tx.send("second").await.unwrap(); + /// drop(tx); + /// + /// assert_eq!(rx.recv().await, Ok("first")); + /// assert_eq!(rx.recv().await, Ok("second")); + /// assert_eq!(rx.recv().await, Err(mpsc::RecvError::Disconnected)); + /// # } + /// ``` + pub async fn recv(&mut self) -> Result { + poll_fn(|cx| self.poll_recv(cx)).await + } + + pub(crate) fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { + for registered in [false, true] { + match self.pull() { + Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), + Poll::Ready(Err(TryRecvError::Disconnected)) => { + drop(self.shared.rx_waker.take()); + return Poll::Ready(Err(RecvError::Disconnected)); + } + Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {} + } + if !registered { + self.shared.rx_waker.register(cx.waker()); + } + } + Poll::Pending + } + + #[cfg(test)] + pub(crate) fn set_head(&mut self, head: usize) { + self.head = head; + } +} diff --git a/asyncband/src/mpsc/bounded/semaphore.rs b/asyncband/src/mpsc/bounded/semaphore.rs new file mode 100644 index 00000000..e0ec4925 --- /dev/null +++ b/asyncband/src/mpsc/bounded/semaphore.rs @@ -0,0 +1,286 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! The channel's capacity: a counting semaphore with a close operation and a fair wait queue. +//! +//! This channel-local semaphore keeps its permit counter and channel state in one atomic. The +//! general-purpose semaphore has neither a close operation nor acquisition errors. +//! +//! `state` is the available permit count, plus two sentinel values at the top of the range: +//! +//! * `CLOSED`: the receiver is gone. No permits are issued or returned, and waiters drain with an +//! error. +//! * `WAITING`: the wait queue may be non-empty. Releases then take the locked path and grant the +//! permit directly to the oldest waiter instead of returning it to the counter, so capacity is +//! handed out in registration order and new arrivals cannot steal an already granted slot. The +//! counter is zero while this sentinel stands: waiters only register after observing exhaustion, +//! and grants bypass the counter. +//! +//! With neither sentinel installed, acquire and release are single lock-free operations on `state`. +//! Wait-queue mutations always hold the queue lock; a registration installs `WAITING` before its +//! final capacity recheck, which switches any racing release to the locked path and strands no +//! permit without a wake. + +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use super::SendError; +use super::TrySendError; +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaitList; +use crate::internal::waitlist::WaiterId; +use crate::internal::waker_batch::WakerBatch; + +pub struct Semaphore { + state: AtomicUsize, + waiters: Mutex>, +} + +const CLOSED: usize = usize::MAX; +const WAITING: usize = usize::MAX - 1; + +pub struct Waiter { + granted: bool, + waker: Option, +} + +impl Semaphore { + pub fn new(available: usize) -> Self { + Self { + state: AtomicUsize::new(available), + waiters: Mutex::new(WaitList::new()), + } + } + + pub fn try_acquire(&self) -> Result<(), TrySendError<()>> { + let mut state = self.state.load(Ordering::Acquire); + loop { + if state == CLOSED { + return Err(TrySendError::Disconnected(())); + } + if state == WAITING || state == 0 { + return Err(TrySendError::Full(())); + } + match self.state.compare_exchange_weak( + state, + state - 1, + Ordering::Acquire, + Ordering::Acquire, + ) { + Ok(_) => return Ok(()), + Err(actual) => state = actual, + } + } + } + + /// Acquires one permit asynchronously, waiting in registration order when the semaphore + /// is exhausted. + pub fn acquire(&self) -> Acquire<'_> { + Acquire { + semaphore: self, + waiter: None, + } + } + + pub fn is_closed(&self) -> bool { + self.state.load(Ordering::Acquire) == CLOSED + } + + // Installs WAITING over an exhausted counter. A permit that arrived first wins the compare + // exchange, and the caller's recheck under the queue lock picks it up instead. + fn set_waiting(&self) { + let _ = self + .state + .compare_exchange(0, WAITING, Ordering::AcqRel, Ordering::Acquire); + } + + // Removes WAITING, keeping whatever count a racing grant restoration left behind. + fn clear_waiting(&self) { + let _ = self + .state + .compare_exchange(WAITING, 0, Ordering::Release, Ordering::Relaxed); + } + + pub fn release(&self) { + // Fast path: with no waiting sender and no close in sight, the permit goes straight + // back to the counter. + let mut state = self.state.load(Ordering::Relaxed); + loop { + if state == WAITING || state == CLOSED { + break; + } + match self.state.compare_exchange_weak( + state, + state + 1, + Ordering::Release, + Ordering::Relaxed, + ) { + Ok(_) => return, + Err(actual) => state = actual, + } + } + let wake = self.release_locked(&mut self.waiters.lock()); + if let Some(waker) = wake { + waker.wake(); + } + } + + pub fn release_locked(&self, waiters: &mut WaitList) -> Option { + if self.is_closed() { + return None; + } + if let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { + // Grant ownership before waking; new arrivals cannot steal this capacity. + waiter.granted = true; + let waker = waiter.waker.take(); + if waiters.is_empty() { + self.clear_waiting(); + } + return waker; + } + // The queue is empty: return the permit to the counter. An outstanding grant already + // owns its capacity. Adding to a plain count is safe because only lock-holding + // operations install a sentinel, and this operation holds the lock; WAITING itself + // must be displaced rather than incremented, because WAITING + 1 is CLOSED. + if self.state.load(Ordering::Relaxed) == WAITING { + let _displaced = + self.state + .compare_exchange(WAITING, 1, Ordering::Release, Ordering::Relaxed); + debug_assert_eq!(_displaced, Ok(WAITING)); + } else { + self.state.fetch_add(1, Ordering::Release); + } + None + } + + pub fn close(&self) -> WakerBatch { + let mut waiters = self.waiters.lock(); + self.state.store(CLOSED, Ordering::Release); + let mut wakers = WakerBatch::new(); + while let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } + } + wakers + } +} + +/// An in-flight [`Semaphore::acquire`] operation. +/// +/// Dropping the operation removes its wait-queue registration; a capacity grant that already +/// reached the registration is released to the next waiter or returned to the counter. +pub struct Acquire<'a> { + semaphore: &'a Semaphore, + waiter: Option, +} + +impl Acquire<'_> { + pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll>> { + let semaphore = self.semaphore; + let mut cloned_waker = None; + let result = loop { + if self.waiter.is_none() { + match semaphore.try_acquire() { + Ok(()) => break Ok(()), + Err(TrySendError::Disconnected(())) => break Err(SendError::new(())), + Err(TrySendError::Full(())) => {} + } + } + let mut waiters = semaphore.waiters.lock(); + if semaphore.is_closed() { + // Drop removes any remaining registration, including an unused grant. + break Err(SendError::new(())); + } + if let Some(index) = self.waiter { + let waiter = waiters.waiter_mut(index); + if waiter.granted { + let waiter = waiters.remove_unlinked_waiter(index); + self.waiter = None; + drop(waiters); + drop(waiter); + break Ok(()); + } + if waiter + .waker + .as_ref() + .is_some_and(|w| w.will_wake(cx.waker())) + { + return Poll::Pending; + } + if let Some(waker) = cloned_waker.take() { + let old = waiter.waker.replace(waker); + drop(waiters); + drop(old); + return Poll::Pending; + } + } else { + // Install WAITING before the final capacity recheck: if a permit arrived + // first, the installation loses the compare exchange and the recheck picks + // the permit up; otherwise a racing release switches to the locked path, so + // no permit can be stranded without a wake. Waiting senders already in the + // queue take priority over this recheck. + semaphore.set_waiting(); + if waiters.is_empty() && semaphore.try_acquire().is_ok() { + semaphore.clear_waiting(); + break Ok(()); + } + if let Some(waker) = cloned_waker.take() { + self.waiter = Some(waiters.push_back(Waiter { + granted: false, + waker: Some(waker), + })); + return Poll::Pending; + } + } + drop(waiters); + // Clone outside the lock, then recheck capacity and closure before registering. + cloned_waker = Some(cx.waker().clone()); + }; + // The permit owns capacity before an unused cloned waker can panic. + drop(cloned_waker); + Poll::Ready(result) + } +} + +impl Drop for Acquire<'_> { + fn drop(&mut self) { + let Some(index) = self.waiter else { return }; + let semaphore = self.semaphore; + let (waiter, wake) = { + let mut waiters = semaphore.waiters.lock(); + waiters.unlink_waiter(index, |_| true); + let waiter = waiters.remove_unlinked_waiter(index); + let wake = if waiter.granted { + semaphore.release_locked(&mut waiters) + } else { + if waiters.is_empty() { + semaphore.clear_waiting(); + } + None + }; + (waiter, wake) + }; + if let Some(waker) = wake { + waker.wake(); + } + drop(waiter); + } +} diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs new file mode 100644 index 00000000..e2a83142 --- /dev/null +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -0,0 +1,198 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::fmt; +use std::future::poll_fn; +use std::sync::Arc; +use std::sync::atomic::Ordering; + +use super::SendError; +use super::Shared; +use super::TrySendError; + +/// The sending endpoint of a bounded mpsc channel. +/// +/// Instances are created by the [`bounded`](super::bounded) function. +pub struct BoundedSender { + shared: Arc>, +} + +impl Clone for BoundedSender { + fn clone(&self) -> Self { + self.shared.senders.fetch_add(1, Ordering::Relaxed); + BoundedSender { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for BoundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedSender").finish_non_exhaustive() + } +} + +impl Drop for BoundedSender { + fn drop(&mut self) { + if self.shared.senders.fetch_sub(1, Ordering::AcqRel) == 1 { + self.shared.rx_waker.wake(); + } + } +} + +impl BoundedSender { + pub(crate) fn new(shared: Arc>) -> Self { + Self { shared } + } + + /// Sends a message, waiting until the channel has capacity when necessary. + /// + /// If the receiver has been dropped, the returned error contains `value`. + /// + /// # Cancel safety + /// + /// Dropping a pending `send` loses its place waiting for capacity and drops `value`; a call + /// that has returned `Pending` has not sent the message. Use [`Self::try_send`] when the + /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for + /// capacity before constructing the message. + pub async fn send(&self, value: T) -> Result<(), SendError> { + match self.reserve().await { + Ok(permit) => permit.send(value), + Err(_) => Err(SendError::new(value)), + } + } + + /// Reserves capacity for one message before constructing it. + /// + /// A successful reservation returns a [`Permit`]. Dropping the permit releases capacity; + /// [`Permit::send`] publishes a value without waiting for space. Reservations do not establish + /// message order: other producers may send while a permit is held. + /// + /// Returns `SendError(())` if the receiver has been dropped. A permit obtained earlier does + /// not keep the receiver alive; sending with it can still return the unsent value on + /// disconnect. + /// + /// # Cancel safety + /// + /// Dropping a pending reservation loses its place in the wait queue. If capacity has already + /// been granted, it is released to the next waiter or made available to a new sender. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, mut rx) = asyncband::mpsc::bounded(1); + /// let permit = tx.reserve().await.unwrap(); + /// let message = String::from("constructed after capacity became available"); + /// permit.send(message).unwrap(); + /// assert_eq!( + /// rx.recv().await.unwrap(), + /// "constructed after capacity became available" + /// ); + /// # } + /// ``` + pub async fn reserve(&self) -> Result, SendError<()>> { + let mut acquire = self.shared.tx_permits.acquire(); + poll_fn(|cx| acquire.poll(cx)).await?; + Ok(Permit { sender: Some(self) }) + } + + /// Reserves capacity for one message without waiting. + /// + /// Returns [`TrySendError::Full`] if queued messages and outstanding permits occupy the + /// buffer, or [`TrySendError::Disconnected`] if the receiver has been dropped. + pub fn try_reserve(&self) -> Result, TrySendError<()>> { + self.shared.tx_permits.try_acquire()?; + Ok(Permit { sender: Some(self) }) + } + + /// Attempts to send a message without waiting for capacity. + /// + /// A full buffer returns [`TrySendError::Full`], while a dropped receiver returns + /// [`TrySendError::Disconnected`]. Both errors return ownership of the unsent value. + /// + /// # Examples + /// + /// ``` + /// use asyncband::mpsc::TrySendError; + /// use asyncband::mpsc::bounded; + /// + /// let (tx, mut rx) = bounded(1); + /// tx.try_send(10).unwrap(); + /// assert_eq!(tx.try_send(20), Err(TrySendError::Full(20))); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// tx.try_send(20).unwrap(); + /// drop(rx); + /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); + /// ``` + pub fn try_send(&self, value: T) -> Result<(), TrySendError> { + match self.try_reserve() { + Ok(permit) => permit + .send(value) + .map_err(|error| TrySendError::Disconnected(error.into_inner())), + Err(TrySendError::Full(())) => Err(TrySendError::Full(value)), + Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), + } + } + + #[cfg(test)] + pub(crate) fn shared(&self) -> &Arc> { + &self.shared + } +} + +/// Capacity reserved for one message on a bounded channel. +/// +/// Created by [`BoundedSender::reserve`] or [`BoundedSender::try_reserve`]. Holding a permit +/// reduces available capacity but does not prevent other messages from being received. Dropping +/// it without sending releases capacity and notifies a waiting sender. +#[must_use = "dropping the permit releases its reserved capacity"] +pub struct Permit<'a, T> { + sender: Option<&'a BoundedSender>, +} + +impl fmt::Debug for Permit<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Permit").finish_non_exhaustive() + } +} + +impl Permit<'_, T> { + /// Publishes a message using this permit, without waiting for capacity. + /// + /// If the receiver has been dropped, the returned error contains the unsent value. + pub fn send(mut self, value: T) -> Result<(), SendError> { + let shared = &self.sender.unwrap().shared; + // SAFETY: This permit owns one capacity unit. Claiming a slot and writing it is a + // synchronous operation with no user callbacks or await points between the two. + unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; + // Publication owns the capacity before a wake callback can panic. + self.sender = None; + shared.rx_waker.wake(); + Ok(()) + } +} + +impl Drop for Permit<'_, T> { + fn drop(&mut self) { + if let Some(sender) = self.sender { + sender.shared.tx_permits.release(); + } + } +} From 8070c87bbc237b45d57025bef652116d8be7ad17 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 13:19:57 +0800 Subject: [PATCH 24/34] refactor(mpsc): split bounded buffer storage by message size Give the slotted and zero-sized queues separate variants of one storage enum instead of sharing a flat struct with size checks. The zero-sized channel no longer carries a dead ticket, close flag, or slot pointer; the slotted variant is boxed so the enum does not reserve its room either way. Review experiment: the dispatch tag is nearly free, but the extra dependent load behind the box costs the single-threaded round trip. Numbers accompany the pull request discussion; drop this commit if the readability does not pay for them. --- asyncband/src/mpsc/bounded/buffer.rs | 197 +++++++++++++-------- asyncband/src/mpsc/bounded/buffer_tests.rs | 21 ++- 2 files changed, 143 insertions(+), 75 deletions(-) diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index ef4cfdbf..8c1d4e3c 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -42,13 +42,26 @@ use crate::internal::cache_padded::CachePadded; const EMPTY: u8 = 0; const READY: u8 = 1; const CLOSED: u8 = 2; -const ZERO_SIZED_CLOSED: usize = 1 << (usize::BITS - 1); +const CLOSED_BIT: usize = 1 << (usize::BITS - 1); pub struct Buffer { + storage: Storage, +} + +/// Slotted and zero-sized queues share no storage beyond a message count. Splitting them into +/// variants frees the zero-sized queue from a dead ticket, close flag, and slot allocation. +/// Boxing the slotted variant keeps that saving: an unboxed enum reserves room for the larger +/// variant either way. The variant tag sits beside the words every operation already loads, +/// so the dispatch branch is as predictable as the size check it replaces. +enum Storage { + Slots(Box>), + ZeroSized(ZeroSized), +} + +struct Slots { slots: Box<[Slot]>, tail: CachePadded, closed: AtomicBool, - zero_sized: ZeroSized, } struct Slot { @@ -88,7 +101,7 @@ impl ZeroSized { fn push(&self) -> bool { let mut queued = self.queued.load(Ordering::Acquire); loop { - if queued & ZERO_SIZED_CLOSED != 0 { + if queued & CLOSED_BIT != 0 { return false; } // The capacity limit keeps the count far below the closed flag bit. @@ -106,7 +119,7 @@ impl ZeroSized { /// Accounts for one consumed message, returning `false` when the queue was observed empty. fn pop(&self) -> bool { - let queued = self.queued.load(Ordering::Acquire) & !ZERO_SIZED_CLOSED; + let queued = self.queued.load(Ordering::Acquire) & !CLOSED_BIT; if queued == 0 { return false; } @@ -118,27 +131,22 @@ impl ZeroSized { /// Stops publication and returns the queue length transferred to the drain. fn close(&self) -> usize { - self.queued.fetch_or(ZERO_SIZED_CLOSED, Ordering::AcqRel) & !ZERO_SIZED_CLOSED + self.queued.fetch_or(CLOSED_BIT, Ordering::AcqRel) & !CLOSED_BIT } } -impl Buffer { - pub fn new(capacity: usize) -> Self { - let slots = if size_of::() == 0 { - Box::default() - } else { - (0..capacity.next_power_of_two()) - .map(|_| Slot { - state: AtomicU8::new(EMPTY), - value: UnsafeCell::new(MaybeUninit::uninit()), - }) - .collect() - }; +impl Slots { + fn new(capacity: usize) -> Self { + let slots = (0..capacity.next_power_of_two()) + .map(|_| Slot { + state: AtomicU8::new(EMPTY), + value: UnsafeCell::new(MaybeUninit::uninit()), + }) + .collect(); Self { slots, tail: CachePadded::new(AtomicUsize::new(0)), closed: AtomicBool::new(false), - zero_sized: ZeroSized::new(), } } @@ -157,28 +165,11 @@ impl Buffer { Ok(self.tail.fetch_add(1, Ordering::AcqRel)) } - /// Writes and publishes one message. Closing may instead return the unsent value. + /// Writes and publishes one message into a claimed position. /// /// # Safety /// - /// Own one capacity permit before calling; release it only after a failed push or after - /// the consumer reads the published value. No user code runs between claim and publication. - pub unsafe fn push(&self, value: T) -> Result<(), T> { - if size_of::() == 0 { - return if self.zero_sized.push() { - mem::forget(value); - Ok(()) - } else { - Err(value) - }; - } - let Ok(position) = self.claim() else { - return Err(value); - }; - // SAFETY: The caller owns capacity and the atomic increment assigned this position. - unsafe { self.publish(position, value) } - } - + /// Own the position from a claim, backed by a capacity permit, and publish it at most once. unsafe fn publish(&self, position: usize, value: T) -> Result<(), T> { let slot = self.slot(position); // SAFETY: Capacity prevents wrapping over unread slots. AcqRel tail increments carry prior @@ -203,17 +194,8 @@ impl Buffer { /// /// # Safety /// - /// Only the exclusive consumer may call this, using its persistent cursor. Release one - /// capacity permit after each successful pop, after the value has been read completely. - pub unsafe fn pop(&self, head: &mut usize) -> Poll> { - if size_of::() == 0 { - return if self.zero_sized.pop() { - // SAFETY: A queued value proves that this ZST is inhabited and owns one value. - Poll::Ready(Some(unsafe { Self::read_zero_sized() })) - } else { - Poll::Ready(None) - }; - } + /// Only the exclusive consumer may call this, using its persistent cursor. + unsafe fn pop(&self, head: &mut usize) -> Poll> { let slot = self.slot(*head); if slot.state.load(Ordering::Acquire) == READY { // SAFETY: Publication initialized the value, and only this consumer can read it. @@ -229,19 +211,83 @@ impl Buffer { } } + /// Stops new claims and returns the physical slot count the drain must cover: a producer + /// may have passed the open check but not obtained its ticket yet, and such a late claim + /// must also find a CLOSED slot. + fn close(&self) -> usize { + self.closed.store(true, Ordering::Release); + self.slots.len() + } +} + +impl Buffer { + pub fn new(capacity: usize) -> Self { + let storage = if size_of::() == 0 { + Storage::ZeroSized(ZeroSized::new()) + } else { + Storage::Slots(Box::new(Slots::new(capacity))) + }; + Self { storage } + } + + /// Writes and publishes one message. Closing may instead return the unsent value. + /// + /// # Safety + /// + /// Own one capacity permit before calling; release it only after a failed push or after + /// the consumer reads the published value. No user code runs between claim and publication. + pub unsafe fn push(&self, value: T) -> Result<(), T> { + match &self.storage { + Storage::Slots(slots) => { + let Ok(position) = slots.claim() else { + return Err(value); + }; + // SAFETY: The caller owns capacity and the ticket assigned this position. + unsafe { slots.publish(position, value) } + } + Storage::ZeroSized(zero_sized) => { + debug_assert_eq!(size_of::(), 0); + if zero_sized.push() { + mem::forget(value); + Ok(()) + } else { + Err(value) + } + } + } + } + + /// Pending means a producer claimed the head but has not published it yet. + /// + /// # Safety + /// + /// Only the exclusive consumer may call this, using its persistent cursor. Release one + /// capacity permit after each successful pop, after the value has been read completely. + pub unsafe fn pop(&self, head: &mut usize) -> Poll> { + match &self.storage { + // SAFETY: The caller's guarantee forwards unchanged. + Storage::Slots(slots) => unsafe { slots.pop(head) }, + Storage::ZeroSized(zero_sized) => { + debug_assert_eq!(size_of::(), 0); + if zero_sized.pop() { + // SAFETY: A queued value proves that this ZST is inhabited and owns one value. + Poll::Ready(Some(unsafe { read_zero_sized() })) + } else { + Poll::Ready(None) + } + } + } + } + /// Stops new claims and returns ownership of published values to a drain guard. /// /// # Safety /// /// Only the exclusive consumer may close the buffer, once, using its current cursor. pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { - self.closed.store(true, Ordering::Release); - let remaining = if size_of::() == 0 { - self.zero_sized.close() - } else { - // Cover every physical slot: a producer may have passed the open check but not - // obtained its ticket yet. Such a late claim must also find a CLOSED slot. - self.slots.len() + let remaining = match &self.storage { + Storage::Slots(slots) => slots.close(), + Storage::ZeroSized(zero_sized) => zero_sized.close(), }; Drain { buffer: self, @@ -250,13 +296,24 @@ impl Buffer { } } - unsafe fn read_zero_sized() -> T { - // SAFETY: The caller owns an initialized, inhabited ZST. Reading it accesses no bytes; - // dangling supplies a non-null, correctly aligned pointer, as in a ZST Vec. - unsafe { NonNull::::dangling().as_ptr().read() } + #[cfg(test)] + fn slots(&self) -> &Slots { + match &self.storage { + Storage::Slots(slots) => slots, + Storage::ZeroSized(_) => unreachable!("zero-sized messages have no slots"), + } } } +/// # Safety +/// +/// The caller must own an initialized, inhabited ZST value. +unsafe fn read_zero_sized() -> T { + // SAFETY: Reading it accesses no bytes; dangling supplies a non-null, correctly aligned + // pointer, as in a ZST Vec. + unsafe { NonNull::::dangling().as_ptr().read() } +} + pub struct Drain<'a, T> { buffer: &'a Buffer, position: usize, @@ -271,18 +328,20 @@ impl Iterator for Drain<'_, T> { let position = self.position; self.remaining -= 1; self.position = self.position.wrapping_add(1); - if size_of::() == 0 { + match &self.buffer.storage { + Storage::Slots(slots) => { + let slot = slots.slot(position); + if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { + // SAFETY: The drain won ownership of a published value. The cursor and + // state already advanced, so a panicking destructor cannot read twice. + return Some(unsafe { (*slot.value.get()).assume_init_read() }); + } + // An unpublished slot stays owned by its producer, which will observe CLOSED + // and recover its value. The shared Arc keeps this allocation alive until then. + } // SAFETY: Closing transferred this many initialized ZST values to the drain. - return Some(unsafe { Buffer::::read_zero_sized() }); - } - let slot = self.buffer.slot(position); - if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { - // SAFETY: The drain won ownership of a published value. The cursor and state - // already advanced, so a panicking destructor cannot cause a second read. - return Some(unsafe { (*slot.value.get()).assume_init_read() }); + Storage::ZeroSized(_) => return Some(unsafe { read_zero_sized() }), } - // An unpublished slot stays owned by its producer, which will observe CLOSED and - // recover its value. The shared Arc keeps this allocation alive until that finishes. } None } diff --git a/asyncband/src/mpsc/bounded/buffer_tests.rs b/asyncband/src/mpsc/bounded/buffer_tests.rs index 6ceba628..f6ec4fc8 100644 --- a/asyncband/src/mpsc/bounded/buffer_tests.rs +++ b/asyncband/src/mpsc/bounded/buffer_tests.rs @@ -38,7 +38,7 @@ fn publish_claimed( value: T, ) -> Result<(), T> { // SAFETY: The test claimed this position while holding the same capacity permit. - unsafe { tx.shared().buffer.publish(position, value) }?; + unsafe { tx.shared().buffer.slots().publish(position, value) }?; // Publication owns the capacity now; forgetting skips the permit's release on drop. std::mem::forget(permit); tx.shared().rx_waker.wake(); @@ -51,12 +51,16 @@ fn a_claimed_head_waits_for_publication_across_laps() { for initial in [0, usize::MAX - 1] { let (tx, mut rx) = bounded(capacity); // Start an empty ring near ticket overflow instead of running usize::MAX sends. - tx.shared().buffer.tail.store(initial, Ordering::Relaxed); + tx.shared() + .buffer + .slots() + .tail + .store(initial, Ordering::Relaxed); rx.set_head(initial); let mut cx = Context::from_waker(Waker::noop()); for lap in 0..8 { let permit = tx.try_reserve().unwrap(); - let position = tx.shared().buffer.claim().unwrap(); + let position = tx.shared().buffer.slots().claim().unwrap(); for offset in 1..capacity { tx.try_send(lap * capacity + offset).unwrap(); } @@ -84,9 +88,14 @@ fn a_claim_delayed_past_close_returns_its_value() { }; let allocation = Arc::downgrade(tx.shared()); // Pause after claim's open check, then resume its atomic ticket allocation after close. - assert!(!tx.shared().buffer.closed.load(Ordering::Acquire)); + assert!(!tx.shared().buffer.slots().closed.load(Ordering::Acquire)); drop(rx); - let position = tx.shared().buffer.tail.fetch_add(1, Ordering::AcqRel); + let position = tx + .shared() + .buffer + .slots() + .tail + .fetch_add(1, Ordering::AcqRel); let unsent = publish_claimed(&tx, permit, position, value).unwrap_err(); assert_eq!(unsent.bytes, [7; 1024]); drop(unsent); @@ -125,7 +134,7 @@ fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { let paused = &paused; let publisher = scope.spawn(move || { let permit = sender.try_reserve().unwrap(); - let position = sender.shared().buffer.claim().unwrap(); + let position = sender.shared().buffer.slots().claim().unwrap(); let value = Payload { bytes: [1; 1024], drops: drops.clone(), From 4c922a9368324994c442b80fb86d3821499da5a2 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 13:43:54 +0800 Subject: [PATCH 25/34] test(mpsc): inline the bounded buffer tests The file boundary bought nothing over a test module, and the split hid how little code the unsafe core actually has. --- asyncband/src/mpsc/bounded/buffer.rs | 222 ++++++++++++++++++- asyncband/src/mpsc/bounded/buffer_tests.rs | 235 --------------------- 2 files changed, 220 insertions(+), 237 deletions(-) delete mode 100644 asyncband/src/mpsc/bounded/buffer_tests.rs diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index 8c1d4e3c..ff0bdebd 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -369,5 +369,223 @@ impl Drop for Drain<'_, T> { } #[cfg(test)] -#[path = "buffer_tests.rs"] -mod tests; +mod tests { + use std::sync::Arc; + use std::sync::Barrier; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::task::Context; + use std::task::Poll; + use std::task::Waker; + use std::thread; + + use crate::mpsc::BoundedSender; + use crate::mpsc::Permit; + use crate::mpsc::TryRecvError; + use crate::mpsc::bounded; + + // Exercise the scheduling window inside synchronous send, while retaining a real capacity + // permit. Ordinary callers cannot split a claim from its publication. + fn publish_claimed( + tx: &BoundedSender, + permit: Permit<'_, T>, + position: usize, + value: T, + ) -> Result<(), T> { + // SAFETY: The test claimed this position while holding the same capacity permit. + unsafe { tx.shared().buffer.slots().publish(position, value) }?; + // Publication owns the capacity now; forgetting skips the permit's release on drop. + std::mem::forget(permit); + tx.shared().rx_waker.wake(); + Ok(()) + } + + #[test] + fn a_claimed_head_waits_for_publication_across_laps() { + for capacity in [1, 3, 7] { + for initial in [0, usize::MAX - 1] { + let (tx, mut rx) = bounded(capacity); + // Start an empty ring near ticket overflow instead of running usize::MAX sends. + tx.shared() + .buffer + .slots() + .tail + .store(initial, Ordering::Relaxed); + rx.set_head(initial); + let mut cx = Context::from_waker(Waker::noop()); + for lap in 0..8 { + let permit = tx.try_reserve().unwrap(); + let position = tx.shared().buffer.slots().claim().unwrap(); + for offset in 1..capacity { + tx.try_send(lap * capacity + offset).unwrap(); + } + // A full ring must differ from an empty one even with no head value ready. + assert!(rx.poll_recv(&mut cx).is_pending()); + publish_claimed(&tx, permit, position, lap * capacity).unwrap(); + for offset in 0..capacity { + assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } + } + } + } + + #[test] + fn a_claim_delayed_past_close_returns_its_value() { + let (tx, rx) = bounded(3); + let permit = tx.try_reserve().unwrap(); + let drops = Arc::new(AtomicUsize::new(0)); + let value = Payload { + bytes: [7; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let allocation = Arc::downgrade(tx.shared()); + // Pause after claim's open check, then resume its atomic ticket allocation after close. + assert!(!tx.shared().buffer.slots().closed.load(Ordering::Acquire)); + drop(rx); + let position = tx + .shared() + .buffer + .slots() + .tail + .fetch_add(1, Ordering::AcqRel); + let unsent = publish_claimed(&tx, permit, position, value).unwrap_err(); + assert_eq!(unsent.bytes, [7; 1024]); + drop(unsent); + assert_eq!(drops.load(Ordering::Relaxed), 1); + drop(tx); + assert!(allocation.upgrade().is_none()); + } + + #[derive(Debug)] + #[repr(align(128))] + struct Payload { + bytes: [u8; 1024], + drops: Arc, + // Queued messages must not keep the shared allocation alive through a sender cycle. + _sender: BoundedSender, + } + + impl Drop for Payload { + fn drop(&mut self) { + self.drops.fetch_add(1, Ordering::Relaxed); + } + } + + #[test] + fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { + let (tx, rx) = bounded(2); + let allocation = Arc::downgrade(tx.shared()); + let drops = Arc::new(AtomicUsize::new(0)); + let paused = Barrier::new(2); + let (resume_tx, resume_rx) = std::sync::mpsc::channel(); + let (closed_tx, closed_rx) = std::sync::mpsc::channel(); + + thread::scope(|scope| { + let sender = &tx; + let drops = &drops; + let paused = &paused; + let publisher = scope.spawn(move || { + let permit = sender.try_reserve().unwrap(); + let position = sender.shared().buffer.slots().claim().unwrap(); + let value = Payload { + bytes: [1; 1024], + drops: drops.clone(), + _sender: sender.clone(), + }; + paused.wait(); + resume_rx.recv().unwrap(); + let unsent = publish_claimed(sender, permit, position, value).unwrap_err(); + assert_eq!(unsent.bytes, [1; 1024]); + drop(unsent); + }); + paused.wait(); + tx.try_send(Payload { + bytes: [2; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }) + .unwrap(); + let closer = scope.spawn(move || { + drop(rx); + closed_tx.send(()).unwrap(); + }); + #[cfg(not(miri))] + let closed = closed_rx.recv_timeout(std::time::Duration::from_secs(10)); + #[cfg(miri)] + let closed = closed_rx.recv(); + let dropped_before_resume = drops.load(Ordering::Relaxed); + // Unblock the publisher before asserting so a failed close cannot strand the scope. + resume_tx.send(()).unwrap(); + publisher.join().unwrap(); + closer.join().unwrap(); + assert!(closed.is_ok(), "close waited for the paused publisher"); + assert_eq!(dropped_before_resume, 1); + }); + + assert_eq!(drops.load(Ordering::Relaxed), 2); + drop(tx); + assert!(allocation.upgrade().is_none()); + } + + #[test] + fn publication_racing_with_close_drops_every_payload_once() { + for _ in 0..if cfg!(miri) { 8 } else { 128 } { + let (tx, rx) = bounded(3); + let allocation = Arc::downgrade(tx.shared()); + let drops = Arc::new(AtomicUsize::new(0)); + let start = Barrier::new(4); + thread::scope(|scope| { + for byte in 0..3 { + let permit = tx.try_reserve().unwrap(); + let value = Payload { + bytes: [byte; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let start = &start; + scope.spawn(move || { + start.wait(); + if let Err(error) = permit.send(value) { + let value = error.into_inner(); + assert_eq!(value.bytes, [byte; 1024]); + drop(value); + } + }); + } + start.wait(); + drop(rx); + }); + assert_eq!(drops.load(Ordering::Relaxed), 3); + drop(tx); + assert!(allocation.upgrade().is_none()); + } + } + + #[test] + fn an_old_permit_can_publish_after_other_producers_wrap_the_ring() { + let (tx, mut rx) = bounded(3); + let old = tx.try_reserve().unwrap(); + for lap in 0..16 { + for offset in 0..2 { + tx.try_send([lap * 2 + offset; 1024]).unwrap(); + } + for offset in 0..2 { + assert_eq!(rx.try_recv(), Ok([lap * 2 + offset; 1024])); + } + } + thread::scope(|scope| { + scope + .spawn(move || old.send([42; 1024]).unwrap()) + .join() + .unwrap(); + }); + assert_eq!( + rx.poll_recv(&mut Context::from_waker(Waker::noop())), + Poll::Ready(Ok([42; 1024])) + ); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } +} diff --git a/asyncband/src/mpsc/bounded/buffer_tests.rs b/asyncband/src/mpsc/bounded/buffer_tests.rs deleted file mode 100644 index f6ec4fc8..00000000 --- a/asyncband/src/mpsc/bounded/buffer_tests.rs +++ /dev/null @@ -1,235 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; -use std::sync::Barrier; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use std::task::Waker; -use std::thread; - -use crate::mpsc::BoundedSender; -use crate::mpsc::Permit; -use crate::mpsc::TryRecvError; -use crate::mpsc::bounded; - -// Exercise the scheduling window inside synchronous send, while retaining a real capacity -// permit. Ordinary callers cannot split a claim from its publication. -fn publish_claimed( - tx: &BoundedSender, - permit: Permit<'_, T>, - position: usize, - value: T, -) -> Result<(), T> { - // SAFETY: The test claimed this position while holding the same capacity permit. - unsafe { tx.shared().buffer.slots().publish(position, value) }?; - // Publication owns the capacity now; forgetting skips the permit's release on drop. - std::mem::forget(permit); - tx.shared().rx_waker.wake(); - Ok(()) -} - -#[test] -fn a_claimed_head_waits_for_publication_across_laps() { - for capacity in [1, 3, 7] { - for initial in [0, usize::MAX - 1] { - let (tx, mut rx) = bounded(capacity); - // Start an empty ring near ticket overflow instead of running usize::MAX sends. - tx.shared() - .buffer - .slots() - .tail - .store(initial, Ordering::Relaxed); - rx.set_head(initial); - let mut cx = Context::from_waker(Waker::noop()); - for lap in 0..8 { - let permit = tx.try_reserve().unwrap(); - let position = tx.shared().buffer.slots().claim().unwrap(); - for offset in 1..capacity { - tx.try_send(lap * capacity + offset).unwrap(); - } - // A full ring must differ from an empty one even if no head value is ready yet. - assert!(rx.poll_recv(&mut cx).is_pending()); - publish_claimed(&tx, permit, position, lap * capacity).unwrap(); - for offset in 0..capacity { - assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); - } - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - } - } - } -} - -#[test] -fn a_claim_delayed_past_close_returns_its_value() { - let (tx, rx) = bounded(3); - let permit = tx.try_reserve().unwrap(); - let drops = Arc::new(AtomicUsize::new(0)); - let value = Payload { - bytes: [7; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }; - let allocation = Arc::downgrade(tx.shared()); - // Pause after claim's open check, then resume its atomic ticket allocation after close. - assert!(!tx.shared().buffer.slots().closed.load(Ordering::Acquire)); - drop(rx); - let position = tx - .shared() - .buffer - .slots() - .tail - .fetch_add(1, Ordering::AcqRel); - let unsent = publish_claimed(&tx, permit, position, value).unwrap_err(); - assert_eq!(unsent.bytes, [7; 1024]); - drop(unsent); - assert_eq!(drops.load(Ordering::Relaxed), 1); - drop(tx); - assert!(allocation.upgrade().is_none()); -} - -#[derive(Debug)] -#[repr(align(128))] -struct Payload { - bytes: [u8; 1024], - drops: Arc, - // Queued messages must not keep the shared allocation alive through a sender cycle. - _sender: BoundedSender, -} - -impl Drop for Payload { - fn drop(&mut self) { - self.drops.fetch_add(1, Ordering::Relaxed); - } -} - -#[test] -fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { - let (tx, rx) = bounded(2); - let allocation = Arc::downgrade(tx.shared()); - let drops = Arc::new(AtomicUsize::new(0)); - let paused = Barrier::new(2); - let (resume_tx, resume_rx) = std::sync::mpsc::channel(); - let (closed_tx, closed_rx) = std::sync::mpsc::channel(); - - thread::scope(|scope| { - let sender = &tx; - let drops = &drops; - let paused = &paused; - let publisher = scope.spawn(move || { - let permit = sender.try_reserve().unwrap(); - let position = sender.shared().buffer.slots().claim().unwrap(); - let value = Payload { - bytes: [1; 1024], - drops: drops.clone(), - _sender: sender.clone(), - }; - paused.wait(); - resume_rx.recv().unwrap(); - let unsent = publish_claimed(sender, permit, position, value).unwrap_err(); - assert_eq!(unsent.bytes, [1; 1024]); - drop(unsent); - }); - paused.wait(); - tx.try_send(Payload { - bytes: [2; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }) - .unwrap(); - let closer = scope.spawn(move || { - drop(rx); - closed_tx.send(()).unwrap(); - }); - #[cfg(not(miri))] - let closed = closed_rx.recv_timeout(std::time::Duration::from_secs(10)); - #[cfg(miri)] - let closed = closed_rx.recv(); - let dropped_before_resume = drops.load(Ordering::Relaxed); - // Unblock the publisher before asserting so a failed close cannot strand the scope. - resume_tx.send(()).unwrap(); - publisher.join().unwrap(); - closer.join().unwrap(); - assert!(closed.is_ok(), "close waited for the paused publisher"); - assert_eq!(dropped_before_resume, 1); - }); - - assert_eq!(drops.load(Ordering::Relaxed), 2); - drop(tx); - assert!(allocation.upgrade().is_none()); -} - -#[test] -fn publication_racing_with_close_drops_every_payload_once() { - for _ in 0..if cfg!(miri) { 8 } else { 128 } { - let (tx, rx) = bounded(3); - let allocation = Arc::downgrade(tx.shared()); - let drops = Arc::new(AtomicUsize::new(0)); - let start = Barrier::new(4); - thread::scope(|scope| { - for byte in 0..3 { - let permit = tx.try_reserve().unwrap(); - let value = Payload { - bytes: [byte; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }; - let start = &start; - scope.spawn(move || { - start.wait(); - if let Err(error) = permit.send(value) { - let value = error.into_inner(); - assert_eq!(value.bytes, [byte; 1024]); - drop(value); - } - }); - } - start.wait(); - drop(rx); - }); - assert_eq!(drops.load(Ordering::Relaxed), 3); - drop(tx); - assert!(allocation.upgrade().is_none()); - } -} - -#[test] -fn an_old_permit_can_publish_after_other_producers_wrap_the_ring() { - let (tx, mut rx) = bounded(3); - let old = tx.try_reserve().unwrap(); - for lap in 0..16 { - for offset in 0..2 { - tx.try_send([lap * 2 + offset; 1024]).unwrap(); - } - for offset in 0..2 { - assert_eq!(rx.try_recv(), Ok([lap * 2 + offset; 1024])); - } - } - thread::scope(|scope| { - scope - .spawn(move || old.send([42; 1024]).unwrap()) - .join() - .unwrap(); - }); - assert_eq!( - rx.poll_recv(&mut Context::from_waker(Waker::noop())), - Poll::Ready(Ok([42; 1024])) - ); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); -} From 7093466b498419bee46627e17fbd081836b0613a Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 14:16:54 +0800 Subject: [PATCH 26/34] fixup Signed-off-by: tison --- asyncband/src/mpsc/bounded/buffer.rs | 25 +++++------ asyncband/src/mpsc/bounded/mod.rs | 61 +++++++++++++++++++------- asyncband/src/mpsc/bounded/receiver.rs | 11 +---- asyncband/src/mpsc/bounded/sender.rs | 19 +------- 4 files changed, 58 insertions(+), 58 deletions(-) diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index ff0bdebd..2180e1c7 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -393,10 +393,10 @@ mod tests { value: T, ) -> Result<(), T> { // SAFETY: The test claimed this position while holding the same capacity permit. - unsafe { tx.shared().buffer.slots().publish(position, value) }?; + unsafe { tx.shared.buffer.slots().publish(position, value) }?; // Publication owns the capacity now; forgetting skips the permit's release on drop. std::mem::forget(permit); - tx.shared().rx_waker.wake(); + tx.shared.rx_waker.wake(); Ok(()) } @@ -406,7 +406,7 @@ mod tests { for initial in [0, usize::MAX - 1] { let (tx, mut rx) = bounded(capacity); // Start an empty ring near ticket overflow instead of running usize::MAX sends. - tx.shared() + tx.shared .buffer .slots() .tail @@ -415,7 +415,7 @@ mod tests { let mut cx = Context::from_waker(Waker::noop()); for lap in 0..8 { let permit = tx.try_reserve().unwrap(); - let position = tx.shared().buffer.slots().claim().unwrap(); + let position = tx.shared.buffer.slots().claim().unwrap(); for offset in 1..capacity { tx.try_send(lap * capacity + offset).unwrap(); } @@ -441,16 +441,11 @@ mod tests { drops: drops.clone(), _sender: tx.clone(), }; - let allocation = Arc::downgrade(tx.shared()); + let allocation = Arc::downgrade(&tx.shared); // Pause after claim's open check, then resume its atomic ticket allocation after close. - assert!(!tx.shared().buffer.slots().closed.load(Ordering::Acquire)); + assert!(!&tx.shared.buffer.slots().closed.load(Ordering::Acquire)); drop(rx); - let position = tx - .shared() - .buffer - .slots() - .tail - .fetch_add(1, Ordering::AcqRel); + let position = tx.shared.buffer.slots().tail.fetch_add(1, Ordering::AcqRel); let unsent = publish_claimed(&tx, permit, position, value).unwrap_err(); assert_eq!(unsent.bytes, [7; 1024]); drop(unsent); @@ -477,7 +472,7 @@ mod tests { #[test] fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { let (tx, rx) = bounded(2); - let allocation = Arc::downgrade(tx.shared()); + let allocation = Arc::downgrade(&tx.shared); let drops = Arc::new(AtomicUsize::new(0)); let paused = Barrier::new(2); let (resume_tx, resume_rx) = std::sync::mpsc::channel(); @@ -489,7 +484,7 @@ mod tests { let paused = &paused; let publisher = scope.spawn(move || { let permit = sender.try_reserve().unwrap(); - let position = sender.shared().buffer.slots().claim().unwrap(); + let position = sender.shared.buffer.slots().claim().unwrap(); let value = Payload { bytes: [1; 1024], drops: drops.clone(), @@ -534,7 +529,7 @@ mod tests { fn publication_racing_with_close_drops_every_payload_once() { for _ in 0..if cfg!(miri) { 8 } else { 128 } { let (tx, rx) = bounded(3); - let allocation = Arc::downgrade(tx.shared()); + let allocation = Arc::downgrade(&tx.shared); let drops = Arc::new(AtomicUsize::new(0)); let start = Barrier::new(4); thread::scope(|scope| { diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 914f7560..b2dcf367 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -35,10 +35,6 @@ mod receiver; mod semaphore; mod sender; -pub use self::receiver::BoundedReceiver; -pub use self::sender::BoundedSender; -pub use self::sender::Permit; - /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// /// [`BoundedSender::send`] waits for capacity when the buffer is full. Receiving a message releases @@ -53,32 +49,65 @@ pub use self::sender::Permit; /// Panics if `buffer` is zero or exceeds the maximum capacity of `usize::MAX >> 1`. #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { - assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); + /// The largest capacity accepted by [`bounded`]. + /// + /// The shared permit counter reserves two sentinel values above the usable range, and the + /// zero-sized queue length packs a closed flag into its top bit. This bound also keeps the + /// rounded-up slot storage from overflowing a power of two. + const MAX_CAPACITY: usize = usize::MAX >> 1; + + assert!( + buffer > 0, + "mpsc bounded channel capacity {buffer} must be nonzero", + ); assert!( buffer <= MAX_CAPACITY, - "mpsc bounded channel capacity {buffer} exceeds the maximum of {MAX_CAPACITY}" + "mpsc bounded channel capacity {buffer} exceeds the maximum of {MAX_CAPACITY}", ); + let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), rx_waker: AtomicWaker::new(), buffer: Buffer::new(buffer), }); - let sender = BoundedSender::new(shared.clone()); - let receiver = BoundedReceiver::new(shared); + let sender = BoundedSender { + shared: shared.clone(), + }; + let receiver = BoundedReceiver { shared, head: 0 }; (sender, receiver) } -pub struct Shared { +/// The sending endpoint of a bounded mpsc channel. +/// +/// Instances are created by the [`bounded`] function. +pub struct BoundedSender { + shared: Arc>, +} + +/// The receiving endpoint of a bounded mpsc channel. +/// +/// Instances are created by the [`bounded`] function. Dropping the receiver discards queued values. +/// The backing allocation remains alive until all endpoints are dropped, so a concurrent sender +/// can safely finish returning an unsent value. +pub struct BoundedReceiver { + shared: Arc>, + head: usize, +} + +/// Capacity reserved for one message on a bounded channel. +/// +/// Created by [`BoundedSender::reserve`] or [`BoundedSender::try_reserve`]. Holding a permit +/// reduces available capacity but does not prevent other messages from being received. Dropping +/// it without sending releases capacity and notifies a waiting sender. +#[must_use = "dropping the permit releases its reserved capacity"] +pub struct Permit<'a, T> { + sender: Option<&'a BoundedSender>, +} + +struct Shared { senders: AtomicUsize, tx_permits: CachePadded, rx_waker: AtomicWaker, buffer: Buffer, } - -/// The largest capacity accepted by [`bounded`]. -/// -/// The shared permit counter reserves two sentinel values above the usable range, and the -/// zero-sized queue length packs a closed flag into its top bit. This bound also keeps the -/// rounded-up slot storage from overflowing a power of two. -const MAX_CAPACITY: usize = usize::MAX >> 1; diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs index 9a3cbd73..8c675d9c 100644 --- a/asyncband/src/mpsc/bounded/receiver.rs +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -25,18 +25,9 @@ use std::task::Poll; use super::RecvError; use super::Shared; use super::TryRecvError; +use super::BoundedReceiver; use crate::internal::wake_all; -/// The receiving endpoint of a bounded mpsc channel. -/// -/// Instances are created by the [`bounded`](super::bounded) function. -/// Dropping the receiver discards queued values. The backing allocation remains alive until -/// all endpoints are dropped, so a concurrent sender can safely finish returning an unsent value. -pub struct BoundedReceiver { - shared: Arc>, - head: usize, -} - impl fmt::Debug for BoundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BoundedReceiver").finish_non_exhaustive() diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index e2a83142..2e1e573f 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -20,17 +20,12 @@ use std::future::poll_fn; use std::sync::Arc; use std::sync::atomic::Ordering; +use super::BoundedSender; +use super::Permit; use super::SendError; use super::Shared; use super::TrySendError; -/// The sending endpoint of a bounded mpsc channel. -/// -/// Instances are created by the [`bounded`](super::bounded) function. -pub struct BoundedSender { - shared: Arc>, -} - impl Clone for BoundedSender { fn clone(&self) -> Self { self.shared.senders.fetch_add(1, Ordering::Relaxed); @@ -157,16 +152,6 @@ impl BoundedSender { } } -/// Capacity reserved for one message on a bounded channel. -/// -/// Created by [`BoundedSender::reserve`] or [`BoundedSender::try_reserve`]. Holding a permit -/// reduces available capacity but does not prevent other messages from being received. Dropping -/// it without sending releases capacity and notifies a waiting sender. -#[must_use = "dropping the permit releases its reserved capacity"] -pub struct Permit<'a, T> { - sender: Option<&'a BoundedSender>, -} - impl fmt::Debug for Permit<'_, T> { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("Permit").finish_non_exhaustive() From 7707e0275d02fe8c07c2f19b3a5821b4b26ecb68 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 14:39:28 +0800 Subject: [PATCH 27/34] fixup Signed-off-by: tison --- asyncband/src/mpsc/bounded/receiver.rs | 14 ++++---------- asyncband/src/mpsc/bounded/sender.rs | 11 ----------- 2 files changed, 4 insertions(+), 21 deletions(-) diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs index 8c675d9c..88db33a1 100644 --- a/asyncband/src/mpsc/bounded/receiver.rs +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -17,15 +17,13 @@ use std::fmt; use std::future::poll_fn; -use std::sync::Arc; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use super::BoundedReceiver; use super::RecvError; -use super::Shared; use super::TryRecvError; -use super::BoundedReceiver; use crate::internal::wake_all; impl fmt::Debug for BoundedReceiver { @@ -48,10 +46,6 @@ impl Drop for BoundedReceiver { } impl BoundedReceiver { - pub(crate) fn new(shared: Arc>) -> Self { - Self { shared, head: 0 } - } - /// Attempts to receive the next queued value without waiting for a new message. /// /// Receiving a value frees one buffer slot. An empty channel returns [`TryRecvError::Empty`] @@ -129,9 +123,9 @@ impl BoundedReceiver { /// /// # Cancel safety /// - /// Dropping a pending `recv` does not remove a message from the channel. A later receive - /// operation can still observe the next queued value, so `recv` may safely be raced with other - /// futures in a selection construct. + /// Dropping a pending `recv` does not remove a message from the channel. A later `recv` call + /// can still observe the next queued value, so `recv` may safely be raced with other futures + /// in a selection construct. /// /// # Examples /// diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index 2e1e573f..7cbe8f77 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -17,13 +17,11 @@ use std::fmt; use std::future::poll_fn; -use std::sync::Arc; use std::sync::atomic::Ordering; use super::BoundedSender; use super::Permit; use super::SendError; -use super::Shared; use super::TrySendError; impl Clone for BoundedSender { @@ -50,10 +48,6 @@ impl Drop for BoundedSender { } impl BoundedSender { - pub(crate) fn new(shared: Arc>) -> Self { - Self { shared } - } - /// Sends a message, waiting until the channel has capacity when necessary. /// /// If the receiver has been dropped, the returned error contains `value`. @@ -145,11 +139,6 @@ impl BoundedSender { Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), } } - - #[cfg(test)] - pub(crate) fn shared(&self) -> &Arc> { - &self.shared - } } impl fmt::Debug for Permit<'_, T> { From 25775fb2f3cfa3cce25cf9e5254ec4095fabc725 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 15:48:30 +0800 Subject: [PATCH 28/34] fix(mpsc): preserve wakeups and capacity ownership --- asyncband/src/internal/atomic_waker.rs | 83 +++-- asyncband/src/mpsc/bounded/buffer.rs | 304 +++++++++--------- asyncband/src/mpsc/bounded/mod.rs | 44 +-- asyncband/src/mpsc/bounded/receiver.rs | 77 +++-- asyncband/src/mpsc/bounded/semaphore.rs | 74 +++-- asyncband/src/mpsc/bounded/sender.rs | 64 ++-- .../tests/mpsc_test/callbacks.rs | 59 ++++ .../tests/mpsc_test/concurrency.rs | 41 +++ tests-integration/tests/mpsc_test/main.rs | 2 +- 9 files changed, 448 insertions(+), 300 deletions(-) diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs index a8c0448b..7eb99e5c 100644 --- a/asyncband/src/internal/atomic_waker.rs +++ b/asyncband/src/internal/atomic_waker.rs @@ -57,7 +57,7 @@ const WAKING: usize = 0b10; /// REGISTERING ------------AcqRel CAS----------------> WAITING /// /// wake: WAITING ----------------AcqRel fetch_or-----------> WAKING -/// WAKING -----------------Release store-------------> WAITING +/// WAKING -----------------Release swap--------------> WAITING /// /// race: REGISTERING ------------AcqRel fetch_or-----------> REGISTERING | WAKING /// REGISTERING | WAKING ---AcqRel swap---------------> WAITING @@ -221,19 +221,7 @@ impl AtomicWaker { WAITING => { // SAFETY: changing WAITING to WAKING grants this thread exclusive access to the // waker slot until the state is returned to WAITING. - let waker = unsafe { (*self.waker.get()).take() }; - - // ORDERING: Release publishes the emptied slot before another operation acquires - // it. The fetch_or above already performed the required Acquire operation. A - // plain store suffices: only this claim moves the state out of WAKING, because - // registration enters from WAITING and concurrent wakes keep the bit set. Debug - // builds pay for a swap to assert that invariant. - if cfg!(debug_assertions) { - debug_assert_eq!(self.state.swap(WAITING, Ordering::Release), WAKING); - } else { - self.state.store(WAITING, Ordering::Release); - } - waker + unsafe { self.take_locked() } } state => { // The thread registering a waker observes WAKING and completes this notification, @@ -245,6 +233,25 @@ impl AtomicWaker { } } } + + /// Removes the waker after this thread has acquired the WAKING state. + /// + /// # Safety + /// + /// The caller must have changed `state` from WAITING to WAKING and must be the only thread + /// accessing `waker`. + #[inline] + unsafe fn take_locked(&self) -> Option { + // SAFETY: The caller owns the waker slot until returning the state to WAITING. + let waker = unsafe { (*self.waker.get()).take() }; + + // ORDERING: Release publishes the emptied slot. The RMW also preserves the release + // sequence of coalesced wakes, including ones after this thread acquired WAKING. A + // plain store would sever those publications from the next registration's Acquire. + let previous = self.state.swap(WAITING, Ordering::Release); + debug_assert_eq!(previous, WAKING); + waker + } } #[cfg(test)] @@ -350,26 +357,40 @@ mod tests { #[test] fn failed_wake_synchronizes_with_next_registration() { - for _ in 0..1_000 { - let did_publish = AtomicBool::new(false); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(Waker::noop()); + struct Publication(UnsafeCell); + // SAFETY: The notifier's write precedes its wake. The read follows a registration that + // acquires that wake's publication; the relaxed scheduling flag adds no synchronization. + unsafe impl Sync for Publication {} - std::thread::scope(|scope| { - let wake = scope.spawn(|| { - did_publish.store(true, Ordering::Relaxed); - atomic_waker.take() - }); - - let local_waker = atomic_waker.take(); - atomic_waker.register(Waker::noop()); + let publication = Arc::new(Publication(UnsafeCell::new(0))); + let did_wake = AtomicBool::new(false); + let atomic_waker = AtomicWaker::new(); + atomic_waker.register(Waker::noop()); + assert_eq!( + atomic_waker.state.fetch_or(WAKING, Ordering::AcqRel), + WAITING + ); - let publication_is_visible = did_publish.load(Ordering::Relaxed); - let concurrent_thread_took_waker = wake.join().unwrap().is_some(); - assert!(publication_is_visible || concurrent_thread_took_waker); - drop(local_waker); + std::thread::scope(|scope| { + let wake = scope.spawn(|| { + // SAFETY: The reader waits for this write's publication through AtomicWaker. + unsafe { *publication.0.get() = 42 }; + assert!(atomic_waker.take().is_none()); + did_wake.store(true, Ordering::Relaxed); }); - } + while !did_wake.load(Ordering::Relaxed) { + std::thread::yield_now(); + } + + // SAFETY: This thread acquired WAKING above. The coalesced notifier never touches + // the slot. Complete the first wake only after the second has published its update. + drop(unsafe { atomic_waker.take_locked() }); + atomic_waker.register(Waker::noop()); + // SAFETY: Registration acquires the coalesced wake through the release sequence. + // Miri detects a data race here if restoring WAITING severs that sequence. + assert_eq!(unsafe { *publication.0.get() }, 42); + wake.join().unwrap(); + }); } #[cfg(panic = "unwind")] diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index 2180e1c7..6a484337 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -44,21 +44,98 @@ const READY: u8 = 1; const CLOSED: u8 = 2; const CLOSED_BIT: usize = 1 << (usize::BITS - 1); -pub struct Buffer { - storage: Storage, -} - -/// Slotted and zero-sized queues share no storage beyond a message count. Splitting them into -/// variants frees the zero-sized queue from a dead ticket, close flag, and slot allocation. -/// Boxing the slotted variant keeps that saving: an unboxed enum reserves room for the larger -/// variant either way. The variant tag sits beside the words every operation already loads, -/// so the dispatch branch is as predictable as the size check it replaces. -enum Storage { +/// Zero-sized messages need only a count. Boxing the slotted variant keeps its cursor and +/// per-slot storage out of the zero-sized queue's allocation. +pub enum Buffer { Slots(Box>), ZeroSized(ZeroSized), } -struct Slots { +impl Buffer { + pub fn new(capacity: usize) -> Self { + if size_of::() == 0 { + Self::ZeroSized(ZeroSized::new()) + } else { + Self::Slots(Box::new(Slots::new(capacity))) + } + } + + /// Writes and publishes one message. Closing may instead return the unsent value. + /// + /// # Safety + /// + /// Own one capacity permit before calling; release it only after a failed push or after + /// the consumer reads the published value. No user code runs between claim and publication. + pub unsafe fn push(&self, value: T) -> Result<(), T> { + match self { + Self::Slots(slots) => { + let Ok(position) = slots.claim() else { + return Err(value); + }; + // SAFETY: The caller owns capacity and the ticket assigned this position. + unsafe { slots.publish(position, value) } + } + Self::ZeroSized(zero_sized) => { + debug_assert_eq!(size_of::(), 0); + if zero_sized.push() { + mem::forget(value); + Ok(()) + } else { + Err(value) + } + } + } + } + + /// Pending means a producer claimed the head but has not published it yet. + /// + /// # Safety + /// + /// Only the exclusive consumer may call this, using its persistent cursor. Release one + /// capacity permit after each successful pop, after the value has been read completely. + pub unsafe fn pop(&self, head: &mut usize) -> Poll> { + match self { + // SAFETY: The caller's guarantee forwards unchanged. + Self::Slots(slots) => unsafe { slots.pop(head) }, + Self::ZeroSized(zero_sized) => { + debug_assert_eq!(size_of::(), 0); + if zero_sized.pop() { + // SAFETY: A queued value proves that this ZST is inhabited and owns one value. + Poll::Ready(Some(unsafe { read_zero_sized() })) + } else { + Poll::Ready(None) + } + } + } + } + + /// Stops new claims and returns ownership of published values to a drain guard. + /// + /// # Safety + /// + /// Only the exclusive consumer may close the buffer, once, using its current cursor. + pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { + let remaining = match self { + Self::Slots(slots) => slots.close(), + Self::ZeroSized(zero_sized) => zero_sized.close(), + }; + Drain { + buffer: self, + position: head, + remaining, + } + } + + #[cfg(test)] + fn slots(&self) -> &Slots { + match self { + Self::Slots(slots) => slots, + Self::ZeroSized(_) => unreachable!("zero-sized messages have no slots"), + } + } +} + +pub struct Slots { slots: Box<[Slot]>, tail: CachePadded, closed: AtomicBool, @@ -79,62 +156,6 @@ unsafe impl Sync for Slot {} impl std::panic::UnwindSafe for Slot {} impl std::panic::RefUnwindSafe for Slot {} -/// Queue storage for zero-sized messages, which need no slots, positions, or per-slot flags. -/// Counting them separately also allows every nonzero usize capacity without allocating -/// publication metadata for nonexistent bytes. -/// -/// The queue is entirely its length. The count packs a closed flag into its top bit so that -/// publication and close stay atomic: a publication that raced ahead of the flag is included -/// in the drained count, and every later one observes the flag and fails. -struct ZeroSized { - queued: AtomicUsize, -} - -impl ZeroSized { - fn new() -> Self { - Self { - queued: AtomicUsize::new(0), - } - } - - /// Accounts for one published message, returning `false` once the queue is closed. - fn push(&self) -> bool { - let mut queued = self.queued.load(Ordering::Acquire); - loop { - if queued & CLOSED_BIT != 0 { - return false; - } - // The capacity limit keeps the count far below the closed flag bit. - match self.queued.compare_exchange_weak( - queued, - queued + 1, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => return true, - Err(actual) => queued = actual, - } - } - } - - /// Accounts for one consumed message, returning `false` when the queue was observed empty. - fn pop(&self) -> bool { - let queued = self.queued.load(Ordering::Acquire) & !CLOSED_BIT; - if queued == 0 { - return false; - } - // Only the consumer decrements, and producers can only add: the count observed above - // is a lower bound, so this cannot wrap. - self.queued.fetch_sub(1, Ordering::AcqRel); - true - } - - /// Stops publication and returns the queue length transferred to the drain. - fn close(&self) -> usize { - self.queued.fetch_or(CLOSED_BIT, Ordering::AcqRel) & !CLOSED_BIT - } -} - impl Slots { fn new(capacity: usize) -> Self { let slots = (0..capacity.next_power_of_two()) @@ -220,88 +241,58 @@ impl Slots { } } -impl Buffer { - pub fn new(capacity: usize) -> Self { - let storage = if size_of::() == 0 { - Storage::ZeroSized(ZeroSized::new()) - } else { - Storage::Slots(Box::new(Slots::new(capacity))) - }; - Self { storage } - } +/// Queue storage for zero-sized messages, which need no slots, positions, or per-slot flags. +/// Counting them separately supports the full channel capacity limit without per-slot metadata. +/// +/// The queue is entirely its length. The count packs a closed flag into its top bit so that +/// publication and close stay atomic: a publication that raced ahead of the flag is included +/// in the drained count, and every later one observes the flag and fails. +pub struct ZeroSized { + queued: AtomicUsize, +} - /// Writes and publishes one message. Closing may instead return the unsent value. - /// - /// # Safety - /// - /// Own one capacity permit before calling; release it only after a failed push or after - /// the consumer reads the published value. No user code runs between claim and publication. - pub unsafe fn push(&self, value: T) -> Result<(), T> { - match &self.storage { - Storage::Slots(slots) => { - let Ok(position) = slots.claim() else { - return Err(value); - }; - // SAFETY: The caller owns capacity and the ticket assigned this position. - unsafe { slots.publish(position, value) } - } - Storage::ZeroSized(zero_sized) => { - debug_assert_eq!(size_of::(), 0); - if zero_sized.push() { - mem::forget(value); - Ok(()) - } else { - Err(value) - } - } +impl ZeroSized { + fn new() -> Self { + Self { + queued: AtomicUsize::new(0), } } - /// Pending means a producer claimed the head but has not published it yet. - /// - /// # Safety - /// - /// Only the exclusive consumer may call this, using its persistent cursor. Release one - /// capacity permit after each successful pop, after the value has been read completely. - pub unsafe fn pop(&self, head: &mut usize) -> Poll> { - match &self.storage { - // SAFETY: The caller's guarantee forwards unchanged. - Storage::Slots(slots) => unsafe { slots.pop(head) }, - Storage::ZeroSized(zero_sized) => { - debug_assert_eq!(size_of::(), 0); - if zero_sized.pop() { - // SAFETY: A queued value proves that this ZST is inhabited and owns one value. - Poll::Ready(Some(unsafe { read_zero_sized() })) - } else { - Poll::Ready(None) - } + /// Accounts for one published message, returning `false` once the queue is closed. + fn push(&self) -> bool { + let mut queued = self.queued.load(Ordering::Acquire); + loop { + if queued & CLOSED_BIT != 0 { + return false; + } + // The capacity limit keeps the count far below the closed flag bit. + match self.queued.compare_exchange_weak( + queued, + queued + 1, + Ordering::AcqRel, + Ordering::Acquire, + ) { + Ok(_) => return true, + Err(actual) => queued = actual, } } } - /// Stops new claims and returns ownership of published values to a drain guard. - /// - /// # Safety - /// - /// Only the exclusive consumer may close the buffer, once, using its current cursor. - pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { - let remaining = match &self.storage { - Storage::Slots(slots) => slots.close(), - Storage::ZeroSized(zero_sized) => zero_sized.close(), - }; - Drain { - buffer: self, - position: head, - remaining, + /// Accounts for one consumed message, returning `false` when the queue was observed empty. + fn pop(&self) -> bool { + let queued = self.queued.load(Ordering::Acquire) & !CLOSED_BIT; + if queued == 0 { + return false; } + // Only the consumer decrements, and producers can only add: the count observed above + // is a lower bound, so this cannot wrap. + self.queued.fetch_sub(1, Ordering::AcqRel); + true } - #[cfg(test)] - fn slots(&self) -> &Slots { - match &self.storage { - Storage::Slots(slots) => slots, - Storage::ZeroSized(_) => unreachable!("zero-sized messages have no slots"), - } + /// Stops publication and returns the queue length transferred to the drain. + fn close(&self) -> usize { + self.queued.fetch_or(CLOSED_BIT, Ordering::AcqRel) & !CLOSED_BIT } } @@ -328,8 +319,8 @@ impl Iterator for Drain<'_, T> { let position = self.position; self.remaining -= 1; self.position = self.position.wrapping_add(1); - match &self.buffer.storage { - Storage::Slots(slots) => { + match self.buffer { + Buffer::Slots(slots) => { let slot = slots.slot(position); if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { // SAFETY: The drain won ownership of a published value. The cursor and @@ -340,7 +331,7 @@ impl Iterator for Drain<'_, T> { // and recover its value. The shared Arc keeps this allocation alive until then. } // SAFETY: Closing transferred this many initialized ZST values to the drain. - Storage::ZeroSized(_) => return Some(unsafe { read_zero_sized() }), + Buffer::ZeroSized(_) => return Some(unsafe { read_zero_sized() }), } } None @@ -370,6 +361,8 @@ impl Drop for Drain<'_, T> { #[cfg(test)] mod tests { + use std::future::Future; + use std::pin::pin; use std::sync::Arc; use std::sync::Barrier; use std::sync::atomic::AtomicUsize; @@ -393,10 +386,10 @@ mod tests { value: T, ) -> Result<(), T> { // SAFETY: The test claimed this position while holding the same capacity permit. - unsafe { tx.shared.buffer.slots().publish(position, value) }?; + unsafe { tx.shared().buffer.slots().publish(position, value) }?; // Publication owns the capacity now; forgetting skips the permit's release on drop. std::mem::forget(permit); - tx.shared.rx_waker.wake(); + tx.shared().rx_waker.wake(); Ok(()) } @@ -406,7 +399,7 @@ mod tests { for initial in [0, usize::MAX - 1] { let (tx, mut rx) = bounded(capacity); // Start an empty ring near ticket overflow instead of running usize::MAX sends. - tx.shared + tx.shared() .buffer .slots() .tail @@ -415,12 +408,12 @@ mod tests { let mut cx = Context::from_waker(Waker::noop()); for lap in 0..8 { let permit = tx.try_reserve().unwrap(); - let position = tx.shared.buffer.slots().claim().unwrap(); + let position = tx.shared().buffer.slots().claim().unwrap(); for offset in 1..capacity { tx.try_send(lap * capacity + offset).unwrap(); } // A full ring must differ from an empty one even with no head value ready. - assert!(rx.poll_recv(&mut cx).is_pending()); + assert!(pin!(rx.recv()).poll(&mut cx).is_pending()); publish_claimed(&tx, permit, position, lap * capacity).unwrap(); for offset in 0..capacity { assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); @@ -441,11 +434,16 @@ mod tests { drops: drops.clone(), _sender: tx.clone(), }; - let allocation = Arc::downgrade(&tx.shared); + let allocation = Arc::downgrade(tx.shared()); // Pause after claim's open check, then resume its atomic ticket allocation after close. - assert!(!&tx.shared.buffer.slots().closed.load(Ordering::Acquire)); + assert!(!tx.shared().buffer.slots().closed.load(Ordering::Acquire)); drop(rx); - let position = tx.shared.buffer.slots().tail.fetch_add(1, Ordering::AcqRel); + let position = tx + .shared() + .buffer + .slots() + .tail + .fetch_add(1, Ordering::AcqRel); let unsent = publish_claimed(&tx, permit, position, value).unwrap_err(); assert_eq!(unsent.bytes, [7; 1024]); drop(unsent); @@ -472,7 +470,7 @@ mod tests { #[test] fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { let (tx, rx) = bounded(2); - let allocation = Arc::downgrade(&tx.shared); + let allocation = Arc::downgrade(tx.shared()); let drops = Arc::new(AtomicUsize::new(0)); let paused = Barrier::new(2); let (resume_tx, resume_rx) = std::sync::mpsc::channel(); @@ -484,7 +482,7 @@ mod tests { let paused = &paused; let publisher = scope.spawn(move || { let permit = sender.try_reserve().unwrap(); - let position = sender.shared.buffer.slots().claim().unwrap(); + let position = sender.shared().buffer.slots().claim().unwrap(); let value = Payload { bytes: [1; 1024], drops: drops.clone(), @@ -529,7 +527,7 @@ mod tests { fn publication_racing_with_close_drops_every_payload_once() { for _ in 0..if cfg!(miri) { 8 } else { 128 } { let (tx, rx) = bounded(3); - let allocation = Arc::downgrade(&tx.shared); + let allocation = Arc::downgrade(tx.shared()); let drops = Arc::new(AtomicUsize::new(0)); let start = Barrier::new(4); thread::scope(|scope| { @@ -578,7 +576,7 @@ mod tests { .unwrap(); }); assert_eq!( - rx.poll_recv(&mut Context::from_waker(Waker::noop())), + pin!(rx.recv()).poll(&mut Context::from_waker(Waker::noop())), Poll::Ready(Ok([42; 1024])) ); assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index b2dcf367..59e7298c 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -23,10 +23,6 @@ use std::sync::atomic::AtomicUsize; use self::buffer::Buffer; use self::semaphore::Semaphore; -use super::RecvError; -use super::SendError; -use super::TryRecvError; -use super::TrySendError; use crate::internal::atomic_waker::AtomicWaker; use crate::internal::cache_padded::CachePadded; @@ -35,6 +31,10 @@ mod receiver; mod semaphore; mod sender; +pub use self::receiver::BoundedReceiver; +pub use self::sender::BoundedSender; +pub use self::sender::Permit; + /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// /// [`BoundedSender::send`] waits for capacity when the buffer is full. Receiving a message releases @@ -71,38 +71,10 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { rx_waker: AtomicWaker::new(), buffer: Buffer::new(buffer), }); - let sender = BoundedSender { - shared: shared.clone(), - }; - let receiver = BoundedReceiver { shared, head: 0 }; - (sender, receiver) -} - -/// The sending endpoint of a bounded mpsc channel. -/// -/// Instances are created by the [`bounded`] function. -pub struct BoundedSender { - shared: Arc>, -} - -/// The receiving endpoint of a bounded mpsc channel. -/// -/// Instances are created by the [`bounded`] function. Dropping the receiver discards queued values. -/// The backing allocation remains alive until all endpoints are dropped, so a concurrent sender -/// can safely finish returning an unsent value. -pub struct BoundedReceiver { - shared: Arc>, - head: usize, -} - -/// Capacity reserved for one message on a bounded channel. -/// -/// Created by [`BoundedSender::reserve`] or [`BoundedSender::try_reserve`]. Holding a permit -/// reduces available capacity but does not prevent other messages from being received. Dropping -/// it without sending releases capacity and notifies a waiting sender. -#[must_use = "dropping the permit releases its reserved capacity"] -pub struct Permit<'a, T> { - sender: Option<&'a BoundedSender>, + ( + BoundedSender::new(shared.clone()), + BoundedReceiver::new(shared), + ) } struct Shared { diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs index 88db33a1..377bfa13 100644 --- a/asyncband/src/mpsc/bounded/receiver.rs +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -17,14 +17,26 @@ use std::fmt; use std::future::poll_fn; +use std::sync::Arc; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; -use super::BoundedReceiver; -use super::RecvError; -use super::TryRecvError; +use super::Shared; use crate::internal::wake_all; +use crate::mpsc::RecvError; +use crate::mpsc::TryRecvError; + +/// The receiving endpoint of a bounded mpsc channel. +/// +/// Instances are created by the [`bounded`](crate::mpsc::bounded) function. Dropping the receiver +/// discards queued values. +/// The backing allocation remains alive until all endpoints are dropped, so a concurrent sender +/// can safely finish returning an unsent value. +pub struct BoundedReceiver { + shared: Arc>, + head: usize, +} impl fmt::Debug for BoundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { @@ -46,6 +58,10 @@ impl Drop for BoundedReceiver { } impl BoundedReceiver { + pub(super) fn new(shared: Arc>) -> Self { + Self { shared, head: 0 } + } + /// Attempts to receive the next queued value without waiting for a new message. /// /// Receiving a value frees one buffer slot. An empty channel returns [`TryRecvError::Empty`] @@ -90,31 +106,6 @@ impl BoundedReceiver { } } - /// One attempt to take the head value: a message, an empty-or-disconnected classification, - /// or `Pending` while a claimed head waits for its publication. - fn pull(&mut self) -> Poll> { - let mut disconnected = false; - loop { - // SAFETY: Only this receiver owns head. Capacity is released after the buffer - // finishes reading and advances the cursor, so no producer can overwrite the value. - match unsafe { self.shared.buffer.pop(&mut self.head) } { - Poll::Ready(Some(value)) => { - self.shared.tx_permits.release(); - return Poll::Ready(Ok(value)); - } - Poll::Ready(None) if disconnected => { - return Poll::Ready(Err(TryRecvError::Disconnected)); - } - Poll::Ready(None) if self.shared.senders.load(Ordering::Acquire) == 0 => { - // Acquire the last sender's completed publications before checking again. - disconnected = true; - } - Poll::Ready(None) => return Poll::Ready(Err(TryRecvError::Empty)), - Poll::Pending => return Poll::Pending, - } - } - } - /// Waits for and receives the next value, freeing one buffer slot. /// /// If no value is queued, this method waits until a sender adds one or the last sender is @@ -148,7 +139,32 @@ impl BoundedReceiver { poll_fn(|cx| self.poll_recv(cx)).await } - pub(crate) fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { + /// One attempt to take the head value: a message, an empty-or-disconnected classification, + /// or `Pending` while a claimed head waits for its publication. + fn pull(&mut self) -> Poll> { + let mut disconnected = false; + loop { + // SAFETY: Only this receiver owns head. Capacity is released after the buffer + // finishes reading and advances the cursor, so no producer can overwrite the value. + match unsafe { self.shared.buffer.pop(&mut self.head) } { + Poll::Ready(Some(value)) => { + self.shared.tx_permits.release(); + return Poll::Ready(Ok(value)); + } + Poll::Ready(None) if disconnected => { + return Poll::Ready(Err(TryRecvError::Disconnected)); + } + Poll::Ready(None) if self.shared.senders.load(Ordering::Acquire) == 0 => { + // Acquire the last sender's completed publications before checking again. + disconnected = true; + } + Poll::Ready(None) => return Poll::Ready(Err(TryRecvError::Empty)), + Poll::Pending => return Poll::Pending, + } + } + } + + fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { for registered in [false, true] { match self.pull() { Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), @@ -164,9 +180,8 @@ impl BoundedReceiver { } Poll::Pending } - #[cfg(test)] - pub(crate) fn set_head(&mut self, head: usize) { + pub(super) fn set_head(&mut self, head: usize) { self.head = head; } } diff --git a/asyncband/src/mpsc/bounded/semaphore.rs b/asyncband/src/mpsc/bounded/semaphore.rs index e0ec4925..8e60ed39 100644 --- a/asyncband/src/mpsc/bounded/semaphore.rs +++ b/asyncband/src/mpsc/bounded/semaphore.rs @@ -31,9 +31,9 @@ //! and grants bypass the counter. //! //! With neither sentinel installed, acquire and release are single lock-free operations on `state`. -//! Wait-queue mutations always hold the queue lock; a registration installs `WAITING` before its -//! final capacity recheck, which switches any racing release to the locked path and strands no -//! permit without a wake. +//! Wait-queue mutations always hold the queue lock. A registration must install or observe +//! `WAITING` before joining the queue, so every subsequent release takes the locked path. If a +//! release wins that transition, acquisition retries instead of registering against a plain count. use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -41,12 +41,12 @@ use std::task::Context; use std::task::Poll; use std::task::Waker; -use super::SendError; -use super::TrySendError; use crate::internal::mutex::Mutex; use crate::internal::waitlist::WaitList; use crate::internal::waitlist::WaiterId; use crate::internal::waker_batch::WakerBatch; +use crate::mpsc::SendError; +use crate::mpsc::TrySendError; pub struct Semaphore { state: AtomicUsize, @@ -56,7 +56,7 @@ pub struct Semaphore { const CLOSED: usize = usize::MAX; const WAITING: usize = usize::MAX - 1; -pub struct Waiter { +struct Waiter { granted: bool, waker: Option, } @@ -69,7 +69,7 @@ impl Semaphore { } } - pub fn try_acquire(&self) -> Result<(), TrySendError<()>> { + pub fn try_acquire(&self) -> Result, TrySendError<()>> { let mut state = self.state.load(Ordering::Acquire); loop { if state == CLOSED { @@ -84,7 +84,7 @@ impl Semaphore { Ordering::Acquire, Ordering::Acquire, ) { - Ok(_) => return Ok(()), + Ok(_) => return Ok(Capacity { semaphore: self }), Err(actual) => state = actual, } } @@ -103,12 +103,14 @@ impl Semaphore { self.state.load(Ordering::Acquire) == CLOSED } - // Installs WAITING over an exhausted counter. A permit that arrived first wins the compare - // exchange, and the caller's recheck under the queue lock picks it up instead. - fn set_waiting(&self) { - let _ = self - .state - .compare_exchange(0, WAITING, Ordering::AcqRel, Ordering::Acquire); + // Called with the queue locked. A failed installation requires retrying acquisition: a + // racing sender may consume the returned capacity before a separate recheck can see it. + fn set_waiting(&self) -> bool { + matches!( + self.state + .compare_exchange(0, WAITING, Ordering::AcqRel, Ordering::Acquire), + Ok(_) | Err(WAITING) + ) } // Removes WAITING, keeping whatever count a racing grant restoration left behind. @@ -142,7 +144,7 @@ impl Semaphore { } } - pub fn release_locked(&self, waiters: &mut WaitList) -> Option { + fn release_locked(&self, waiters: &mut WaitList) -> Option { if self.is_closed() { return None; } @@ -183,6 +185,24 @@ impl Semaphore { } } +/// Owns one capacity unit until publication transfers it to a queued message. +#[must_use = "dropping the guard releases its capacity"] +pub struct Capacity<'a> { + semaphore: &'a Semaphore, +} + +impl Capacity<'_> { + pub fn forget(self) { + std::mem::forget(self); + } +} + +impl Drop for Capacity<'_> { + fn drop(&mut self) { + self.semaphore.release(); + } +} + /// An in-flight [`Semaphore::acquire`] operation. /// /// Dropping the operation removes its wait-queue registration; a capacity grant that already @@ -192,14 +212,14 @@ pub struct Acquire<'a> { waiter: Option, } -impl Acquire<'_> { - pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll>> { +impl<'a> Acquire<'a> { + pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { let semaphore = self.semaphore; let mut cloned_waker = None; let result = loop { if self.waiter.is_none() { match semaphore.try_acquire() { - Ok(()) => break Ok(()), + Ok(capacity) => break Ok(capacity), Err(TrySendError::Disconnected(())) => break Err(SendError::new(())), Err(TrySendError::Full(())) => {} } @@ -214,9 +234,10 @@ impl Acquire<'_> { if waiter.granted { let waiter = waiters.remove_unlinked_waiter(index); self.waiter = None; + let capacity = Capacity { semaphore }; drop(waiters); drop(waiter); - break Ok(()); + break Ok(capacity); } if waiter .waker @@ -232,15 +253,9 @@ impl Acquire<'_> { return Poll::Pending; } } else { - // Install WAITING before the final capacity recheck: if a permit arrived - // first, the installation loses the compare exchange and the recheck picks - // the permit up; otherwise a racing release switches to the locked path, so - // no permit can be stranded without a wake. Waiting senders already in the - // queue take priority over this recheck. - semaphore.set_waiting(); - if waiters.is_empty() && semaphore.try_acquire().is_ok() { - semaphore.clear_waiting(); - break Ok(()); + if !semaphore.set_waiting() { + drop(waiters); + continue; } if let Some(waker) = cloned_waker.take() { self.waiter = Some(waiters.push_back(Waiter { @@ -254,7 +269,8 @@ impl Acquire<'_> { // Clone outside the lock, then recheck capacity and closure before registering. cloned_waker = Some(cx.waker().clone()); }; - // The permit owns capacity before an unused cloned waker can panic. + // A successful result already owns a guard, so a panicking waker destructor returns + // capacity even before the caller has constructed its public permit. drop(cloned_waker); Poll::Ready(result) } diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index 7cbe8f77..1e610c87 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -17,12 +17,20 @@ use std::fmt; use std::future::poll_fn; +use std::sync::Arc; use std::sync::atomic::Ordering; -use super::BoundedSender; -use super::Permit; -use super::SendError; -use super::TrySendError; +use super::Shared; +use super::semaphore::Capacity; +use crate::mpsc::SendError; +use crate::mpsc::TrySendError; + +/// The sending endpoint of a bounded mpsc channel. +/// +/// Instances are created by the [`bounded`](crate::mpsc::bounded) function. +pub struct BoundedSender { + shared: Arc>, +} impl Clone for BoundedSender { fn clone(&self) -> Self { @@ -48,6 +56,10 @@ impl Drop for BoundedSender { } impl BoundedSender { + pub(super) fn new(shared: Arc>) -> Self { + Self { shared } + } + /// Sends a message, waiting until the channel has capacity when necessary. /// /// If the receiver has been dropped, the returned error contains `value`. @@ -97,8 +109,11 @@ impl BoundedSender { /// ``` pub async fn reserve(&self) -> Result, SendError<()>> { let mut acquire = self.shared.tx_permits.acquire(); - poll_fn(|cx| acquire.poll(cx)).await?; - Ok(Permit { sender: Some(self) }) + let capacity = poll_fn(|cx| acquire.poll(cx)).await?; + Ok(Permit { + shared: &self.shared, + capacity, + }) } /// Reserves capacity for one message without waiting. @@ -106,8 +121,11 @@ impl BoundedSender { /// Returns [`TrySendError::Full`] if queued messages and outstanding permits occupy the /// buffer, or [`TrySendError::Disconnected`] if the receiver has been dropped. pub fn try_reserve(&self) -> Result, TrySendError<()>> { - self.shared.tx_permits.try_acquire()?; - Ok(Permit { sender: Some(self) }) + let capacity = self.shared.tx_permits.try_acquire()?; + Ok(Permit { + shared: &self.shared, + capacity, + }) } /// Attempts to send a message without waiting for capacity. @@ -139,6 +157,22 @@ impl BoundedSender { Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), } } + + #[cfg(test)] + pub(super) fn shared(&self) -> &Arc> { + &self.shared + } +} + +/// Capacity reserved for one message on a bounded channel. +/// +/// Created by [`BoundedSender::reserve`] or [`BoundedSender::try_reserve`]. Holding a permit +/// reduces available capacity but does not prevent other messages from being received. Dropping +/// it without sending releases capacity and notifies a waiting sender. +#[must_use = "dropping the permit releases its reserved capacity"] +pub struct Permit<'a, T> { + shared: &'a Shared, + capacity: Capacity<'a>, } impl fmt::Debug for Permit<'_, T> { @@ -151,22 +185,14 @@ impl Permit<'_, T> { /// Publishes a message using this permit, without waiting for capacity. /// /// If the receiver has been dropped, the returned error contains the unsent value. - pub fn send(mut self, value: T) -> Result<(), SendError> { - let shared = &self.sender.unwrap().shared; + pub fn send(self, value: T) -> Result<(), SendError> { + let Self { shared, capacity } = self; // SAFETY: This permit owns one capacity unit. Claiming a slot and writing it is a // synchronous operation with no user callbacks or await points between the two. unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; // Publication owns the capacity before a wake callback can panic. - self.sender = None; + capacity.forget(); shared.rx_waker.wake(); Ok(()) } } - -impl Drop for Permit<'_, T> { - fn drop(&mut self) { - if let Some(sender) = self.sender { - sender.shared.tx_permits.release(); - } - } -} diff --git a/tests-integration/tests/mpsc_test/callbacks.rs b/tests-integration/tests/mpsc_test/callbacks.rs index 3f1bed5e..a069f001 100644 --- a/tests-integration/tests/mpsc_test/callbacks.rs +++ b/tests-integration/tests/mpsc_test/callbacks.rs @@ -104,6 +104,65 @@ fn bounded_send_rechecks_capacity_freed_by_waker_clone() { }); } +#[cfg(panic = "unwind")] +#[test] +fn bounded_send_returns_capacity_when_an_unused_waker_panics_on_drop() { + use std::task::RawWaker; + use std::task::RawWakerVTable; + + struct Callbacks { + receiver: Mutex>, + drop_panics: AtomicBool, + } + + unsafe fn clone(data: *const ()) -> RawWaker { + let pointer = data.cast::(); + // SAFETY: The input waker owns a live Arc. The returned clone gains its own reference. + unsafe { + assert_eq!((*pointer).receiver.lock().unwrap().try_recv(), Ok(1)); + Arc::increment_strong_count(pointer); + } + RawWaker::new(data, &VTABLE) + } + + unsafe fn release(data: *const ()) { + // SAFETY: Consumes this waker's Arc reference, including if the callback unwinds. + let callbacks = unsafe { Arc::from_raw(data.cast::()) }; + assert!( + !callbacks.drop_panics.swap(false, Ordering::Relaxed), + "unused cloned waker panicked on drop" + ); + } + + // A raw vtable is needed to run callbacks for cloning and dropping each waker reference. + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, release, |_| {}, release); + + let (tx, rx) = mpsc::bounded(1); + tx.try_send(1).unwrap(); + let callbacks = Arc::new(Callbacks { + receiver: Mutex::new(rx), + drop_panics: AtomicBool::new(true), + }); + let data = Arc::into_raw(callbacks.clone()).cast(); + // SAFETY: Every waker owns an Arc reference. All callbacks preserve ownership and use only + // synchronized state; wake_by_ref does not touch the reference count. + let waker = unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }; + + let mut send = Box::pin(tx.send(2)); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + poll_with(send.as_mut(), &waker) + })) + .is_err() + ); + drop(send); + let mut receiver = callbacks.receiver.lock().unwrap(); + assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); + tx.try_send(3) + .expect("unwinding must return acquired capacity"); + assert_eq!(receiver.try_recv(), Ok(3)); +} + #[test] fn receive_rechecks_messages_sent_by_waker_clone() { assert_completes_without_deadlock(|| { diff --git a/tests-integration/tests/mpsc_test/concurrency.rs b/tests-integration/tests/mpsc_test/concurrency.rs index ce097fc4..f62339ba 100644 --- a/tests-integration/tests/mpsc_test/concurrency.rs +++ b/tests-integration/tests/mpsc_test/concurrency.rs @@ -63,6 +63,47 @@ fn bounded_receive_racing_with_send_registration_cannot_lose_wakeup() { } } +#[test] +fn bounded_competing_reservation_cannot_strand_a_waiting_send() { + for _ in 0..if cfg!(miri) { 32 } else { 512 } { + let (tx, mut rx) = mpsc::bounded(1); + tx.try_send(1).unwrap(); + let start = Barrier::new(3); + let received = Barrier::new(2); + let (waker, notified) = WakeCounter::new(); + let mut send = Box::pin(tx.send(2)); + + let poll = thread::scope(|scope| { + let receive = scope.spawn(|| { + start.wait(); + assert_eq!(rx.try_recv(), Ok(1)); + received.wait(); + }); + let competitor = scope.spawn(|| { + start.wait(); + received.wait(); + // Try to consume the returned capacity while the other sender registers. + tx.try_reserve().ok() + }); + start.wait(); + let poll = poll_with(send.as_mut(), &waker); + receive.join().unwrap(); + // Keep any competing reservation until registration has completed. Its release + // must notify a pending sender even if it won capacity during that registration. + drop(competitor.join().unwrap()); + poll + }); + + if poll.is_pending() { + assert!(notified.count() > 0, "available capacity stranded a sender"); + assert_eq!(poll_once(send.as_mut()), Poll::Ready(Ok(()))); + } else { + assert_eq!(poll, Poll::Ready(Ok(()))); + } + assert_eq!(rx.try_recv(), Ok(2)); + } +} + #[test] fn bounded_try_recv_does_not_report_empty_after_completed_sends() { const PRODUCERS: usize = 4; diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index a7e60077..9d9647a5 100644 --- a/tests-integration/tests/mpsc_test/main.rs +++ b/tests-integration/tests/mpsc_test/main.rs @@ -176,7 +176,7 @@ fn receives_wake_for_messages_and_the_last_sender_drop() { } #[test] -#[should_panic(expected = "mpsc bounded channel requires buffer > 0")] +#[should_panic(expected = "must be nonzero")] fn bounded_rejects_zero_capacity() { let _ = mpsc::bounded::(0); } From c5b75eeff0600edf92a39f3121013d02a13a1a85 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 16:16:23 +0800 Subject: [PATCH 29/34] refactor(mpsc): unify bounded message storage --- CHANGELOG.md | 4 +- asyncband/src/mpsc/bounded/buffer.rs | 273 ++++-------------- asyncband/src/mpsc/bounded/mod.rs | 9 +- asyncband/src/mpsc/bounded/semaphore.rs | 6 - asyncband/src/mpsc/bounded/sender.rs | 10 +- tests-integration/tests/mpsc_test/main.rs | 7 - .../tests/mpsc_test/reservation.rs | 24 +- 7 files changed, 91 insertions(+), 242 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3b133146..74e099ec 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,10 +16,10 @@ All notable changes to this project will be documented in this file. ### Improvements -* Reject bounded MPSC capacities above `usize::MAX >> 1` up front with an explicit panic message instead of an opaque arithmetic overflow; zero-sized messages need no slot storage and remain limited only by the permit counter. +* Reject bounded MPSC capacities above `usize::MAX >> 1` up front with an explicit panic message instead of an opaque arithmetic overflow. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. -* Improve bounded MPSC throughput: acquiring and releasing capacity no longer takes an internal lock while no sender is waiting, and zero-sized messages no longer take the buffer lock. +* Allow bounded MPSC producers to publish messages concurrently; acquiring and releasing capacity no longer takes an internal lock while no sender is waiting. ## v0.7.2 diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index 6a484337..d1698803 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -28,9 +28,7 @@ //! allocation, so a publisher's slot stays alive even when receiver drop closes it concurrently. use std::cell::UnsafeCell; -use std::mem; use std::mem::MaybeUninit; -use std::ptr::NonNull; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicU8; use std::sync::atomic::AtomicUsize; @@ -42,21 +40,25 @@ use crate::internal::cache_padded::CachePadded; const EMPTY: u8 = 0; const READY: u8 = 1; const CLOSED: u8 = 2; -const CLOSED_BIT: usize = 1 << (usize::BITS - 1); -/// Zero-sized messages need only a count. Boxing the slotted variant keeps its cursor and -/// per-slot storage out of the zero-sized queue's allocation. -pub enum Buffer { - Slots(Box>), - ZeroSized(ZeroSized), +pub struct Buffer { + slots: Box<[Slot]>, + tail: CachePadded, + closed: AtomicBool, } impl Buffer { pub fn new(capacity: usize) -> Self { - if size_of::() == 0 { - Self::ZeroSized(ZeroSized::new()) - } else { - Self::Slots(Box::new(Slots::new(capacity))) + let slots = (0..capacity.next_power_of_two()) + .map(|_| Slot { + state: AtomicU8::new(EMPTY), + value: UnsafeCell::new(MaybeUninit::uninit()), + }) + .collect(); + Self { + slots, + tail: CachePadded::new(AtomicUsize::new(0)), + closed: AtomicBool::new(false), } } @@ -67,24 +69,11 @@ impl Buffer { /// Own one capacity permit before calling; release it only after a failed push or after /// the consumer reads the published value. No user code runs between claim and publication. pub unsafe fn push(&self, value: T) -> Result<(), T> { - match self { - Self::Slots(slots) => { - let Ok(position) = slots.claim() else { - return Err(value); - }; - // SAFETY: The caller owns capacity and the ticket assigned this position. - unsafe { slots.publish(position, value) } - } - Self::ZeroSized(zero_sized) => { - debug_assert_eq!(size_of::(), 0); - if zero_sized.push() { - mem::forget(value); - Ok(()) - } else { - Err(value) - } - } - } + let Ok(position) = self.claim() else { + return Err(value); + }; + // SAFETY: The caller owns capacity and the ticket assigned this position. + unsafe { self.publish(position, value) } } /// Pending means a producer claimed the head but has not published it yet. @@ -94,18 +83,18 @@ impl Buffer { /// Only the exclusive consumer may call this, using its persistent cursor. Release one /// capacity permit after each successful pop, after the value has been read completely. pub unsafe fn pop(&self, head: &mut usize) -> Poll> { - match self { - // SAFETY: The caller's guarantee forwards unchanged. - Self::Slots(slots) => unsafe { slots.pop(head) }, - Self::ZeroSized(zero_sized) => { - debug_assert_eq!(size_of::(), 0); - if zero_sized.pop() { - // SAFETY: A queued value proves that this ZST is inhabited and owns one value. - Poll::Ready(Some(unsafe { read_zero_sized() })) - } else { - Poll::Ready(None) - } - } + let slot = self.slot(*head); + if slot.state.load(Ordering::Acquire) == READY { + // SAFETY: Publication initialized the value, and only this consumer can read it. + // Capacity is still held until this method has returned the value to its caller. + let value = unsafe { (*slot.value.get()).assume_init_read() }; + slot.state.store(EMPTY, Ordering::Release); + *head = head.wrapping_add(1); + Poll::Ready(Some(value)) + } else if self.tail.load(Ordering::Acquire) == *head { + Poll::Ready(None) + } else { + Poll::Pending } } @@ -115,59 +104,13 @@ impl Buffer { /// /// Only the exclusive consumer may close the buffer, once, using its current cursor. pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { - let remaining = match self { - Self::Slots(slots) => slots.close(), - Self::ZeroSized(zero_sized) => zero_sized.close(), - }; + self.closed.store(true, Ordering::Release); + // Cover every physical slot: a producer may have passed the open check but not yet + // claimed its ticket. Such a late claim must also find a CLOSED slot. Drain { buffer: self, position: head, - remaining, - } - } - - #[cfg(test)] - fn slots(&self) -> &Slots { - match self { - Self::Slots(slots) => slots, - Self::ZeroSized(_) => unreachable!("zero-sized messages have no slots"), - } - } -} - -pub struct Slots { - slots: Box<[Slot]>, - tail: CachePadded, - closed: AtomicBool, -} - -struct Slot { - state: AtomicU8, - value: UnsafeCell>, -} - -// SAFETY: Capacity and the tail ticket give a producer exclusive ownership of an empty slot. -// Release publication transfers its value to the exclusive consumer. Closing an unpublished -// slot leaves its value with the producer; closing a READY slot transfers it to the drain. -unsafe impl Sync for Slot {} - -// No reference to a stored value escapes. Every value is removed from the slot's ownership -// before running a callback or destructor that might panic. -impl std::panic::UnwindSafe for Slot {} -impl std::panic::RefUnwindSafe for Slot {} - -impl Slots { - fn new(capacity: usize) -> Self { - let slots = (0..capacity.next_power_of_two()) - .map(|_| Slot { - state: AtomicU8::new(EMPTY), - value: UnsafeCell::new(MaybeUninit::uninit()), - }) - .collect(); - Self { - slots, - tail: CachePadded::new(AtomicUsize::new(0)), - closed: AtomicBool::new(false), + remaining: self.slots.len(), } } @@ -210,100 +153,22 @@ impl Slots { } } } - - /// Pending means a producer claimed the head but has not published it yet. - /// - /// # Safety - /// - /// Only the exclusive consumer may call this, using its persistent cursor. - unsafe fn pop(&self, head: &mut usize) -> Poll> { - let slot = self.slot(*head); - if slot.state.load(Ordering::Acquire) == READY { - // SAFETY: Publication initialized the value, and only this consumer can read it. - // Capacity is still held until this method has returned the value to its caller. - let value = unsafe { (*slot.value.get()).assume_init_read() }; - slot.state.store(EMPTY, Ordering::Release); - *head = head.wrapping_add(1); - Poll::Ready(Some(value)) - } else if self.tail.load(Ordering::Acquire) == *head { - Poll::Ready(None) - } else { - Poll::Pending - } - } - - /// Stops new claims and returns the physical slot count the drain must cover: a producer - /// may have passed the open check but not obtained its ticket yet, and such a late claim - /// must also find a CLOSED slot. - fn close(&self) -> usize { - self.closed.store(true, Ordering::Release); - self.slots.len() - } } -/// Queue storage for zero-sized messages, which need no slots, positions, or per-slot flags. -/// Counting them separately supports the full channel capacity limit without per-slot metadata. -/// -/// The queue is entirely its length. The count packs a closed flag into its top bit so that -/// publication and close stay atomic: a publication that raced ahead of the flag is included -/// in the drained count, and every later one observes the flag and fails. -pub struct ZeroSized { - queued: AtomicUsize, +struct Slot { + state: AtomicU8, + value: UnsafeCell>, } -impl ZeroSized { - fn new() -> Self { - Self { - queued: AtomicUsize::new(0), - } - } - - /// Accounts for one published message, returning `false` once the queue is closed. - fn push(&self) -> bool { - let mut queued = self.queued.load(Ordering::Acquire); - loop { - if queued & CLOSED_BIT != 0 { - return false; - } - // The capacity limit keeps the count far below the closed flag bit. - match self.queued.compare_exchange_weak( - queued, - queued + 1, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => return true, - Err(actual) => queued = actual, - } - } - } - - /// Accounts for one consumed message, returning `false` when the queue was observed empty. - fn pop(&self) -> bool { - let queued = self.queued.load(Ordering::Acquire) & !CLOSED_BIT; - if queued == 0 { - return false; - } - // Only the consumer decrements, and producers can only add: the count observed above - // is a lower bound, so this cannot wrap. - self.queued.fetch_sub(1, Ordering::AcqRel); - true - } - - /// Stops publication and returns the queue length transferred to the drain. - fn close(&self) -> usize { - self.queued.fetch_or(CLOSED_BIT, Ordering::AcqRel) & !CLOSED_BIT - } -} +// SAFETY: Capacity and the tail ticket give a producer exclusive ownership of an empty slot. +// Release publication transfers its value to the exclusive consumer. Closing an unpublished +// slot leaves its value with the producer; closing a READY slot transfers it to the drain. +unsafe impl Sync for Slot {} -/// # Safety -/// -/// The caller must own an initialized, inhabited ZST value. -unsafe fn read_zero_sized() -> T { - // SAFETY: Reading it accesses no bytes; dangling supplies a non-null, correctly aligned - // pointer, as in a ZST Vec. - unsafe { NonNull::::dangling().as_ptr().read() } -} +// No reference to a stored value escapes. Every value is removed from the slot's ownership +// before running a callback or destructor that might panic. +impl std::panic::UnwindSafe for Slot {} +impl std::panic::RefUnwindSafe for Slot {} pub struct Drain<'a, T> { buffer: &'a Buffer, @@ -319,20 +184,14 @@ impl Iterator for Drain<'_, T> { let position = self.position; self.remaining -= 1; self.position = self.position.wrapping_add(1); - match self.buffer { - Buffer::Slots(slots) => { - let slot = slots.slot(position); - if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { - // SAFETY: The drain won ownership of a published value. The cursor and - // state already advanced, so a panicking destructor cannot read twice. - return Some(unsafe { (*slot.value.get()).assume_init_read() }); - } - // An unpublished slot stays owned by its producer, which will observe CLOSED - // and recover its value. The shared Arc keeps this allocation alive until then. - } - // SAFETY: Closing transferred this many initialized ZST values to the drain. - Buffer::ZeroSized(_) => return Some(unsafe { read_zero_sized() }), + let slot = self.buffer.slot(position); + if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { + // SAFETY: The drain won ownership of a published value. The cursor and + // state already advanced, so a panicking destructor cannot read twice. + return Some(unsafe { (*slot.value.get()).assume_init_read() }); } + // An unpublished slot stays owned by its producer, which will observe CLOSED + // and recover its value. The shared Arc keeps this allocation alive until then. } None } @@ -362,6 +221,7 @@ impl Drop for Drain<'_, T> { #[cfg(test)] mod tests { use std::future::Future; + use std::mem; use std::pin::pin; use std::sync::Arc; use std::sync::Barrier; @@ -386,9 +246,9 @@ mod tests { value: T, ) -> Result<(), T> { // SAFETY: The test claimed this position while holding the same capacity permit. - unsafe { tx.shared().buffer.slots().publish(position, value) }?; - // Publication owns the capacity now; forgetting skips the permit's release on drop. - std::mem::forget(permit); + unsafe { tx.shared().buffer.publish(position, value) }?; + // Publication owns the capacity now, so the permit must not release it on drop. + mem::forget(permit); tx.shared().rx_waker.wake(); Ok(()) } @@ -399,16 +259,12 @@ mod tests { for initial in [0, usize::MAX - 1] { let (tx, mut rx) = bounded(capacity); // Start an empty ring near ticket overflow instead of running usize::MAX sends. - tx.shared() - .buffer - .slots() - .tail - .store(initial, Ordering::Relaxed); + tx.shared().buffer.tail.store(initial, Ordering::Relaxed); rx.set_head(initial); let mut cx = Context::from_waker(Waker::noop()); for lap in 0..8 { let permit = tx.try_reserve().unwrap(); - let position = tx.shared().buffer.slots().claim().unwrap(); + let position = tx.shared().buffer.claim().unwrap(); for offset in 1..capacity { tx.try_send(lap * capacity + offset).unwrap(); } @@ -436,14 +292,9 @@ mod tests { }; let allocation = Arc::downgrade(tx.shared()); // Pause after claim's open check, then resume its atomic ticket allocation after close. - assert!(!tx.shared().buffer.slots().closed.load(Ordering::Acquire)); + assert!(!tx.shared().buffer.closed.load(Ordering::Acquire)); drop(rx); - let position = tx - .shared() - .buffer - .slots() - .tail - .fetch_add(1, Ordering::AcqRel); + let position = tx.shared().buffer.tail.fetch_add(1, Ordering::AcqRel); let unsent = publish_claimed(&tx, permit, position, value).unwrap_err(); assert_eq!(unsent.bytes, [7; 1024]); drop(unsent); @@ -482,7 +333,7 @@ mod tests { let paused = &paused; let publisher = scope.spawn(move || { let permit = sender.try_reserve().unwrap(); - let position = sender.shared().buffer.slots().claim().unwrap(); + let position = sender.shared().buffer.claim().unwrap(); let value = Payload { bytes: [1; 1024], drops: drops.clone(), diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 59e7298c..167e2b22 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -41,8 +41,8 @@ pub use self::sender::Permit; /// one slot for a waiting sender. Capacity is granted in the order that pending sends and /// reservations enter the wait queue; new senders cannot take an already granted slot. /// -/// Storage for nonzero-sized messages is preallocated and rounded up to a power of two; the -/// channel's capacity remains exactly `buffer`. Zero-sized messages need no per-slot storage. +/// Message slots are preallocated and rounded up to a power of two; the channel's capacity +/// remains exactly `buffer`. Every slot needs state metadata, including for zero-sized messages. /// /// # Panics /// @@ -51,9 +51,8 @@ pub use self::sender::Permit; pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { /// The largest capacity accepted by [`bounded`]. /// - /// The shared permit counter reserves two sentinel values above the usable range, and the - /// zero-sized queue length packs a closed flag into its top bit. This bound also keeps the - /// rounded-up slot storage from overflowing a power of two. + /// The shared permit counter reserves two sentinel values above the usable range. This + /// bound also keeps the rounded-up slot storage from overflowing a power of two. const MAX_CAPACITY: usize = usize::MAX >> 1; assert!( diff --git a/asyncband/src/mpsc/bounded/semaphore.rs b/asyncband/src/mpsc/bounded/semaphore.rs index 8e60ed39..e7f165e6 100644 --- a/asyncband/src/mpsc/bounded/semaphore.rs +++ b/asyncband/src/mpsc/bounded/semaphore.rs @@ -191,12 +191,6 @@ pub struct Capacity<'a> { semaphore: &'a Semaphore, } -impl Capacity<'_> { - pub fn forget(self) { - std::mem::forget(self); - } -} - impl Drop for Capacity<'_> { fn drop(&mut self) { self.semaphore.release(); diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index 1e610c87..d72e403c 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -17,6 +17,7 @@ use std::fmt; use std::future::poll_fn; +use std::mem; use std::sync::Arc; use std::sync::atomic::Ordering; @@ -186,13 +187,12 @@ impl Permit<'_, T> { /// /// If the receiver has been dropped, the returned error contains the unsent value. pub fn send(self, value: T) -> Result<(), SendError> { - let Self { shared, capacity } = self; // SAFETY: This permit owns one capacity unit. Claiming a slot and writing it is a // synchronous operation with no user callbacks or await points between the two. - unsafe { shared.buffer.push(value) }.map_err(SendError::new)?; - // Publication owns the capacity before a wake callback can panic. - capacity.forget(); - shared.rx_waker.wake(); + unsafe { self.shared.buffer.push(value) }.map_err(SendError::new)?; + // The queued message now owns capacity, even if the wake callback panics. + mem::forget(self.capacity); + self.shared.rx_waker.wake(); Ok(()) } } diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index 9d9647a5..0dad066b 100644 --- a/tests-integration/tests/mpsc_test/main.rs +++ b/tests-integration/tests/mpsc_test/main.rs @@ -186,10 +186,3 @@ fn bounded_rejects_zero_capacity() { fn bounded_rejects_capacity_above_the_maximum() { let _ = mpsc::bounded::((usize::MAX >> 1) + 1); } - -#[test] -fn bounded_zero_sized_messages_need_no_slot_storage() { - let (tx, mut rx) = mpsc::bounded::<()>(usize::MAX >> 1); - tx.try_send(()).unwrap(); - assert_eq!(rx.try_recv(), Ok(())); -} diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index 6f470832..631a261d 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -16,6 +16,7 @@ // under the License. use std::cell::Cell; +use std::mem; use std::panic::AssertUnwindSafe; use std::panic::catch_unwind; use std::sync::Arc; @@ -57,16 +58,26 @@ fn held_permits_consume_capacity_without_claiming_message_order() { } #[test] -fn zero_sized_messages_support_the_full_capacity_range() { - for capacity in [(usize::MAX >> 1) - 1, usize::MAX >> 1] { +fn zero_sized_messages_preserve_capacity_across_reservation_and_close() { + for capacity in [1, 3, 64] { let (tx, mut rx) = mpsc::bounded::<()>(capacity); let permit = tx.try_reserve().unwrap(); - tx.try_send(()).unwrap(); - assert_eq!(rx.try_recv(), Ok(())); + for _ in 1..capacity { + tx.try_send(()).unwrap(); + } + assert_eq!(tx.try_send(()), Err(TrySendError::Full(()))); + for _ in 1..capacity { + assert_eq!(rx.try_recv(), Ok(())); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); drop(permit); tx.try_reserve().unwrap().send(()).unwrap(); + assert_eq!(rx.try_recv(), Ok(())); // Closing restores buffered capacity before outstanding permits are dropped. let held = tx.try_reserve().unwrap(); + for _ in 1..capacity { + tx.try_send(()).unwrap(); + } drop(rx); drop(held); assert!(matches!( @@ -82,6 +93,7 @@ fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { use std::sync::atomic::Ordering; static DROPS: AtomicUsize = AtomicUsize::new(0); + #[repr(align(128))] struct Message; impl Drop for Message { fn drop(&mut self) { @@ -89,7 +101,7 @@ fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { } } - let (tx, mut rx) = mpsc::bounded(usize::MAX >> 1); + let (tx, mut rx) = mpsc::bounded(3); for _ in 0..3 { assert!(tx.try_send(Message).is_ok()); } @@ -159,7 +171,7 @@ fn closing_after_a_grant_returns_the_unsent_message() { fn receiver_drop_does_not_wait_for_held_or_forgotten_permits() { let (tx, mut rx) = mpsc::bounded(3); let held = tx.try_reserve().unwrap(); - std::mem::forget(tx.try_reserve().unwrap()); + mem::forget(tx.try_reserve().unwrap()); tx.try_send(String::from("ready")).unwrap(); assert_eq!(rx.try_recv().unwrap(), "ready"); tx.try_send(String::from("discarded on close")).unwrap(); From 22c85e4441351ec8d999e10b877af057eb9a3fdb Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 17:12:46 +0800 Subject: [PATCH 30/34] refactor(mpsc): simplify publication and receiver wakeups --- LICENSE | 12 - asyncband/src/internal/atomic_waker.rs | 491 ------------------ asyncband/src/internal/mod.rs | 3 - asyncband/src/mpsc/bounded/buffer.rs | 280 ++-------- asyncband/src/mpsc/bounded/mod.rs | 6 +- asyncband/src/mpsc/bounded/receiver.rs | 58 ++- asyncband/src/mpsc/bounded/sender.rs | 14 +- .../tests/mpsc_test/concurrency.rs | 49 ++ .../tests/mpsc_test/reservation.rs | 19 +- 9 files changed, 166 insertions(+), 766 deletions(-) delete mode 100644 asyncband/src/internal/atomic_waker.rs diff --git a/LICENSE b/LICENSE index 942ddcee..3bd8900a 100644 --- a/LICENSE +++ b/LICENSE @@ -377,18 +377,6 @@ the Apache-2.0 option for the incorporated portions. Asyncband does not provide the upstream crate's synchronized receive operations and simplifies the incorporated implementation accordingly. -Portions of asyncband/src/internal/atomic_waker.rs are derived from futures-rs -0.3.34 at the following exact revision and source path: - - https://github.com/rust-lang/futures-rs/blob/705e6b5c0f06535b1aac1cb1989a172b3d45be8c/futures-core/src/task/__internal/atomic_waker.rs - -futures-rs is licensed under Apache-2.0 or MIT. Apache Asyncband uses the -Apache-2.0 option for the incorporated portions. The upstream source carries -the following copyright notices: - - Copyright (c) 2016 Alex Crichton - Copyright (c) 2017 The Tokio Authors - The polling loop in asyncband/src/blocking/executor.rs is adapted from Pollster 1.0.1 at the following exact revision and source path: diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs deleted file mode 100644 index 7eb99e5c..00000000 --- a/asyncband/src/internal/atomic_waker.rs +++ /dev/null @@ -1,491 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -// This file contains a state machine derived from futures-rs 0.3.34 and panic-recovery behavior -// informed by Tokio 1.53.1. -// Asyncband uses the Apache-2.0 license option for code incorporated from futures-rs. -// The incorporated code has been modified for use in Apache Asyncband. -// Upstream sources: -// https://github.com/rust-lang/futures-rs/blob/705e6b5c0f06535b1aac1cb1989a172b3d45be8c/futures-core/src/task/__internal/atomic_waker.rs -// https://github.com/tokio-rs/tokio/blob/75fef53d0a8590c2d1dbb63672aa7b7d1ef51155/tokio/src/sync/task/atomic_waker.rs - -use std::cell::UnsafeCell; -use std::panic::AssertUnwindSafe; -use std::panic::RefUnwindSafe; -use std::panic::UnwindSafe; -use std::panic::catch_unwind; -use std::panic::resume_unwind; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Waker; - -const WAITING: usize = 0; -const REGISTERING: usize = 0b01; -const WAKING: usize = 0b10; - -/// A single-registerer, multi-notifier cell for task wake-up. -/// -/// The atomic state both grants exclusive access to `waker` and records one coalesced wake request. -/// The operation that moves the state out of `WAITING` remains the only slot owner until it returns -/// the state to `WAITING`. -/// -/// * `WAITING`: the slot is unlocked and may contain a registered waker. -/// * `REGISTERING`: `register` exclusively owns the slot and no concurrent wake is pending. -/// * `WAKING`: `wake` exclusively owns the slot. A racing `register` self-wakes without touching -/// the slot. -/// * `REGISTERING | WAKING`: `register` still owns the slot and must complete a concurrent wake -/// before returning to `WAITING`. -/// -/// Valid state transitions are: -/// -/// ```text -/// register: WAITING ----------------Acquire CAS---------------> REGISTERING -/// REGISTERING ------------AcqRel CAS----------------> WAITING -/// -/// wake: WAITING ----------------AcqRel fetch_or-----------> WAKING -/// WAKING -----------------Release swap--------------> WAITING -/// -/// race: REGISTERING ------------AcqRel fetch_or-----------> REGISTERING | WAKING -/// REGISTERING | WAKING ---AcqRel swap---------------> WAITING -/// ``` -/// -/// Additional calls to `wake` while `WAKING` is set are coalesced. A wake completed before a -/// registration starts is not remembered, so callers must register before rechecking the condition -/// that determines whether to return `Pending`. -/// -/// Every transition that acquires slot ownership has an Acquire operation paired with the previous -/// owner's Release transition to `WAITING`. The Release half of `wake` also publishes the caller's -/// preceding condition update; a racing `register` acquires that publication before it returns. -pub struct AtomicWaker { - state: AtomicUsize, - waker: UnsafeCell>, -} - -// SAFETY: `state` grants exclusive access to `waker`, and losing concurrent registrations do not -// touch the slot. `Waker` itself is `Send + Sync`. -unsafe impl Sync for AtomicWaker {} - -// `Waker` callbacks may unwind, but no panic leaves a state bit owned by the unwinding operation. A -// failed clone leaves the old slot intact and completes any raced wake, while wake and drop -// callbacks run after that operation's critical section has been released. -impl RefUnwindSafe for AtomicWaker {} -impl UnwindSafe for AtomicWaker {} - -impl AtomicWaker { - #[inline] - pub const fn new() -> Self { - Self { - state: AtomicUsize::new(WAITING), - waker: UnsafeCell::new(None), - } - } - - /// Registers `waker`, replacing a previously registered task if it differs. - /// - /// Calls to this method must not overlap. It may run concurrently with any number of calls to - /// [`wake`](Self::wake). - #[inline] - pub fn register(&self, waker: &Waker) { - // ORDERING: On success, Acquire pairs with the Release operation that last returned the - // state to WAITING and transfers exclusive ownership of the waker slot to this thread. On - // failure, Acquire matters when this reads WAKING from a notifier's AcqRel fetch_or: it - // receives the condition update that preceded that wake before this method returns. - match self - .state - .compare_exchange(WAITING, REGISTERING, Ordering::Acquire, Ordering::Acquire) - .unwrap_or_else(|state| state) - { - WAITING => { - // SAFETY: changing WAITING to REGISTERING grants this thread exclusive access to - // the waker slot until the state is returned to WAITING. - unsafe { self.register_locked(waker) } - } - WAKING => { - // A concurrent wake owns the slot. Self-waking ensures that this registration is - // not lost even though it cannot replace the slot right now. - waker.wake_by_ref(); - } - state => { - // Concurrent registration violates this type's contract. Ignoring the losing - // registration preserves memory safety and lets the winner provide notification. - debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); - } - } - } - - /// Registers a waker after this thread has acquired the REGISTERING state. - /// - /// # Safety - /// - /// The caller must have changed `state` from WAITING to REGISTERING and must be the only - /// thread accessing `waker`. - #[inline] - unsafe fn register_locked(&self, waker: &Waker) { - // Avoid both cloning and dropping the common case where an executor polls the receiver - // repeatedly with the same task waker. - let needs_replacement = match unsafe { &*self.waker.get() } { - Some(current) => !current.will_wake(waker), - None => true, - }; - - let mut clone_panic = None; - let old_waker = if needs_replacement { - match catch_unwind(AssertUnwindSafe(|| waker.clone())) { - Ok(new_waker) => unsafe { (*self.waker.get()).replace(new_waker) }, - Err(payload) => { - clone_panic = Some(payload); - None - } - } - } else { - None - }; - - // ORDERING: Release publishes a newly registered waker when the CAS succeeds. If it fails, - // Acquire receives the concurrent notifier's Release publication before the wake is - // completed below. AcqRel is the weakest success ordering that permits an Acquire failure - // ordering, although its Acquire half is not otherwise relied upon on the success path. - let concurrent_wake = match self.state.compare_exchange( - REGISTERING, - WAITING, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => None, - Err(state) => { - debug_assert_eq!(state, REGISTERING | WAKING); - - // SAFETY: REGISTERING remains set, so this thread still owns the waker slot. - let registered = unsafe { (*self.waker.get()).take() }; - - // ORDERING: Acquire receives all coalesced wake publications. Release publishes - // the empty slot and makes it available to the next register or wake operation. - self.state.swap(WAITING, Ordering::AcqRel); - registered - } - }; - - if let Some(payload) = clone_panic { - // Preserve the original clone panic while still completing a wake that raced with it. - if let Some(waker) = concurrent_wake { - let _ = catch_unwind(AssertUnwindSafe(|| waker.wake())); - } - resume_unwind(payload); - } - - // User waker code runs only after the state machine is back in WAITING, so a panic cannot - // leave the cell locked. If the wake raced with a replacement, notify both tasks: the - // concurrent call may have targeted the old registration, while future progress relies on - // the new one. A panic from the superseded waker must not prevent the new task from waking. - if let Some(waker) = concurrent_wake { - if let Some(old_waker) = old_waker { - let _ = catch_unwind(AssertUnwindSafe(|| old_waker.wake())); - } - waker.wake(); - } else { - // Drop a replaced waker only after releasing the state lock. - drop(old_waker); - } - } - - /// Wakes and removes the most recently registered waker, if any. - #[inline] - pub fn wake(&self) { - if let Some(waker) = self.take() { - waker.wake(); - } - } - - /// Removes the registered waker if this call acquires the slot. A concurrent registration or - /// wake may instead take responsibility for notifying it. - #[inline] - pub fn take(&self) -> Option { - // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the - // previous owner. Release publishes the condition update that the caller performed before - // calling wake, including when a registering thread already owns the slot. - match self.state.fetch_or(WAKING, Ordering::AcqRel) { - WAITING => { - // SAFETY: changing WAITING to WAKING grants this thread exclusive access to the - // waker slot until the state is returned to WAITING. - unsafe { self.take_locked() } - } - state => { - // The thread registering a waker observes WAKING and completes this notification, - // or another waking thread has already taken responsibility for it. - debug_assert!( - state == REGISTERING || state == REGISTERING | WAKING || state == WAKING - ); - None - } - } - } - - /// Removes the waker after this thread has acquired the WAKING state. - /// - /// # Safety - /// - /// The caller must have changed `state` from WAITING to WAKING and must be the only thread - /// accessing `waker`. - #[inline] - unsafe fn take_locked(&self) -> Option { - // SAFETY: The caller owns the waker slot until returning the state to WAITING. - let waker = unsafe { (*self.waker.get()).take() }; - - // ORDERING: Release publishes the emptied slot. The RMW also preserves the release - // sequence of coalesced wakes, including ones after this thread acquired WAKING. A - // plain store would sever those publications from the next registration's Acquire. - let previous = self.state.swap(WAITING, Ordering::Release); - debug_assert_eq!(previous, WAKING); - waker - } -} - -#[cfg(test)] -mod tests { - use std::ptr; - use std::sync::Arc; - use std::sync::atomic::AtomicBool; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - use std::task::RawWaker; - use std::task::RawWakerVTable; - use std::task::Wake; - - use super::*; - - struct WakeCounter(AtomicUsize); - - impl Wake for WakeCounter { - fn wake(self: Arc) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } - - #[cfg(panic = "unwind")] - fn clone_panicking_waker() -> Waker { - static VTABLE: RawWakerVTable = RawWakerVTable::new( - |_| panic!("clone failed"), - |_| unreachable!(), - |_| unreachable!(), - |_| {}, - ); - - unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &VTABLE)) } - } - - #[test] - fn wake_notifies_once() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.register(&waker); - atomic_waker.wake(); - atomic_waker.wake(); - - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn reregistering_same_task_does_not_clone_waker() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.register(&waker); - let registered_refs = Arc::strong_count(&counter); - atomic_waker.register(&waker); - - assert_eq!(Arc::strong_count(&counter), registered_refs); - } - - #[test] - fn wake_before_register_is_not_remembered() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.wake(); - atomic_waker.register(&waker); - - assert_eq!(counter.0.load(Ordering::Relaxed), 0); - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn wake_during_replacement_notifies_old_and_new_tasks() { - let old_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let old_waker = Waker::from(old_counter.clone()); - let new_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let new_waker = Waker::from(new_counter.clone()); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&old_waker); - - assert_eq!( - atomic_waker.state.compare_exchange( - WAITING, - REGISTERING, - Ordering::AcqRel, - Ordering::Acquire, - ), - Ok(WAITING) - ); - std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); - - // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching - // the slot. Calling the helper completes the interrupted registration. - unsafe { atomic_waker.register_locked(&new_waker) }; - - assert_eq!(old_counter.0.load(Ordering::Relaxed), 1); - assert_eq!(new_counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn failed_wake_synchronizes_with_next_registration() { - struct Publication(UnsafeCell); - // SAFETY: The notifier's write precedes its wake. The read follows a registration that - // acquires that wake's publication; the relaxed scheduling flag adds no synchronization. - unsafe impl Sync for Publication {} - - let publication = Arc::new(Publication(UnsafeCell::new(0))); - let did_wake = AtomicBool::new(false); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(Waker::noop()); - assert_eq!( - atomic_waker.state.fetch_or(WAKING, Ordering::AcqRel), - WAITING - ); - - std::thread::scope(|scope| { - let wake = scope.spawn(|| { - // SAFETY: The reader waits for this write's publication through AtomicWaker. - unsafe { *publication.0.get() = 42 }; - assert!(atomic_waker.take().is_none()); - did_wake.store(true, Ordering::Relaxed); - }); - while !did_wake.load(Ordering::Relaxed) { - std::thread::yield_now(); - } - - // SAFETY: This thread acquired WAKING above. The coalesced notifier never touches - // the slot. Complete the first wake only after the second has published its update. - drop(unsafe { atomic_waker.take_locked() }); - atomic_waker.register(Waker::noop()); - // SAFETY: Registration acquires the coalesced wake through the release sequence. - // Miri detects a data race here if restoring WAITING severs that sequence. - assert_eq!(unsafe { *publication.0.get() }, 42); - wake.join().unwrap(); - }); - } - - #[cfg(panic = "unwind")] - #[test] - fn clone_panic_does_not_poison_state() { - let atomic_waker = AtomicWaker::new(); - - assert!( - catch_unwind(|| { - atomic_waker.register(&clone_panicking_waker()); - }) - .is_err() - ); - - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - atomic_waker.register(&Waker::from(counter.clone())); - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[cfg(panic = "unwind")] - #[test] - fn clone_panic_completes_concurrent_wake() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&Waker::from(counter.clone())); - - assert_eq!( - atomic_waker.state.compare_exchange( - WAITING, - REGISTERING, - Ordering::Acquire, - Ordering::Acquire, - ), - Ok(WAITING) - ); - std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); - - // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching - // the state. Calling the helper completes the interrupted registration. - assert!( - catch_unwind(|| unsafe { - atomic_waker.register_locked(&clone_panicking_waker()); - }) - .is_err() - ); - - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - - let next_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - atomic_waker.register(&Waker::from(next_counter.clone())); - atomic_waker.wake(); - assert_eq!(next_counter.0.load(Ordering::Relaxed), 1); - } - - #[cfg(panic = "unwind")] - #[test] - fn drop_panic_does_not_poison_state() { - unsafe fn clone_drop_panicker(data: *const ()) -> RawWaker { - RawWaker::new(data, &DROP_PANICKING_VTABLE) - } - - unsafe fn wake_drop_panicker(_: *const ()) {} - - unsafe fn drop_drop_panicker(data: *const ()) { - // SAFETY: the test keeps the pointed-to AtomicBool alive until every derived waker has - // been dropped. - let should_panic = unsafe { &*data.cast::() }; - if should_panic.swap(false, Ordering::Relaxed) { - panic!("drop failed"); - } - } - - static DROP_PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( - clone_drop_panicker, - wake_drop_panicker, - wake_drop_panicker, - drop_drop_panicker, - ); - - let should_panic = AtomicBool::new(true); - let old_waker = unsafe { - Waker::from_raw(RawWaker::new( - ptr::from_ref(&should_panic).cast(), - &DROP_PANICKING_VTABLE, - )) - }; - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let new_waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&old_waker); - - assert!(catch_unwind(AssertUnwindSafe(|| atomic_waker.register(&new_waker))).is_err()); - - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } -} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 6271cade..6d065bd1 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -49,9 +49,6 @@ pub(crate) fn wake_all(mut wakers: impl Iterator) { } } -#[cfg(feature = "mpsc")] -pub(crate) mod atomic_waker; - #[cfg(feature = "mpsc")] pub(crate) mod cache_padded; diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs index d1698803..41fb1bcc 100644 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ b/asyncband/src/mpsc/bounded/buffer.rs @@ -69,11 +69,30 @@ impl Buffer { /// Own one capacity permit before calling; release it only after a failed push or after /// the consumer reads the published value. No user code runs between claim and publication. pub unsafe fn push(&self, value: T) -> Result<(), T> { - let Ok(position) = self.claim() else { + if self.closed.load(Ordering::Acquire) { return Err(value); - }; - // SAFETY: The caller owns capacity and the ticket assigned this position. - unsafe { self.publish(position, value) } + } + // Closing may race after this check. It marks every physical slot CLOSED, so even a + // delayed claimant will recover its own value instead of publishing into a dead queue. + let position = self.tail.fetch_add(1, Ordering::AcqRel); + let slot = self.slot(position); + + // SAFETY: Capacity prevents wrapping over unread slots. AcqRel tail increments carry prior + // claimants' capacity-acquire edges even when this producer held its permit for a long + // time. The previous consumer has therefore finished reading before this write. + unsafe { (*slot.value.get()).write(value) }; + match slot + .state + .compare_exchange(EMPTY, READY, Ordering::Release, Ordering::Acquire) + { + Ok(_) => Ok(()), + Err(state) => { + debug_assert_eq!(state, CLOSED); + // SAFETY: Close saw an unpublished slot and did not read it. Failed publication + // leaves exclusive ownership with this producer, including during receiver drop. + Err(unsafe { (*slot.value.get()).assume_init_read() }) + } + } } /// Pending means a producer claimed the head but has not published it yet. @@ -119,40 +138,6 @@ impl Buffer { // semaphore still enforces the exact requested capacity, including non-powers of two. &self.slots[position & (self.slots.len() - 1)] } - - fn claim(&self) -> Result { - if self.closed.load(Ordering::Acquire) { - return Err(()); - } - // Closing may race after this check. It marks every physical slot CLOSED, so even a - // delayed claimant will recover its own value instead of publishing into a dead queue. - Ok(self.tail.fetch_add(1, Ordering::AcqRel)) - } - - /// Writes and publishes one message into a claimed position. - /// - /// # Safety - /// - /// Own the position from a claim, backed by a capacity permit, and publish it at most once. - unsafe fn publish(&self, position: usize, value: T) -> Result<(), T> { - let slot = self.slot(position); - // SAFETY: Capacity prevents wrapping over unread slots. AcqRel tail increments carry prior - // claimants' capacity-acquire edges even when this producer held its permit for a long - // time. The previous consumer has therefore finished reading before this write. - unsafe { (*slot.value.get()).write(value) }; - match slot - .state - .compare_exchange(EMPTY, READY, Ordering::Release, Ordering::Acquire) - { - Ok(_) => Ok(()), - Err(state) => { - debug_assert_eq!(state, CLOSED); - // SAFETY: Close saw an unpublished slot and did not read it. Failed publication - // leaves exclusive ownership with this producer, including during receiver drop. - Err(unsafe { (*slot.value.get()).assume_init_read() }) - } - } - } } struct Slot { @@ -220,216 +205,33 @@ impl Drop for Drain<'_, T> { #[cfg(test)] mod tests { - use std::future::Future; - use std::mem; - use std::pin::pin; - use std::sync::Arc; - use std::sync::Barrier; - use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; - use std::task::Context; use std::task::Poll; - use std::task::Waker; - use std::thread; - use crate::mpsc::BoundedSender; - use crate::mpsc::Permit; - use crate::mpsc::TryRecvError; - use crate::mpsc::bounded; - - // Exercise the scheduling window inside synchronous send, while retaining a real capacity - // permit. Ordinary callers cannot split a claim from its publication. - fn publish_claimed( - tx: &BoundedSender, - permit: Permit<'_, T>, - position: usize, - value: T, - ) -> Result<(), T> { - // SAFETY: The test claimed this position while holding the same capacity permit. - unsafe { tx.shared().buffer.publish(position, value) }?; - // Publication owns the capacity now, so the permit must not release it on drop. - mem::forget(permit); - tx.shared().rx_waker.wake(); - Ok(()) - } + use super::Buffer; #[test] - fn a_claimed_head_waits_for_publication_across_laps() { + fn fifo_survives_cursor_wraparound() { for capacity in [1, 3, 7] { - for initial in [0, usize::MAX - 1] { - let (tx, mut rx) = bounded(capacity); - // Start an empty ring near ticket overflow instead of running usize::MAX sends. - tx.shared().buffer.tail.store(initial, Ordering::Relaxed); - rx.set_head(initial); - let mut cx = Context::from_waker(Waker::noop()); - for lap in 0..8 { - let permit = tx.try_reserve().unwrap(); - let position = tx.shared().buffer.claim().unwrap(); - for offset in 1..capacity { - tx.try_send(lap * capacity + offset).unwrap(); - } - // A full ring must differ from an empty one even with no head value ready. - assert!(pin!(rx.recv()).poll(&mut cx).is_pending()); - publish_claimed(&tx, permit, position, lap * capacity).unwrap(); - for offset in 0..capacity { - assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); - } - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + let buffer = Buffer::new(capacity); + let mut head = usize::MAX - 1; + // Start near overflow instead of requiring usize::MAX messages to reach it. + buffer.tail.store(head, Ordering::Relaxed); + for lap in 0..4 { + for offset in 0..capacity { + // SAFETY: This test owns all capacity and queues at most capacity values. + unsafe { buffer.push(lap * capacity + offset) }.unwrap(); } - } - } - } - - #[test] - fn a_claim_delayed_past_close_returns_its_value() { - let (tx, rx) = bounded(3); - let permit = tx.try_reserve().unwrap(); - let drops = Arc::new(AtomicUsize::new(0)); - let value = Payload { - bytes: [7; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }; - let allocation = Arc::downgrade(tx.shared()); - // Pause after claim's open check, then resume its atomic ticket allocation after close. - assert!(!tx.shared().buffer.closed.load(Ordering::Acquire)); - drop(rx); - let position = tx.shared().buffer.tail.fetch_add(1, Ordering::AcqRel); - let unsent = publish_claimed(&tx, permit, position, value).unwrap_err(); - assert_eq!(unsent.bytes, [7; 1024]); - drop(unsent); - assert_eq!(drops.load(Ordering::Relaxed), 1); - drop(tx); - assert!(allocation.upgrade().is_none()); - } - - #[derive(Debug)] - #[repr(align(128))] - struct Payload { - bytes: [u8; 1024], - drops: Arc, - // Queued messages must not keep the shared allocation alive through a sender cycle. - _sender: BoundedSender, - } - - impl Drop for Payload { - fn drop(&mut self) { - self.drops.fetch_add(1, Ordering::Relaxed); - } - } - - #[test] - fn closing_reclaims_ready_values_without_waiting_for_a_paused_publisher() { - let (tx, rx) = bounded(2); - let allocation = Arc::downgrade(tx.shared()); - let drops = Arc::new(AtomicUsize::new(0)); - let paused = Barrier::new(2); - let (resume_tx, resume_rx) = std::sync::mpsc::channel(); - let (closed_tx, closed_rx) = std::sync::mpsc::channel(); - - thread::scope(|scope| { - let sender = &tx; - let drops = &drops; - let paused = &paused; - let publisher = scope.spawn(move || { - let permit = sender.try_reserve().unwrap(); - let position = sender.shared().buffer.claim().unwrap(); - let value = Payload { - bytes: [1; 1024], - drops: drops.clone(), - _sender: sender.clone(), - }; - paused.wait(); - resume_rx.recv().unwrap(); - let unsent = publish_claimed(sender, permit, position, value).unwrap_err(); - assert_eq!(unsent.bytes, [1; 1024]); - drop(unsent); - }); - paused.wait(); - tx.try_send(Payload { - bytes: [2; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }) - .unwrap(); - let closer = scope.spawn(move || { - drop(rx); - closed_tx.send(()).unwrap(); - }); - #[cfg(not(miri))] - let closed = closed_rx.recv_timeout(std::time::Duration::from_secs(10)); - #[cfg(miri)] - let closed = closed_rx.recv(); - let dropped_before_resume = drops.load(Ordering::Relaxed); - // Unblock the publisher before asserting so a failed close cannot strand the scope. - resume_tx.send(()).unwrap(); - publisher.join().unwrap(); - closer.join().unwrap(); - assert!(closed.is_ok(), "close waited for the paused publisher"); - assert_eq!(dropped_before_resume, 1); - }); - - assert_eq!(drops.load(Ordering::Relaxed), 2); - drop(tx); - assert!(allocation.upgrade().is_none()); - } - - #[test] - fn publication_racing_with_close_drops_every_payload_once() { - for _ in 0..if cfg!(miri) { 8 } else { 128 } { - let (tx, rx) = bounded(3); - let allocation = Arc::downgrade(tx.shared()); - let drops = Arc::new(AtomicUsize::new(0)); - let start = Barrier::new(4); - thread::scope(|scope| { - for byte in 0..3 { - let permit = tx.try_reserve().unwrap(); - let value = Payload { - bytes: [byte; 1024], - drops: drops.clone(), - _sender: tx.clone(), - }; - let start = &start; - scope.spawn(move || { - start.wait(); - if let Err(error) = permit.send(value) { - let value = error.into_inner(); - assert_eq!(value.bytes, [byte; 1024]); - drop(value); - } - }); + for offset in 0..capacity { + // SAFETY: This is the only consumer, and the next batch starts after draining. + assert_eq!( + unsafe { buffer.pop(&mut head) }, + Poll::Ready(Some(lap * capacity + offset)) + ); } - start.wait(); - drop(rx); - }); - assert_eq!(drops.load(Ordering::Relaxed), 3); - drop(tx); - assert!(allocation.upgrade().is_none()); - } - } - - #[test] - fn an_old_permit_can_publish_after_other_producers_wrap_the_ring() { - let (tx, mut rx) = bounded(3); - let old = tx.try_reserve().unwrap(); - for lap in 0..16 { - for offset in 0..2 { - tx.try_send([lap * 2 + offset; 1024]).unwrap(); - } - for offset in 0..2 { - assert_eq!(rx.try_recv(), Ok([lap * 2 + offset; 1024])); + // SAFETY: The test still owns the exclusive consumer cursor. + assert_eq!(unsafe { buffer.pop(&mut head) }, Poll::Ready(None)); } } - thread::scope(|scope| { - scope - .spawn(move || old.send([42; 1024]).unwrap()) - .join() - .unwrap(); - }); - assert_eq!( - pin!(rx.recv()).poll(&mut Context::from_waker(Waker::noop())), - Poll::Ready(Ok([42; 1024])) - ); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } } diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 167e2b22..5f1dd31b 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -22,8 +22,8 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use self::buffer::Buffer; +use self::receiver::ReceiverWaker; use self::semaphore::Semaphore; -use crate::internal::atomic_waker::AtomicWaker; use crate::internal::cache_padded::CachePadded; mod buffer; @@ -67,7 +67,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { let shared = Arc::new(Shared { senders: AtomicUsize::new(1), tx_permits: CachePadded::new(Semaphore::new(buffer)), - rx_waker: AtomicWaker::new(), + rx_waker: ReceiverWaker::new(), buffer: Buffer::new(buffer), }); ( @@ -79,6 +79,6 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { struct Shared { senders: AtomicUsize, tx_permits: CachePadded, - rx_waker: AtomicWaker, + rx_waker: ReceiverWaker, buffer: Buffer, } diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs index 377bfa13..27f5f00d 100644 --- a/asyncband/src/mpsc/bounded/receiver.rs +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -18,11 +18,16 @@ use std::fmt; use std::future::poll_fn; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::Ordering; +use std::sync::atomic::fence; use std::task::Context; use std::task::Poll; +use std::task::Waker; use super::Shared; +use crate::internal::cache_padded::CachePadded; +use crate::internal::mutex::Mutex; use crate::internal::wake_all; use crate::mpsc::RecvError; use crate::mpsc::TryRecvError; @@ -180,8 +185,55 @@ impl BoundedReceiver { } Poll::Pending } - #[cfg(test)] - pub(super) fn set_head(&mut self, head: usize) { - self.head = head; +} + +// The receiver checks the queue after registering; publishers check this flag after publication. +// Paired SeqCst fences prevent both sides from missing the other's transition. The stable false +// flag avoids modifying the waker's cache line for every message while the receiver is running. +pub struct ReceiverWaker { + waiting: CachePadded, + waker: Mutex>, +} + +impl ReceiverWaker { + pub fn new() -> Self { + Self { + waiting: CachePadded::new(AtomicBool::new(false)), + waker: Mutex::new(None), + } + } + + pub fn register(&self, waker: &Waker) { + let mut current = self.waker.lock(); + let old = if current.as_ref().is_some_and(|old| old.will_wake(waker)) { + None + } else { + // Only the receiver registers. User clone callbacks run outside the lock. + drop(current); + let waker = waker.clone(); + current = self.waker.lock(); + current.replace(waker) + }; + self.waiting.store(true, Ordering::Relaxed); + fence(Ordering::SeqCst); + drop(current); + drop(old); + } + + pub fn wake(&self) { + fence(Ordering::SeqCst); + if !self.waiting.load(Ordering::Relaxed) || !self.waiting.swap(false, Ordering::Relaxed) { + return; + } + if let Some(waker) = self.take() { + waker.wake(); + } + } + + pub fn take(&self) -> Option { + let mut current = self.waker.lock(); + // Clearing under the lock also takes responsibility for a newer registration. + self.waiting.store(false, Ordering::Relaxed); + current.take() } } diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index d72e403c..96e3f194 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -72,8 +72,13 @@ impl BoundedSender { /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for /// capacity before constructing the message. pub async fn send(&self, value: T) -> Result<(), SendError> { - match self.reserve().await { - Ok(permit) => permit.send(value), + let mut acquire = self.shared.tx_permits.acquire(); + match poll_fn(|cx| acquire.poll(cx)).await { + Ok(capacity) => Permit { + shared: &self.shared, + capacity, + } + .send(value), Err(_) => Err(SendError::new(value)), } } @@ -158,11 +163,6 @@ impl BoundedSender { Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), } } - - #[cfg(test)] - pub(super) fn shared(&self) -> &Arc> { - &self.shared - } } /// Capacity reserved for one message on a bounded channel. diff --git a/tests-integration/tests/mpsc_test/concurrency.rs b/tests-integration/tests/mpsc_test/concurrency.rs index f62339ba..788fbcbf 100644 --- a/tests-integration/tests/mpsc_test/concurrency.rs +++ b/tests-integration/tests/mpsc_test/concurrency.rs @@ -17,6 +17,7 @@ use std::future::Future; use std::pin::Pin; +use std::sync::Arc; use std::sync::Barrier; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -35,6 +36,54 @@ use tokio_test::assert_ok; use super::support::WakeCounter; use super::support::poll_with; +#[test] +fn publication_racing_with_close_drops_every_payload_once() { + #[derive(Debug)] + #[repr(align(128))] + struct Payload { + bytes: [u8; 1024], + drops: Arc<[AtomicUsize; 3]>, + // Queued messages must not retain the channel through a sender cycle. + _sender: mpsc::BoundedSender, + } + + impl Drop for Payload { + fn drop(&mut self) { + self.drops[self.bytes[0] as usize].fetch_add(1, Ordering::Relaxed); + } + } + + for _ in 0..if cfg!(miri) { 8 } else { 128 } { + let (tx, rx) = mpsc::bounded(3); + let drops = Arc::new(std::array::from_fn(|_| AtomicUsize::new(0))); + let start = Barrier::new(4); + thread::scope(|scope| { + for byte in 0..3 { + let permit = tx.try_reserve().unwrap(); + let value = Payload { + bytes: [byte; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let start = &start; + scope.spawn(move || { + start.wait(); + if let Err(error) = permit.send(value) { + let value = error.into_inner(); + assert_eq!(value.bytes, [byte; 1024]); + drop(value); + } + }); + } + start.wait(); + drop(rx); + }); + for count in drops.iter() { + assert_eq!(count.load(Ordering::Relaxed), 1); + } + } +} + #[test] fn bounded_receive_racing_with_send_registration_cannot_lose_wakeup() { for _ in 0..128 { diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index 631a261d..5ba54834 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -37,15 +37,18 @@ fn held_permits_consume_capacity_without_claiming_message_order() { for capacity in [1, 3, 64] { let (tx, mut rx) = mpsc::bounded(capacity); let permit = tx.try_reserve().unwrap(); - for value in 1..capacity { - tx.try_send(value).unwrap(); - } - assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); - assert_eq!(tx.try_send(0), Err(TrySendError::Full(0))); - for value in 1..capacity { - assert_eq!(rx.try_recv(), Ok(value)); + // The held permit stays usable while other messages repeatedly reuse the buffer. + for lap in 0..8 { + for offset in 1..capacity { + tx.try_send(lap * capacity + offset).unwrap(); + } + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + assert_eq!(tx.try_send(0), Err(TrySendError::Full(0))); + for offset in 1..capacity { + assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); permit.send(0).unwrap(); assert_eq!(rx.try_recv(), Ok(0)); // Repeated reservation and cancellation must restore the exact original capacity. From 898a3c47355b194b1471fef4d119e54172350f78 Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 17:27:36 +0800 Subject: [PATCH 31/34] refactor(mpsc): use one mutex for bounded channels --- CHANGELOG.md | 1 - asyncband/src/internal/cache_padded.rs | 55 ---- asyncband/src/internal/mod.rs | 3 - asyncband/src/mpsc/bounded/buffer.rs | 237 -------------- asyncband/src/mpsc/bounded/mod.rs | 94 ++++-- asyncband/src/mpsc/bounded/receiver.rs | 175 ++++------- asyncband/src/mpsc/bounded/semaphore.rs | 296 ------------------ asyncband/src/mpsc/bounded/sender.rs | 188 +++++++++-- .../tests/mpsc_test/reservation.rs | 2 +- 9 files changed, 282 insertions(+), 769 deletions(-) delete mode 100644 asyncband/src/internal/cache_padded.rs delete mode 100644 asyncband/src/mpsc/bounded/buffer.rs delete mode 100644 asyncband/src/mpsc/bounded/semaphore.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 74e099ec..72b1515b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -19,7 +19,6 @@ All notable changes to this project will be documented in this file. * Reject bounded MPSC capacities above `usize::MAX >> 1` up front with an explicit panic message instead of an opaque arithmetic overflow. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. -* Allow bounded MPSC producers to publish messages concurrently; acquiring and releasing capacity no longer takes an internal lock while no sender is waiting. ## v0.7.2 diff --git a/asyncband/src/internal/cache_padded.rs b/asyncband/src/internal/cache_padded.rs deleted file mode 100644 index 7268030d..00000000 --- a/asyncband/src/internal/cache_padded.rs +++ /dev/null @@ -1,55 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -// Use conservative architecture estimates, not a guarantee about every CPU's cache line. -// Keep 128 bytes for large ARM/PowerPC lines and adjacent-line prefetching on x86-64, -// 256 bytes for s390x, and at least 64 bytes elsewhere. -#[cfg_attr(target_arch = "s390x", repr(align(256)))] -#[cfg_attr( - any( - target_arch = "aarch64", - target_arch = "arm64ec", - target_arch = "powerpc64", - target_arch = "x86_64", - ), - repr(align(128)) -)] -#[cfg_attr( - not(any( - target_arch = "s390x", - target_arch = "aarch64", - target_arch = "arm64ec", - target_arch = "powerpc64", - target_arch = "x86_64", - )), - repr(align(64)) -)] -pub struct CachePadded(T); - -impl CachePadded { - pub const fn new(value: T) -> Self { - Self(value) - } -} - -impl std::ops::Deref for CachePadded { - type Target = T; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 6d065bd1..0252aa83 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -49,9 +49,6 @@ pub(crate) fn wake_all(mut wakers: impl Iterator) { } } -#[cfg(feature = "mpsc")] -pub(crate) mod cache_padded; - #[cfg(any( feature = "barrier", feature = "broadcast", diff --git a/asyncband/src/mpsc/bounded/buffer.rs b/asyncband/src/mpsc/bounded/buffer.rs deleted file mode 100644 index 41fb1bcc..00000000 --- a/asyncband/src/mpsc/bounded/buffer.rs +++ /dev/null @@ -1,237 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! Capacity, position, and publication are separate ownership transitions: -//! -//! - A permit owns capacity, but holds no position until synchronous `push` claims a ticket. -//! - The ticket gives one producer a slot. `READY` publishes its initialized value to the receiver. -//! - The receiver finishes reading before returning capacity. AcqRel ticket increments carry that -//! reuse ordering even to a producer that acquired its permit on an earlier lap. -//! - Close competes with publication on the slot state. The drain owns `READY` values; a producer -//! that encounters `CLOSED` owns its unpublished value. Neither waits for the other to resume. -//! -//! Only the non-cloneable receiver advances the read cursor. All endpoints retain the shared -//! allocation, so a publisher's slot stays alive even when receiver drop closes it concurrently. - -use std::cell::UnsafeCell; -use std::mem::MaybeUninit; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicU8; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Poll; - -use crate::internal::cache_padded::CachePadded; - -const EMPTY: u8 = 0; -const READY: u8 = 1; -const CLOSED: u8 = 2; - -pub struct Buffer { - slots: Box<[Slot]>, - tail: CachePadded, - closed: AtomicBool, -} - -impl Buffer { - pub fn new(capacity: usize) -> Self { - let slots = (0..capacity.next_power_of_two()) - .map(|_| Slot { - state: AtomicU8::new(EMPTY), - value: UnsafeCell::new(MaybeUninit::uninit()), - }) - .collect(); - Self { - slots, - tail: CachePadded::new(AtomicUsize::new(0)), - closed: AtomicBool::new(false), - } - } - - /// Writes and publishes one message. Closing may instead return the unsent value. - /// - /// # Safety - /// - /// Own one capacity permit before calling; release it only after a failed push or after - /// the consumer reads the published value. No user code runs between claim and publication. - pub unsafe fn push(&self, value: T) -> Result<(), T> { - if self.closed.load(Ordering::Acquire) { - return Err(value); - } - // Closing may race after this check. It marks every physical slot CLOSED, so even a - // delayed claimant will recover its own value instead of publishing into a dead queue. - let position = self.tail.fetch_add(1, Ordering::AcqRel); - let slot = self.slot(position); - - // SAFETY: Capacity prevents wrapping over unread slots. AcqRel tail increments carry prior - // claimants' capacity-acquire edges even when this producer held its permit for a long - // time. The previous consumer has therefore finished reading before this write. - unsafe { (*slot.value.get()).write(value) }; - match slot - .state - .compare_exchange(EMPTY, READY, Ordering::Release, Ordering::Acquire) - { - Ok(_) => Ok(()), - Err(state) => { - debug_assert_eq!(state, CLOSED); - // SAFETY: Close saw an unpublished slot and did not read it. Failed publication - // leaves exclusive ownership with this producer, including during receiver drop. - Err(unsafe { (*slot.value.get()).assume_init_read() }) - } - } - } - - /// Pending means a producer claimed the head but has not published it yet. - /// - /// # Safety - /// - /// Only the exclusive consumer may call this, using its persistent cursor. Release one - /// capacity permit after each successful pop, after the value has been read completely. - pub unsafe fn pop(&self, head: &mut usize) -> Poll> { - let slot = self.slot(*head); - if slot.state.load(Ordering::Acquire) == READY { - // SAFETY: Publication initialized the value, and only this consumer can read it. - // Capacity is still held until this method has returned the value to its caller. - let value = unsafe { (*slot.value.get()).assume_init_read() }; - slot.state.store(EMPTY, Ordering::Release); - *head = head.wrapping_add(1); - Poll::Ready(Some(value)) - } else if self.tail.load(Ordering::Acquire) == *head { - Poll::Ready(None) - } else { - Poll::Pending - } - } - - /// Stops new claims and returns ownership of published values to a drain guard. - /// - /// # Safety - /// - /// Only the exclusive consumer may close the buffer, once, using its current cursor. - pub unsafe fn close(&self, head: usize) -> Drain<'_, T> { - self.closed.store(true, Ordering::Release); - // Cover every physical slot: a producer may have passed the open check but not yet - // claimed its ticket. Such a late claim must also find a CLOSED slot. - Drain { - buffer: self, - position: head, - remaining: self.slots.len(), - } - } - - fn slot(&self, position: usize) -> &Slot { - // Power-of-two storage preserves indexing when the full-width ticket wraps. The - // semaphore still enforces the exact requested capacity, including non-powers of two. - &self.slots[position & (self.slots.len() - 1)] - } -} - -struct Slot { - state: AtomicU8, - value: UnsafeCell>, -} - -// SAFETY: Capacity and the tail ticket give a producer exclusive ownership of an empty slot. -// Release publication transfers its value to the exclusive consumer. Closing an unpublished -// slot leaves its value with the producer; closing a READY slot transfers it to the drain. -unsafe impl Sync for Slot {} - -// No reference to a stored value escapes. Every value is removed from the slot's ownership -// before running a callback or destructor that might panic. -impl std::panic::UnwindSafe for Slot {} -impl std::panic::RefUnwindSafe for Slot {} - -pub struct Drain<'a, T> { - buffer: &'a Buffer, - position: usize, - remaining: usize, -} - -impl Iterator for Drain<'_, T> { - type Item = T; - - fn next(&mut self) -> Option { - while self.remaining != 0 { - let position = self.position; - self.remaining -= 1; - self.position = self.position.wrapping_add(1); - let slot = self.buffer.slot(position); - if slot.state.swap(CLOSED, Ordering::AcqRel) == READY { - // SAFETY: The drain won ownership of a published value. The cursor and - // state already advanced, so a panicking destructor cannot read twice. - return Some(unsafe { (*slot.value.get()).assume_init_read() }); - } - // An unpublished slot stays owned by its producer, which will observe CLOSED - // and recover its value. The shared Arc keeps this allocation alive until then. - } - None - } -} - -impl Drop for Drain<'_, T> { - fn drop(&mut self) { - struct Remaining<'a, 'b, T>(&'a mut Drain<'b, T>); - - impl Drop for Remaining<'_, '_, T> { - fn drop(&mut self) { - for value in self.0.by_ref() { - drop(value); - } - } - } - - // A guard inside Drop is necessary: Drop itself is not called again if a payload's - // destructor panics while this normal drain is running. - let remaining = Remaining(self); - for value in remaining.0.by_ref() { - drop(value); - } - } -} - -#[cfg(test)] -mod tests { - use std::sync::atomic::Ordering; - use std::task::Poll; - - use super::Buffer; - - #[test] - fn fifo_survives_cursor_wraparound() { - for capacity in [1, 3, 7] { - let buffer = Buffer::new(capacity); - let mut head = usize::MAX - 1; - // Start near overflow instead of requiring usize::MAX messages to reach it. - buffer.tail.store(head, Ordering::Relaxed); - for lap in 0..4 { - for offset in 0..capacity { - // SAFETY: This test owns all capacity and queues at most capacity values. - unsafe { buffer.push(lap * capacity + offset) }.unwrap(); - } - for offset in 0..capacity { - // SAFETY: This is the only consumer, and the next batch starts after draining. - assert_eq!( - unsafe { buffer.pop(&mut head) }, - Poll::Ready(Some(lap * capacity + offset)) - ); - } - // SAFETY: The test still owns the exclusive consumer cursor. - assert_eq!(unsafe { buffer.pop(&mut head) }, Poll::Ready(None)); - } - } - } -} diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 5f1dd31b..b1b75449 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -18,17 +18,16 @@ //! A bounded multi-producer, single-consumer queue for sending values between asynchronous //! tasks with backpressure control. +use std::collections::VecDeque; use std::sync::Arc; -use std::sync::atomic::AtomicUsize; +use std::task::Waker; -use self::buffer::Buffer; -use self::receiver::ReceiverWaker; -use self::semaphore::Semaphore; -use crate::internal::cache_padded::CachePadded; +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaitList; +use crate::mpsc::TryRecvError; +use crate::mpsc::TrySendError; -mod buffer; mod receiver; -mod semaphore; mod sender; pub use self::receiver::BoundedReceiver; @@ -41,18 +40,18 @@ pub use self::sender::Permit; /// one slot for a waiting sender. Capacity is granted in the order that pending sends and /// reservations enter the wait queue; new senders cannot take an already granted slot. /// -/// Message slots are preallocated and rounded up to a power of two; the channel's capacity -/// remains exactly `buffer`. Every slot needs state metadata, including for zero-sized messages. +/// Message storage is preallocated for `buffer` values. Queued messages and outstanding +/// reservations together occupy at most `buffer` capacity units. +/// +/// Operations briefly acquire an internal mutex; no lock is held across an await point or while +/// invoking waker callbacks or message destructors. The `try_*` methods do not wait for capacity +/// or messages, but may wait to acquire this mutex. /// /// # Panics /// /// Panics if `buffer` is zero or exceeds the maximum capacity of `usize::MAX >> 1`. #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { - /// The largest capacity accepted by [`bounded`]. - /// - /// The shared permit counter reserves two sentinel values above the usable range. This - /// bound also keeps the rounded-up slot storage from overflowing a power of two. const MAX_CAPACITY: usize = usize::MAX >> 1; assert!( @@ -64,21 +63,68 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { "mpsc bounded channel capacity {buffer} exceeds the maximum of {MAX_CAPACITY}", ); - let shared = Arc::new(Shared { - senders: AtomicUsize::new(1), - tx_permits: CachePadded::new(Semaphore::new(buffer)), - rx_waker: ReceiverWaker::new(), - buffer: Buffer::new(buffer), - }); + let shared = Arc::new(Mutex::new(State { + queue: VecDeque::with_capacity(buffer), + available: buffer, + senders: 1, + receiver_open: true, + receiver_waker: None, + waiters: WaitList::new(), + })); ( BoundedSender::new(shared.clone()), BoundedReceiver::new(shared), ) } -struct Shared { - senders: AtomicUsize, - tx_permits: CachePadded, - rx_waker: ReceiverWaker, - buffer: Buffer, +// While open, capacity belongs to available, a queued message, a Permit, or a granted waiter. +// All transitions hold one mutex. Waker callbacks and message destruction run after unlocking. +struct State { + queue: VecDeque, + available: usize, + senders: usize, + receiver_open: bool, + receiver_waker: Option, + waiters: WaitList, +} + +impl State { + fn acquire(&mut self) -> Result<(), TrySendError<()>> { + if !self.receiver_open { + Err(TrySendError::Disconnected(())) + } else if self.available == 0 { + Err(TrySendError::Full(())) + } else { + self.available -= 1; + Ok(()) + } + } + + fn release(&mut self) -> Option { + if !self.receiver_open { + return None; + } + if let Some((_, waiter)) = self.waiters.unlink_first_waiter(|_| true) { + // The detached node owns capacity until its future claims or cancels the grant. + waiter.granted = true; + return waiter.waker.take(); + } + self.available += 1; + None + } + + fn pop(&mut self) -> Result<(T, Option), TryRecvError> { + if let Some(value) = self.queue.pop_front() { + Ok((value, self.release())) + } else if self.senders == 0 { + Err(TryRecvError::Disconnected) + } else { + Err(TryRecvError::Empty) + } + } +} + +struct Waiter { + granted: bool, + waker: Option, } diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs index 27f5f00d..c97315a4 100644 --- a/asyncband/src/mpsc/bounded/receiver.rs +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -17,30 +17,24 @@ use std::fmt; use std::future::poll_fn; +use std::mem; use std::sync::Arc; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; -use std::sync::atomic::fence; use std::task::Context; use std::task::Poll; -use std::task::Waker; -use super::Shared; -use crate::internal::cache_padded::CachePadded; +use super::State; use crate::internal::mutex::Mutex; use crate::internal::wake_all; +use crate::internal::waker_batch::WakerBatch; use crate::mpsc::RecvError; use crate::mpsc::TryRecvError; /// The receiving endpoint of a bounded mpsc channel. /// /// Instances are created by the [`bounded`](crate::mpsc::bounded) function. Dropping the receiver -/// discards queued values. -/// The backing allocation remains alive until all endpoints are dropped, so a concurrent sender -/// can safely finish returning an unsent value. +/// discards queued values and disconnects pending sends and reservations. pub struct BoundedReceiver { - shared: Arc>, - head: usize, + shared: Arc>>, } impl fmt::Debug for BoundedReceiver { @@ -51,20 +45,29 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - // SAFETY: Receiver ownership provides exclusive access to the consumption cursor. - // The drain first prevents new claims. Its destructor completes cleanup on unwinding. - let drain = unsafe { self.shared.buffer.close(self.head) }; - let wakers = self.shared.tx_permits.close(); - let receiver_waker = self.shared.rx_waker.take(); + let (queue, receiver_waker, wakers) = { + let mut state = self.shared.lock(); + state.receiver_open = false; + let queue = mem::take(&mut state.queue); + let receiver_waker = state.receiver_waker.take(); + let mut wakers = WakerBatch::new(); + while let Some((_, waiter)) = state.waiters.unlink_first_waiter(|_| true) { + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } + } + (queue, receiver_waker, wakers) + }; + // Local ownership also drains the queue if a wake or waker destructor unwinds. wake_all(wakers.into_iter()); drop(receiver_waker); - drop(drain); + drop(queue); } } impl BoundedReceiver { - pub(super) fn new(shared: Arc>) -> Self { - Self { shared, head: 0 } + pub(super) fn new(shared: Arc>>) -> Self { + Self { shared } } /// Attempts to receive the next queued value without waiting for a new message. @@ -73,9 +76,6 @@ impl BoundedReceiver { /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has /// been dropped and all queued values have been consumed. /// - /// If a producer is still completing a synchronous publication at the queue head, this - /// method waits for that publication. Use [`Self::recv`] to wait asynchronously instead. - /// /// # Examples /// /// ``` @@ -93,22 +93,11 @@ impl BoundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - let mut spins = 0; - loop { - match self.pull() { - Poll::Ready(result) => return result, - Poll::Pending => { - // A synchronous publisher already owns the head. Reporting Empty here - // could hide a later send that has completed. Async recv parks instead. - if spins < 32 { - std::hint::spin_loop(); - spins += 1; - } else { - std::thread::yield_now(); - } - } - } + let (value, wake) = self.shared.lock().pop()?; + if let Some(waker) = wake { + waker.wake(); } + Ok(value) } /// Waits for and receives the next value, freeing one buffer slot. @@ -144,96 +133,42 @@ impl BoundedReceiver { poll_fn(|cx| self.poll_recv(cx)).await } - /// One attempt to take the head value: a message, an empty-or-disconnected classification, - /// or `Pending` while a claimed head waits for its publication. - fn pull(&mut self) -> Poll> { - let mut disconnected = false; + fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { + let mut cloned_waker = None; loop { - // SAFETY: Only this receiver owns head. Capacity is released after the buffer - // finishes reading and advances the cursor, so no producer can overwrite the value. - match unsafe { self.shared.buffer.pop(&mut self.head) } { - Poll::Ready(Some(value)) => { - self.shared.tx_permits.release(); + let mut state = self.shared.lock(); + match state.pop() { + Ok((value, wake)) => { + drop(state); + if let Some(waker) = wake { + waker.wake(); + } return Poll::Ready(Ok(value)); } - Poll::Ready(None) if disconnected => { - return Poll::Ready(Err(TryRecvError::Disconnected)); - } - Poll::Ready(None) if self.shared.senders.load(Ordering::Acquire) == 0 => { - // Acquire the last sender's completed publications before checking again. - disconnected = true; - } - Poll::Ready(None) => return Poll::Ready(Err(TryRecvError::Empty)), - Poll::Pending => return Poll::Pending, - } - } - } - - fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - for registered in [false, true] { - match self.pull() { - Poll::Ready(Ok(value)) => return Poll::Ready(Ok(value)), - Poll::Ready(Err(TryRecvError::Disconnected)) => { - drop(self.shared.rx_waker.take()); + Err(TryRecvError::Disconnected) => { + let old = state.receiver_waker.take(); + drop(state); + drop(old); return Poll::Ready(Err(RecvError::Disconnected)); } - Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {} + Err(TryRecvError::Empty) => {} } - if !registered { - self.shared.rx_waker.register(cx.waker()); + if state + .receiver_waker + .as_ref() + .is_some_and(|w| w.will_wake(cx.waker())) + { + return Poll::Pending; } + if let Some(waker) = cloned_waker.take() { + let old = state.receiver_waker.replace(waker); + drop(state); + drop(old); + return Poll::Pending; + } + drop(state); + // Clone can reenter the channel, so check the queue again after acquiring the lock. + cloned_waker = Some(cx.waker().clone()); } - Poll::Pending - } -} - -// The receiver checks the queue after registering; publishers check this flag after publication. -// Paired SeqCst fences prevent both sides from missing the other's transition. The stable false -// flag avoids modifying the waker's cache line for every message while the receiver is running. -pub struct ReceiverWaker { - waiting: CachePadded, - waker: Mutex>, -} - -impl ReceiverWaker { - pub fn new() -> Self { - Self { - waiting: CachePadded::new(AtomicBool::new(false)), - waker: Mutex::new(None), - } - } - - pub fn register(&self, waker: &Waker) { - let mut current = self.waker.lock(); - let old = if current.as_ref().is_some_and(|old| old.will_wake(waker)) { - None - } else { - // Only the receiver registers. User clone callbacks run outside the lock. - drop(current); - let waker = waker.clone(); - current = self.waker.lock(); - current.replace(waker) - }; - self.waiting.store(true, Ordering::Relaxed); - fence(Ordering::SeqCst); - drop(current); - drop(old); - } - - pub fn wake(&self) { - fence(Ordering::SeqCst); - if !self.waiting.load(Ordering::Relaxed) || !self.waiting.swap(false, Ordering::Relaxed) { - return; - } - if let Some(waker) = self.take() { - waker.wake(); - } - } - - pub fn take(&self) -> Option { - let mut current = self.waker.lock(); - // Clearing under the lock also takes responsibility for a newer registration. - self.waiting.store(false, Ordering::Relaxed); - current.take() } } diff --git a/asyncband/src/mpsc/bounded/semaphore.rs b/asyncband/src/mpsc/bounded/semaphore.rs deleted file mode 100644 index e7f165e6..00000000 --- a/asyncband/src/mpsc/bounded/semaphore.rs +++ /dev/null @@ -1,296 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -//! The channel's capacity: a counting semaphore with a close operation and a fair wait queue. -//! -//! This channel-local semaphore keeps its permit counter and channel state in one atomic. The -//! general-purpose semaphore has neither a close operation nor acquisition errors. -//! -//! `state` is the available permit count, plus two sentinel values at the top of the range: -//! -//! * `CLOSED`: the receiver is gone. No permits are issued or returned, and waiters drain with an -//! error. -//! * `WAITING`: the wait queue may be non-empty. Releases then take the locked path and grant the -//! permit directly to the oldest waiter instead of returning it to the counter, so capacity is -//! handed out in registration order and new arrivals cannot steal an already granted slot. The -//! counter is zero while this sentinel stands: waiters only register after observing exhaustion, -//! and grants bypass the counter. -//! -//! With neither sentinel installed, acquire and release are single lock-free operations on `state`. -//! Wait-queue mutations always hold the queue lock. A registration must install or observe -//! `WAITING` before joining the queue, so every subsequent release takes the locked path. If a -//! release wins that transition, acquisition retries instead of registering against a plain count. - -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Context; -use std::task::Poll; -use std::task::Waker; - -use crate::internal::mutex::Mutex; -use crate::internal::waitlist::WaitList; -use crate::internal::waitlist::WaiterId; -use crate::internal::waker_batch::WakerBatch; -use crate::mpsc::SendError; -use crate::mpsc::TrySendError; - -pub struct Semaphore { - state: AtomicUsize, - waiters: Mutex>, -} - -const CLOSED: usize = usize::MAX; -const WAITING: usize = usize::MAX - 1; - -struct Waiter { - granted: bool, - waker: Option, -} - -impl Semaphore { - pub fn new(available: usize) -> Self { - Self { - state: AtomicUsize::new(available), - waiters: Mutex::new(WaitList::new()), - } - } - - pub fn try_acquire(&self) -> Result, TrySendError<()>> { - let mut state = self.state.load(Ordering::Acquire); - loop { - if state == CLOSED { - return Err(TrySendError::Disconnected(())); - } - if state == WAITING || state == 0 { - return Err(TrySendError::Full(())); - } - match self.state.compare_exchange_weak( - state, - state - 1, - Ordering::Acquire, - Ordering::Acquire, - ) { - Ok(_) => return Ok(Capacity { semaphore: self }), - Err(actual) => state = actual, - } - } - } - - /// Acquires one permit asynchronously, waiting in registration order when the semaphore - /// is exhausted. - pub fn acquire(&self) -> Acquire<'_> { - Acquire { - semaphore: self, - waiter: None, - } - } - - pub fn is_closed(&self) -> bool { - self.state.load(Ordering::Acquire) == CLOSED - } - - // Called with the queue locked. A failed installation requires retrying acquisition: a - // racing sender may consume the returned capacity before a separate recheck can see it. - fn set_waiting(&self) -> bool { - matches!( - self.state - .compare_exchange(0, WAITING, Ordering::AcqRel, Ordering::Acquire), - Ok(_) | Err(WAITING) - ) - } - - // Removes WAITING, keeping whatever count a racing grant restoration left behind. - fn clear_waiting(&self) { - let _ = self - .state - .compare_exchange(WAITING, 0, Ordering::Release, Ordering::Relaxed); - } - - pub fn release(&self) { - // Fast path: with no waiting sender and no close in sight, the permit goes straight - // back to the counter. - let mut state = self.state.load(Ordering::Relaxed); - loop { - if state == WAITING || state == CLOSED { - break; - } - match self.state.compare_exchange_weak( - state, - state + 1, - Ordering::Release, - Ordering::Relaxed, - ) { - Ok(_) => return, - Err(actual) => state = actual, - } - } - let wake = self.release_locked(&mut self.waiters.lock()); - if let Some(waker) = wake { - waker.wake(); - } - } - - fn release_locked(&self, waiters: &mut WaitList) -> Option { - if self.is_closed() { - return None; - } - if let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { - // Grant ownership before waking; new arrivals cannot steal this capacity. - waiter.granted = true; - let waker = waiter.waker.take(); - if waiters.is_empty() { - self.clear_waiting(); - } - return waker; - } - // The queue is empty: return the permit to the counter. An outstanding grant already - // owns its capacity. Adding to a plain count is safe because only lock-holding - // operations install a sentinel, and this operation holds the lock; WAITING itself - // must be displaced rather than incremented, because WAITING + 1 is CLOSED. - if self.state.load(Ordering::Relaxed) == WAITING { - let _displaced = - self.state - .compare_exchange(WAITING, 1, Ordering::Release, Ordering::Relaxed); - debug_assert_eq!(_displaced, Ok(WAITING)); - } else { - self.state.fetch_add(1, Ordering::Release); - } - None - } - - pub fn close(&self) -> WakerBatch { - let mut waiters = self.waiters.lock(); - self.state.store(CLOSED, Ordering::Release); - let mut wakers = WakerBatch::new(); - while let Some((_, waiter)) = waiters.unlink_first_waiter(|_| true) { - if let Some(waker) = waiter.waker.take() { - wakers.push(waker); - } - } - wakers - } -} - -/// Owns one capacity unit until publication transfers it to a queued message. -#[must_use = "dropping the guard releases its capacity"] -pub struct Capacity<'a> { - semaphore: &'a Semaphore, -} - -impl Drop for Capacity<'_> { - fn drop(&mut self) { - self.semaphore.release(); - } -} - -/// An in-flight [`Semaphore::acquire`] operation. -/// -/// Dropping the operation removes its wait-queue registration; a capacity grant that already -/// reached the registration is released to the next waiter or returned to the counter. -pub struct Acquire<'a> { - semaphore: &'a Semaphore, - waiter: Option, -} - -impl<'a> Acquire<'a> { - pub fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { - let semaphore = self.semaphore; - let mut cloned_waker = None; - let result = loop { - if self.waiter.is_none() { - match semaphore.try_acquire() { - Ok(capacity) => break Ok(capacity), - Err(TrySendError::Disconnected(())) => break Err(SendError::new(())), - Err(TrySendError::Full(())) => {} - } - } - let mut waiters = semaphore.waiters.lock(); - if semaphore.is_closed() { - // Drop removes any remaining registration, including an unused grant. - break Err(SendError::new(())); - } - if let Some(index) = self.waiter { - let waiter = waiters.waiter_mut(index); - if waiter.granted { - let waiter = waiters.remove_unlinked_waiter(index); - self.waiter = None; - let capacity = Capacity { semaphore }; - drop(waiters); - drop(waiter); - break Ok(capacity); - } - if waiter - .waker - .as_ref() - .is_some_and(|w| w.will_wake(cx.waker())) - { - return Poll::Pending; - } - if let Some(waker) = cloned_waker.take() { - let old = waiter.waker.replace(waker); - drop(waiters); - drop(old); - return Poll::Pending; - } - } else { - if !semaphore.set_waiting() { - drop(waiters); - continue; - } - if let Some(waker) = cloned_waker.take() { - self.waiter = Some(waiters.push_back(Waiter { - granted: false, - waker: Some(waker), - })); - return Poll::Pending; - } - } - drop(waiters); - // Clone outside the lock, then recheck capacity and closure before registering. - cloned_waker = Some(cx.waker().clone()); - }; - // A successful result already owns a guard, so a panicking waker destructor returns - // capacity even before the caller has constructed its public permit. - drop(cloned_waker); - Poll::Ready(result) - } -} - -impl Drop for Acquire<'_> { - fn drop(&mut self) { - let Some(index) = self.waiter else { return }; - let semaphore = self.semaphore; - let (waiter, wake) = { - let mut waiters = semaphore.waiters.lock(); - waiters.unlink_waiter(index, |_| true); - let waiter = waiters.remove_unlinked_waiter(index); - let wake = if waiter.granted { - semaphore.release_locked(&mut waiters) - } else { - if waiters.is_empty() { - semaphore.clear_waiting(); - } - None - }; - (waiter, wake) - }; - if let Some(waker) = wake { - waker.wake(); - } - drop(waiter); - } -} diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index 96e3f194..4a174c1b 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -19,10 +19,13 @@ use std::fmt; use std::future::poll_fn; use std::mem; use std::sync::Arc; -use std::sync::atomic::Ordering; +use std::task::Context; +use std::task::Poll; -use super::Shared; -use super::semaphore::Capacity; +use super::State; +use super::Waiter; +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaiterId; use crate::mpsc::SendError; use crate::mpsc::TrySendError; @@ -30,12 +33,12 @@ use crate::mpsc::TrySendError; /// /// Instances are created by the [`bounded`](crate::mpsc::bounded) function. pub struct BoundedSender { - shared: Arc>, + shared: Arc>>, } impl Clone for BoundedSender { fn clone(&self) -> Self { - self.shared.senders.fetch_add(1, Ordering::Relaxed); + self.shared.lock().senders += 1; BoundedSender { shared: self.shared.clone(), } @@ -50,14 +53,23 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { - if self.shared.senders.fetch_sub(1, Ordering::AcqRel) == 1 { - self.shared.rx_waker.wake(); + let wake = { + let mut state = self.shared.lock(); + state.senders -= 1; + if state.senders == 0 { + state.receiver_waker.take() + } else { + None + } + }; + if let Some(waker) = wake { + waker.wake(); } } } impl BoundedSender { - pub(super) fn new(shared: Arc>) -> Self { + pub(super) fn new(shared: Arc>>) -> Self { Self { shared } } @@ -72,13 +84,24 @@ impl BoundedSender { /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for /// capacity before constructing the message. pub async fn send(&self, value: T) -> Result<(), SendError> { - let mut acquire = self.shared.tx_permits.acquire(); - match poll_fn(|cx| acquire.poll(cx)).await { - Ok(capacity) => Permit { - shared: &self.shared, - capacity, + { + let mut state = self.shared.lock(); + match state.acquire() { + Ok(()) => { + state.queue.push_back(value); + let wake = state.receiver_waker.take(); + drop(state); + if let Some(waker) = wake { + waker.wake(); + } + return Ok(()); + } + Err(TrySendError::Disconnected(())) => return Err(SendError::new(value)), + Err(TrySendError::Full(())) => {} } - .send(value), + } + match self.reserve().await { + Ok(permit) => permit.send(value), Err(_) => Err(SendError::new(value)), } } @@ -114,12 +137,11 @@ impl BoundedSender { /// # } /// ``` pub async fn reserve(&self) -> Result, SendError<()>> { - let mut acquire = self.shared.tx_permits.acquire(); - let capacity = poll_fn(|cx| acquire.poll(cx)).await?; - Ok(Permit { + let mut reservation = Reservation { shared: &self.shared, - capacity, - }) + waiter: None, + }; + poll_fn(|cx| reservation.poll(cx)).await } /// Reserves capacity for one message without waiting. @@ -127,10 +149,9 @@ impl BoundedSender { /// Returns [`TrySendError::Full`] if queued messages and outstanding permits occupy the /// buffer, or [`TrySendError::Disconnected`] if the receiver has been dropped. pub fn try_reserve(&self) -> Result, TrySendError<()>> { - let capacity = self.shared.tx_permits.try_acquire()?; + self.shared.lock().acquire()?; Ok(Permit { shared: &self.shared, - capacity, }) } @@ -155,10 +176,17 @@ impl BoundedSender { /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - match self.try_reserve() { - Ok(permit) => permit - .send(value) - .map_err(|error| TrySendError::Disconnected(error.into_inner())), + let mut state = self.shared.lock(); + match state.acquire() { + Ok(()) => { + state.queue.push_back(value); + let wake = state.receiver_waker.take(); + drop(state); + if let Some(waker) = wake { + waker.wake(); + } + Ok(()) + } Err(TrySendError::Full(())) => Err(TrySendError::Full(value)), Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), } @@ -172,8 +200,7 @@ impl BoundedSender { /// it without sending releases capacity and notifies a waiting sender. #[must_use = "dropping the permit releases its reserved capacity"] pub struct Permit<'a, T> { - shared: &'a Shared, - capacity: Capacity<'a>, + shared: &'a Mutex>, } impl fmt::Debug for Permit<'_, T> { @@ -187,12 +214,109 @@ impl Permit<'_, T> { /// /// If the receiver has been dropped, the returned error contains the unsent value. pub fn send(self, value: T) -> Result<(), SendError> { - // SAFETY: This permit owns one capacity unit. Claiming a slot and writing it is a - // synchronous operation with no user callbacks or await points between the two. - unsafe { self.shared.buffer.push(value) }.map_err(SendError::new)?; + let mut state = self.shared.lock(); + if !state.receiver_open { + return Err(SendError::new(value)); + } + state.queue.push_back(value); // The queued message now owns capacity, even if the wake callback panics. - mem::forget(self.capacity); - self.shared.rx_waker.wake(); + mem::forget(self); + let wake = state.receiver_waker.take(); + drop(state); + if let Some(waker) = wake { + waker.wake(); + } Ok(()) } } + +impl Drop for Permit<'_, T> { + fn drop(&mut self) { + let wake = self.shared.lock().release(); + if let Some(waker) = wake { + waker.wake(); + } + } +} + +struct Reservation<'a, T> { + shared: &'a Mutex>, + waiter: Option, +} + +impl<'a, T> Reservation<'a, T> { + fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { + let mut cloned_waker = None; + let result = loop { + let mut state = self.shared.lock(); + if !state.receiver_open { + // Drop removes any remaining registration, including an unused grant. + break Err(SendError::new(())); + } + if let Some(index) = self.waiter { + let waiter = state.waiters.waiter_mut(index); + if waiter.granted { + let waiter = state.waiters.remove_unlinked_waiter(index); + self.waiter = None; + let permit = Permit { + shared: self.shared, + }; + drop(state); + drop(waiter); + break Ok(permit); + } + if waiter + .waker + .as_ref() + .is_some_and(|w| w.will_wake(cx.waker())) + { + return Poll::Pending; + } + if let Some(waker) = cloned_waker.take() { + let old = waiter.waker.replace(waker); + drop(state); + drop(old); + return Poll::Pending; + } + } else if state.available != 0 { + state.available -= 1; + break Ok(Permit { + shared: self.shared, + }); + } else if let Some(waker) = cloned_waker.take() { + self.waiter = Some(state.waiters.push_back(Waiter { + granted: false, + waker: Some(waker), + })); + return Poll::Pending; + } + drop(state); + // Clone outside the lock, then recheck capacity and closure before registering. + cloned_waker = Some(cx.waker().clone()); + }; + // The permit already owns capacity if dropping an unused clone unwinds. + drop(cloned_waker); + Poll::Ready(result) + } +} + +impl Drop for Reservation<'_, T> { + fn drop(&mut self) { + let Some(index) = self.waiter else { return }; + let (waiter, wake) = { + let mut state = self.shared.lock(); + state.waiters.unlink_waiter(index, |_| true); + let waiter = state.waiters.remove_unlinked_waiter(index); + let wake = if waiter.granted { + state.release() + } else { + None + }; + (waiter, wake) + }; + if let Some(waker) = wake { + waker.wake(); + } + drop(waiter); + } +} diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index 5ba54834..7c71c813 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -76,7 +76,7 @@ fn zero_sized_messages_preserve_capacity_across_reservation_and_close() { drop(permit); tx.try_reserve().unwrap().send(()).unwrap(); assert_eq!(rx.try_recv(), Ok(())); - // Closing restores buffered capacity before outstanding permits are dropped. + // An outstanding permit can be dropped after the receiver closes. let held = tx.try_reserve().unwrap(); for _ in 1..capacity { tx.try_send(()).unwrap(); From 010a3340484aa6bc27e46a6524bd55366202f52d Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 17:52:14 +0800 Subject: [PATCH 32/34] refactor(mpsc): clarify bounded state and capacity limits --- CHANGELOG.md | 1 - asyncband/src/mpsc/bounded/mod.rs | 33 ++++++++---------- asyncband/src/mpsc/bounded/receiver.rs | 18 +++++----- asyncband/src/mpsc/bounded/sender.rs | 42 ++++++++++------------- tests-integration/tests/mpsc_test/main.rs | 10 ++++-- 5 files changed, 49 insertions(+), 55 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 72b1515b..7113dc6b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -16,7 +16,6 @@ All notable changes to this project will be documented in this file. ### Improvements -* Reject bounded MPSC capacities above `usize::MAX >> 1` up front with an explicit panic message instead of an opaque arithmetic overflow. * Finish releasing buffered bounded MPSC messages even if one message destructor panics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index b1b75449..a9ae677b 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -49,27 +49,21 @@ pub use self::sender::Permit; /// /// # Panics /// -/// Panics if `buffer` is zero or exceeds the maximum capacity of `usize::MAX >> 1`. +/// Panics if `buffer` is zero or the preallocated message storage exceeds the allocation size +/// limit. #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { - const MAX_CAPACITY: usize = usize::MAX >> 1; - assert!( buffer > 0, "mpsc bounded channel capacity {buffer} must be nonzero", ); - assert!( - buffer <= MAX_CAPACITY, - "mpsc bounded channel capacity {buffer} exceeds the maximum of {MAX_CAPACITY}", - ); - let shared = Arc::new(Mutex::new(State { queue: VecDeque::with_capacity(buffer), available: buffer, senders: 1, - receiver_open: true, - receiver_waker: None, - waiters: WaitList::new(), + receiver: true, + recv_waker: None, + send_waiters: WaitList::new(), })); ( BoundedSender::new(shared.clone()), @@ -83,14 +77,15 @@ struct State { queue: VecDeque, available: usize, senders: usize, - receiver_open: bool, - receiver_waker: Option, - waiters: WaitList, + // True while the receiving endpoint is alive. + receiver: bool, + recv_waker: Option, + send_waiters: WaitList, } impl State { fn acquire(&mut self) -> Result<(), TrySendError<()>> { - if !self.receiver_open { + if !self.receiver { Err(TrySendError::Disconnected(())) } else if self.available == 0 { Err(TrySendError::Full(())) @@ -101,12 +96,12 @@ impl State { } fn release(&mut self) -> Option { - if !self.receiver_open { + if !self.receiver { return None; } - if let Some((_, waiter)) = self.waiters.unlink_first_waiter(|_| true) { + if let Some((_, waiter)) = self.send_waiters.unlink_first_waiter(|_| true) { // The detached node owns capacity until its future claims or cancels the grant. - waiter.granted = true; + waiter.grant = true; return waiter.waker.take(); } self.available += 1; @@ -125,6 +120,6 @@ impl State { } struct Waiter { - granted: bool, + grant: bool, waker: Option, } diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs index c97315a4..09997183 100644 --- a/asyncband/src/mpsc/bounded/receiver.rs +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -45,22 +45,22 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - let (queue, receiver_waker, wakers) = { + let (queue, recv_waker, wakers) = { let mut state = self.shared.lock(); - state.receiver_open = false; + state.receiver = false; let queue = mem::take(&mut state.queue); - let receiver_waker = state.receiver_waker.take(); + let recv_waker = state.recv_waker.take(); let mut wakers = WakerBatch::new(); - while let Some((_, waiter)) = state.waiters.unlink_first_waiter(|_| true) { + while let Some((_, waiter)) = state.send_waiters.unlink_first_waiter(|_| true) { if let Some(waker) = waiter.waker.take() { wakers.push(waker); } } - (queue, receiver_waker, wakers) + (queue, recv_waker, wakers) }; // Local ownership also drains the queue if a wake or waker destructor unwinds. wake_all(wakers.into_iter()); - drop(receiver_waker); + drop(recv_waker); drop(queue); } } @@ -146,7 +146,7 @@ impl BoundedReceiver { return Poll::Ready(Ok(value)); } Err(TryRecvError::Disconnected) => { - let old = state.receiver_waker.take(); + let old = state.recv_waker.take(); drop(state); drop(old); return Poll::Ready(Err(RecvError::Disconnected)); @@ -154,14 +154,14 @@ impl BoundedReceiver { Err(TryRecvError::Empty) => {} } if state - .receiver_waker + .recv_waker .as_ref() .is_some_and(|w| w.will_wake(cx.waker())) { return Poll::Pending; } if let Some(waker) = cloned_waker.take() { - let old = state.receiver_waker.replace(waker); + let old = state.recv_waker.replace(waker); drop(state); drop(old); return Poll::Pending; diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index 4a174c1b..92018a6b 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -57,7 +57,7 @@ impl Drop for BoundedSender { let mut state = self.shared.lock(); state.senders -= 1; if state.senders == 0 { - state.receiver_waker.take() + state.recv_waker.take() } else { None } @@ -89,7 +89,7 @@ impl BoundedSender { match state.acquire() { Ok(()) => { state.queue.push_back(value); - let wake = state.receiver_waker.take(); + let wake = state.recv_waker.take(); drop(state); if let Some(waker) = wake { waker.wake(); @@ -137,11 +137,11 @@ impl BoundedSender { /// # } /// ``` pub async fn reserve(&self) -> Result, SendError<()>> { - let mut reservation = Reservation { + let mut reserve = Reserve { shared: &self.shared, waiter: None, }; - poll_fn(|cx| reservation.poll(cx)).await + poll_fn(|cx| reserve.poll(cx)).await } /// Reserves capacity for one message without waiting. @@ -180,7 +180,7 @@ impl BoundedSender { match state.acquire() { Ok(()) => { state.queue.push_back(value); - let wake = state.receiver_waker.take(); + let wake = state.recv_waker.take(); drop(state); if let Some(waker) = wake { waker.wake(); @@ -215,13 +215,13 @@ impl Permit<'_, T> { /// If the receiver has been dropped, the returned error contains the unsent value. pub fn send(self, value: T) -> Result<(), SendError> { let mut state = self.shared.lock(); - if !state.receiver_open { + if !state.receiver { return Err(SendError::new(value)); } state.queue.push_back(value); // The queued message now owns capacity, even if the wake callback panics. mem::forget(self); - let wake = state.receiver_waker.take(); + let wake = state.recv_waker.take(); drop(state); if let Some(waker) = wake { waker.wake(); @@ -239,24 +239,24 @@ impl Drop for Permit<'_, T> { } } -struct Reservation<'a, T> { +struct Reserve<'a, T> { shared: &'a Mutex>, waiter: Option, } -impl<'a, T> Reservation<'a, T> { +impl<'a, T> Reserve<'a, T> { fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { let mut cloned_waker = None; let result = loop { let mut state = self.shared.lock(); - if !state.receiver_open { + if !state.receiver { // Drop removes any remaining registration, including an unused grant. break Err(SendError::new(())); } if let Some(index) = self.waiter { - let waiter = state.waiters.waiter_mut(index); - if waiter.granted { - let waiter = state.waiters.remove_unlinked_waiter(index); + let waiter = state.send_waiters.waiter_mut(index); + if waiter.grant { + let waiter = state.send_waiters.remove_unlinked_waiter(index); self.waiter = None; let permit = Permit { shared: self.shared, @@ -284,8 +284,8 @@ impl<'a, T> Reservation<'a, T> { shared: self.shared, }); } else if let Some(waker) = cloned_waker.take() { - self.waiter = Some(state.waiters.push_back(Waiter { - granted: false, + self.waiter = Some(state.send_waiters.push_back(Waiter { + grant: false, waker: Some(waker), })); return Poll::Pending; @@ -300,18 +300,14 @@ impl<'a, T> Reservation<'a, T> { } } -impl Drop for Reservation<'_, T> { +impl Drop for Reserve<'_, T> { fn drop(&mut self) { let Some(index) = self.waiter else { return }; let (waiter, wake) = { let mut state = self.shared.lock(); - state.waiters.unlink_waiter(index, |_| true); - let waiter = state.waiters.remove_unlinked_waiter(index); - let wake = if waiter.granted { - state.release() - } else { - None - }; + state.send_waiters.unlink_waiter(index, |_| true); + let waiter = state.send_waiters.remove_unlinked_waiter(index); + let wake = if waiter.grant { state.release() } else { None }; (waiter, wake) }; if let Some(waker) = wake { diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index 0dad066b..f2ab74f5 100644 --- a/tests-integration/tests/mpsc_test/main.rs +++ b/tests-integration/tests/mpsc_test/main.rs @@ -182,7 +182,11 @@ fn bounded_rejects_zero_capacity() { } #[test] -#[should_panic(expected = "exceeds the maximum")] -fn bounded_rejects_capacity_above_the_maximum() { - let _ = mpsc::bounded::((usize::MAX >> 1) + 1); +fn bounded_supports_full_usize_capacity_for_zero_sized_messages() { + let (tx, mut rx) = mpsc::bounded::<()>(usize::MAX); + // Returning capacity at this boundary must not overflow the counter. + drop(tx.try_reserve().unwrap()); + tx.try_reserve().unwrap().send(()).unwrap(); + assert_eq!(rx.try_recv(), Ok(())); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } From 824bac1bd0dc8da667e359ed35320fd145ece3cb Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 18:32:44 +0800 Subject: [PATCH 33/34] docs(mpsc): clarify reservation errors and method links --- asyncband/src/mpsc/bounded/sender.rs | 6 ++--- asyncband/src/mpsc/error.rs | 34 ++++++++++++++++------------ 2 files changed, 22 insertions(+), 18 deletions(-) diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index 92018a6b..eab76a0f 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -80,9 +80,9 @@ impl BoundedSender { /// # Cancel safety /// /// Dropping a pending `send` loses its place waiting for capacity and drops `value`; a call - /// that has returned `Pending` has not sent the message. Use [`Self::try_send`] when the - /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for - /// capacity before constructing the message. + /// that has returned `Pending` has not sent the message. Use [`try_send`](Self::try_send) when + /// the caller must retain ownership if capacity is unavailable, or [`reserve`](Self::reserve) + /// to wait for capacity before constructing the message. pub async fn send(&self, value: T) -> Result<(), SendError> { { let mut state = self.shared.lock(); diff --git a/asyncband/src/mpsc/error.rs b/asyncband/src/mpsc/error.rs index 586e4e90..69e545e7 100644 --- a/asyncband/src/mpsc/error.rs +++ b/asyncband/src/mpsc/error.rs @@ -18,29 +18,29 @@ use std::any::type_name; use std::fmt; -/// A send failed because the receiving endpoint has been dropped. +/// A send or capacity reservation failed because the receiver has been dropped. /// -/// Returned from [`UnboundedSender::send`] or [`BoundedSender::send`] if the -/// corresponding [`UnboundedReceiver`] or [`BoundedReceiver`] has already been -/// dropped. +/// Returned by [`UnboundedSender::send`], [`BoundedSender::send`], [`Permit::send`], and +/// [`reserve`]. /// -/// The rejected message remains available through [`SendError::as_inner`] or -/// [`SendError::into_inner`]. +/// A failed send retains the unsent message. A failed reservation carries `()` because no +/// message has been provided yet. Access the value with [`as_inner`](Self::as_inner) or +/// [`into_inner`](Self::into_inner). /// /// [`UnboundedSender::send`]: crate::mpsc::UnboundedSender::send /// [`BoundedSender::send`]: crate::mpsc::BoundedSender::send -/// [`UnboundedReceiver`]: crate::mpsc::UnboundedReceiver -/// [`BoundedReceiver`]: crate::mpsc::BoundedReceiver +/// [`reserve`]: crate::mpsc::BoundedSender::reserve +/// [`Permit::send`]: crate::mpsc::Permit::send #[derive(Clone, PartialEq, Eq)] pub struct SendError(T); impl SendError { - /// Get a reference to the message that failed to be sent. + /// Gets a reference to the unsent message, or `()` for a failed reservation. pub fn as_inner(&self) -> &T { &self.0 } - /// Consumes the error and returns the message that failed to be sent. + /// Consumes the error and returns the unsent message, or `()` for a failed reservation. pub fn into_inner(self) -> T { self.0 } @@ -65,24 +65,28 @@ impl fmt::Debug for SendError { impl std::error::Error for SendError {} -/// A non-blocking send could not accept its message. +/// An attempt to send or reserve capacity failed. +/// +/// Returned by [`try_send`](crate::mpsc::BoundedSender::try_send) and +/// [`try_reserve`](crate::mpsc::BoundedSender::try_reserve). A failed send retains the unsent +/// message; a failed reservation carries `()` because no message has been provided yet. #[derive(Clone, PartialEq, Eq)] pub enum TrySendError { - /// The channel is full, so the message cannot be sent without waiting for capacity. + /// No capacity is available for sending or reserving a message. Full(T), - /// The receiver has been dropped, so the message can never be received. + /// The receiver has been dropped. Disconnected(T), } impl TrySendError { - /// Gets a reference to the message that failed to be sent. + /// Gets a reference to the unsent message, or `()` for a failed reservation. pub fn as_inner(&self) -> &T { match self { TrySendError::Full(msg) | TrySendError::Disconnected(msg) => msg, } } - /// Consumes the error and returns the message that failed to be sent. + /// Consumes the error and returns the unsent message, or `()` for a failed reservation. pub fn into_inner(self) -> T { match self { TrySendError::Full(msg) | TrySendError::Disconnected(msg) => msg, From 6cb4d6b74940a3f64883e638d27945e2f1cac56e Mon Sep 17 00:00:00 2001 From: tison Date: Wed, 9 Sep 2026 18:32:57 +0800 Subject: [PATCH 34/34] refactor(mpsc): simplify bounded send and wake paths Reuse try_send for ready sends and clone wakers before locking channel state. Remove redundant ZST lifecycle coverage while keeping the maximum-capacity boundary test. --- asyncband/src/mpsc/bounded/receiver.rs | 43 +++----- asyncband/src/mpsc/bounded/sender.rs | 103 +++++++----------- .../tests/mpsc_test/reservation.rs | 56 ---------- 3 files changed, 54 insertions(+), 148 deletions(-) diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs index 09997183..8b46e059 100644 --- a/asyncband/src/mpsc/bounded/receiver.rs +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -134,41 +134,28 @@ impl BoundedReceiver { } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - let mut cloned_waker = None; - loop { - let mut state = self.shared.lock(); - match state.pop() { - Ok((value, wake)) => { - drop(state); - if let Some(waker) = wake { - waker.wake(); - } - return Poll::Ready(Ok(value)); - } - Err(TryRecvError::Disconnected) => { - let old = state.recv_waker.take(); - drop(state); - drop(old); - return Poll::Ready(Err(RecvError::Disconnected)); + let waker = cx.waker().clone(); + let mut state = self.shared.lock(); + match state.pop() { + Ok((value, wake)) => { + drop(state); + if let Some(waker) = wake { + waker.wake(); } - Err(TryRecvError::Empty) => {} + Poll::Ready(Ok(value)) } - if state - .recv_waker - .as_ref() - .is_some_and(|w| w.will_wake(cx.waker())) - { - return Poll::Pending; + Err(TryRecvError::Disconnected) => { + let old = state.recv_waker.take(); + drop(state); + drop(old); + Poll::Ready(Err(RecvError::Disconnected)) } - if let Some(waker) = cloned_waker.take() { + Err(TryRecvError::Empty) => { let old = state.recv_waker.replace(waker); drop(state); drop(old); - return Poll::Pending; + Poll::Pending } - drop(state); - // Clone can reenter the channel, so check the queue again after acquiring the lock. - cloned_waker = Some(cx.waker().clone()); } } } diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs index eab76a0f..392fa159 100644 --- a/asyncband/src/mpsc/bounded/sender.rs +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -84,22 +84,11 @@ impl BoundedSender { /// the caller must retain ownership if capacity is unavailable, or [`reserve`](Self::reserve) /// to wait for capacity before constructing the message. pub async fn send(&self, value: T) -> Result<(), SendError> { - { - let mut state = self.shared.lock(); - match state.acquire() { - Ok(()) => { - state.queue.push_back(value); - let wake = state.recv_waker.take(); - drop(state); - if let Some(waker) = wake { - waker.wake(); - } - return Ok(()); - } - Err(TrySendError::Disconnected(())) => return Err(SendError::new(value)), - Err(TrySendError::Full(())) => {} - } - } + let value = match self.try_send(value) { + Ok(()) => return Ok(()), + Err(TrySendError::Disconnected(value)) => return Err(SendError::new(value)), + Err(TrySendError::Full(value)) => value, + }; match self.reserve().await { Ok(permit) => permit.send(value), Err(_) => Err(SendError::new(value)), @@ -246,57 +235,43 @@ struct Reserve<'a, T> { impl<'a, T> Reserve<'a, T> { fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { - let mut cloned_waker = None; - let result = loop { - let mut state = self.shared.lock(); - if !state.receiver { - // Drop removes any remaining registration, including an unused grant. - break Err(SendError::new(())); - } - if let Some(index) = self.waiter { - let waiter = state.send_waiters.waiter_mut(index); - if waiter.grant { - let waiter = state.send_waiters.remove_unlinked_waiter(index); - self.waiter = None; - let permit = Permit { - shared: self.shared, - }; - drop(state); - drop(waiter); - break Ok(permit); - } - if waiter - .waker - .as_ref() - .is_some_and(|w| w.will_wake(cx.waker())) - { - return Poll::Pending; - } - if let Some(waker) = cloned_waker.take() { - let old = waiter.waker.replace(waker); - drop(state); - drop(old); - return Poll::Pending; - } - } else if state.available != 0 { - state.available -= 1; - break Ok(Permit { + let waker = cx.waker().clone(); + let mut state = self.shared.lock(); + if !state.receiver { + return Poll::Ready(Err(SendError::new(()))); + } + if let Some(index) = self.waiter { + let waiter = state.send_waiters.waiter_mut(index); + if waiter.grant { + let waiter = state.send_waiters.remove_unlinked_waiter(index); + self.waiter = None; + let permit = Permit { shared: self.shared, - }); - } else if let Some(waker) = cloned_waker.take() { - self.waiter = Some(state.send_waiters.push_back(Waiter { - grant: false, - waker: Some(waker), - })); - return Poll::Pending; + }; + drop(state); + drop(waiter); + drop(waker); + return Poll::Ready(Ok(permit)); } + let old = waiter.waker.replace(waker); drop(state); - // Clone outside the lock, then recheck capacity and closure before registering. - cloned_waker = Some(cx.waker().clone()); - }; - // The permit already owns capacity if dropping an unused clone unwinds. - drop(cloned_waker); - Poll::Ready(result) + drop(old); + return Poll::Pending; + } + if state.available != 0 { + state.available -= 1; + let permit = Permit { + shared: self.shared, + }; + drop(state); + drop(waker); + return Poll::Ready(Ok(permit)); + } + self.waiter = Some(state.send_waiters.push_back(Waiter { + grant: false, + waker: Some(waker), + })); + Poll::Pending } } diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs index 7c71c813..7d2a0741 100644 --- a/tests-integration/tests/mpsc_test/reservation.rs +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -60,62 +60,6 @@ fn held_permits_consume_capacity_without_claiming_message_order() { } } -#[test] -fn zero_sized_messages_preserve_capacity_across_reservation_and_close() { - for capacity in [1, 3, 64] { - let (tx, mut rx) = mpsc::bounded::<()>(capacity); - let permit = tx.try_reserve().unwrap(); - for _ in 1..capacity { - tx.try_send(()).unwrap(); - } - assert_eq!(tx.try_send(()), Err(TrySendError::Full(()))); - for _ in 1..capacity { - assert_eq!(rx.try_recv(), Ok(())); - } - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - drop(permit); - tx.try_reserve().unwrap().send(()).unwrap(); - assert_eq!(rx.try_recv(), Ok(())); - // An outstanding permit can be dropped after the receiver closes. - let held = tx.try_reserve().unwrap(); - for _ in 1..capacity { - tx.try_send(()).unwrap(); - } - drop(rx); - drop(held); - assert!(matches!( - tx.try_reserve(), - Err(TrySendError::Disconnected(())) - )); - } -} - -#[test] -fn zero_sized_messages_are_dropped_once_when_received_or_discarded() { - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - - static DROPS: AtomicUsize = AtomicUsize::new(0); - #[repr(align(128))] - struct Message; - impl Drop for Message { - fn drop(&mut self) { - DROPS.fetch_add(1, Ordering::Relaxed); - } - } - - let (tx, mut rx) = mpsc::bounded(3); - for _ in 0..3 { - assert!(tx.try_send(Message).is_ok()); - } - drop(rx.try_recv().unwrap()); - assert_eq!(DROPS.load(Ordering::Relaxed), 1); - drop(rx); - assert_eq!(DROPS.load(Ordering::Relaxed), 3); - drop(tx.try_send(Message).err().unwrap().into_inner()); - assert_eq!(DROPS.load(Ordering::Relaxed), 4); -} - #[test] fn released_capacity_is_granted_to_the_oldest_waiter() { let (tx, mut rx) = mpsc::bounded(1);