From cbc3381412308f20f358cf3bbe29ecd8cb6ecd6e Mon Sep 17 00:00:00 2001 From: mxsm Date: Sun, 30 Aug 2026 21:28:30 +0800 Subject: [PATCH 01/16] refactor(mpsc): replace the legacy queue backend --- CHANGELOG.md | 1 + asyncband/src/mpsc/bounded.rs | 58 ++--- asyncband/src/mpsc/mod.rs | 1 + asyncband/src/mpsc/queue.rs | 352 +++++++++++++++++++++++++++ asyncband/src/mpsc/unbounded.rs | 57 +++-- tests-integration/tests/mpsc_test.rs | 38 +++ 6 files changed, 449 insertions(+), 58 deletions(-) create mode 100644 asyncband/src/mpsc/queue.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 214b359c..3ca7d5ed 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -37,3 +37,4 @@ All notable changes to this project will be documented in this file. * Remove the `slab` dependency in favor of a focused internal waiter arena. * Describe disconnected channel states consistently in channel error messages. +* Replace the legacy standard-library MPSC backend with an owned queue core and remove the receiver types' manual `Sync` implementations. diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index c4d4bcf2..157c90b0 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -32,6 +32,8 @@ use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; +use super::queue::BoundedQueue; +use super::queue::PushError; use crate::internal::atomic_waker::AtomicWaker; use crate::internal::semaphore::Acquire; use crate::internal::semaphore::Semaphore; @@ -46,23 +48,20 @@ use crate::internal::semaphore::Semaphore; pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); let state = Arc::new(BoundedState { + queue: BoundedQueue::new(buffer), senders: AtomicUsize::new(1), tx_permits: Semaphore::new(0), rx_waker: AtomicWaker::new(), }); - let (sender, receiver) = std::sync::mpsc::sync_channel(buffer); let sender = BoundedSender { state: state.clone(), - sender: Some(sender), - }; - let receiver = BoundedReceiver { - state: state.clone(), - receiver: Some(receiver), }; + let receiver = BoundedReceiver { state }; (sender, receiver) } -struct BoundedState { +struct BoundedState { + queue: BoundedQueue, senders: AtomicUsize, tx_permits: Semaphore, rx_waker: AtomicWaker, @@ -72,8 +71,7 @@ struct BoundedState { /// /// Instances are created by the [`bounded`] function. pub struct BoundedSender { - state: Arc, - sender: Option>, + state: Arc>, } impl Clone for BoundedSender { @@ -81,7 +79,6 @@ impl Clone for BoundedSender { self.state.senders.fetch_add(1, Ordering::Release); BoundedSender { state: self.state.clone(), - sender: self.sender.clone(), } } } @@ -94,9 +91,6 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { - // Dropping the final underlying sender disconnects the channel. - drop(self.sender.take()); - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { 1 => { // Wake the receiver so it can observe the channel's disconnected state. @@ -193,18 +187,14 @@ impl BoundedSender { /// # } /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - // SAFETY: The sender is guaranteed to be non-null before dropped. - let sender = self.sender.as_ref().unwrap(); - match sender.try_send(value) { + match self.state.queue.try_push(value) { Ok(()) => { self.state.rx_waker.wake(); Ok(()) } - Err(std::sync::mpsc::TrySendError::Full(value)) => Err(TrySendError::Full(value)), - Err(std::sync::mpsc::TrySendError::Disconnected(value)) => { - Err(TrySendError::Disconnected(value)) - } + Err(PushError::Full(value)) => Err(TrySendError::Full(value)), + Err(PushError::Disconnected(value)) => Err(TrySendError::Disconnected(value)), } } } @@ -213,14 +203,9 @@ impl BoundedSender { /// /// Instances are created by the [`bounded`] function. pub struct BoundedReceiver { - state: Arc, - receiver: Option>, + state: Arc>, } -/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`. -/// That is, `BoundedReceiver` can only be accessed by one thread at a time. -unsafe impl Sync for BoundedReceiver {} - impl fmt::Debug for BoundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BoundedReceiver").finish_non_exhaustive() @@ -229,7 +214,7 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - drop(self.receiver.take()); + self.state.queue.disconnect_receiver(); self.state.tx_permits.notify_all(); } } @@ -270,15 +255,20 @@ impl BoundedReceiver { /// # } /// ``` pub fn try_recv(&mut self) -> Result { - // SAFETY: The receiver is guaranteed to be non-null before dropped. - let receiver = self.receiver.as_ref().unwrap(); - match receiver.try_recv() { - Ok(v) => { + if let Some(value) = self.state.queue.pop() { + self.state.tx_permits.release_if_nonempty(1); + Ok(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. + if let Some(value) = self.state.queue.pop() { self.state.tx_permits.release_if_nonempty(1); - Ok(v) + Ok(value) + } else { + Err(TryRecvError::Disconnected) } - Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected), - Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty), + } else { + Err(TryRecvError::Empty) } } diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs index 45bc7406..a1fd4d9d 100644 --- a/asyncband/src/mpsc/mod.rs +++ b/asyncband/src/mpsc/mod.rs @@ -19,6 +19,7 @@ mod bounded; mod error; +mod queue; mod unbounded; pub use self::bounded::BoundedReceiver; diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs new file mode 100644 index 00000000..b0838cd1 --- /dev/null +++ b/asyncband/src/mpsc/queue.rs @@ -0,0 +1,352 @@ +// 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::collections::VecDeque; +use std::hint::spin_loop; +use std::mem; +use std::mem::MaybeUninit; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::sync::atomic::fence; + +use crate::internal::mutex::Mutex; + +pub(super) struct UnboundedQueue { + inner: Mutex>, +} + +struct UnboundedInner { + // Storage and receiver liveness share one lock so a send is linearized either before receiver + // disconnection, with its value in the queue, or after it, with the value returned to sender. + messages: VecDeque, + receiver_alive: bool, +} + +pub(super) struct UnboundedConsumer { + local: Mutex>, +} + +pub(super) enum PushError { + Full(T), + Disconnected(T), +} + +impl UnboundedQueue { + pub(super) const fn new() -> Self { + Self { + inner: Mutex::new(UnboundedInner { + messages: VecDeque::new(), + receiver_alive: true, + }), + } + } + + pub(super) fn push(&self, value: T) -> Result<(), PushError> { + let mut inner = self.inner.lock(); + if !inner.receiver_alive { + return Err(PushError::Disconnected(value)); + } + inner.messages.push_back(value); + Ok(()) + } + + pub(super) fn pop(&self, consumer: &UnboundedConsumer) -> Option { + let mut local = consumer.local.lock(); + if local.is_empty() { + let mut inner = self.inner.lock(); + mem::swap(&mut *local, &mut inner.messages); + } + local.pop_front() + } + + pub(super) fn disconnect_receiver(&self, consumer: &UnboundedConsumer) { + let (local, shared) = { + let mut local = consumer.local.lock(); + let mut inner = self.inner.lock(); + inner.receiver_alive = false; + (mem::take(&mut *local), mem::take(&mut inner.messages)) + }; + drop((local, shared)); + } +} + +impl UnboundedConsumer { + pub(super) const fn new() -> Self { + Self { + local: Mutex::new(VecDeque::new()), + } + } +} + +pub(super) struct BoundedQueue { + slots: Box<[Slot]>, + head: CachePadded, + tail: CachePadded, + capacity: usize, + one_lap: usize, + mark_bit: usize, +} + +#[repr(align(64))] +struct CachePadded(T); + +impl std::ops::Deref for CachePadded { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +struct Slot { + stamp: AtomicUsize, + value: UnsafeCell>, +} + +// SAFETY: A successful tail CAS 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. +unsafe impl Sync for Slot {} + +impl BoundedQueue { + pub(super) 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), + value: UnsafeCell::new(MaybeUninit::uninit()), + }) + .collect(); + Self { + slots, + head: CachePadded(AtomicUsize::new(0)), + tail: CachePadded(AtomicUsize::new(0)), + capacity, + one_lap, + mark_bit, + } + } + + pub(super) fn try_push(&self, value: T) -> Result<(), PushError> { + let mut tail = self.tail.load(Ordering::Relaxed); + let mut backoff = 0; + loop { + if tail & self.mark_bit != 0 { + return Err(PushError::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); + 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(PushError::Full(value)); + } + tail = self.tail.load(Ordering::Relaxed); + } else { + tail = self.tail.load(Ordering::Relaxed); + } + Self::spin(&mut backoff); + } + } + + pub(super) fn pop(&self) -> Option { + 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 Some(value); + } + + if stamp == head { + fence(Ordering::SeqCst); + if self.tail.load(Ordering::Relaxed) & !self.mark_bit == head { + return None; + } + } + if backoff == 8 { + return None; + } + Self::spin(&mut backoff); + head = self.head.load(Ordering::Relaxed); + } + } + + pub(super) fn disconnect_receiver(&self) { + let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; + self.discard_until(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) + } + } + + fn discard_until(&self, tail: usize) { + let mut head = self.head.load(Ordering::Relaxed); + let mut backoff = 0; + while head != tail { + let index = head & (self.mark_bit - 1); + 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. + unsafe { (*slot.value.get()).assume_init_drop() }; + head = next_head; + 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); + } +} + +impl Drop for BoundedQueue { + fn drop(&mut self) { + let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; + self.discard_until(tail); + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::thread; + + use super::BoundedQueue; + use super::PushError; + use super::UnboundedConsumer; + use super::UnboundedQueue; + + #[test] + fn bounded_queue_preserves_capacity_and_fifo_order() { + let queue = BoundedQueue::new(3); + for value in 0..3 { + assert!(queue.try_push(value).is_ok()); + } + assert!(matches!(queue.try_push(3), Err(PushError::Full(3)))); + for value in 0..3 { + assert_eq!(queue.pop(), Some(value)); + } + assert_eq!(queue.pop(), None); + + for value in 3..12 { + assert!(queue.try_push(value).is_ok()); + assert_eq!(queue.pop(), Some(value)); + } + } + + #[test] + fn bounded_queue_coordinates_multiple_producers() { + let queue = Arc::new(BoundedQueue::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(PushError::Full(returned)) => { + value = returned; + thread::yield_now(); + } + Err(PushError::Disconnected(_)) => panic!("queue disconnected"), + } + } + } + }) + }) + .collect(); + + let mut values = Vec::new(); + while values.len() < 64 { + if let Some(value) = 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::>()); + } + + #[test] + fn unbounded_queue_batches_without_reordering() { + let queue = UnboundedQueue::new(); + let consumer = UnboundedConsumer::new(); + assert!(queue.push(1).is_ok()); + assert!(queue.push(2).is_ok()); + assert_eq!(queue.pop(&consumer), Some(1)); + assert!(queue.push(3).is_ok()); + assert_eq!(queue.pop(&consumer), Some(2)); + assert_eq!(queue.pop(&consumer), Some(3)); + assert_eq!(queue.pop(&consumer), None); + } +} diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index 57552d2a..542e1c7c 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -29,6 +29,9 @@ use std::task::Poll; use super::RecvError; use super::SendError; use super::TryRecvError; +use super::queue::PushError; +use super::queue::UnboundedConsumer; +use super::queue::UnboundedQueue; use crate::internal::atomic_waker::AtomicWaker; /// Creates an unbounded mpsc channel for communicating between asynchronous @@ -42,22 +45,22 @@ use crate::internal::atomic_waker::AtomicWaker; /// process to run out of memory. In this case, the process will be aborted. pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { + queue: UnboundedQueue::new(), senders: AtomicUsize::new(1), rx_waker: AtomicWaker::new(), }); - let (sender, receiver) = std::sync::mpsc::channel(); let sender = UnboundedSender { state: state.clone(), - sender: Some(sender), }; let receiver = UnboundedReceiver { - state: state.clone(), - receiver, + state, + consumer: UnboundedConsumer::new(), }; (sender, receiver) } -struct UnboundedState { +struct UnboundedState { + queue: UnboundedQueue, senders: AtomicUsize, rx_waker: AtomicWaker, } @@ -66,8 +69,7 @@ struct UnboundedState { /// /// Instances are created by the [`unbounded`] function. pub struct UnboundedSender { - state: Arc, - sender: Option>, + state: Arc>, } impl Clone for UnboundedSender { @@ -75,7 +77,6 @@ impl Clone for UnboundedSender { self.state.senders.fetch_add(1, Ordering::Release); UnboundedSender { state: self.state.clone(), - sender: self.sender.clone(), } } } @@ -88,9 +89,6 @@ impl fmt::Debug for UnboundedSender { impl Drop for UnboundedSender { fn drop(&mut self) { - // Dropping the final underlying sender disconnects the channel. - drop(self.sender.take()); - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { 1 => { // Wake the receiver so it can observe the channel's disconnected state. @@ -113,9 +111,11 @@ impl UnboundedSender { /// If the receiver has been dropped, this function returns an error. The error includes /// the value passed to `send`. pub fn send(&self, value: T) -> Result<(), SendError> { - // SAFETY: The sender is guaranteed to be non-null before dropped. - let sender = self.sender.as_ref().unwrap(); - sender.send(value).map_err(|err| SendError::new(err.0))?; + match self.state.queue.push(value) { + Ok(()) => {} + Err(PushError::Disconnected(value)) => return Err(SendError::new(value)), + Err(PushError::Full(_)) => unreachable!("unbounded queue cannot be full"), + } self.state.rx_waker.wake(); @@ -127,20 +127,22 @@ impl UnboundedSender { /// /// Instances are created by the [`unbounded`] function. pub struct UnboundedReceiver { - state: Arc, - receiver: std::sync::mpsc::Receiver, + state: Arc>, + consumer: UnboundedConsumer, } -/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`. -/// That is, `UnboundedReceiver` can only be accessed by one thread at a time. -unsafe impl Sync for UnboundedReceiver {} - impl fmt::Debug for UnboundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("UnboundedReceiver").finish_non_exhaustive() } } +impl Drop for UnboundedReceiver { + fn drop(&mut self) { + self.state.queue.disconnect_receiver(&self.consumer); + } +} + impl UnboundedReceiver { /// Tries to receive the next value for this receiver. /// @@ -176,10 +178,17 @@ impl UnboundedReceiver { /// # } /// ``` pub fn try_recv(&mut self) -> Result { - match self.receiver.try_recv() { - Ok(v) => Ok(v), - Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected), - Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty), + if let Some(value) = self.state.queue.pop(&self.consumer) { + Ok(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. + self.state + .queue + .pop(&self.consumer) + .ok_or(TryRecvError::Disconnected) + } else { + Err(TryRecvError::Empty) } } diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index ed132d36..8ef59ff7 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -209,6 +209,44 @@ fn try_recv_reports_disconnection_while_empty_unbounded() { assert_eq!(Err(TryRecvError::Disconnected), rx.try_recv()); } +#[test] +fn cancelled_receive_does_not_consume_a_later_message() { + let (unbounded_tx, mut unbounded_rx) = mpsc::unbounded(); + { + let mut receive = Box::pin(unbounded_rx.recv()); + assert!(poll_once(receive.as_mut()).is_pending()); + } + unbounded_tx.send(1).unwrap(); + assert_eq!(unbounded_rx.try_recv(), Ok(1)); + + let (bounded_tx, mut bounded_rx) = mpsc::bounded(1); + { + let mut receive = Box::pin(bounded_rx.recv()); + assert!(poll_once(receive.as_mut()).is_pending()); + } + bounded_tx.try_send(2).unwrap(); + assert_eq!(bounded_rx.try_recv(), Ok(2)); +} + +#[test] +fn buffered_messages_are_drained_before_disconnection() { + let (unbounded_tx, mut unbounded_rx) = mpsc::unbounded(); + unbounded_tx.send(1).unwrap(); + unbounded_tx.send(2).unwrap(); + drop(unbounded_tx); + assert_eq!(unbounded_rx.try_recv(), Ok(1)); + assert_eq!(unbounded_rx.try_recv(), Ok(2)); + assert_eq!(unbounded_rx.try_recv(), Err(TryRecvError::Disconnected)); + + let (bounded_tx, mut bounded_rx) = mpsc::bounded(2); + bounded_tx.try_send(3).unwrap(); + bounded_tx.try_send(4).unwrap(); + drop(bounded_tx); + assert_eq!(bounded_rx.try_recv(), Ok(3)); + assert_eq!(bounded_rx.try_recv(), Ok(4)); + assert_eq!(bounded_rx.try_recv(), Err(TryRecvError::Disconnected)); +} + #[tokio::test] async fn send_recv_bounded() { let (tx, mut rx) = mpsc::bounded(1); From 13a100eb1bd57b47007963f42030c667395361a1 Mon Sep 17 00:00:00 2001 From: mxsm Date: Wed, 2 Sep 2026 13:02:08 +0800 Subject: [PATCH 02/16] fix(mpsc): restore queue endpoint traits --- asyncband/src/mpsc/queue.rs | 69 ++++++++++++++++++++++++++ tests-integration/tests/traits_test.rs | 37 ++++++++++++++ 2 files changed, 106 insertions(+) diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs index b0838cd1..52c75be0 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/queue.rs @@ -41,6 +41,9 @@ pub(super) struct UnboundedConsumer { local: Mutex>, } +// The consumer never relies on a pinned location for its local queue or queued values. +impl Unpin for UnboundedConsumer {} + pub(super) enum PushError { Full(T), Disconnected(T), @@ -123,6 +126,11 @@ struct Slot { // consumer reads only after acquiring that stamp and publishes the following lap 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 BoundedQueue { pub(super) fn new(capacity: usize) -> Self { assert!(capacity <= usize::MAX / 4, "mpsc capacity is too large"); @@ -273,6 +281,8 @@ impl Drop for BoundedQueue { #[cfg(test)] mod tests { use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; use std::thread; use super::BoundedQueue; @@ -337,6 +347,65 @@ mod tests { assert_eq!(values, (0..64).collect::>()); } + #[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); + } + } + + // 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 = BoundedQueue::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. + let popped = queue.pop(); + assert!(popped.is_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.disconnect_receiver(); + + // `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" + ); + } + } + #[test] fn unbounded_queue_batches_without_reordering() { let queue = UnboundedQueue::new(); diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index 34ff5068..d6340bbc 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -16,6 +16,9 @@ // under the License. use std::cell::Cell; +use std::marker::PhantomPinned; +use std::panic::RefUnwindSafe; +use std::panic::UnwindSafe; use asyncband::barrier::Barrier; use asyncband::broadcast; @@ -180,6 +183,40 @@ fn public_types_are_unpin() { assert_unpin::(); } +#[test] +fn mpsc_endpoints_keep_legacy_traits_regardless_of_payload() { + fn assert_send() {} + fn assert_sync() {} + fn assert_unpin() {} + fn assert_unwind_safe() {} + fn assert_ref_unwind_safe() {} + + macro_rules! assert_endpoint_traits { + ($endpoint:ident, $payload:ty) => { + assert_send::>(); + assert_sync::>(); + assert_unpin::>(); + assert_unwind_safe::>(); + assert_ref_unwind_safe::>(); + }; + } + + macro_rules! assert_payload_traits { + ($endpoint:ident) => { + assert_endpoint_traits!($endpoint, i32); + assert_endpoint_traits!($endpoint, Cell); + assert_endpoint_traits!($endpoint, &'static mut i32); + assert_endpoint_traits!($endpoint, PhantomPinned); + }; + } + + // Four endpoint types × four payloads × five traits = 80 compile-time assertions. + assert_payload_traits!(BoundedSender); + assert_payload_traits!(BoundedReceiver); + assert_payload_traits!(UnboundedSender); + assert_payload_traits!(UnboundedReceiver); +} + #[test] fn unbounded_manual_manager_traits_do_not_depend_on_the_object() { fn assert_copy() {} From 08d3f19eea383f9851df8680b85c6621b0381d66 Mon Sep 17 00:00:00 2001 From: mxsm Date: Fri, 4 Sep 2026 16:46:24 -0700 Subject: [PATCH 03/16] perf(mpsc): choose queue padding by target architecture --- asyncband/src/mpsc/queue.rs | 50 ++++++++++++++++++++++++++----------- 1 file changed, 36 insertions(+), 14 deletions(-) diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs index 52c75be0..d67f2867 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/queue.rs @@ -26,7 +26,7 @@ use std::sync::atomic::fence; use crate::internal::mutex::Mutex; -pub(super) struct UnboundedQueue { +pub struct UnboundedQueue { inner: Mutex>, } @@ -37,20 +37,20 @@ struct UnboundedInner { receiver_alive: bool, } -pub(super) struct UnboundedConsumer { +pub struct UnboundedConsumer { local: Mutex>, } // The consumer never relies on a pinned location for its local queue or queued values. impl Unpin for UnboundedConsumer {} -pub(super) enum PushError { +pub enum PushError { Full(T), Disconnected(T), } impl UnboundedQueue { - pub(super) const fn new() -> Self { + pub const fn new() -> Self { Self { inner: Mutex::new(UnboundedInner { messages: VecDeque::new(), @@ -59,7 +59,7 @@ impl UnboundedQueue { } } - pub(super) fn push(&self, value: T) -> Result<(), PushError> { + pub fn push(&self, value: T) -> Result<(), PushError> { let mut inner = self.inner.lock(); if !inner.receiver_alive { return Err(PushError::Disconnected(value)); @@ -68,7 +68,7 @@ impl UnboundedQueue { Ok(()) } - pub(super) fn pop(&self, consumer: &UnboundedConsumer) -> Option { + pub fn pop(&self, consumer: &UnboundedConsumer) -> Option { let mut local = consumer.local.lock(); if local.is_empty() { let mut inner = self.inner.lock(); @@ -77,7 +77,7 @@ impl UnboundedQueue { local.pop_front() } - pub(super) fn disconnect_receiver(&self, consumer: &UnboundedConsumer) { + pub fn disconnect_receiver(&self, consumer: &UnboundedConsumer) { let (local, shared) = { let mut local = consumer.local.lock(); let mut inner = self.inner.lock(); @@ -89,14 +89,14 @@ impl UnboundedQueue { } impl UnboundedConsumer { - pub(super) const fn new() -> Self { + pub const fn new() -> Self { Self { local: Mutex::new(VecDeque::new()), } } } -pub(super) struct BoundedQueue { +pub struct BoundedQueue { slots: Box<[Slot]>, head: CachePadded, tail: CachePadded, @@ -105,7 +105,29 @@ pub(super) struct BoundedQueue { mark_bit: usize, } -#[repr(align(64))] +// 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)) +)] struct CachePadded(T); impl std::ops::Deref for CachePadded { @@ -132,7 +154,7 @@ impl std::panic::UnwindSafe for Slot {} impl std::panic::RefUnwindSafe for Slot {} impl BoundedQueue { - pub(super) fn new(capacity: usize) -> Self { + 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; @@ -152,7 +174,7 @@ impl BoundedQueue { } } - pub(super) fn try_push(&self, value: T) -> Result<(), PushError> { + pub fn try_push(&self, value: T) -> Result<(), PushError> { let mut tail = self.tail.load(Ordering::Relaxed); let mut backoff = 0; loop { @@ -193,7 +215,7 @@ impl BoundedQueue { } } - pub(super) fn pop(&self) -> Option { + pub fn pop(&self) -> Option { let mut head = self.head.load(Ordering::Relaxed); let mut backoff = 0; loop { @@ -225,7 +247,7 @@ impl BoundedQueue { } } - pub(super) fn disconnect_receiver(&self) { + pub fn disconnect_receiver(&self) { let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; self.discard_until(tail); } From 0f039429c1046cf5061f5b2659479dda07075084 Mon Sep 17 00:00:00 2001 From: mxsm Date: Fri, 4 Sep 2026 16:52:03 -0700 Subject: [PATCH 04/16] docs: omit unverified MPSC performance improvement --- CHANGELOG.md | 4 ---- 1 file changed, 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0bac052f..0543a046 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,10 +4,6 @@ All notable changes to this project will be documented in this file. ## Unreleased -### Improvements - -* Reduce unbounded MPSC queue contention by receiving messages in local batches while preserving FIFO ordering and endpoint compatibility. - ## v0.7.2 ### Improvements From 918214a0956c5a00882d24fcbcd1031e673d2411 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 7 Sep 2026 23:25:08 +0800 Subject: [PATCH 05/16] fix(mpsc): distinguish pending publication from an empty queue --- asyncband/src/mpsc/bounded.rs | 56 ++++++++++++++++++---------- asyncband/src/mpsc/queue.rs | 44 +++++++++++++++++----- tests-integration/tests/mpsc_test.rs | 48 ++++++++++++++++++++++++ 3 files changed, 119 insertions(+), 29 deletions(-) diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index dc707d78..2d3ffe0d 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -27,6 +27,7 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use std::task::ready; use super::RecvError; use super::SendError; @@ -219,7 +220,10 @@ impl Drop for BoundedReceiver { } impl BoundedReceiver { - /// Attempts to receive the next queued value without waiting. + /// 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 @@ -242,21 +246,29 @@ impl BoundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - if let Some(value) = self.state.queue.pop() { - self.state.tx_permits.release_if_nonempty(1); - Ok(value) + loop { + if let Poll::Ready(result) = self.try_recv_once() { + return result; + } + std::thread::yield_now(); + } + } + + fn try_recv_once(&mut self) -> Poll> { + let value = if let Some(value) = ready!(self.state.queue.pop()) { + 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. - if let Some(value) = self.state.queue.pop() { - self.state.tx_permits.release_if_nonempty(1); - Ok(value) - } else { - Err(TryRecvError::Disconnected) - } + let Some(value) = ready!(self.state.queue.pop()) else { + return Poll::Ready(Err(TryRecvError::Disconnected)); + }; + value } else { - Err(TryRecvError::Empty) - } + return Poll::Ready(Err(TryRecvError::Empty)); + }; + self.state.tx_permits.release_if_nonempty(1); + Poll::Ready(Ok(value)) } /// Waits for and receives the next value, freeing one buffer slot. @@ -293,16 +305,20 @@ impl BoundedReceiver { } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => { + 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.rx_waker.register(cx.waker()); - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => Poll::Pending, + 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)) => Poll::Pending, } } } diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs index d67f2867..f4ce7826 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/queue.rs @@ -23,6 +23,7 @@ use std::mem::MaybeUninit; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::sync::atomic::fence; +use std::task::Poll; use crate::internal::mutex::Mutex; @@ -215,7 +216,9 @@ impl BoundedQueue { } } - pub fn pop(&self) -> Option { + // 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. + pub fn pop(&self) -> Poll> { let mut head = self.head.load(Ordering::Relaxed); let mut backoff = 0; loop { @@ -230,17 +233,17 @@ impl BoundedQueue { slot.stamp .store(head.wrapping_add(self.one_lap), Ordering::Release); self.head.store(next_head, Ordering::SeqCst); - return Some(value); + return Poll::Ready(Some(value)); } if stamp == head { fence(Ordering::SeqCst); if self.tail.load(Ordering::Relaxed) & !self.mark_bit == head { - return None; + return Poll::Ready(None); } } if backoff == 8 { - return None; + return Poll::Pending; } Self::spin(&mut backoff); head = self.head.load(Ordering::Relaxed); @@ -305,6 +308,7 @@ mod tests { use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; + use std::task::Poll; use std::thread; use super::BoundedQueue; @@ -320,16 +324,38 @@ mod tests { } assert!(matches!(queue.try_push(3), Err(PushError::Full(3)))); for value in 0..3 { - assert_eq!(queue.pop(), Some(value)); + assert_eq!(queue.pop(), Poll::Ready(Some(value))); } - assert_eq!(queue.pop(), None); + assert_eq!(queue.pop(), Poll::Ready(None)); for value in 3..12 { assert!(queue.try_push(value).is_ok()); - assert_eq!(queue.pop(), Some(value)); + assert_eq!(queue.pop(), Poll::Ready(Some(value))); } } + #[test] + fn bounded_queue_does_not_report_empty_behind_an_unpublished_head() { + let queue = BoundedQueue::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); + let receive = 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); + assert_eq!(queue.pop(), Poll::Ready(Some(1))); + assert_eq!(queue.pop(), Poll::Ready(Some(2))); + assert_eq!(queue.pop(), Poll::Ready(None)); + } + #[test] fn bounded_queue_coordinates_multiple_producers() { let queue = Arc::new(BoundedQueue::new(4)); @@ -356,7 +382,7 @@ mod tests { let mut values = Vec::new(); while values.len() < 64 { - if let Some(value) = queue.pop() { + if let Poll::Ready(Some(value)) = queue.pop() { values.push(value); } else { thread::yield_now(); @@ -397,7 +423,7 @@ mod tests { // Free slot 0, then reuse it on the next lap at position 8. let popped = queue.pop(); - assert!(popped.is_some()); + 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()); diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 5c696f28..087b40f5 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -15,7 +15,10 @@ // specific language governing permissions and limitations // under the License. +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::task::Poll; +use std::thread; use asyncband::mpsc; use asyncband::mpsc::RecvError; @@ -261,6 +264,51 @@ fn bounded_try_send_respects_capacity_and_order() { } } +#[test] +fn bounded_try_recv_does_not_report_empty_after_completed_sends() { + const PRODUCERS: usize = 4; + const MESSAGES_PER_PRODUCER: usize = 16_384; + let (tx, mut rx) = mpsc::bounded(64); + let completed = AtomicUsize::new(0); + let mut premature_empty = 0; + + thread::scope(|scope| { + for producer in 0..PRODUCERS { + let tx = tx.clone(); + let completed = &completed; + scope.spawn(move || { + for sequence in 0..MESSAGES_PER_PRODUCER { + loop { + match tx.try_send((producer, sequence)) { + Ok(()) => break, + Err(TrySendError::Full(_)) => thread::yield_now(), + Err(TrySendError::Disconnected(_)) => panic!("receiver is still alive"), + } + } + completed.fetch_add(1, Ordering::Release); + } + }); + } + + let mut received = 0; + while received < PRODUCERS * MESSAGES_PER_PRODUCER { + // Once more sends have completed than messages received, Empty cannot be correct. + let has_completed_send = completed.load(Ordering::Acquire) > received; + match rx.try_recv() { + Ok(_) => received += 1, + Err(TryRecvError::Empty) => { + premature_empty += usize::from(has_completed_send); + thread::yield_now(); + } + Err(TryRecvError::Disconnected) => panic!("original sender is still alive"), + } + } + }); + + // Drain and join before asserting so a failure cannot strand a producer on a full channel. + assert_eq!(premature_empty, 0); +} + #[tokio::test] async fn try_send_after_disconnection_bounded() { let (tx, rx) = mpsc::bounded(1); From 296eeeb274265049ad3fbf77f334c6f0846059d0 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 7 Sep 2026 23:26:13 +0800 Subject: [PATCH 06/16] refactor(mpsc): make bounded consumer safety requirements explicit --- asyncband/src/mpsc/bounded.rs | 9 ++++-- asyncband/src/mpsc/queue.rs | 56 ++++++++++++++++++++++++----------- 2 files changed, 45 insertions(+), 20 deletions(-) diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index 2d3ffe0d..5d88827c 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -214,7 +214,8 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - self.state.queue.disconnect_receiver(); + // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows. + unsafe { self.state.queue.disconnect_receiver() }; self.state.tx_permits.notify_all(); } } @@ -255,12 +256,14 @@ impl BoundedReceiver { } fn try_recv_once(&mut self) -> Poll> { - let value = if let Some(value) = ready!(self.state.queue.pop()) { + // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows. + let value = if let Some(value) = ready!(unsafe { self.state.queue.pop() }) { 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. - let Some(value) = ready!(self.state.queue.pop()) else { + // SAFETY: The exclusive receiver borrow still guarantees a single consumer. + let Some(value) = ready!(unsafe { self.state.queue.pop() }) else { return Poll::Ready(Err(TryRecvError::Disconnected)); }; value diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs index f4ce7826..add73003 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/queue.rs @@ -216,9 +216,13 @@ impl BoundedQueue { } } - // 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. - pub fn pop(&self) -> Poll> { + /// 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. + /// + /// # Safety + /// + /// The caller must serialize all calls to `pop` and `disconnect_receiver` for this queue. + pub unsafe fn pop(&self) -> Poll> { let mut head = self.head.load(Ordering::Relaxed); let mut backoff = 0; loop { @@ -250,9 +254,15 @@ impl BoundedQueue { } } - pub fn disconnect_receiver(&self) { + /// Closes the queue and drops all remaining values. + /// + /// # Safety + /// + /// The caller must serialize all calls to `pop` and `disconnect_receiver` for this queue. + pub unsafe fn disconnect_receiver(&self) { let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; - self.discard_until(tail); + // SAFETY: The caller guarantees exclusive consumer access. + unsafe { self.discard_until(tail) }; } fn advance(&self, position: usize) -> usize { @@ -265,7 +275,8 @@ impl BoundedQueue { } } - fn discard_until(&self, tail: usize) { + // 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); let mut backoff = 0; while head != tail { @@ -299,7 +310,8 @@ impl BoundedQueue { impl Drop for BoundedQueue { fn drop(&mut self) { let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; - self.discard_until(tail); + // SAFETY: The queue is closed and its exclusive borrow rules out concurrent access. + unsafe { self.discard_until(tail) }; } } @@ -324,13 +336,16 @@ mod tests { } assert!(matches!(queue.try_push(3), Err(PushError::Full(3)))); for value in 0..3 { - assert_eq!(queue.pop(), Poll::Ready(Some(value))); + // SAFETY: This thread is the only consumer. + assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); } - assert_eq!(queue.pop(), Poll::Ready(None)); + // 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()); - assert_eq!(queue.pop(), Poll::Ready(Some(value))); + // SAFETY: This thread is the only consumer. + assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); } } @@ -345,15 +360,19 @@ mod tests { // 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); - let receive = queue.pop(); + // 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); - assert_eq!(queue.pop(), Poll::Ready(Some(1))); - assert_eq!(queue.pop(), Poll::Ready(Some(2))); - assert_eq!(queue.pop(), Poll::Ready(None)); + // SAFETY: This thread is the only consumer. + unsafe { + assert_eq!(queue.pop(), Poll::Ready(Some(1))); + assert_eq!(queue.pop(), Poll::Ready(Some(2))); + assert_eq!(queue.pop(), Poll::Ready(None)); + } } #[test] @@ -382,7 +401,8 @@ mod tests { let mut values = Vec::new(); while values.len() < 64 { - if let Poll::Ready(Some(value)) = queue.pop() { + // 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(); @@ -422,7 +442,8 @@ mod tests { } // Free slot 0, then reuse it on the next lap at position 8. - let popped = queue.pop(); + // 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); @@ -432,7 +453,8 @@ mod tests { assert_eq!(queue.head.load(Ordering::Relaxed), 1); assert_eq!(queue.tail.load(Ordering::Relaxed), queue.one_lap + 1); - queue.disconnect_receiver(); + // SAFETY: This thread is the only consumer and no pop is in progress. + unsafe { queue.disconnect_receiver() }; // `discard_until` must dispose every value exactly once, including position 8. for (value, counter) in drops.iter().enumerate() { From 9e3fb911b7f4d4c078d5368e22e7ac416e540538 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 7 Sep 2026 23:28:09 +0800 Subject: [PATCH 07/16] perf(mpsc): access receiver-local batches without locking --- asyncband/src/internal/mutex.rs | 6 ++++++ asyncband/src/mpsc/queue.rs | 24 +++++++++++++----------- asyncband/src/mpsc/unbounded.rs | 6 +++--- 3 files changed, 22 insertions(+), 14 deletions(-) diff --git a/asyncband/src/internal/mutex.rs b/asyncband/src/internal/mutex.rs index ce5229f8..581f4a4f 100644 --- a/asyncband/src/internal/mutex.rs +++ b/asyncband/src/internal/mutex.rs @@ -34,4 +34,10 @@ impl Mutex { pub fn lock(&self) -> std::sync::MutexGuard<'_, T> { self.0.lock().unwrap_or_else(PoisonError::into_inner) } + + // Exclusive access is currently used only by the mpsc receiver. + #[cfg_attr(not(feature = "mpsc"), allow(dead_code))] + pub fn get_mut(&mut self) -> &mut T { + self.0.get_mut().unwrap_or_else(PoisonError::into_inner) + } } diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs index add73003..c78a65bd 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/queue.rs @@ -39,6 +39,8 @@ struct UnboundedInner { } pub struct UnboundedConsumer { + // Preserve Sync for Send-only values; exclusive consumer access never needs to lock this + // mutex. local: Mutex>, } @@ -69,21 +71,21 @@ impl UnboundedQueue { Ok(()) } - pub fn pop(&self, consumer: &UnboundedConsumer) -> Option { - let mut local = consumer.local.lock(); + pub fn pop(&self, consumer: &mut UnboundedConsumer) -> Option { + let local = consumer.local.get_mut(); if local.is_empty() { let mut inner = self.inner.lock(); - mem::swap(&mut *local, &mut inner.messages); + mem::swap(local, &mut inner.messages); } local.pop_front() } - pub fn disconnect_receiver(&self, consumer: &UnboundedConsumer) { + pub fn disconnect_receiver(&self, consumer: &mut UnboundedConsumer) { let (local, shared) = { - let mut local = consumer.local.lock(); + let local = consumer.local.get_mut(); let mut inner = self.inner.lock(); inner.receiver_alive = false; - (mem::take(&mut *local), mem::take(&mut inner.messages)) + (mem::take(local), mem::take(&mut inner.messages)) }; drop((local, shared)); } @@ -479,13 +481,13 @@ mod tests { #[test] fn unbounded_queue_batches_without_reordering() { let queue = UnboundedQueue::new(); - let consumer = UnboundedConsumer::new(); + let mut consumer = UnboundedConsumer::new(); assert!(queue.push(1).is_ok()); assert!(queue.push(2).is_ok()); - assert_eq!(queue.pop(&consumer), Some(1)); + assert_eq!(queue.pop(&mut consumer), Some(1)); assert!(queue.push(3).is_ok()); - assert_eq!(queue.pop(&consumer), Some(2)); - assert_eq!(queue.pop(&consumer), Some(3)); - assert_eq!(queue.pop(&consumer), None); + assert_eq!(queue.pop(&mut consumer), Some(2)); + assert_eq!(queue.pop(&mut consumer), Some(3)); + assert_eq!(queue.pop(&mut consumer), None); } } diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index 1ec6cd09..1bf9936c 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -131,7 +131,7 @@ impl fmt::Debug for UnboundedReceiver { impl Drop for UnboundedReceiver { fn drop(&mut self) { - self.state.queue.disconnect_receiver(&self.consumer); + self.state.queue.disconnect_receiver(&mut self.consumer); } } @@ -159,14 +159,14 @@ impl UnboundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - if let Some(value) = self.state.queue.pop(&self.consumer) { + if let Some(value) = self.state.queue.pop(&mut self.consumer) { Ok(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. self.state .queue - .pop(&self.consumer) + .pop(&mut self.consumer) .ok_or(TryRecvError::Disconnected) } else { Err(TryRecvError::Empty) From 84bfb24bf09e817f067942b1654f0e05dafc6248 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 7 Sep 2026 23:30:23 +0800 Subject: [PATCH 08/16] fix(mpsc): release oversized unbounded batches after draining --- asyncband/src/mpsc/queue.rs | 69 ++++++++++++++++++++++++++++++++- asyncband/src/mpsc/unbounded.rs | 3 ++ 2 files changed, 71 insertions(+), 1 deletion(-) diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs index c78a65bd..2fa21a6b 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/queue.rs @@ -27,6 +27,10 @@ use std::task::Poll; use crate::internal::mutex::Mutex; +// Keep small batches reusable without retaining an arbitrarily large historical burst. Count +// inline storage bytes: boxed payloads own separate allocations, and zero-sized values own none. +const UNBOUNDED_CACHE_BYTES: usize = 64 * 1024; + pub struct UnboundedQueue { inner: Mutex>, } @@ -77,7 +81,15 @@ impl UnboundedQueue { let mut inner = self.inner.lock(); mem::swap(local, &mut inner.messages); } - local.pop_front() + if local.len() == 1 + && local.capacity().saturating_mul(mem::size_of::()) > UNBOUNDED_CACHE_BYTES + { + // Retire the allocation on the last value, outside the shared lock. Keep the ordinary + // pop as a tail expression so large inline values need no intermediate storage. + mem::take(local).pop_front() + } else { + local.pop_front() + } } pub fn disconnect_receiver(&self, consumer: &mut UnboundedConsumer) { @@ -490,4 +502,59 @@ mod tests { assert_eq!(queue.pop(&mut consumer), Some(3)); assert_eq!(queue.pop(&mut consumer), None); } + + #[test] + fn unbounded_queue_releases_large_batches_after_the_last_value() { + let queue = UnboundedQueue::new(); + let mut consumer = UnboundedConsumer::new(); + for value in 0..128u8 { + assert!(queue.push([value; 1024]).is_ok()); + } + + for value in 0..64u8 { + assert_eq!(queue.pop(&mut consumer), Some([value; 1024])); + } + // A partially consumed batch must survive while producers start filling the next batch. + assert!(queue.push([128; 1024]).is_ok()); + for value in 64..128u8 { + assert_eq!(queue.pop(&mut consumer), Some([value; 1024])); + } + + // Reclaim on the final successful receive, without requiring an extra empty poll. + assert_eq!(consumer.local.get_mut().capacity(), 0); + assert_eq!(queue.pop(&mut consumer), Some([128; 1024])); + assert_eq!(queue.pop(&mut consumer), None); + } + + #[test] + fn unbounded_queue_keeps_small_batches_for_reuse() { + let queue = UnboundedQueue::new(); + let mut consumer = UnboundedConsumer::new(); + for value in 0..32usize { + assert!(queue.push(value).is_ok()); + } + assert_eq!(queue.pop(&mut consumer), Some(0)); + let capacity = consumer.local.get_mut().capacity(); + for value in 1..32 { + assert_eq!(queue.pop(&mut consumer), Some(value)); + } + assert_eq!(consumer.local.get_mut().capacity(), capacity); + + // An empty poll returns the allocation to producers instead of discarding a small cache. + assert_eq!(queue.pop(&mut consumer), None); + assert_eq!(queue.inner.lock().messages.capacity(), capacity); + } + + #[test] + fn unbounded_queue_drains_zero_sized_values() { + let queue = UnboundedQueue::new(); + let mut consumer = UnboundedConsumer::new(); + for _ in 0..32 { + assert!(queue.push(()).is_ok()); + } + for _ in 0..32 { + assert_eq!(queue.pop(&mut consumer), Some(())); + } + assert_eq!(queue.pop(&mut consumer), None); + } } diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index 1bf9936c..c4be5ab2 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -39,6 +39,9 @@ use crate::internal::atomic_waker::AtomicWaker; /// While the receiver is alive, each send appends its value immediately. Pending messages can /// therefore grow with producer demand and are limited only by successful memory allocation. Use a /// bounded channel or external admission control when producers may outpace the receiver. +/// +/// After all messages have been received, large backing allocations are released; small buffers +/// may be retained for reuse. Partially consumed batches can retain their original allocation. pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { queue: UnboundedQueue::new(), From 375cbda63bfcf6475eec352c38daa414dec67ad3 Mon Sep 17 00:00:00 2001 From: tison Date: Mon, 7 Sep 2026 23:33:35 +0800 Subject: [PATCH 09/16] bench(mpsc): compare recurring bursts and payload storage --- benchmarks/ecosystem/mpsc/adapters.rs | 67 ++++++++++++------------ benchmarks/ecosystem/mpsc/unbounded.rs | 70 ++++++++++++++++++++++++++ 2 files changed, 104 insertions(+), 33 deletions(-) diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs index 7e78c217..1fd4dbb0 100644 --- a/benchmarks/ecosystem/mpsc/adapters.rs +++ b/benchmarks/ecosystem/mpsc/adapters.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::Debug; use std::task::Context; use crate::support::poll_ready; @@ -37,15 +38,15 @@ pub trait BoundedMpsc: Send + Sync + 'static { fn recv_blocking(receiver: &mut Self::Receiver) -> usize; } -pub trait UnboundedMpsc: Send + Sync + 'static { +pub trait UnboundedMpsc: Send + Sync + 'static { type Sender: Clone + Send + 'static; type Receiver: Send + 'static; fn channel() -> (Self::Sender, Self::Receiver); - fn send(sender: &Self::Sender, value: usize); - fn try_recv(receiver: &mut Self::Receiver) -> usize; - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize; - fn recv_blocking(receiver: &mut Self::Receiver) -> usize; + fn send(sender: &Self::Sender, value: T); + fn try_recv(receiver: &mut Self::Receiver) -> T; + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T; + fn recv_blocking(receiver: &mut Self::Receiver) -> T; } impl BoundedMpsc for Asyncband { @@ -180,102 +181,102 @@ impl BoundedMpsc for Flume { } } -impl UnboundedMpsc for Asyncband { - type Receiver = asyncband::mpsc::UnboundedReceiver; - type Sender = asyncband::mpsc::UnboundedSender; +impl UnboundedMpsc for Asyncband { + type Receiver = asyncband::mpsc::UnboundedReceiver; + type Sender = asyncband::mpsc::UnboundedSender; fn channel() -> (Self::Sender, Self::Receiver) { asyncband::mpsc::unbounded() } - fn send(sender: &Self::Sender, value: usize) { + fn send(sender: &Self::Sender, value: T) { sender.send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().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() } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl UnboundedMpsc for Tokio { - type Receiver = tokio::sync::mpsc::UnboundedReceiver; - type Sender = tokio::sync::mpsc::UnboundedSender; +impl UnboundedMpsc for Tokio { + type Receiver = tokio::sync::mpsc::UnboundedReceiver; + type Sender = tokio::sync::mpsc::UnboundedSender; fn channel() -> (Self::Sender, Self::Receiver) { tokio::sync::mpsc::unbounded_channel() } - fn send(sender: &Self::Sender, value: usize) { + fn send(sender: &Self::Sender, value: T) { sender.send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().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() } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl UnboundedMpsc for AsyncChannel { - type Receiver = async_channel::Receiver; - type Sender = async_channel::Sender; +impl UnboundedMpsc for AsyncChannel { + type Receiver = async_channel::Receiver; + type Sender = async_channel::Sender; fn channel() -> (Self::Sender, Self::Receiver) { async_channel::unbounded() } - fn send(sender: &Self::Sender, value: usize) { + fn 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 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() } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl UnboundedMpsc for Flume { - type Receiver = flume::Receiver; - type Sender = flume::Sender; +impl UnboundedMpsc for Flume { + type Receiver = flume::Receiver; + type Sender = flume::Sender; fn channel() -> (Self::Sender, Self::Receiver) { flume::unbounded() } - fn send(sender: &Self::Sender, value: usize) { + fn send(sender: &Self::Sender, value: T) { sender.send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().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() } - 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/unbounded.rs b/benchmarks/ecosystem/mpsc/unbounded.rs index 6c8ac549..6933a7ca 100644 --- a/benchmarks/ecosystem/mpsc/unbounded.rs +++ b/benchmarks/ecosystem/mpsc/unbounded.rs @@ -63,3 +63,73 @@ fn concurrent(bencher: Bencher, producer_count: usize) { .with_inputs(|| ConcurrentBatch::>::new(producer_count)) .bench_local_refs(|batch| batch.run()); } + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = [32, 1024, 65_536], + sample_count = 20, + sample_size = 1, +)] +fn burst_drain(bencher: Bencher, messages: usize) { + repeated_bursts::(bencher, messages, 0, || usize::MAX); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + consts = [64, 1024], + args = [32, 1024, 65_536], + sample_count = 20, + sample_size = 1, +)] +fn burst_drain_inline, const SIZE: usize>( + bencher: Bencher, + messages: usize, +) { + repeated_bursts::(bencher, messages, 0, || [1; SIZE]); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = [32, 1024, 65_536], + sample_count = 20, + sample_size = 1, +)] +fn burst_drain_boxed>>(bencher: Bencher, messages: usize) { + // Include payload allocation and destruction to compare the complete boxed-message lifecycle. + repeated_bursts::(bencher, messages, 0, || Box::new([1; 1024])); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = [1024, 65_536], + sample_count = 20, + sample_size = 1, +)] +fn burst_with_backlog(bencher: Bencher, messages: usize) { + repeated_bursts::(bencher, messages, messages / 2, || usize::MAX); +} + +fn repeated_bursts, T, F: Fn() -> T>( + bencher: Bencher, + messages: usize, + backlog: usize, + make_value: F, +) { + // Keep one channel alive across samples so allocation reuse and reclamation are measured. + let (sender, mut receiver) = C::channel(); + for _ in 0..backlog { + C::send(&sender, make_value()); + } + let mut run = || { + for _ in 0..messages { + C::send(&sender, black_box(make_value())); + } + for _ in 0..messages { + black_box(C::try_recv(&mut receiver)); + } + }; + // Measure recurring bursts after the initial allocation, including a deliberately retained + // backlog where requested. Do not require an extra empty receive to trigger reclamation. + run(); + bencher.counter(ItemsCount::new(messages)).bench_local(run); +} From 9b8a75a39d4256db047653f0d7adde49f8352332 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 00:08:38 +0800 Subject: [PATCH 10/16] refactor(internal): share cache-line padding between primitives --- asyncband/src/internal/cache_padded.rs | 55 ++++++++++++++++++++++++++ asyncband/src/internal/mod.rs | 3 ++ asyncband/src/mpsc/queue.rs | 38 ++---------------- 3 files changed, 61 insertions(+), 35 deletions(-) create mode 100644 asyncband/src/internal/cache_padded.rs 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 c6363791..4f33dff8 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -70,6 +70,9 @@ pub(crate) mod atomic_waker; #[allow(dead_code)] pub(crate) mod arena; +#[allow(dead_code)] +pub(crate) mod cache_padded; + #[cfg(any(feature = "latch", feature = "once"))] pub(crate) mod countdown; diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs index 2fa21a6b..e314dad3 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/queue.rs @@ -25,6 +25,7 @@ use std::sync::atomic::Ordering; use std::sync::atomic::fence; use std::task::Poll; +use crate::internal::cache_padded::CachePadded; use crate::internal::mutex::Mutex; // Keep small batches reusable without retaining an arbitrarily large historical burst. Count @@ -120,39 +121,6 @@ pub struct BoundedQueue { mark_bit: usize, } -// 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)) -)] -struct CachePadded(T); - -impl std::ops::Deref for CachePadded { - type Target = T; - - fn deref(&self) -> &Self::Target { - &self.0 - } -} - struct Slot { stamp: AtomicUsize, value: UnsafeCell>, @@ -181,8 +149,8 @@ impl BoundedQueue { .collect(); Self { slots, - head: CachePadded(AtomicUsize::new(0)), - tail: CachePadded(AtomicUsize::new(0)), + head: CachePadded::new(AtomicUsize::new(0)), + tail: CachePadded::new(AtomicUsize::new(0)), capacity, one_lap, mark_bit, From 5b161be98b67de96e208871fca431361c8e3bf97 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 00:12:26 +0800 Subject: [PATCH 11/16] refactor(internal): allow unused mutex operations at the module boundary --- asyncband/src/internal/mod.rs | 2 ++ asyncband/src/internal/mutex.rs | 2 -- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 4f33dff8..7acb2f00 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -95,6 +95,8 @@ pub(crate) mod value_cell; feature = "waitgroup", feature = "watch", ))] +// Some primitives use only shared access, leaving `Mutex::get_mut` unused. +#[allow(dead_code)] pub(crate) mod mutex; #[cfg(any( diff --git a/asyncband/src/internal/mutex.rs b/asyncband/src/internal/mutex.rs index 581f4a4f..93baf8f2 100644 --- a/asyncband/src/internal/mutex.rs +++ b/asyncband/src/internal/mutex.rs @@ -35,8 +35,6 @@ impl Mutex { self.0.lock().unwrap_or_else(PoisonError::into_inner) } - // Exclusive access is currently used only by the mpsc receiver. - #[cfg_attr(not(feature = "mpsc"), allow(dead_code))] pub fn get_mut(&mut self) -> &mut T { self.0.get_mut().unwrap_or_else(PoisonError::into_inner) } From 8470215ec653f450acce07f28b39eeb242e85699 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 00:38:45 +0800 Subject: [PATCH 12/16] fix(mpsc): release receiver wakers on disconnection --- CHANGELOG.md | 4 +++ asyncband/src/internal/atomic_waker.rs | 4 ++- asyncband/src/mpsc/bounded.rs | 3 ++ asyncband/src/mpsc/unbounded.rs | 3 ++ tests-integration/tests/mpsc_test.rs | 49 ++++++++++++++++++++++++++ 5 files changed, 62 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0543a046..fa53b69d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file. ## Unreleased +### Bug fixes + +* Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. + ## v0.7.2 ### Improvements diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs index 3616691a..a6dcd20f 100644 --- a/asyncband/src/internal/atomic_waker.rs +++ b/asyncband/src/internal/atomic_waker.rs @@ -210,8 +210,10 @@ impl AtomicWaker { } } + /// Removes the registered waker if this call acquires the slot. A concurrent registration or + /// wake may instead take responsibility for notifying it. #[inline] - fn take(&self) -> Option { + 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. diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index 5d88827c..8687b996 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -214,9 +214,12 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { + // A registered waker may own a sender; release it to break that ownership cycle. + let receiver_waker = self.state.rx_waker.take(); // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows. unsafe { self.state.queue.disconnect_receiver() }; self.state.tx_permits.notify_all(); + drop(receiver_waker); } } diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index c4be5ab2..4a0060b6 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -134,7 +134,10 @@ impl fmt::Debug for UnboundedReceiver { impl Drop for UnboundedReceiver { fn drop(&mut self) { + // A registered waker may own a sender; release it to break that ownership cycle. + let receiver_waker = self.state.rx_waker.take(); self.state.queue.disconnect_receiver(&mut self.consumer); + drop(receiver_waker); } } diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 087b40f5..2aa45cc3 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -15,9 +15,14 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; +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::Wake; +use std::task::Waker; use std::thread; use asyncband::mpsc; @@ -35,6 +40,50 @@ fn expect_ready(poll: Poll) -> T { } } +struct HoldSender { + _sender: S, +} + +// This waker must own the sender so its final drop can break the tested reference cycle. +#[allow(clippy::manual_noop_waker)] +impl Wake for HoldSender { + fn wake(self: Arc) {} +} + +#[test] +fn bounded_receiver_drop_releases_registered_waker() { + let (tx, mut rx) = mpsc::bounded::<()>(1); + let holder = Arc::new(HoldSender { _sender: tx }); + let retained = Arc::downgrade(&holder); + let waker = Waker::from(holder); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + drop(rx); + assert!(retained.upgrade().is_none()); +} + +#[test] +fn unbounded_receiver_drop_releases_registered_waker() { + let (tx, mut rx) = mpsc::unbounded::<()>(); + let holder = Arc::new(HoldSender { _sender: tx }); + let retained = Arc::downgrade(&holder); + let waker = Waker::from(holder); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + drop(rx); + assert!(retained.upgrade().is_none()); +} + #[test] fn unbounded_collects_from_multiple_producers() { let (tx, mut rx) = mpsc::unbounded(); From b01bd86b170351ddacfafecba58625d882b7c8eb Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 00:42:43 +0800 Subject: [PATCH 13/16] perf(mpsc): coordinate unbounded messages and notifications together --- CHANGELOG.md | 4 + asyncband/src/mpsc/queue.rs | 152 ------------------ asyncband/src/mpsc/unbounded.rs | 227 +++++++++++++++++++++------ tests-integration/tests/mpsc_test.rs | 130 +++++++++++++++ 4 files changed, 310 insertions(+), 203 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index fa53b69d..35e870e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,10 @@ All notable changes to this project will be documented in this file. * Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. +### Improvements + +* Reduce unbounded MPSC synchronization overhead by transferring messages in batches and coordinating receiver notifications with queued messages; release large empty batch allocations while retaining small buffers for reuse. + ## v0.7.2 ### Improvements diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/queue.rs index e314dad3..68c0550b 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/queue.rs @@ -16,9 +16,7 @@ // under the License. use std::cell::UnsafeCell; -use std::collections::VecDeque; use std::hint::spin_loop; -use std::mem; use std::mem::MaybeUninit; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -26,92 +24,12 @@ use std::sync::atomic::fence; use std::task::Poll; use crate::internal::cache_padded::CachePadded; -use crate::internal::mutex::Mutex; - -// Keep small batches reusable without retaining an arbitrarily large historical burst. Count -// inline storage bytes: boxed payloads own separate allocations, and zero-sized values own none. -const UNBOUNDED_CACHE_BYTES: usize = 64 * 1024; - -pub struct UnboundedQueue { - inner: Mutex>, -} - -struct UnboundedInner { - // Storage and receiver liveness share one lock so a send is linearized either before receiver - // disconnection, with its value in the queue, or after it, with the value returned to sender. - messages: VecDeque, - receiver_alive: bool, -} - -pub struct UnboundedConsumer { - // Preserve Sync for Send-only values; exclusive consumer access never needs to lock this - // mutex. - local: Mutex>, -} - -// The consumer never relies on a pinned location for its local queue or queued values. -impl Unpin for UnboundedConsumer {} pub enum PushError { Full(T), Disconnected(T), } -impl UnboundedQueue { - pub const fn new() -> Self { - Self { - inner: Mutex::new(UnboundedInner { - messages: VecDeque::new(), - receiver_alive: true, - }), - } - } - - pub fn push(&self, value: T) -> Result<(), PushError> { - let mut inner = self.inner.lock(); - if !inner.receiver_alive { - return Err(PushError::Disconnected(value)); - } - inner.messages.push_back(value); - Ok(()) - } - - pub fn pop(&self, consumer: &mut UnboundedConsumer) -> Option { - let local = consumer.local.get_mut(); - if local.is_empty() { - let mut inner = self.inner.lock(); - mem::swap(local, &mut inner.messages); - } - if local.len() == 1 - && local.capacity().saturating_mul(mem::size_of::()) > UNBOUNDED_CACHE_BYTES - { - // Retire the allocation on the last value, outside the shared lock. Keep the ordinary - // pop as a tail expression so large inline values need no intermediate storage. - mem::take(local).pop_front() - } else { - local.pop_front() - } - } - - pub fn disconnect_receiver(&self, consumer: &mut UnboundedConsumer) { - let (local, shared) = { - let local = consumer.local.get_mut(); - let mut inner = self.inner.lock(); - inner.receiver_alive = false; - (mem::take(local), mem::take(&mut inner.messages)) - }; - drop((local, shared)); - } -} - -impl UnboundedConsumer { - pub const fn new() -> Self { - Self { - local: Mutex::new(VecDeque::new()), - } - } -} - pub struct BoundedQueue { slots: Box<[Slot]>, head: CachePadded, @@ -307,8 +225,6 @@ mod tests { use super::BoundedQueue; use super::PushError; - use super::UnboundedConsumer; - use super::UnboundedQueue; #[test] fn bounded_queue_preserves_capacity_and_fifo_order() { @@ -457,72 +373,4 @@ mod tests { ); } } - - #[test] - fn unbounded_queue_batches_without_reordering() { - let queue = UnboundedQueue::new(); - let mut consumer = UnboundedConsumer::new(); - assert!(queue.push(1).is_ok()); - assert!(queue.push(2).is_ok()); - assert_eq!(queue.pop(&mut consumer), Some(1)); - assert!(queue.push(3).is_ok()); - assert_eq!(queue.pop(&mut consumer), Some(2)); - assert_eq!(queue.pop(&mut consumer), Some(3)); - assert_eq!(queue.pop(&mut consumer), None); - } - - #[test] - fn unbounded_queue_releases_large_batches_after_the_last_value() { - let queue = UnboundedQueue::new(); - let mut consumer = UnboundedConsumer::new(); - for value in 0..128u8 { - assert!(queue.push([value; 1024]).is_ok()); - } - - for value in 0..64u8 { - assert_eq!(queue.pop(&mut consumer), Some([value; 1024])); - } - // A partially consumed batch must survive while producers start filling the next batch. - assert!(queue.push([128; 1024]).is_ok()); - for value in 64..128u8 { - assert_eq!(queue.pop(&mut consumer), Some([value; 1024])); - } - - // Reclaim on the final successful receive, without requiring an extra empty poll. - assert_eq!(consumer.local.get_mut().capacity(), 0); - assert_eq!(queue.pop(&mut consumer), Some([128; 1024])); - assert_eq!(queue.pop(&mut consumer), None); - } - - #[test] - fn unbounded_queue_keeps_small_batches_for_reuse() { - let queue = UnboundedQueue::new(); - let mut consumer = UnboundedConsumer::new(); - for value in 0..32usize { - assert!(queue.push(value).is_ok()); - } - assert_eq!(queue.pop(&mut consumer), Some(0)); - let capacity = consumer.local.get_mut().capacity(); - for value in 1..32 { - assert_eq!(queue.pop(&mut consumer), Some(value)); - } - assert_eq!(consumer.local.get_mut().capacity(), capacity); - - // An empty poll returns the allocation to producers instead of discarding a small cache. - assert_eq!(queue.pop(&mut consumer), None); - assert_eq!(queue.inner.lock().messages.capacity(), capacity); - } - - #[test] - fn unbounded_queue_drains_zero_sized_values() { - let queue = UnboundedQueue::new(); - let mut consumer = UnboundedConsumer::new(); - for _ in 0..32 { - assert!(queue.push(()).is_ok()); - } - for _ in 0..32 { - assert_eq!(queue.pop(&mut consumer), Some(())); - } - assert_eq!(queue.pop(&mut consumer), None); - } } diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index 4a0060b6..92be6990 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -18,21 +18,21 @@ //! An unbounded multi-producer, single-consumer queue for sending values between asynchronous //! tasks. +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::Waker; use super::RecvError; use super::SendError; use super::TryRecvError; -use super::queue::PushError; -use super::queue::UnboundedConsumer; -use super::queue::UnboundedQueue; -use crate::internal::atomic_waker::AtomicWaker; +use crate::internal::mutex::Mutex; /// Creates an unbounded mpsc channel whose send operation never waits for capacity. /// @@ -44,24 +44,35 @@ use crate::internal::atomic_waker::AtomicWaker; /// may be retained for reuse. Partially consumed batches can retain their original allocation. pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { - queue: UnboundedQueue::new(), senders: AtomicUsize::new(1), - rx_waker: AtomicWaker::new(), + inbox: Mutex::new(Inbox { + messages: VecDeque::new(), + receiver_alive: true, + rx_waker: None, + }), }); let sender = UnboundedSender { state: state.clone(), }; let receiver = UnboundedReceiver { state, - consumer: UnboundedConsumer::new(), + batch: Mutex::new(VecDeque::new()), }; (sender, receiver) } struct UnboundedState { - queue: UnboundedQueue, + // Endpoint cloning and ordinary drops do not contend with message traffic. senders: AtomicUsize, - rx_waker: AtomicWaker, + inbox: Mutex>, +} + +// Queue contents, receiver liveness, and its wake registration share one lock. Registering a +// wait and checking its condition cannot race with sending or receiver disconnection. +struct Inbox { + messages: VecDeque, + receiver_alive: bool, + rx_waker: Option, } /// The sending endpoint of an unbounded mpsc channel. @@ -74,7 +85,7 @@ pub struct UnboundedSender { impl Clone for UnboundedSender { fn clone(&self) -> Self { self.state.senders.fetch_add(1, Ordering::Release); - UnboundedSender { + Self { state: self.state.clone(), } } @@ -88,13 +99,10 @@ impl fmt::Debug for UnboundedSender { impl Drop for UnboundedSender { fn drop(&mut self) { - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { - 1 => { - // Wake the receiver so it can observe the channel's disconnected state. - self.state.rx_waker.wake(); - } - _ => { - // there are still other senders left, do nothing + if self.state.senders.fetch_sub(1, Ordering::AcqRel) == 1 { + let waker = self.state.inbox.lock().rx_waker.take(); + if let Some(waker) = waker { + waker.wake(); } } } @@ -106,14 +114,17 @@ impl UnboundedSender { /// This operation is synchronous because the channel has no capacity limit. If the receiver has /// been dropped, the returned error contains `value`. pub fn send(&self, value: T) -> Result<(), SendError> { - match self.state.queue.push(value) { - Ok(()) => {} - Err(PushError::Disconnected(value)) => return Err(SendError::new(value)), - Err(PushError::Full(_)) => unreachable!("unbounded queue cannot be full"), + let waker = { + let mut state = self.state.inbox.lock(); + if !state.receiver_alive { + return Err(SendError::new(value)); + } + state.messages.push_back(value); + state.rx_waker.take() + }; + if let Some(waker) = waker { + waker.wake(); } - - self.state.rx_waker.wake(); - Ok(()) } } @@ -123,7 +134,8 @@ impl UnboundedSender { /// Instances are created by the [`unbounded`] function. pub struct UnboundedReceiver { state: Arc>, - consumer: UnboundedConsumer, + // Only accessed through `get_mut`; the mutex preserves Sync for Send-only payloads. + batch: Mutex>, } impl fmt::Debug for UnboundedReceiver { @@ -134,10 +146,14 @@ impl fmt::Debug for UnboundedReceiver { impl Drop for UnboundedReceiver { fn drop(&mut self) { - // A registered waker may own a sender; release it to break that ownership cycle. - let receiver_waker = self.state.rx_waker.take(); - self.state.queue.disconnect_receiver(&mut self.consumer); - drop(receiver_waker); + let batch = mem::take(self.batch.get_mut()); + let (shared, waker) = { + let mut state = self.state.inbox.lock(); + state.receiver_alive = false; + (mem::take(&mut state.messages), state.rx_waker.take()) + }; + // Destructors may send again. A waker may also own a sender and form an ownership cycle. + drop((batch, shared, waker)); } } @@ -165,18 +181,21 @@ impl UnboundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - if let Some(value) = self.state.queue.pop(&mut self.consumer) { - Ok(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. - self.state - .queue - .pop(&mut self.consumer) - .ok_or(TryRecvError::Disconnected) - } else { - Err(TryRecvError::Empty) + let batch = self.batch.get_mut(); + if batch.is_empty() { + let mut state = self.state.inbox.lock(); + if state.messages.is_empty() { + // Holding the inbox lock excludes a final send between the empty observation and + // the sender-count check; disconnection needs no second queue read. + return Err(if self.state.senders.load(Ordering::Acquire) == 0 { + TryRecvError::Disconnected + } else { + TryRecvError::Empty + }); + } + mem::swap(batch, &mut state.messages); } + Ok(pop_batch(batch)) } /// Waits for and receives the next value. @@ -213,18 +232,124 @@ impl UnboundedReceiver { } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => { - self.state.rx_waker.register(cx.waker()); - - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => Poll::Pending, - } + let batch = self.batch.get_mut(); + if !batch.is_empty() { + return Poll::Ready(Ok(pop_batch(batch))); + } + // Waker clone/drop callbacks can reenter this channel. Clone outside the lock, then + // recheck the condition before registering; keep replaced wakers outside the lock too. + let mut new_waker = None; + loop { + let mut state = self.state.inbox.lock(); + if !state.messages.is_empty() { + mem::swap(batch, &mut state.messages); + drop(state); + return Poll::Ready(Ok(pop_batch(batch))); + } + if self.state.senders.load(Ordering::Acquire) == 0 { + return Poll::Ready(Err(RecvError::Disconnected)); } + if state + .rx_waker + .as_ref() + .is_some_and(|waker| waker.will_wake(cx.waker())) + { + return Poll::Pending; + } + if let Some(waker) = new_waker.take() { + let old_waker = state.rx_waker.replace(waker); + drop(state); + drop(old_waker); + return Poll::Pending; + } + drop(state); + new_waker = Some(cx.waker().clone()); + } + } +} + +// No operation relies on a pinned location for the receiver batch or its values. +impl Unpin for UnboundedReceiver {} + +// Retain small buffers for reuse, measuring inline storage rather than separately boxed payloads. +const BATCH_CACHE_BYTES: usize = 64 * 1024; + +fn pop_batch(batch: &mut VecDeque) -> T { + if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::()) > BATCH_CACHE_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() + } else { + batch.pop_front() + } + .expect("receiver batch must not be empty") +} + +#[cfg(test)] +mod tests { + use super::unbounded; + use crate::mpsc::TryRecvError; + + #[test] + fn batches_preserve_order_across_refills() { + let (tx, mut rx) = unbounded(); + tx.send(1).unwrap(); + tx.send(2).unwrap(); + assert_eq!(rx.try_recv(), Ok(1)); + tx.send(3).unwrap(); + assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv(), Ok(3)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } + + #[test] + fn releases_large_batches_after_the_last_value() { + let (tx, mut rx) = unbounded(); + for value in 0..128u8 { + tx.send([value; 1024]).unwrap(); + } + for value in 0..64u8 { + assert_eq!(rx.try_recv(), Ok([value; 1024])); + } + // A partially consumed batch survives while producers start filling the next batch. + tx.send([128; 1024]).unwrap(); + for value in 64..128u8 { + assert_eq!(rx.try_recv(), Ok([value; 1024])); + } + // Reclaim on the final successful receive, without requiring an extra empty poll. + assert_eq!(rx.batch.get_mut().capacity(), 0); + assert_eq!(rx.try_recv(), Ok([128; 1024])); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } + + #[test] + fn reuses_small_batches_on_refill() { + let (tx, mut rx) = unbounded(); + for value in 0..32usize { + tx.send(value).unwrap(); + } + assert_eq!(rx.try_recv(), Ok(0)); + let capacity = rx.batch.get_mut().capacity(); + for value in 1..32 { + assert_eq!(rx.try_recv(), Ok(value)); + } + assert_eq!(rx.batch.get_mut().capacity(), capacity); + + tx.send(32).unwrap(); + assert_eq!(rx.try_recv(), Ok(32)); + assert_eq!(rx.state.inbox.lock().messages.capacity(), capacity); + } + + #[test] + fn drains_zero_sized_values() { + let (tx, mut rx) = unbounded(); + for _ in 0..32 { + tx.send(()).unwrap(); + } + for _ in 0..32 { + assert_eq!(rx.try_recv(), Ok(())); } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } } diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 2aa45cc3..7b84b933 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -21,9 +21,12 @@ use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use std::task::RawWaker; +use std::task::RawWakerVTable; use std::task::Wake; use std::task::Waker; use std::thread; +use std::time::Duration; use asyncband::mpsc; use asyncband::mpsc::RecvError; @@ -84,6 +87,133 @@ fn unbounded_receiver_drop_releases_registered_waker() { assert!(retained.upgrade().is_none()); } +fn assert_completes_without_deadlock(test: impl FnOnce() + Send + 'static) { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + test(); + finished_tx.send(()).unwrap(); + }); + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("waker callback did not finish"); + worker.join().unwrap(); +} + +#[test] +fn unbounded_wake_callback_can_send() { + struct SendOnWake(mpsc::UnboundedSender); + + impl Wake for SendOnWake { + fn wake(self: Arc) { + self.0.send(2).unwrap(); + } + } + + assert_completes_without_deadlock(|| { + let (tx, mut rx) = mpsc::unbounded(); + let waker = Waker::from(Arc::new(SendOnWake(tx.clone()))); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + tx.send(1).unwrap(); + assert_eq!(rx.try_recv(), Ok(1)); + assert_eq!(rx.try_recv(), Ok(2)); + }); +} + +#[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) {} + } + + 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); + } + } + + 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(), + })); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + if disconnected { + drop(rx); + } else { + // Replacing the waker can enqueue a message during this poll. Either immediate + // completion or a notified Pending is valid, but the message must not be lost. + let poll = poll_once(Box::pin(rx.recv()).as_mut()); + if poll.is_pending() { + assert_eq!(rx.try_recv(), Ok(7)); + } else { + assert_eq!(poll, Poll::Ready(Ok(7))); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); + } + assert_eq!(drops.load(Ordering::Relaxed), 1); + } + }); +} + +#[test] +fn unbounded_waker_clone_rechecks_messages_sent_during_registration() { + unsafe fn clone_sender(data: *const ()) -> RawWaker { + let sender = data.cast::>(); + // SAFETY: Each raw waker owns an Arc to this sender; cloning borrows the live sender and + // then adds the strong reference owned by the returned waker. + unsafe { + (*sender).send(7).unwrap(); + Arc::increment_strong_count(sender); + } + RawWaker::new(data, &VTABLE) + } + + unsafe fn drop_sender(data: *const ()) { + // SAFETY: Consumes exactly the Arc reference owned by this raw waker. + drop(unsafe { Arc::from_raw(data.cast::>()) }); + } + + static VTABLE: RawWakerVTable = + RawWakerVTable::new(clone_sender, drop_sender, |_| {}, drop_sender); + + assert_completes_without_deadlock(|| { + let (tx, mut rx) = mpsc::unbounded::(); + let data = Arc::into_raw(Arc::new(tx)).cast(); + // SAFETY: The vtable manages one Arc reference per waker and the sender is Send + Sync. + let waker = unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }; + assert_eq!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)), + Poll::Ready(Ok(7)) + ); + }); +} + #[test] fn unbounded_collects_from_multiple_producers() { let (tx, mut rx) = mpsc::unbounded(); From 67af39d2c364be3a68a626c1e26ab9974c878e4c Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 00:44:50 +0800 Subject: [PATCH 14/16] refactor(mpsc): isolate bounded ring and sender wait responsibilities --- asyncband/src/mpsc/bounded.rs | 44 ++++++++----------- .../src/mpsc/{queue.rs => bounded/ring.rs} | 36 +++++++-------- asyncband/src/mpsc/mod.rs | 1 - 3 files changed, 35 insertions(+), 46 deletions(-) rename asyncband/src/mpsc/{queue.rs => bounded/ring.rs} (94%) diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index 8687b996..550bf165 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -19,9 +19,7 @@ //! tasks with backpressure control. use std::fmt; -use std::future::Future; use std::future::poll_fn; -use std::pin::pin; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -29,16 +27,17 @@ use std::task::Context; use std::task::Poll; use std::task::ready; +use self::ring::Ring; use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; -use super::queue::BoundedQueue; -use super::queue::PushError; use crate::internal::atomic_waker::AtomicWaker; use crate::internal::semaphore::Acquire; use crate::internal::semaphore::Semaphore; +mod ring; + /// 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 @@ -51,9 +50,9 @@ use crate::internal::semaphore::Semaphore; pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); let state = Arc::new(BoundedState { - queue: BoundedQueue::new(buffer), + buffer: Ring::new(buffer), senders: AtomicUsize::new(1), - tx_permits: Semaphore::new(0), + send_waiters: Semaphore::new(0), rx_waker: AtomicWaker::new(), }); let sender = BoundedSender { @@ -64,9 +63,10 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { } struct BoundedState { - queue: BoundedQueue, + buffer: Ring, senders: AtomicUsize, - tx_permits: Semaphore, + // Notifications grant retries; only the ring determines whether buffer capacity is available. + send_waiters: Semaphore, rx_waker: AtomicWaker, } @@ -137,7 +137,7 @@ impl BoundedSender { }; loop { - let poll = pin!(&mut self.acquire).poll(cx); + let poll = self.acquire.poll_once(cx.waker()); value = match self.sender.try_send(value) { Ok(()) => return Poll::Ready(Ok(())), @@ -148,7 +148,7 @@ impl BoundedSender { }; if poll.is_ready() { - self.acquire = self.sender.state.tx_permits.poll_acquire(1); + self.acquire = self.sender.state.send_waiters.poll_acquire(1); } else { self.value = Some(value); return Poll::Pending; @@ -157,7 +157,7 @@ impl BoundedSender { } } - let acquire = self.state.tx_permits.poll_acquire(1); + let acquire = self.state.send_waiters.poll_acquire(1); let mut send = SendState { sender: self, value: Some(value), @@ -187,15 +187,9 @@ impl BoundedSender { /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - match self.state.queue.try_push(value) { - Ok(()) => { - self.state.rx_waker.wake(); - - Ok(()) - } - Err(PushError::Full(value)) => Err(TrySendError::Full(value)), - Err(PushError::Disconnected(value)) => Err(TrySendError::Disconnected(value)), - } + self.state.buffer.try_push(value)?; + self.state.rx_waker.wake(); + Ok(()) } } @@ -217,8 +211,8 @@ impl Drop for BoundedReceiver { // A registered waker may own a sender; release it to break that ownership cycle. let receiver_waker = self.state.rx_waker.take(); // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows. - unsafe { self.state.queue.disconnect_receiver() }; - self.state.tx_permits.notify_all(); + unsafe { self.state.buffer.disconnect_receiver() }; + self.state.send_waiters.notify_all(); drop(receiver_waker); } } @@ -260,20 +254,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.queue.pop() }) { + let value = if let Some(value) = ready!(unsafe { self.state.buffer.pop() }) { 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.queue.pop() }) else { + let Some(value) = ready!(unsafe { self.state.buffer.pop() }) else { return Poll::Ready(Err(TryRecvError::Disconnected)); }; value } else { return Poll::Ready(Err(TryRecvError::Empty)); }; - self.state.tx_permits.release_if_nonempty(1); + self.state.send_waiters.release_if_nonempty(1); Poll::Ready(Ok(value)) } diff --git a/asyncband/src/mpsc/queue.rs b/asyncband/src/mpsc/bounded/ring.rs similarity index 94% rename from asyncband/src/mpsc/queue.rs rename to asyncband/src/mpsc/bounded/ring.rs index 68c0550b..4eb95085 100644 --- a/asyncband/src/mpsc/queue.rs +++ b/asyncband/src/mpsc/bounded/ring.rs @@ -24,13 +24,9 @@ use std::sync::atomic::fence; use std::task::Poll; use crate::internal::cache_padded::CachePadded; +use crate::mpsc::TrySendError; -pub enum PushError { - Full(T), - Disconnected(T), -} - -pub struct BoundedQueue { +pub struct Ring { slots: Box<[Slot]>, head: CachePadded, tail: CachePadded, @@ -54,7 +50,7 @@ unsafe impl Sync for Slot {} impl std::panic::UnwindSafe for Slot {} impl std::panic::RefUnwindSafe for Slot {} -impl BoundedQueue { +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(); @@ -75,12 +71,12 @@ impl BoundedQueue { } } - pub fn try_push(&self, value: T) -> Result<(), PushError> { + 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(PushError::Disconnected(value)); + return Err(TrySendError::Disconnected(value)); } let index = tail & (self.mark_bit - 1); @@ -106,7 +102,7 @@ impl BoundedQueue { } 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(PushError::Full(value)); + return Err(TrySendError::Full(value)); } tail = self.tail.load(Ordering::Relaxed); } else { @@ -207,7 +203,7 @@ impl BoundedQueue { } } -impl Drop for BoundedQueue { +impl Drop for Ring { fn drop(&mut self) { let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; // SAFETY: The queue is closed and its exclusive borrow rules out concurrent access. @@ -223,16 +219,16 @@ mod tests { use std::task::Poll; use std::thread; - use super::BoundedQueue; - use super::PushError; + use super::Ring; + use super::TrySendError; #[test] fn bounded_queue_preserves_capacity_and_fifo_order() { - let queue = BoundedQueue::new(3); + let queue = Ring::new(3); for value in 0..3 { assert!(queue.try_push(value).is_ok()); } - assert!(matches!(queue.try_push(3), Err(PushError::Full(3)))); + 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))); @@ -249,7 +245,7 @@ mod tests { #[test] fn bounded_queue_does_not_report_empty_behind_an_unpublished_head() { - let queue = BoundedQueue::new(2); + 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. @@ -275,7 +271,7 @@ mod tests { #[test] fn bounded_queue_coordinates_multiple_producers() { - let queue = Arc::new(BoundedQueue::new(4)); + let queue = Arc::new(Ring::new(4)); let producers: Vec<_> = (0..2) .map(|producer| { let queue = queue.clone(); @@ -285,11 +281,11 @@ mod tests { loop { match queue.try_push(value) { Ok(()) => break, - Err(PushError::Full(returned)) => { + Err(TrySendError::Full(returned)) => { value = returned; thread::yield_now(); } - Err(PushError::Disconnected(_)) => panic!("queue disconnected"), + Err(TrySendError::Disconnected(_)) => panic!("queue disconnected"), } } } @@ -332,7 +328,7 @@ mod tests { AtomicUsize::new(0), AtomicUsize::new(0), ]; - let queue = BoundedQueue::new(3); + let queue = Ring::new(3); // Positions: 0, 1, 2 (then tail wraps to 8). for counter in &drops[..3] { diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs index 30ef65e9..040c98b2 100644 --- a/asyncband/src/mpsc/mod.rs +++ b/asyncband/src/mpsc/mod.rs @@ -23,7 +23,6 @@ mod bounded; mod error; -mod queue; mod unbounded; pub use self::bounded::BoundedReceiver; From 81535127848326deeb33b191d380927fd09ae50d Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 00:52:07 +0800 Subject: [PATCH 15/16] bench(mpsc): measure sustained traffic and sender lifecycle costs --- benchmarks/ecosystem/mpsc/bounded.rs | 20 ++++++++ benchmarks/ecosystem/mpsc/support.rs | 65 ++++++++++++++++++++++++++ benchmarks/ecosystem/mpsc/unbounded.rs | 20 ++++++++ 3 files changed, 105 insertions(+) diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index bda71079..80495e98 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -29,6 +29,7 @@ use super::support::BOUNDED_CAPACITY; use super::support::Bounded; use super::support::ConcurrentBatch; use super::support::PRODUCER_COUNTS; +use super::support::RepeatedBatch; use crate::support::bench_context; #[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] @@ -64,3 +65,22 @@ fn concurrent(bencher: Bencher, producer_count: usize) { .with_inputs(|| ConcurrentBatch::>::new(producer_count)) .bench_local_refs(|batch| batch.run()); } + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = PRODUCER_COUNTS, + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn sustained(bencher: Bencher, producer_count: usize) { + let mut batch = RepeatedBatch::>::new(producer_count); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[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()))); +} diff --git a/benchmarks/ecosystem/mpsc/support.rs b/benchmarks/ecosystem/mpsc/support.rs index aea28858..11edb58f 100644 --- a/benchmarks/ecosystem/mpsc/support.rs +++ b/benchmarks/ecosystem/mpsc/support.rs @@ -18,6 +18,8 @@ use std::marker::PhantomData; use std::sync::Arc; use std::sync::Barrier; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::thread; use std::thread::JoinHandle; @@ -133,3 +135,66 @@ impl Drop for ConcurrentBatch { } } } + +// Reuse worker threads and channel storage so steady-state samples exclude thread creation. +pub struct RepeatedBatch { + receiver: C::Receiver, + start: Arc, + stop: Arc, + workers: Vec>, +} + +impl RepeatedBatch { + pub fn new(producer_count: usize) -> Self { + assert_eq!(BATCH_MESSAGES % producer_count, 0); + let (sender, receiver) = C::channel(); + let start = Arc::new(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) + .map(|producer| { + let sender = sender.clone(); + let start = start.clone(); + let stop = stop.clone(); + thread::spawn(move || { + loop { + start.wait(); + if stop.load(Ordering::Acquire) { + break; + } + let first = producer * messages_per_producer; + for offset in 0..messages_per_producer { + C::send(&sender, black_box(first + offset)); + } + } + }) + }) + .collect(); + drop(sender); + Self { + receiver, + start, + stop, + workers, + } + } + + pub fn run(&mut self) -> usize { + self.start.wait(); + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(C::recv(&mut self.receiver)); + } + black_box(checksum) + } +} + +impl Drop for RepeatedBatch { + fn drop(&mut self) { + self.stop.store(true, Ordering::Release); + self.start.wait(); + for worker in self.workers.drain(..) { + worker.join().expect("benchmark producer panicked"); + } + } +} diff --git a/benchmarks/ecosystem/mpsc/unbounded.rs b/benchmarks/ecosystem/mpsc/unbounded.rs index 6933a7ca..b08ed700 100644 --- a/benchmarks/ecosystem/mpsc/unbounded.rs +++ b/benchmarks/ecosystem/mpsc/unbounded.rs @@ -27,6 +27,7 @@ use super::adapters::UnboundedMpsc; use super::support::BATCH_MESSAGES; use super::support::ConcurrentBatch; use super::support::PRODUCER_COUNTS; +use super::support::RepeatedBatch; use super::support::Unbounded; use crate::support::bench_context; @@ -133,3 +134,22 @@ fn repeated_bursts, T, F: Fn() -> T>( run(); bencher.counter(ItemsCount::new(messages)).bench_local(run); } + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = PRODUCER_COUNTS, + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn sustained(bencher: Bencher, producer_count: usize) { + let mut batch = RepeatedBatch::>::new(producer_count); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +fn clone_drop_sender(bencher: Bencher) { + let (sender, _receiver) = C::channel(); + bencher.bench_local(|| drop(black_box(sender.clone()))); +} From 9c5be38f5edb9c4866cdc05110b8c5408acc0ad4 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 01:00:25 +0800 Subject: [PATCH 16/16] refactor(internal): gate cache padding on its current consumer --- asyncband/src/internal/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 7acb2f00..84309fc2 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -70,7 +70,7 @@ pub(crate) mod atomic_waker; #[allow(dead_code)] pub(crate) mod arena; -#[allow(dead_code)] +#[cfg(feature = "mpsc")] pub(crate) mod cache_padded; #[cfg(any(feature = "latch", feature = "once"))]