You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@asyncband.apache.org by ti...@apache.org on 2026/09/08 04:25:56 UTC
(asyncband) branch main updated: perf(mpsc): separate bounded waiting and unbounded reclamation (#290)
This is an automated email from the ASF dual-hosted git repository.
tisonkun pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/asyncband.git
The following commit(s) were added to refs/heads/main by this push:
new 64d7a94 perf(mpsc): separate bounded waiting and unbounded reclamation (#290)
64d7a94 is described below
commit 64d7a94ddb5ab9f2e130f8f1cc8f392301e73ebe
Author: tison <wa...@gmail.com>
AuthorDate: Tue Sep 8 12:25:51 2026 +0800
perf(mpsc): separate bounded waiting and unbounded reclamation (#290)
---
CHANGELOG.md | 6 +-
asyncband/src/internal/mod.rs | 16 +-
asyncband/src/internal/semaphore.rs | 34 --
asyncband/src/mpsc/{bounded.rs => bounded/mod.rs} | 127 +++--
asyncband/src/mpsc/bounded/ring.rs | 265 ++++------
asyncband/src/mpsc/bounded/ring_tests.rs | 200 +++++++
asyncband/src/mpsc/bounded/waiters.rs | 158 ++++++
asyncband/src/mpsc/unbounded/buffer.rs | 93 ++++
asyncband/src/mpsc/unbounded/buffer_tests.rs | 114 ++++
.../src/mpsc/{unbounded.rs => unbounded/mod.rs} | 119 +----
benchmarks/Cargo.toml | 2 +-
benchmarks/ecosystem/mpsc/adapters.rs | 56 +-
benchmarks/ecosystem/mpsc/bounded.rs | 35 ++
benchmarks/ecosystem/mpsc/support.rs | 108 +++-
benchmarks/ecosystem/mpsc/unbounded.rs | 133 +++++
tests-integration/tests/mpsc_test.rs | 574 ---------------------
tests-integration/tests/mpsc_test/backpressure.rs | 107 ++++
tests-integration/tests/mpsc_test/callbacks.rs | 328 ++++++++++++
tests-integration/tests/mpsc_test/concurrency.rs | 289 +++++++++++
tests-integration/tests/mpsc_test/main.rs | 175 +++++++
tests-integration/tests/mpsc_test/support.rs | 106 ++++
xtask/src/main.rs | 1 +
22 files changed, 2087 insertions(+), 959 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 35e870e..43ffd8e 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,10 +7,14 @@ All notable changes to this project will be documented in this file.
### Bug fixes
* Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender.
+* Notify all blocked bounded MPSC senders on receiver disconnection even when a buffered message destructor panics.
+* Avoid deadlocks when a bounded MPSC sender's waker clone callback receives from the same channel.
### 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.
+* Finish releasing buffered bounded MPSC messages even if one message destructor panics.
+* Reduce bounded MPSC contention when senders or the receiver are not waiting, improving throughput without changing capacity or cancellation semantics.
+* Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy.
## v0.7.2
diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs
index 84309fc..045a1c2 100644
--- a/asyncband/src/internal/mod.rs
+++ b/asyncband/src/internal/mod.rs
@@ -49,7 +49,9 @@ pub(crate) fn wake_all(mut wakers: impl Iterator<Item = Waker>) {
}
}
-#[cfg(feature = "mpsc")]
+// MPSC owns its receiver wait protocol; the general-purpose waker currently has no production
+// users.
+#[cfg(test)]
pub(crate) mod atomic_waker;
#[cfg(any(
@@ -99,15 +101,9 @@ pub(crate) mod value_cell;
#[allow(dead_code)]
pub(crate) mod mutex;
-#[cfg(any(
- feature = "mpsc",
- feature = "mutex",
- feature = "rwlock",
- feature = "semaphore",
-))]
-// `mpsc` uses `poll_acquire`, `release_if_nonempty`, and `notify_all`; mutexes and rwlocks use
-// `acquire`, `try_acquire`, and `release`; the public semaphore also uses the accounting methods.
-// Each single-primitive build intentionally leaves the other groups unused.
+#[cfg(any(feature = "mutex", feature = "rwlock", feature = "semaphore"))]
+// Mutexes and rwlocks use the acquire/release operations; the public semaphore also exposes
+// permit accounting. Each single-primitive build leaves part of this shared API unused.
#[allow(dead_code)]
pub(crate) mod semaphore;
diff --git a/asyncband/src/internal/semaphore.rs b/asyncband/src/internal/semaphore.rs
index f537d23..3828361 100644
--- a/asyncband/src/internal/semaphore.rs
+++ b/asyncband/src/internal/semaphore.rs
@@ -39,7 +39,6 @@ use crate::internal::mutex::Mutex;
use crate::internal::waitlist::WaitList;
use crate::internal::waitlist::WaiterId;
use crate::internal::wake_all;
-use crate::internal::waker_batch::WakerBatch;
/// The internal semaphore that provides low-level async primitives.
#[derive(Debug)]
@@ -206,39 +205,6 @@ impl Semaphore {
}
}
- /// Adds `n` permits to the semaphore if there is any waiter.
- pub fn release_if_nonempty(&self, n: usize) {
- let waiters = self.waiters.lock();
- if !waiters.is_empty() {
- self.insert_permits_with_lock(n, waiters);
- }
- }
-
- /// Adds as many permits until there is no waiter.
- pub fn notify_all(&self) {
- let mut waiters = self.waiters.lock();
- let mut wakers = WakerBatch::new();
- loop {
- match waiters.unlink_first_waiter(|node| {
- node.permits = 0;
- true
- }) {
- None => break,
- Some((id, waiter)) => {
- let remove_now = waiter.waker.is_none();
- if let Some(waker) = waiter.waker.take() {
- wakers.push(waker);
- }
- if remove_now {
- waiters.remove_unlinked_waiter(id);
- }
- }
- }
- }
- drop(waiters);
- wake_all(wakers.into_iter());
- }
-
fn insert_permits_with_lock(
&self,
mut rem: usize,
diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded/mod.rs
similarity index 77%
rename from asyncband/src/mpsc/bounded.rs
rename to asyncband/src/mpsc/bounded/mod.rs
index 550bf16..455e709 100644
--- a/asyncband/src/mpsc/bounded.rs
+++ b/asyncband/src/mpsc/bounded/mod.rs
@@ -28,15 +28,16 @@ use std::task::Poll;
use std::task::ready;
use self::ring::Ring;
+use self::waiters::SendWaiters;
use super::RecvError;
use super::SendError;
use super::TryRecvError;
use super::TrySendError;
-use crate::internal::atomic_waker::AtomicWaker;
-use crate::internal::semaphore::Acquire;
-use crate::internal::semaphore::Semaphore;
+// Ring owns capacity, publication, and waiting for the head slot. SendWaiters only schedules
+// retries after receiving frees capacity; a notification does not reserve a slot.
mod ring;
+mod waiters;
/// Creates a bounded mpsc channel with room for `buffer` queued messages.
///
@@ -49,11 +50,10 @@ mod ring;
#[track_caller]
pub fn bounded<T>(buffer: usize) -> (BoundedSender<T>, BoundedReceiver<T>) {
assert!(buffer > 0, "mpsc bounded channel requires buffer > 0");
- let state = Arc::new(BoundedState {
+ let state = Arc::new(Shared {
buffer: Ring::new(buffer),
senders: AtomicUsize::new(1),
- send_waiters: Semaphore::new(0),
- rx_waker: AtomicWaker::new(),
+ send_waiters: SendWaiters::new(),
});
let sender = BoundedSender {
state: state.clone(),
@@ -62,19 +62,17 @@ pub fn bounded<T>(buffer: usize) -> (BoundedSender<T>, BoundedReceiver<T>) {
(sender, receiver)
}
-struct BoundedState<T> {
+struct Shared<T> {
buffer: Ring<T>,
senders: AtomicUsize,
- // Notifications grant retries; only the ring determines whether buffer capacity is available.
- send_waiters: Semaphore,
- rx_waker: AtomicWaker,
+ send_waiters: SendWaiters,
}
/// The sending endpoint of a bounded mpsc channel.
///
/// Instances are created by the [`bounded`] function.
pub struct BoundedSender<T> {
- state: Arc<BoundedState<T>>,
+ state: Arc<Shared<T>>,
}
impl<T> Clone for BoundedSender<T> {
@@ -94,14 +92,8 @@ impl<T> fmt::Debug for BoundedSender<T> {
impl<T> Drop for BoundedSender<T> {
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 {
+ self.state.buffer.wake_receiver();
}
}
}
@@ -122,48 +114,38 @@ impl<T> BoundedSender<T> {
Err(TrySendError::Disconnected(value)) => return Err(SendError::new(value)),
Err(TrySendError::Full(value)) => value,
};
-
- struct SendState<'a, T> {
- sender: &'a BoundedSender<T>,
- value: Option<T>,
- acquire: Acquire<'a>,
- }
-
- impl<T> SendState<'_, T> {
- fn poll_send(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), SendError<T>>> {
- let mut value = match self.value.take() {
- Some(value) => value,
- None => return Poll::Ready(Ok(())),
- };
-
- loop {
- let poll = self.acquire.poll_once(cx.waker());
-
- value = match self.sender.try_send(value) {
- Ok(()) => return Poll::Ready(Ok(())),
- Err(TrySendError::Disconnected(value)) => {
- return Poll::Ready(Err(SendError::new(value)));
- }
- Err(TrySendError::Full(value)) => value,
- };
-
- if poll.is_ready() {
- self.acquire = self.sender.state.send_waiters.poll_acquire(1);
- } else {
- self.value = Some(value);
- return Poll::Pending;
- }
+ let mut waiter = self.state.send_waiters.waiter();
+ let mut value = Some(value);
+ poll_fn(|cx| {
+ let message = value.take().expect("send polled after completion");
+ let message = match self.try_send(message) {
+ Ok(()) => {
+ waiter.finish();
+ return Poll::Ready(Ok(()));
+ }
+ Err(TrySendError::Disconnected(message)) => {
+ waiter.finish();
+ return Poll::Ready(Err(SendError::new(message)));
+ }
+ Err(TrySendError::Full(message)) => message,
+ };
+ waiter.register(cx.waker());
+ match self.try_send(message) {
+ Ok(()) => {
+ waiter.finish();
+ Poll::Ready(Ok(()))
+ }
+ Err(TrySendError::Disconnected(message)) => {
+ waiter.finish();
+ Poll::Ready(Err(SendError::new(message)))
+ }
+ Err(TrySendError::Full(message)) => {
+ value = Some(message);
+ Poll::Pending
}
}
- }
-
- let acquire = self.state.send_waiters.poll_acquire(1);
- let mut send = SendState {
- sender: self,
- value: Some(value),
- acquire,
- };
- poll_fn(|cx| send.poll_send(cx)).await
+ })
+ .await
}
/// Attempts to send a message without waiting for capacity.
@@ -187,9 +169,7 @@ impl<T> BoundedSender<T> {
/// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30)));
/// ```
pub fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
- self.state.buffer.try_push(value)?;
- self.state.rx_waker.wake();
- Ok(())
+ self.state.buffer.try_push(value)
}
}
@@ -197,7 +177,7 @@ impl<T> BoundedSender<T> {
///
/// Instances are created by the [`bounded`] function.
pub struct BoundedReceiver<T> {
- state: Arc<BoundedState<T>>,
+ state: Arc<Shared<T>>,
}
impl<T> fmt::Debug for BoundedReceiver<T> {
@@ -208,12 +188,23 @@ impl<T> fmt::Debug for BoundedReceiver<T> {
impl<T> Drop for BoundedReceiver<T> {
fn drop(&mut self) {
+ struct DrainOnDrop<'a, T>(&'a Ring<T>);
+ impl<T> Drop for DrainOnDrop<'_, T> {
+ fn drop(&mut self) {
+ // SAFETY: This guard lives only within the exclusive receiver's drop, after close.
+ unsafe { self.0.drain() };
+ }
+ }
+
+ self.state.buffer.close();
+ let drain = DrainOnDrop(&self.state.buffer);
// 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.buffer.disconnect_receiver() };
+ let receiver_waker = self.state.buffer.take_receiver_waker();
+ // Complete notifications before dropping messages. Either kind of callback may panic;
+ // the drain guard still releases buffered values if a wake or waker drop unwinds.
self.state.send_waiters.notify_all();
drop(receiver_waker);
+ drop(drain);
}
}
@@ -267,7 +258,7 @@ impl<T> BoundedReceiver<T> {
} else {
return Poll::Ready(Err(TryRecvError::Empty));
};
- self.state.send_waiters.release_if_nonempty(1);
+ self.state.send_waiters.notify_one();
Poll::Ready(Ok(value))
}
@@ -311,7 +302,7 @@ impl<T> BoundedReceiver<T> {
Poll::Ready(Err(RecvError::Disconnected))
}
Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => {
- self.state.rx_waker.register(cx.waker());
+ self.state.buffer.register_receiver(cx.waker());
match self.try_recv_once() {
Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)),
diff --git a/asyncband/src/mpsc/bounded/ring.rs b/asyncband/src/mpsc/bounded/ring.rs
index 4eb9508..ace5aa7 100644
--- a/asyncband/src/mpsc/bounded/ring.rs
+++ b/asyncband/src/mpsc/bounded/ring.rs
@@ -18,18 +18,25 @@
use std::cell::UnsafeCell;
use std::hint::spin_loop;
use std::mem::MaybeUninit;
+use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;
use std::sync::atomic::fence;
use std::task::Poll;
+use std::task::Waker;
use crate::internal::cache_padded::CachePadded;
+use crate::internal::mutex::Mutex;
use crate::mpsc::TrySendError;
pub struct Ring<T> {
slots: Box<[Slot<T>]>,
head: CachePadded<AtomicUsize>,
tail: CachePadded<AtomicUsize>,
+ // This flag is usually stable while the consumer advances head. Sharing head
+ // for notifications would make every producer track a constantly invalidated cache line.
+ receiver_waiting: CachePadded<AtomicBool>,
+ receiver: Mutex<Option<Waker>>,
capacity: usize,
one_lap: usize,
mark_bit: usize,
@@ -65,6 +72,8 @@ impl<T> Ring<T> {
slots,
head: CachePadded::new(AtomicUsize::new(0)),
tail: CachePadded::new(AtomicUsize::new(0)),
+ receiver_waiting: CachePadded::new(AtomicBool::new(false)),
+ receiver: Mutex::new(None),
capacity,
one_lap,
mark_bit,
@@ -95,6 +104,16 @@ impl<T> Ring<T> {
// matching stamp proves the consumer completed its previous lap.
unsafe { (*slot.value.get()).write(value) };
slot.stamp.store(tail.wrapping_add(1), Ordering::Release);
+ // Publication precedes the wait check; registration pairs this fence
+ // with a second pop before the receiver is allowed to return Pending.
+ fence(Ordering::SeqCst);
+ if self.receiver_waiting.load(Ordering::Relaxed)
+ && self.receiver_waiting.swap(false, Ordering::Relaxed)
+ {
+ // Claim the notification before locking so concurrent publishers
+ // do not all queue behind the same receiver registration.
+ self.wake_receiver();
+ }
return Ok(());
}
Err(actual) => tail = actual,
@@ -106,7 +125,17 @@ impl<T> Ring<T> {
}
tail = self.tail.load(Ordering::Relaxed);
} else {
- tail = self.tail.load(Ordering::Relaxed);
+ let actual = self.tail.load(Ordering::Relaxed);
+ if actual == tail {
+ // Reserved but unpublished messages also occupy capacity. In particular, a
+ // capacity-one queue must report Full without waiting for its producer to
+ // publish the slot's stamp.
+ fence(Ordering::SeqCst);
+ if self.head.load(Ordering::Relaxed).wrapping_add(self.one_lap) == tail {
+ return Err(TrySendError::Full(value));
+ }
+ }
+ tail = actual;
}
Self::spin(&mut backoff);
}
@@ -117,7 +146,7 @@ impl<T> Ring<T> {
///
/// # Safety
///
- /// The caller must serialize all calls to `pop` and `disconnect_receiver` for this queue.
+ /// The caller must serialize all calls to `pop` and `drain` for this queue.
pub unsafe fn pop(&self) -> Poll<Option<T>> {
let mut head = self.head.load(Ordering::Relaxed);
let mut backoff = 0;
@@ -150,15 +179,71 @@ impl<T> Ring<T> {
}
}
- /// Closes the queue and drops all remaining values.
+ /// Registers the exclusive receiver, which must retry `pop` before returning Pending.
+ pub fn register_receiver(&self, waker: &Waker) {
+ let mut receiver = self.receiver.lock();
+ let old_waker = if receiver.as_ref().is_some_and(|old| old.will_wake(waker)) {
+ None
+ } else {
+ // Only the receiver registers. Producers can take the old waker while we clone,
+ // but cannot install a replacement. Clone/drop callbacks may send into this channel.
+ drop(receiver);
+ let waker = waker.clone();
+ receiver = self.receiver.lock();
+ receiver.replace(waker)
+ };
+ self.receiver_waiting.store(true, Ordering::Relaxed);
+ // Paired with the publisher's fence: either it sees this flag, or the receiver's
+ // subsequent pop sees its stamp. Taking a waker clears the flag under the same lock,
+ // so it cannot erase a newer registration without also taking responsibility for it.
+ fence(Ordering::SeqCst);
+ drop(receiver);
+ drop(old_waker);
+ }
+
+ pub fn take_receiver_waker(&self) -> Option<Waker> {
+ let mut receiver = self.receiver.lock();
+ self.receiver_waiting.store(false, Ordering::Relaxed);
+ receiver.take()
+ }
+
+ pub fn wake_receiver(&self) {
+ if let Some(waker) = self.take_receiver_waker() {
+ waker.wake();
+ }
+ }
+
+ /// Prevents subsequent sends from reserving slots. Already reserved slots still publish.
+ pub fn close(&self) {
+ self.tail.fetch_or(self.mark_bit, Ordering::SeqCst);
+ }
+
+ /// Drops all values after closing, including values whose publication is still in progress.
///
/// # 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;
- // SAFETY: The caller guarantees exclusive consumer access.
- unsafe { self.discard_until(tail) };
+ /// The queue must be closed. The caller must serialize all calls to `pop` and `drain`.
+ pub unsafe fn drain(&self) {
+ struct DrainRemaining<'a, T> {
+ ring: &'a Ring<T>,
+ tail: usize,
+ }
+ impl<T> Drop for DrainRemaining<'_, T> {
+ fn drop(&mut self) {
+ // SAFETY: The guard is scoped to the exclusive consumer's drain of a closed ring.
+ unsafe { self.ring.discard_until(self.tail) };
+ }
+ }
+
+ let tail = self.tail.load(Ordering::Relaxed);
+ debug_assert_ne!(tail & self.mark_bit, 0);
+ let remaining = DrainRemaining {
+ ring: self,
+ tail: tail & !self.mark_bit,
+ };
+ // SAFETY: The caller guarantees exclusive consumer access. The guard finishes draining if
+ // a value's destructor panics, so messages that own senders cannot retain the closed ring.
+ unsafe { self.discard_until(remaining.tail) };
}
fn advance(&self, position: usize) -> usize {
@@ -205,168 +290,12 @@ impl<T> Ring<T> {
impl<T> Drop for Ring<T> {
fn drop(&mut self) {
- let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit;
+ self.close();
// SAFETY: The queue is closed and its exclusive borrow rules out concurrent access.
- unsafe { self.discard_until(tail) };
+ unsafe { self.drain() };
}
}
#[cfg(test)]
-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::Ring;
- use super::TrySendError;
-
- #[test]
- fn bounded_queue_preserves_capacity_and_fifo_order() {
- let queue = Ring::new(3);
- for value in 0..3 {
- assert!(queue.try_push(value).is_ok());
- }
- assert!(matches!(queue.try_push(3), Err(TrySendError::Full(3))));
- for value in 0..3 {
- // SAFETY: This thread is the only consumer.
- assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value)));
- }
- // SAFETY: This thread is the only consumer.
- assert_eq!(unsafe { queue.pop() }, Poll::Ready(None));
-
- for value in 3..12 {
- assert!(queue.try_push(value).is_ok());
- // SAFETY: This thread is the only consumer.
- assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value)));
- }
- }
-
- #[test]
- fn bounded_queue_does_not_report_empty_behind_an_unpublished_head() {
- let queue = Ring::new(2);
-
- // Pause a synthetic producer after reserving and initializing slot 0, before publishing
- // its stamp. Another producer can finish sending into slot 1 in the meantime.
- queue.tail.store(1, Ordering::SeqCst);
- let slot = &queue.slots[0];
- // SAFETY: advancing the tail reserved this initially empty slot for the synthetic producer.
- unsafe { (*slot.value.get()).write(1) };
- let later_send = queue.try_push(2);
- // SAFETY: This thread is the only consumer, even while a producer is unpublished.
- let receive = unsafe { queue.pop() };
-
- // Finish publication before asserting so even a failed assertion can safely drop the queue.
- slot.stamp.store(1, Ordering::Release);
- assert!(later_send.is_ok());
- assert_eq!(receive, Poll::Pending);
- // SAFETY: This thread is the only consumer.
- 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]
- fn bounded_queue_coordinates_multiple_producers() {
- let queue = Arc::new(Ring::new(4));
- let producers: Vec<_> = (0..2)
- .map(|producer| {
- let queue = queue.clone();
- thread::spawn(move || {
- for offset in 0..32 {
- let mut value = producer * 32 + offset;
- loop {
- match queue.try_push(value) {
- Ok(()) => break,
- Err(TrySendError::Full(returned)) => {
- value = returned;
- thread::yield_now();
- }
- Err(TrySendError::Disconnected(_)) => panic!("queue disconnected"),
- }
- }
- }
- })
- })
- .collect();
-
- let mut values = Vec::new();
- while values.len() < 64 {
- // SAFETY: Worker threads only push; this thread is the only consumer.
- if let Poll::Ready(Some(value)) = unsafe { queue.pop() } {
- values.push(value);
- } else {
- thread::yield_now();
- }
- }
- for producer in producers {
- producer.join().unwrap();
- }
- values.sort_unstable();
- assert_eq!(values, (0..64).collect::<Vec<_>>());
- }
-
- #[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 = Ring::new(3);
-
- // Positions: 0, 1, 2 (then tail wraps to 8).
- for counter in &drops[..3] {
- assert!(queue.try_push(DropSpy(counter)).is_ok());
- }
-
- // Free slot 0, then reuse it on the next lap at position 8.
- // SAFETY: This thread is the only consumer.
- let popped = unsafe { queue.pop() };
- assert!(matches!(popped, Poll::Ready(Some(_))));
- drop(popped);
- assert_eq!(drops[0].load(Ordering::Relaxed), 1);
- assert!(queue.try_push(DropSpy(&drops[3])).is_ok());
-
- // The pending range is positions 1 -> 2 -> 8 -> 9, not a contiguous integer range.
- assert_eq!(queue.head.load(Ordering::Relaxed), 1);
- assert_eq!(queue.tail.load(Ordering::Relaxed), queue.one_lap + 1);
-
- // 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() {
- 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"
- );
- }
- }
-}
+#[path = "ring_tests.rs"]
+mod tests;
diff --git a/asyncband/src/mpsc/bounded/ring_tests.rs b/asyncband/src/mpsc/bounded/ring_tests.rs
new file mode 100644
index 0000000..d6d5043
--- /dev/null
+++ b/asyncband/src/mpsc/bounded/ring_tests.rs
@@ -0,0 +1,200 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::task::Poll;
+use std::thread;
+
+use super::Ring;
+use super::TrySendError;
+
+#[test]
+fn bounded_queue_preserves_capacity_and_fifo_order() {
+ let queue = Ring::new(3);
+ for value in 0..3 {
+ assert!(queue.try_push(value).is_ok());
+ }
+ assert!(matches!(queue.try_push(3), Err(TrySendError::Full(3))));
+ for value in 0..3 {
+ // SAFETY: This thread is the only consumer.
+ assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value)));
+ }
+ // SAFETY: This thread is the only consumer.
+ assert_eq!(unsafe { queue.pop() }, Poll::Ready(None));
+
+ for value in 3..12 {
+ assert!(queue.try_push(value).is_ok());
+ // SAFETY: This thread is the only consumer.
+ assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value)));
+ }
+}
+
+#[test]
+fn bounded_queue_does_not_report_empty_behind_an_unpublished_head() {
+ let queue = Ring::new(2);
+
+ // Pause a synthetic producer after reserving and initializing slot 0, before publishing
+ // its stamp. Another producer can finish sending into slot 1 in the meantime.
+ queue.tail.store(1, Ordering::SeqCst);
+ let slot = &queue.slots[0];
+ // SAFETY: advancing the tail reserved this initially empty slot for the synthetic producer.
+ unsafe { (*slot.value.get()).write(1) };
+ let later_send = queue.try_push(2);
+ // SAFETY: This thread is the only consumer, even while a producer is unpublished.
+ let receive = unsafe { queue.pop() };
+
+ // Finish publication before asserting so even a failed assertion can safely drop the queue.
+ slot.stamp.store(1, Ordering::Release);
+ assert!(later_send.is_ok());
+ assert_eq!(receive, Poll::Pending);
+ // SAFETY: This thread is the only consumer.
+ 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]
+fn unpublished_reservations_count_toward_capacity() {
+ let queue = Arc::new(Ring::new(1));
+ queue.tail.store(queue.one_lap, Ordering::SeqCst);
+ let (done, completed) = std::sync::mpsc::channel();
+ let producer = {
+ let queue = queue.clone();
+ thread::spawn(move || done.send(queue.try_push(2)).unwrap())
+ };
+ #[cfg(not(miri))]
+ let result = completed.recv_timeout(std::time::Duration::from_secs(10));
+ // Miri reports a deadlock directly instead of relying on an interpretation-time deadline.
+ #[cfg(miri)]
+ let result = completed.recv();
+ // Finish the synthetic reservation even if the other producer stalled. This lets the
+ // worker and the ring's destructor finish before the failure is reported.
+ let slot = &queue.slots[0];
+ // SAFETY: Advancing the tail above exclusively reserved the initially empty slot.
+ unsafe { (*slot.value.get()).write(1) };
+ slot.stamp.store(1, Ordering::Release);
+ producer.join().unwrap();
+ assert!(matches!(result, Ok(Err(TrySendError::Full(2)))));
+ // SAFETY: Both producers have finished and this thread is the only consumer.
+ assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(1)));
+}
+
+#[test]
+fn bounded_queue_coordinates_multiple_producers() {
+ let queue = Arc::new(Ring::new(4));
+ let producers: Vec<_> = (0..2)
+ .map(|producer| {
+ let queue = queue.clone();
+ thread::spawn(move || {
+ for offset in 0..32 {
+ let mut value = producer * 32 + offset;
+ loop {
+ match queue.try_push(value) {
+ Ok(()) => break,
+ Err(TrySendError::Full(returned)) => {
+ value = returned;
+ thread::yield_now();
+ }
+ Err(TrySendError::Disconnected(_)) => panic!("queue disconnected"),
+ }
+ }
+ }
+ })
+ })
+ .collect();
+
+ let mut values = Vec::new();
+ while values.len() < 64 {
+ // SAFETY: Worker threads only push; this thread is the only consumer.
+ if let Poll::Ready(Some(value)) = unsafe { queue.pop() } {
+ values.push(value);
+ } else {
+ thread::yield_now();
+ }
+ }
+ for producer in producers {
+ producer.join().unwrap();
+ }
+ values.sort_unstable();
+ assert_eq!(values, (0..64).collect::<Vec<_>>());
+}
+
+#[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 = Ring::new(3);
+
+ // Positions: 0, 1, 2 (then tail wraps to 8).
+ for counter in &drops[..3] {
+ assert!(queue.try_push(DropSpy(counter)).is_ok());
+ }
+
+ // Free slot 0, then reuse it on the next lap at position 8.
+ // SAFETY: This thread is the only consumer.
+ let popped = unsafe { queue.pop() };
+ assert!(matches!(popped, Poll::Ready(Some(_))));
+ drop(popped);
+ assert_eq!(drops[0].load(Ordering::Relaxed), 1);
+ assert!(queue.try_push(DropSpy(&drops[3])).is_ok());
+
+ // The pending range is positions 1 -> 2 -> 8 -> 9, not a contiguous integer range.
+ assert_eq!(queue.head.load(Ordering::Relaxed), 1);
+ assert_eq!(queue.tail.load(Ordering::Relaxed), queue.one_lap + 1);
+
+ queue.close();
+ // SAFETY: The queue is closed and this thread is the only consumer.
+ unsafe { queue.drain() };
+
+ // `discard_until` must dispose every value exactly once, including position 8.
+ for (value, counter) in drops.iter().enumerate() {
+ assert_eq!(
+ counter.load(Ordering::Relaxed),
+ 1,
+ "value {value} was dropped an unexpected number of times"
+ );
+ }
+
+ // Queue Drop calls discard_until again; it must see head == tail and not redrop.
+ drop(queue);
+ for (value, counter) in drops.iter().enumerate() {
+ assert_eq!(
+ counter.load(Ordering::Relaxed),
+ 1,
+ "value {value} was dropped more than once"
+ );
+ }
+}
diff --git a/asyncband/src/mpsc/bounded/waiters.rs b/asyncband/src/mpsc/bounded/waiters.rs
new file mode 100644
index 0000000..3a34e4f
--- /dev/null
+++ b/asyncband/src/mpsc/bounded/waiters.rs
@@ -0,0 +1,158 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+use std::task::Waker;
+
+use crate::internal::mutex::Mutex;
+use crate::internal::waitlist::WaitList;
+use crate::internal::waitlist::WaiterId;
+use crate::internal::wake_all;
+use crate::internal::waker_batch::WakerBatch;
+
+// A wake grants a retry, not a capacity permit. Keeping notified nodes until their future
+// consumes the notification lets cancellation pass an unused retry to the next sender.
+pub struct SendWaiters {
+ waiting: AtomicBool,
+ queue: Mutex<WaitList<Option<Waker>>>,
+}
+
+impl SendWaiters {
+ pub fn new() -> Self {
+ Self {
+ waiting: AtomicBool::new(false),
+ queue: Mutex::new(WaitList::new()),
+ }
+ }
+
+ pub fn waiter(&self) -> SendWaiter<'_> {
+ SendWaiter {
+ waiters: self,
+ index: None,
+ }
+ }
+
+ pub fn notify_one(&self) {
+ // The receiver publishes its head with SeqCst before checking this flag. Registration
+ // publishes the flag with SeqCst before rechecking capacity (which uses a SeqCst fence).
+ // Either the receiver sees the registration or the sender sees the released capacity.
+ if !self.waiting.load(Ordering::SeqCst) {
+ return;
+ }
+ let waker = {
+ let mut queue = self.queue.lock();
+ let waker = queue
+ .unlink_first_waiter(|_| true)
+ .and_then(|(_, waker)| waker.take());
+ self.waiting.store(!queue.is_empty(), Ordering::SeqCst);
+ waker
+ };
+ if let Some(waker) = waker {
+ waker.wake();
+ }
+ }
+
+ pub fn notify_all(&self) {
+ let mut wakers = WakerBatch::new();
+ {
+ let mut queue = self.queue.lock();
+ while let Some((_, waker)) = queue.unlink_first_waiter(|_| true) {
+ if let Some(waker) = waker.take() {
+ wakers.push(waker);
+ }
+ }
+ self.waiting.store(false, Ordering::SeqCst);
+ }
+ wake_all(wakers.into_iter());
+ }
+}
+
+pub struct SendWaiter<'a> {
+ waiters: &'a SendWaiters,
+ index: Option<WaiterId>,
+}
+
+impl SendWaiter<'_> {
+ // The caller must retry sending after registration, before returning Pending.
+ pub fn register(&mut self, waker: &Waker) {
+ let mut new_waker = None;
+ loop {
+ let mut queue = self.waiters.queue.lock();
+ if let Some(index) = self.index {
+ if queue
+ .waiter_mut(index)
+ .as_ref()
+ .is_some_and(|current| current.will_wake(waker))
+ {
+ return;
+ }
+ }
+ let Some(waker) = new_waker.take() else {
+ // Waker callbacks may reenter the channel, including clone and drop callbacks.
+ drop(queue);
+ new_waker = Some(waker.clone());
+ continue;
+ };
+ let old_waker = if let Some(index) = self.index {
+ let node = queue.waiter_mut(index);
+ if node.is_some() {
+ node.replace(waker)
+ } else {
+ queue.remove_unlinked_waiter(index);
+ self.index = Some(queue.push_back(Some(waker)));
+ None
+ }
+ } else {
+ self.index = Some(queue.push_back(Some(waker)));
+ None
+ };
+ self.waiters.waiting.store(true, Ordering::SeqCst);
+ drop(queue);
+ drop(old_waker);
+ return;
+ }
+ }
+
+ pub fn finish(&mut self) {
+ if let Some(index) = self.index.take() {
+ drop(self.remove(index));
+ }
+ }
+
+ fn remove(&self, index: WaiterId) -> Option<Waker> {
+ let mut queue = self.waiters.queue.lock();
+ queue.unlink_waiter(index, |_| true);
+ let waker = queue.remove_unlinked_waiter(index);
+ self.waiters
+ .waiting
+ .store(!queue.is_empty(), Ordering::SeqCst);
+ waker
+ }
+}
+
+impl Drop for SendWaiter<'_> {
+ fn drop(&mut self) {
+ if let Some(index) = self.index.take() {
+ let waker = self.remove(index);
+ if waker.is_none() {
+ self.waiters.notify_one();
+ }
+ drop(waker);
+ }
+ }
+}
diff --git a/asyncband/src/mpsc/unbounded/buffer.rs b/asyncband/src/mpsc/unbounded/buffer.rs
new file mode 100644
index 0000000..7c11449
--- /dev/null
+++ b/asyncband/src/mpsc/unbounded/buffer.rs
@@ -0,0 +1,93 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::VecDeque;
+use std::mem;
+
+// Bound the inline storage retained by a partial batch. Boxed payloads belong to individual
+// messages, not these backing allocations. Empty buffers are reused without retaining peak size.
+pub const SEGMENT_BYTES: usize = 32 * 1024;
+
+pub struct Buffer<T> {
+ writable: VecDeque<T>,
+ sealed: VecDeque<VecDeque<T>>,
+ spare: VecDeque<T>,
+}
+
+impl<T> Buffer<T> {
+ pub fn new() -> Self {
+ Self {
+ writable: VecDeque::new(),
+ sealed: VecDeque::new(),
+ spare: VecDeque::new(),
+ }
+ }
+
+ fn segment_capacity() -> usize {
+ if mem::size_of::<T>() == 0 {
+ return usize::MAX;
+ }
+ let limit = (SEGMENT_BYTES / mem::size_of::<T>()).max(1);
+ // Power-of-two limits let VecDeque grow naturally without exceeding the segment budget.
+ 1 << (usize::BITS - 1 - limit.leading_zeros())
+ }
+
+ pub fn push(&mut self, value: T) {
+ if self.writable.len() == Self::segment_capacity() {
+ let next = if self.spare.capacity() == 0 {
+ VecDeque::with_capacity(Self::segment_capacity())
+ } else {
+ mem::take(&mut self.spare)
+ };
+ let sealed = mem::replace(&mut self.writable, next);
+ self.sealed.push_back(sealed);
+ }
+ self.writable.push_back(value);
+ }
+
+ pub fn refill(&mut self, batch: &mut VecDeque<T>) {
+ debug_assert!(batch.is_empty());
+ if let Some(sealed) = self.sealed.pop_front() {
+ // Keep one empty segment for the next producer rollover. Every other consumed
+ // segment is released, so retained payload storage does not track peak occupancy.
+ self.spare = mem::replace(batch, sealed);
+ if self.sealed.is_empty()
+ && self.sealed.capacity() * mem::size_of::<VecDeque<T>>() > 1024
+ {
+ self.sealed = VecDeque::new();
+ }
+ } else if !self.writable.is_empty() {
+ self.spare = VecDeque::new();
+ mem::swap(batch, &mut self.writable);
+ }
+ }
+}
+
+pub fn pop_batch<T>(batch: &mut VecDeque<T>) -> T {
+ if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::<T>()) > SEGMENT_BYTES {
+ // Retire the allocation on the last value, outside the inbox lock. Keep this as a tail
+ // expression to avoid intermediate storage for large inline values.
+ mem::take(batch).pop_front()
+ } else {
+ batch.pop_front()
+ }
+ .expect("receiver batch must not be empty")
+}
+
+#[cfg(test)]
+#[path = "buffer_tests.rs"]
+mod tests;
diff --git a/asyncband/src/mpsc/unbounded/buffer_tests.rs b/asyncband/src/mpsc/unbounded/buffer_tests.rs
new file mode 100644
index 0000000..de40b1e
--- /dev/null
+++ b/asyncband/src/mpsc/unbounded/buffer_tests.rs
@@ -0,0 +1,114 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::collections::VecDeque;
+
+use super::Buffer;
+use super::SEGMENT_BYTES;
+use super::pop_batch;
+
+fn allocated_bytes<T>(buffer: &Buffer<T>, batch: &VecDeque<T>) -> usize {
+ let slots = batch.capacity()
+ + buffer.writable.capacity()
+ + buffer.spare.capacity()
+ + buffer.sealed.iter().map(VecDeque::capacity).sum::<usize>();
+ slots * size_of::<T>()
+}
+
+fn receive<T>(buffer: &mut Buffer<T>, batch: &mut VecDeque<T>) -> T {
+ if batch.is_empty() {
+ buffer.refill(batch);
+ }
+ pop_batch(batch)
+}
+
+#[test]
+fn messages_arriving_during_a_batch_remain_in_fifo_order() {
+ let mut buffer = Buffer::new();
+ let mut batch = VecDeque::new();
+ buffer.push(1);
+ buffer.push(2);
+ assert_eq!(receive(&mut buffer, &mut batch), 1);
+ buffer.push(3);
+ assert_eq!(receive(&mut buffer, &mut batch), 2);
+ assert_eq!(receive(&mut buffer, &mut batch), 3);
+ buffer.refill(&mut batch);
+ assert!(batch.is_empty());
+}
+
+#[test]
+fn a_partial_drain_reclaims_segments_and_preserves_new_sends() {
+ let mut buffer = Buffer::new();
+ let mut batch = VecDeque::new();
+ for value in 0..1024usize {
+ buffer.push([value; 128]);
+ }
+ let peak = allocated_bytes(&buffer, &batch);
+ for value in 0..512 {
+ assert_eq!(receive(&mut buffer, &mut batch), [value; 128]);
+ }
+ assert!(allocated_bytes(&buffer, &batch) <= peak * 3 / 4);
+ // This value must stay behind both the current batch and the sealed segments.
+ buffer.push([1024; 128]);
+ for value in 512..=1024 {
+ assert_eq!(receive(&mut buffer, &mut batch), [value; 128]);
+ }
+ assert!(allocated_bytes(&buffer, &batch) <= 2 * SEGMENT_BYTES);
+ buffer.refill(&mut batch);
+ assert!(batch.is_empty());
+}
+
+#[test]
+fn small_batches_are_reused_on_refill() {
+ let mut buffer = Buffer::new();
+ let mut batch = VecDeque::new();
+ for value in 0..32usize {
+ buffer.push(value);
+ }
+ assert_eq!(receive(&mut buffer, &mut batch), 0);
+ let capacity = batch.capacity();
+ for value in 1..32 {
+ assert_eq!(receive(&mut buffer, &mut batch), value);
+ }
+ assert_eq!(batch.capacity(), capacity);
+ buffer.push(32);
+ assert_eq!(receive(&mut buffer, &mut batch), 32);
+ assert_eq!(buffer.writable.capacity(), capacity);
+}
+
+#[test]
+fn oversized_inline_values_release_the_allocation_on_the_last_receive() {
+ let mut buffer = Buffer::new();
+ let mut batch = VecDeque::new();
+ buffer.push([7u8; SEGMENT_BYTES + 1]);
+ assert_eq!(receive(&mut buffer, &mut batch), [7u8; SEGMENT_BYTES + 1]);
+ assert_eq!(allocated_bytes(&buffer, &batch), 0);
+}
+
+#[test]
+fn zero_sized_values_do_not_require_segments() {
+ let mut buffer = Buffer::new();
+ let mut batch = VecDeque::new();
+ for _ in 0..32 {
+ buffer.push(());
+ }
+ for _ in 0..32 {
+ receive(&mut buffer, &mut batch);
+ }
+ buffer.refill(&mut batch);
+ assert!(batch.is_empty());
+}
diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded/mod.rs
similarity index 71%
rename from asyncband/src/mpsc/unbounded.rs
rename to asyncband/src/mpsc/unbounded/mod.rs
index 92be699..7c69ac4 100644
--- a/asyncband/src/mpsc/unbounded.rs
+++ b/asyncband/src/mpsc/unbounded/mod.rs
@@ -29,24 +29,29 @@ use std::task::Context;
use std::task::Poll;
use std::task::Waker;
+use self::buffer::Buffer;
+use self::buffer::pop_batch;
use super::RecvError;
use super::SendError;
use super::TryRecvError;
use crate::internal::mutex::Mutex;
+// Buffer owns segmentation and reclamation; Inbox serializes enqueueing, registration, and close.
+mod buffer;
+
/// Creates an unbounded mpsc channel whose send operation never waits for capacity.
///
/// 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.
+/// Storage is reclaimed incrementally as messages are received. A bounded amount of empty
+/// storage may be retained for reuse, independently of the channel's previous peak occupancy.
pub fn unbounded<T>() -> (UnboundedSender<T>, UnboundedReceiver<T>) {
- let state = Arc::new(UnboundedState {
+ let state = Arc::new(Shared {
senders: AtomicUsize::new(1),
inbox: Mutex::new(Inbox {
- messages: VecDeque::new(),
+ buffer: Buffer::new(),
receiver_alive: true,
rx_waker: None,
}),
@@ -61,7 +66,7 @@ pub fn unbounded<T>() -> (UnboundedSender<T>, UnboundedReceiver<T>) {
(sender, receiver)
}
-struct UnboundedState<T> {
+struct Shared<T> {
// Endpoint cloning and ordinary drops do not contend with message traffic.
senders: AtomicUsize,
inbox: Mutex<Inbox<T>>,
@@ -70,7 +75,7 @@ struct UnboundedState<T> {
// 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<T> {
- messages: VecDeque<T>,
+ buffer: Buffer<T>,
receiver_alive: bool,
rx_waker: Option<Waker>,
}
@@ -79,7 +84,7 @@ struct Inbox<T> {
///
/// Instances are created by the [`unbounded`] function.
pub struct UnboundedSender<T> {
- state: Arc<UnboundedState<T>>,
+ state: Arc<Shared<T>>,
}
impl<T> Clone for UnboundedSender<T> {
@@ -119,7 +124,7 @@ impl<T> UnboundedSender<T> {
if !state.receiver_alive {
return Err(SendError::new(value));
}
- state.messages.push_back(value);
+ state.buffer.push(value);
state.rx_waker.take()
};
if let Some(waker) = waker {
@@ -133,7 +138,7 @@ impl<T> UnboundedSender<T> {
///
/// Instances are created by the [`unbounded`] function.
pub struct UnboundedReceiver<T> {
- state: Arc<UnboundedState<T>>,
+ state: Arc<Shared<T>>,
// Only accessed through `get_mut`; the mutex preserves Sync for Send-only payloads.
batch: Mutex<VecDeque<T>>,
}
@@ -150,7 +155,10 @@ impl<T> Drop for UnboundedReceiver<T> {
let (shared, waker) = {
let mut state = self.state.inbox.lock();
state.receiver_alive = false;
- (mem::take(&mut state.messages), state.rx_waker.take())
+ (
+ mem::replace(&mut state.buffer, Buffer::new()),
+ state.rx_waker.take(),
+ )
};
// Destructors may send again. A waker may also own a sender and form an ownership cycle.
drop((batch, shared, waker));
@@ -184,7 +192,8 @@ impl<T> UnboundedReceiver<T> {
let batch = self.batch.get_mut();
if batch.is_empty() {
let mut state = self.state.inbox.lock();
- if state.messages.is_empty() {
+ state.buffer.refill(batch);
+ if batch.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 {
@@ -193,7 +202,6 @@ impl<T> UnboundedReceiver<T> {
TryRecvError::Empty
});
}
- mem::swap(batch, &mut state.messages);
}
Ok(pop_batch(batch))
}
@@ -241,8 +249,8 @@ impl<T> UnboundedReceiver<T> {
let mut new_waker = None;
loop {
let mut state = self.state.inbox.lock();
- if !state.messages.is_empty() {
- mem::swap(batch, &mut state.messages);
+ state.buffer.refill(batch);
+ if !batch.is_empty() {
drop(state);
return Poll::Ready(Ok(pop_batch(batch)));
}
@@ -270,86 +278,3 @@ impl<T> UnboundedReceiver<T> {
// No operation relies on a pinned location for the receiver batch or its values.
impl<T> Unpin for UnboundedReceiver<T> {}
-
-// Retain small buffers for reuse, measuring inline storage rather than separately boxed payloads.
-const BATCH_CACHE_BYTES: usize = 64 * 1024;
-
-fn pop_batch<T>(batch: &mut VecDeque<T>) -> T {
- if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::<T>()) > 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/benchmarks/Cargo.toml b/benchmarks/Cargo.toml
index 2104a82..19e5ac3 100644
--- a/benchmarks/Cargo.toml
+++ b/benchmarks/Cargo.toml
@@ -50,7 +50,7 @@ asyncband = { workspace = true, features = [
divan = { workspace = true }
flume = { workspace = true, features = ["async"] }
pollster = { workspace = true }
-tokio = { workspace = true, features = ["sync"] }
+tokio = { workspace = true, features = ["rt-multi-thread", "sync"] }
waitgroup = { workspace = true }
[[bench]]
diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs
index 1fd4dbb..506f167 100644
--- a/benchmarks/ecosystem/mpsc/adapters.rs
+++ b/benchmarks/ecosystem/mpsc/adapters.rs
@@ -16,6 +16,7 @@
// under the License.
use std::fmt::Debug;
+use std::future::Future;
use std::task::Context;
use crate::support::poll_ready;
@@ -26,7 +27,7 @@ pub struct AsyncChannel;
pub struct Flume;
pub trait BoundedMpsc: Send + Sync + 'static {
- type Sender: Clone + Send + 'static;
+ type Sender: Clone + Send + Sync + 'static;
type Receiver: Send + 'static;
fn channel(capacity: usize) -> (Self::Sender, Self::Receiver);
@@ -34,18 +35,21 @@ pub trait BoundedMpsc: Send + Sync + 'static {
fn try_recv(receiver: &mut Self::Receiver) -> usize;
fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>);
fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize;
+ fn send_async(sender: &Self::Sender, value: usize) -> impl Future<Output = ()> + Send;
+ fn recv_async(receiver: &mut Self::Receiver) -> impl Future<Output = usize> + Send;
fn send_blocking(sender: &Self::Sender, value: usize);
fn recv_blocking(receiver: &mut Self::Receiver) -> usize;
}
pub trait UnboundedMpsc<T = usize>: Send + Sync + 'static {
- type Sender: Clone + Send + 'static;
+ type Sender: Clone + Send + Sync + 'static;
type Receiver: Send + 'static;
fn channel() -> (Self::Sender, Self::Receiver);
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_async(receiver: &mut Self::Receiver) -> impl Future<Output = T> + Send;
fn recv_blocking(receiver: &mut Self::Receiver) -> T;
}
@@ -73,6 +77,14 @@ impl BoundedMpsc for Asyncband {
poll_ready(receiver.recv(), context).unwrap()
}
+ async fn send_async(sender: &Self::Sender, value: usize) {
+ sender.send(value).await.unwrap();
+ }
+
+ async fn recv_async(receiver: &mut Self::Receiver) -> usize {
+ receiver.recv().await.unwrap()
+ }
+
fn send_blocking(sender: &Self::Sender, value: usize) {
pollster::block_on(sender.send(value)).unwrap();
}
@@ -106,6 +118,14 @@ impl BoundedMpsc for Tokio {
poll_ready(receiver.recv(), context).unwrap()
}
+ async fn send_async(sender: &Self::Sender, value: usize) {
+ sender.send(value).await.unwrap();
+ }
+
+ async fn recv_async(receiver: &mut Self::Receiver) -> usize {
+ receiver.recv().await.unwrap()
+ }
+
fn send_blocking(sender: &Self::Sender, value: usize) {
pollster::block_on(sender.send(value)).unwrap();
}
@@ -139,6 +159,14 @@ impl BoundedMpsc for AsyncChannel {
poll_ready(receiver.recv(), context).unwrap()
}
+ async fn send_async(sender: &Self::Sender, value: usize) {
+ sender.send(value).await.unwrap();
+ }
+
+ async fn recv_async(receiver: &mut Self::Receiver) -> usize {
+ receiver.recv().await.unwrap()
+ }
+
fn send_blocking(sender: &Self::Sender, value: usize) {
pollster::block_on(sender.send(value)).unwrap();
}
@@ -172,6 +200,14 @@ impl BoundedMpsc for Flume {
poll_ready(receiver.recv_async(), context).unwrap()
}
+ async fn send_async(sender: &Self::Sender, value: usize) {
+ sender.send_async(value).await.unwrap();
+ }
+
+ async fn recv_async(receiver: &mut Self::Receiver) -> usize {
+ receiver.recv_async().await.unwrap()
+ }
+
fn send_blocking(sender: &Self::Sender, value: usize) {
pollster::block_on(sender.send_async(value)).unwrap();
}
@@ -201,6 +237,10 @@ impl<T: Debug + Send + 'static> UnboundedMpsc<T> for Asyncband {
poll_ready(receiver.recv(), context).unwrap()
}
+ async fn recv_async(receiver: &mut Self::Receiver) -> T {
+ receiver.recv().await.unwrap()
+ }
+
fn recv_blocking(receiver: &mut Self::Receiver) -> T {
pollster::block_on(receiver.recv()).unwrap()
}
@@ -226,6 +266,10 @@ impl<T: Debug + Send + 'static> UnboundedMpsc<T> for Tokio {
poll_ready(receiver.recv(), context).unwrap()
}
+ async fn recv_async(receiver: &mut Self::Receiver) -> T {
+ receiver.recv().await.unwrap()
+ }
+
fn recv_blocking(receiver: &mut Self::Receiver) -> T {
pollster::block_on(receiver.recv()).unwrap()
}
@@ -251,6 +295,10 @@ impl<T: Debug + Send + 'static> UnboundedMpsc<T> for AsyncChannel {
poll_ready(receiver.recv(), context).unwrap()
}
+ async fn recv_async(receiver: &mut Self::Receiver) -> T {
+ receiver.recv().await.unwrap()
+ }
+
fn recv_blocking(receiver: &mut Self::Receiver) -> T {
pollster::block_on(receiver.recv()).unwrap()
}
@@ -276,6 +324,10 @@ impl<T: Debug + Send + 'static> UnboundedMpsc<T> for Flume {
poll_ready(receiver.recv_async(), context).unwrap()
}
+ async fn recv_async(receiver: &mut Self::Receiver) -> T {
+ receiver.recv_async().await.unwrap()
+ }
+
fn recv_blocking(receiver: &mut Self::Receiver) -> T {
pollster::block_on(receiver.recv_async()).unwrap()
}
diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs
index 80495e9..d388cc6 100644
--- a/benchmarks/ecosystem/mpsc/bounded.rs
+++ b/benchmarks/ecosystem/mpsc/bounded.rs
@@ -30,6 +30,7 @@ use super::support::Bounded;
use super::support::ConcurrentBatch;
use super::support::PRODUCER_COUNTS;
use super::support::RepeatedBatch;
+use super::support::RepeatedTasks;
use crate::support::bench_context;
#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])]
@@ -84,3 +85,37 @@ fn clone_drop_sender<C: BoundedMpsc>(bencher: Bencher) {
let (sender, _receiver) = C::channel(BOUNDED_CAPACITY);
bencher.bench_local(|| drop(black_box(sender.clone())));
}
+
+#[divan::bench(
+ types = [Asyncband, Tokio, AsyncChannel, Flume],
+ consts = [1, 4096],
+ args = PRODUCER_COUNTS,
+ sample_count = 50,
+ sample_size = 1,
+ counter = ItemsCount::new(BATCH_MESSAGES),
+)]
+fn sustained_capacity<C: BoundedMpsc, const CAPACITY: usize>(
+ bencher: Bencher,
+ producer_count: usize,
+) {
+ let mut batch = RepeatedBatch::<Bounded<C, CAPACITY>>::new(producer_count);
+ batch.run();
+ bencher.bench_local(|| batch.run());
+}
+
+#[divan::bench(
+ types = [Asyncband, Tokio, AsyncChannel, Flume],
+ consts = [1, 64, 4096],
+ args = [(1, 0), (4, 0), (1, 4), (4, 4), (8, 4)],
+ sample_count = 50,
+ sample_size = 1,
+ counter = ItemsCount::new(BATCH_MESSAGES),
+)]
+fn scheduled<C: BoundedMpsc, const CAPACITY: usize>(
+ bencher: Bencher,
+ (producers, workers): (usize, usize),
+) {
+ let mut batch = RepeatedTasks::<Bounded<C, CAPACITY>>::new(producers, workers);
+ batch.run();
+ bencher.bench_local(|| batch.run());
+}
diff --git a/benchmarks/ecosystem/mpsc/support.rs b/benchmarks/ecosystem/mpsc/support.rs
index 11edb58..3e777e0 100644
--- a/benchmarks/ecosystem/mpsc/support.rs
+++ b/benchmarks/ecosystem/mpsc/support.rs
@@ -15,6 +15,7 @@
// specific language governing permissions and limitations
// under the License.
+use std::future::Future;
use std::marker::PhantomData;
use std::sync::Arc;
use std::sync::Barrier;
@@ -33,28 +34,38 @@ pub const BATCH_MESSAGES: usize = 16_384;
pub const PRODUCER_COUNTS: &[usize] = &[1, 2, 4, 8];
pub trait ConcurrentMpsc: Send + Sync + 'static {
- type Sender: Clone + Send + 'static;
+ type Sender: Clone + Send + Sync + 'static;
type Receiver: Send + 'static;
fn channel() -> (Self::Sender, Self::Receiver);
fn send(sender: &Self::Sender, value: usize);
fn recv(receiver: &mut Self::Receiver) -> usize;
+ fn send_async(sender: &Self::Sender, value: usize) -> impl Future<Output = ()> + Send;
+ fn recv_async(receiver: &mut Self::Receiver) -> impl Future<Output = usize> + Send;
}
-pub struct Bounded<C>(PhantomData<C>);
+pub struct Bounded<C, const CAPACITY: usize = BOUNDED_CAPACITY>(PhantomData<C>);
-impl<C: BoundedMpsc> ConcurrentMpsc for Bounded<C> {
+impl<C: BoundedMpsc, const CAPACITY: usize> ConcurrentMpsc for Bounded<C, CAPACITY> {
type Receiver = C::Receiver;
type Sender = C::Sender;
fn channel() -> (Self::Sender, Self::Receiver) {
- C::channel(BOUNDED_CAPACITY)
+ C::channel(CAPACITY)
}
fn send(sender: &Self::Sender, value: usize) {
C::send_blocking(sender, value);
}
+ async fn send_async(sender: &Self::Sender, value: usize) {
+ C::send_async(sender, value).await;
+ }
+
+ async fn recv_async(receiver: &mut Self::Receiver) -> usize {
+ C::recv_async(receiver).await
+ }
+
fn recv(receiver: &mut Self::Receiver) -> usize {
C::recv_blocking(receiver)
}
@@ -74,6 +85,14 @@ impl<C: UnboundedMpsc> ConcurrentMpsc for Unbounded<C> {
C::send(sender, value);
}
+ async fn send_async(sender: &Self::Sender, value: usize) {
+ C::send(sender, value);
+ }
+
+ async fn recv_async(receiver: &mut Self::Receiver) -> usize {
+ C::recv_async(receiver).await
+ }
+
fn recv(receiver: &mut Self::Receiver) -> usize {
C::recv_blocking(receiver)
}
@@ -198,3 +217,84 @@ impl<C: ConcurrentMpsc> Drop for RepeatedBatch<C> {
}
}
}
+
+// Exercise executor wakeups as well as channel traffic. Reuse tasks, threads, and channel storage
+// across samples; a current-thread runtime also exposes polling that monopolizes the executor.
+pub struct RepeatedTasks<C: ConcurrentMpsc> {
+ runtime: tokio::runtime::Runtime,
+ receiver: C::Receiver,
+ start: Arc<tokio::sync::Barrier>,
+ stop: Arc<AtomicBool>,
+ workers: Vec<tokio::task::JoinHandle<()>>,
+}
+
+impl<C: ConcurrentMpsc> RepeatedTasks<C> {
+ pub fn new(producer_count: usize, worker_threads: usize) -> Self {
+ assert_eq!(BATCH_MESSAGES % producer_count, 0);
+ let runtime = if worker_threads == 0 {
+ tokio::runtime::Builder::new_current_thread()
+ .build()
+ .unwrap()
+ } else {
+ tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(worker_threads)
+ .build()
+ .unwrap()
+ };
+ let (sender, receiver) = C::channel();
+ let start = Arc::new(tokio::sync::Barrier::new(producer_count + 1));
+ let stop = Arc::new(AtomicBool::new(false));
+ let messages_per_producer = BATCH_MESSAGES / producer_count;
+ let workers = (0..producer_count)
+ .map(|producer| {
+ let sender = sender.clone();
+ let start = start.clone();
+ let stop = stop.clone();
+ runtime.spawn(async move {
+ loop {
+ start.wait().await;
+ if stop.load(Ordering::Acquire) {
+ break;
+ }
+ let first = producer * messages_per_producer;
+ for offset in 0..messages_per_producer {
+ C::send_async(&sender, black_box(first + offset)).await;
+ }
+ }
+ })
+ })
+ .collect();
+ drop(sender);
+ Self {
+ runtime,
+ receiver,
+ start,
+ stop,
+ workers,
+ }
+ }
+
+ pub fn run(&mut self) -> usize {
+ self.runtime.block_on(async {
+ self.start.wait().await;
+ let mut checksum = 0usize;
+ for _ in 0..BATCH_MESSAGES {
+ checksum = checksum.wrapping_add(C::recv_async(&mut self.receiver).await);
+ }
+ assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2);
+ black_box(checksum)
+ })
+ }
+}
+
+impl<C: ConcurrentMpsc> Drop for RepeatedTasks<C> {
+ fn drop(&mut self) {
+ self.stop.store(true, Ordering::Release);
+ self.runtime.block_on(async {
+ self.start.wait().await;
+ for worker in self.workers.drain(..) {
+ worker.await.expect("benchmark producer panicked");
+ }
+ });
+ }
+}
diff --git a/benchmarks/ecosystem/mpsc/unbounded.rs b/benchmarks/ecosystem/mpsc/unbounded.rs
index b08ed70..b2e79d0 100644
--- a/benchmarks/ecosystem/mpsc/unbounded.rs
+++ b/benchmarks/ecosystem/mpsc/unbounded.rs
@@ -15,6 +15,13 @@
// specific language governing permissions and limitations
// under the License.
+use std::future::Future;
+use std::future::poll_fn;
+use std::pin::pin;
+use std::sync::Arc;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::Ordering;
+
use divan::Bencher;
use divan::black_box;
use divan::counter::ItemsCount;
@@ -28,6 +35,7 @@ use super::support::BATCH_MESSAGES;
use super::support::ConcurrentBatch;
use super::support::PRODUCER_COUNTS;
use super::support::RepeatedBatch;
+use super::support::RepeatedTasks;
use super::support::Unbounded;
use crate::support::bench_context;
@@ -153,3 +161,128 @@ fn clone_drop_sender<C: UnboundedMpsc>(bencher: Bencher) {
let (sender, _receiver) = C::channel();
bencher.bench_local(|| drop(black_box(sender.clone())));
}
+
+#[divan::bench(
+ types = [Asyncband, Tokio, AsyncChannel, Flume],
+ args = [(1, 0), (4, 0), (1, 4), (4, 4), (8, 4)],
+ sample_count = 50,
+ sample_size = 1,
+ counter = ItemsCount::new(BATCH_MESSAGES),
+)]
+fn scheduled<C: UnboundedMpsc>(bencher: Bencher, (producers, workers): (usize, usize)) {
+ let mut batch = RepeatedTasks::<Unbounded<C>>::new(producers, workers);
+ batch.run();
+ bencher.bench_local(|| batch.run());
+}
+
+#[divan::bench(
+ types = [Asyncband, Tokio, AsyncChannel, Flume],
+ args = [(1, 64), (4, 64), (8, 64), (1, 1024), (4, 1024), (8, 1024)],
+ sample_count = 50,
+ sample_size = 1,
+ counter = ItemsCount::new(BATCH_MESSAGES),
+)]
+fn scheduled_bursts_inline<C: UnboundedMpsc<[u8; 1024]>>(
+ bencher: Bencher,
+ (producers, burst_messages): (usize, usize),
+) {
+ assert_eq!(BATCH_MESSAGES % burst_messages, 0);
+ assert_eq!(burst_messages % producers, 0);
+ let runtime = tokio::runtime::Builder::new_multi_thread()
+ .worker_threads(4)
+ .build()
+ .unwrap();
+ let (sender, mut receiver) = C::channel();
+ let start: Vec<_> = (0..producers)
+ .map(|_| Arc::new(tokio::sync::Notify::new()))
+ .collect();
+ let stop = Arc::new(AtomicBool::new(false));
+ let workers: Vec<_> = start
+ .iter()
+ .enumerate()
+ .map(|(producer, start)| {
+ let sender = sender.clone();
+ let start = start.clone();
+ let stop = stop.clone();
+ runtime.spawn(async move {
+ let mut sequence = 0u64;
+ loop {
+ start.notified().await;
+ if stop.load(Ordering::Acquire) {
+ break;
+ }
+ for _ in 0..burst_messages / producers {
+ let mut value = [1; 1024];
+ value[..8].copy_from_slice(&(producer as u64).to_le_bytes());
+ value[8..16].copy_from_slice(&sequence.to_le_bytes());
+ C::send(&sender, black_box(value));
+ sequence += 1;
+ }
+ }
+ })
+ })
+ .collect();
+ drop(sender);
+
+ let mut expected = vec![0u64; producers];
+ let mut run = || {
+ runtime.block_on(async {
+ let mut checksum = 0;
+ for _ in 0..BATCH_MESSAGES / burst_messages {
+ let first = {
+ // Exclude Tokio's cooperative-budget Pending from the initial empty probe.
+ // Retain the same receive future so its channel registration drives the wake.
+ let mut receive =
+ pin!(tokio::task::unconstrained(C::recv_async(&mut receiver)));
+ let mut released = false;
+ poll_fn(|cx| {
+ let result = receive.as_mut().poll(cx);
+ if !released {
+ assert!(
+ result.is_pending(),
+ "each burst must start with an empty wait"
+ );
+ released = true;
+ for producer in &start {
+ producer.notify_one();
+ }
+ }
+ result
+ })
+ .await
+ };
+ checksum += check_inline_message(first, &mut expected);
+ for _ in 1..burst_messages {
+ checksum +=
+ check_inline_message(C::recv_async(&mut receiver).await, &mut expected);
+ }
+ }
+ assert_eq!(checksum, BATCH_MESSAGES);
+ assert!(expected.iter().all(|count| *count == expected[0]));
+ black_box(checksum)
+ })
+ };
+ // Reuse the channel and tasks, including across bursts. Timing includes producer release,
+ // channel wakeups, concurrent payload movement, and storage reclamation, not task creation.
+ run();
+ bencher.bench_local(run);
+
+ stop.store(true, Ordering::Release);
+ for producer in &start {
+ producer.notify_one();
+ }
+ runtime.block_on(async {
+ for worker in workers {
+ worker.await.expect("benchmark producer panicked");
+ }
+ });
+}
+
+fn check_inline_message(value: [u8; 1024], expected: &mut [u64]) -> usize {
+ let value = black_box(value);
+ let producer = u64::from_le_bytes(value[..8].try_into().unwrap()) as usize;
+ let sequence = u64::from_le_bytes(value[8..16].try_into().unwrap());
+ assert_eq!(sequence, expected[producer]);
+ expected[producer] += 1;
+ usize::from(value[1023])
+}
diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs
deleted file mode 100644
index 7b84b93..0000000
--- a/tests-integration/tests/mpsc_test.rs
+++ /dev/null
@@ -1,574 +0,0 @@
-// Licensed to the Apache Software Foundation (ASF) under one
-// or more contributor license agreements. See the NOTICE file
-// distributed with this work for additional information
-// regarding copyright ownership. The ASF licenses this file
-// to you under the Apache License, Version 2.0 (the
-// "License"); you may not use this file except in compliance
-// with the License. You may obtain a copy of the License at
-//
-// http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing,
-// software distributed under the License is distributed on an
-// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-// KIND, either express or implied. See the License for the
-// specific language governing permissions and limitations
-// under the License.
-
-use std::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::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;
-use asyncband::mpsc::TryRecvError;
-use asyncband::mpsc::TrySendError;
-use tests_integration::poll_once;
-use tests_integration::test_runtime;
-use tokio_test::assert_ok;
-
-fn expect_ready<T>(poll: Poll<T>) -> T {
- match poll {
- Poll::Ready(value) => value,
- Poll::Pending => panic!("future should be ready"),
- }
-}
-
-struct HoldSender<S> {
- _sender: S,
-}
-
-// This waker must own the sender so its final drop can break the tested reference cycle.
-#[allow(clippy::manual_noop_waker)]
-impl<S: Send + Sync> Wake for HoldSender<S> {
- fn wake(self: Arc<Self>) {}
-}
-
-#[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());
-}
-
-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<usize>);
-
- impl Wake for SendOnWake {
- fn wake(self: Arc<Self>) {
- 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<usize>,
- disconnected: bool,
- drops: Arc<AtomicUsize>,
- }
-
- // 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<Self>) {}
- }
-
- 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::<mpsc::UnboundedSender<usize>>();
- // 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::<mpsc::UnboundedSender<usize>>()) });
- }
-
- static VTABLE: RawWakerVTable =
- RawWakerVTable::new(clone_sender, drop_sender, |_| {}, drop_sender);
-
- assert_completes_without_deadlock(|| {
- let (tx, mut rx) = mpsc::unbounded::<usize>();
- 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();
-
- test_runtime().block_on(async move {
- for i in 0..8 {
- let tx = tx.clone();
- tokio::spawn(async move {
- tx.send(i).unwrap();
- });
- }
- drop(tx);
-
- let mut sum = 0;
- while let Ok(i) = rx.recv().await {
- sum += i;
- }
- assert_eq!(sum, 28);
- });
-}
-
-#[tokio::test]
-async fn select_streams() {
- let (tx1, mut rx1) = mpsc::unbounded::<i32>();
- let (tx2, mut rx2) = mpsc::unbounded::<i32>();
- let (tx3, mut rx3) = mpsc::bounded(1);
- let (tx4, mut rx4) = mpsc::bounded(1);
-
- tokio::spawn(async move {
- assert_ok!(tx2.send(1));
- tokio::task::yield_now().await;
-
- assert_ok!(tx1.send(2));
- tokio::task::yield_now().await;
-
- assert_ok!(tx2.send(3));
- tokio::task::yield_now().await;
-
- assert_ok!(tx3.send(4).await);
- tokio::task::yield_now().await;
-
- assert_ok!(tx4.send(5).await);
- tokio::task::yield_now().await;
-
- assert_ok!(tx3.send(6).await);
- tokio::task::yield_now().await;
-
- drop((tx1, tx2));
- });
-
- let mut rem = true;
- let mut msgs = vec![];
- let mut rx1_disconnected = false;
- let mut rx2_disconnected = false;
- let mut rx3_disconnected = false;
- let mut rx4_disconnected = false;
-
- while rem {
- rem = !(rx1_disconnected && rx2_disconnected && rx3_disconnected && rx4_disconnected);
-
- tokio::select! {
- result = rx1.recv(), if !rx1_disconnected => {
- match result {
- Ok(x) => msgs.push(x),
- Err(RecvError::Disconnected) => rx1_disconnected = true,
- }
- }
- result = rx2.recv(), if !rx2_disconnected => {
- match result {
- Ok(y) => msgs.push(y),
- Err(RecvError::Disconnected) => rx2_disconnected = true,
- }
- }
- result = rx3.recv(), if !rx3_disconnected => {
- match result {
- Ok(z) => msgs.push(z),
- Err(RecvError::Disconnected) => rx3_disconnected = true,
- }
- }
- result = rx4.recv(), if !rx4_disconnected => {
- match result {
- Ok(w) => msgs.push(w),
- Err(RecvError::Disconnected) => rx4_disconnected = true,
- }
- }
- else => {
- rx1_disconnected = true;
- rx2_disconnected = true;
- rx3_disconnected = true;
- rx4_disconnected = true;
- }
- }
- }
-
- msgs.sort_unstable();
- assert_eq!(&msgs[..], &[1, 2, 3, 4, 5, 6]);
-}
-
-#[tokio::test]
-async fn send_recv_unbounded() {
- let (tx, mut rx) = mpsc::unbounded::<i32>();
-
- // Using `try_send`
- assert_ok!(tx.send(1));
- assert_ok!(tx.send(2));
-
- assert_eq!(rx.recv().await, Ok(1));
- assert_eq!(rx.recv().await, Ok(2));
-
- drop(tx);
-
- assert_eq!(rx.recv().await, Err(RecvError::Disconnected));
-}
-
-#[tokio::test]
-async fn async_send_recv_unbounded() {
- let (tx, mut rx) = mpsc::unbounded();
-
- tokio::spawn(async move {
- assert_ok!(tx.send(1));
- assert_ok!(tx.send(2));
- });
-
- assert_eq!(Ok(1), rx.recv().await);
- assert_eq!(Ok(2), rx.recv().await);
- assert_eq!(Err(RecvError::Disconnected), rx.recv().await);
-}
-
-#[test]
-fn unbounded_try_recv_preserves_order_and_reports_state() {
- let (tx, mut rx) = mpsc::unbounded();
-
- for i in 0..4 {
- tx.send(i).unwrap();
- }
-
- for i in 0..4 {
- assert_eq!(rx.try_recv(), Ok(i));
- }
- assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
- drop(tx);
- assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
-}
-
-#[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);
-
- tx.send(1).await.unwrap();
- assert_eq!(rx.recv().await, Ok(1));
-
- drop(tx);
- assert_eq!(rx.recv().await, Err(RecvError::Disconnected));
-}
-
-#[tokio::test]
-async fn async_send_recv_bounded() {
- let (tx, mut rx) = mpsc::bounded(1);
-
- tx.send(1).await.unwrap();
- // This will block until the receiver is ready to receive.
- tokio::spawn(async move {
- tx.send(2).await.unwrap();
- });
-
- assert_eq!(Ok(1), rx.recv().await);
- assert_eq!(Ok(2), rx.recv().await);
- assert_eq!(Err(RecvError::Disconnected), rx.recv().await);
-}
-
-#[test]
-fn bounded_try_send_respects_capacity_and_order() {
- for capacity in [1, 4, 16] {
- let (tx, mut rx) = mpsc::bounded(capacity);
-
- for i in 0..capacity {
- tx.try_send(i).unwrap();
- }
-
- assert_eq!(tx.try_send(capacity), Err(TrySendError::Full(capacity)));
-
- for i in 0..capacity {
- assert_eq!(rx.try_recv(), Ok(i));
- }
-
- assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
- drop(tx);
- assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
- }
-}
-
-#[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);
-
- tx.try_send(1).unwrap();
- drop(rx);
-
- assert_eq!(tx.try_send(3), Err(TrySendError::Disconnected(3)));
-}
-
-#[tokio::test]
-async fn send_after_disconnection_bounded() {
- let (tx, mut rx) = mpsc::bounded(1);
-
- tx.send(1).await.unwrap();
- assert_eq!(rx.recv().await, Ok(1));
-
- drop(rx);
- let error = tx.send(2).await.unwrap_err();
- assert_eq!(error.into_inner(), 2);
-}
-
-#[test]
-fn bounded_wakes_blocked_senders_one_at_a_time() {
- let (tx, mut rx) = mpsc::bounded(1);
- tx.try_send(0).unwrap();
-
- let first_tx = tx.clone();
- let second_tx = tx.clone();
- let mut first = Box::pin(first_tx.send(1));
- let mut second = Box::pin(second_tx.send(2));
-
- assert!(poll_once(first.as_mut()).is_pending());
- assert!(poll_once(second.as_mut()).is_pending());
-
- assert_eq!(rx.try_recv(), Ok(0));
- assert_eq!(expect_ready(poll_once(first.as_mut())), Ok(()));
- assert!(poll_once(second.as_mut()).is_pending());
-
- assert_eq!(rx.try_recv(), Ok(1));
- assert_eq!(expect_ready(poll_once(second.as_mut())), Ok(()));
- assert_eq!(rx.try_recv(), Ok(2));
-}
-
-#[test]
-fn bounded_cancelled_notified_sender_passes_slot_to_next_sender() {
- let (tx, mut rx) = mpsc::bounded(1);
- tx.try_send(0).unwrap();
-
- let first_tx = tx.clone();
- let second_tx = tx.clone();
- let mut first = Box::pin(first_tx.send(1));
- let mut second = Box::pin(second_tx.send(2));
-
- assert!(poll_once(first.as_mut()).is_pending());
- assert!(poll_once(second.as_mut()).is_pending());
-
- assert_eq!(rx.try_recv(), Ok(0));
- drop(first);
-
- assert_eq!(expect_ready(poll_once(second.as_mut())), Ok(()));
- assert_eq!(rx.try_recv(), Ok(2));
-}
-
-#[test]
-fn bounded_receiver_drop_returns_values_to_all_blocked_senders() {
- let (tx, rx) = mpsc::bounded(1);
- tx.try_send(0).unwrap();
-
- let first_tx = tx.clone();
- let second_tx = tx.clone();
- let mut first = Box::pin(first_tx.send(1));
- let mut second = Box::pin(second_tx.send(2));
-
- assert!(poll_once(first.as_mut()).is_pending());
- assert!(poll_once(second.as_mut()).is_pending());
-
- drop(rx);
-
- let first_error = expect_ready(poll_once(first.as_mut())).unwrap_err();
- let second_error = expect_ready(poll_once(second.as_mut())).unwrap_err();
- assert_eq!(first_error.into_inner(), 1);
- assert_eq!(second_error.into_inner(), 2);
-}
diff --git a/tests-integration/tests/mpsc_test/backpressure.rs b/tests-integration/tests/mpsc_test/backpressure.rs
new file mode 100644
index 0000000..ba21cfc
--- /dev/null
+++ b/tests-integration/tests/mpsc_test/backpressure.rs
@@ -0,0 +1,107 @@
+// 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 asyncband::mpsc;
+use tests_integration::poll_once;
+
+use super::support::WakeCounter;
+use super::support::expect_ready;
+use super::support::poll_with;
+
+#[test]
+fn bounded_wakes_blocked_senders_one_at_a_time() {
+ let (tx, mut rx) = mpsc::bounded(1);
+ tx.try_send(0).unwrap();
+
+ let first_tx = tx.clone();
+ let second_tx = tx.clone();
+ let mut first = Box::pin(first_tx.send(1));
+ let mut second = Box::pin(second_tx.send(2));
+
+ let (first_waker, first_wakes) = WakeCounter::new();
+ let (second_waker, second_wakes) = WakeCounter::new();
+ assert!(poll_with(first.as_mut(), &first_waker).is_pending());
+ assert!(poll_with(second.as_mut(), &second_waker).is_pending());
+
+ assert_eq!(rx.try_recv(), Ok(0));
+ assert_eq!(first_wakes.count(), 1);
+ assert_eq!(second_wakes.count(), 0);
+ assert_eq!(expect_ready(poll_once(first.as_mut())), Ok(()));
+ assert!(poll_with(second.as_mut(), &second_waker).is_pending());
+
+ assert_eq!(rx.try_recv(), Ok(1));
+ assert_eq!(second_wakes.count(), 1);
+ assert_eq!(expect_ready(poll_once(second.as_mut())), Ok(()));
+ assert_eq!(rx.try_recv(), Ok(2));
+}
+
+#[test]
+fn cancelling_a_sender_preserves_capacity_and_notifies_the_next_waiter() {
+ for cancel_after_notification in [false, true] {
+ let (tx, mut rx) = mpsc::bounded(1);
+ tx.try_send(0).unwrap();
+ let mut first = Box::pin(tx.send(1));
+ let mut second = Box::pin(tx.send(2));
+ let (first_waker, first_wakes) = WakeCounter::new();
+ let (second_waker, second_wakes) = WakeCounter::new();
+ assert!(poll_with(first.as_mut(), &first_waker).is_pending());
+ assert!(poll_with(second.as_mut(), &second_waker).is_pending());
+
+ if cancel_after_notification {
+ assert_eq!(rx.try_recv(), Ok(0));
+ assert_eq!(first_wakes.count(), 1);
+ assert_eq!(second_wakes.count(), 0);
+ }
+ drop(first);
+ if !cancel_after_notification {
+ assert_eq!(first_wakes.count(), 0);
+ assert_eq!(second_wakes.count(), 0);
+ assert_eq!(rx.try_recv(), Ok(0));
+ }
+ assert_eq!(second_wakes.count(), 1);
+ assert_eq!(
+ expect_ready(poll_with(second.as_mut(), &second_waker)),
+ Ok(())
+ );
+ assert_eq!(rx.try_recv(), Ok(2));
+ }
+}
+
+#[test]
+fn bounded_receiver_drop_returns_values_to_all_blocked_senders() {
+ let (tx, rx) = mpsc::bounded(1);
+ tx.try_send(0).unwrap();
+
+ let first_tx = tx.clone();
+ let second_tx = tx.clone();
+ let mut first = Box::pin(first_tx.send(1));
+ let mut second = Box::pin(second_tx.send(2));
+
+ let (first_waker, first_wakes) = WakeCounter::new();
+ let (second_waker, second_wakes) = WakeCounter::new();
+ assert!(poll_with(first.as_mut(), &first_waker).is_pending());
+ assert!(poll_with(second.as_mut(), &second_waker).is_pending());
+
+ drop(rx);
+ assert_eq!(first_wakes.count(), 1);
+ assert_eq!(second_wakes.count(), 1);
+
+ let first_error = expect_ready(poll_once(first.as_mut())).unwrap_err();
+ let second_error = expect_ready(poll_once(second.as_mut())).unwrap_err();
+ assert_eq!(first_error.into_inner(), 1);
+ assert_eq!(second_error.into_inner(), 2);
+}
diff --git a/tests-integration/tests/mpsc_test/callbacks.rs b/tests-integration/tests/mpsc_test/callbacks.rs
new file mode 100644
index 0000000..2f654fc
--- /dev/null
+++ b/tests-integration/tests/mpsc_test/callbacks.rs
@@ -0,0 +1,328 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::future::Future;
+use std::sync::Arc;
+use std::sync::Mutex;
+use std::sync::atomic::AtomicBool;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::task::Context;
+use std::task::Poll;
+use std::task::Wake;
+use std::task::Waker;
+
+use asyncband::mpsc;
+use asyncband::mpsc::TryRecvError;
+use tests_integration::poll_once;
+
+use super::support::WakeCounter;
+use super::support::assert_completes_without_deadlock;
+use super::support::expect_ready;
+use super::support::poll_with;
+use super::support::waker_on_clone;
+
+struct HoldSender<S> {
+ _sender: S,
+}
+
+// This waker must own the sender so its final drop can break the tested reference cycle.
+#[allow(clippy::manual_noop_waker)]
+impl<S: Send + Sync> Wake for HoldSender<S> {
+ fn wake(self: Arc<Self>) {}
+}
+
+#[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 bounded_send_rechecks_capacity_freed_by_waker_clone() {
+ assert_completes_without_deadlock(|| {
+ let (tx, rx) = mpsc::bounded(1);
+ tx.try_send(1).unwrap();
+ let receiver = Arc::new(Mutex::new(rx));
+ let received = AtomicBool::new(false);
+ let waker = waker_on_clone({
+ let receiver = receiver.clone();
+ move || {
+ if !received.swap(true, Ordering::Relaxed) {
+ assert_eq!(receiver.lock().unwrap().try_recv(), Ok(1));
+ }
+ }
+ });
+ assert_eq!(
+ poll_with(Box::pin(tx.send(2)).as_mut(), &waker),
+ Poll::Ready(Ok(()))
+ );
+ assert_eq!(receiver.lock().unwrap().try_recv(), Ok(2));
+ });
+}
+
+#[test]
+fn receive_rechecks_messages_sent_by_waker_clone() {
+ assert_completes_without_deadlock(|| {
+ let (tx, mut rx) = mpsc::bounded(1);
+ let waker = waker_on_clone(move || tx.try_send(7).unwrap());
+ assert_eq!(
+ poll_with(Box::pin(rx.recv()).as_mut(), &waker),
+ Poll::Ready(Ok(7))
+ );
+
+ let (tx, mut rx) = mpsc::unbounded();
+ let waker = waker_on_clone(move || tx.send(7).unwrap());
+ assert_eq!(
+ poll_with(Box::pin(rx.recv()).as_mut(), &waker),
+ Poll::Ready(Ok(7))
+ );
+ });
+}
+
+#[test]
+fn wake_callbacks_can_send_into_the_same_channel() {
+ struct SendOnWake(Box<dyn Fn() + Send + Sync>);
+ impl Wake for SendOnWake {
+ fn wake(self: Arc<Self>) {
+ (self.0)();
+ }
+ }
+
+ assert_completes_without_deadlock(|| {
+ let (tx, mut rx) = mpsc::bounded(2);
+ let waker = Waker::from(Arc::new(SendOnWake(Box::new({
+ let tx = tx.clone();
+ move || tx.try_send(2).unwrap()
+ }))));
+ assert!(poll_with(Box::pin(rx.recv()).as_mut(), &waker).is_pending());
+ tx.try_send(1).unwrap();
+ assert_eq!(rx.try_recv(), Ok(1));
+ assert_eq!(rx.try_recv(), Ok(2));
+
+ let (tx, mut rx) = mpsc::unbounded();
+ let waker = Waker::from(Arc::new(SendOnWake(Box::new({
+ let tx = tx.clone();
+ move || tx.send(2).unwrap()
+ }))));
+ assert!(poll_with(Box::pin(rx.recv()).as_mut(), &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<usize>,
+ disconnected: bool,
+ drops: Arc<AtomicUsize>,
+ }
+
+ // 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<Self>) {}
+ }
+
+ 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 (waker, counter) = WakeCounter::new();
+ let poll = poll_with(Box::pin(rx.recv()).as_mut(), &waker);
+ if poll.is_pending() {
+ assert!(counter.count() > 0);
+ assert_eq!(rx.try_recv(), Ok(7));
+ } else {
+ assert_eq!(poll, Poll::Ready(Ok(7)));
+ }
+ assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
+ }
+ assert_eq!(drops.load(Ordering::Relaxed), 1);
+ }
+ });
+}
+
+#[test]
+fn unbounded_disconnect_drops_partial_and_queued_batches_outside_lock() {
+ struct Value(Option<Arc<dyn Fn() + Send + Sync>>);
+ impl Drop for Value {
+ fn drop(&mut self) {
+ if let Some(callback) = &self.0 {
+ callback();
+ }
+ }
+ }
+
+ assert_completes_without_deadlock(|| {
+ let (tx, mut rx) = mpsc::unbounded();
+ let drops = Arc::new(AtomicUsize::new(0));
+ let callback: Arc<dyn Fn() + Send + Sync> = {
+ let tx = tx.clone();
+ let drops = drops.clone();
+ Arc::new(move || {
+ // This exercises both a live receiver and disconnection. The marker has no
+ // callback, so destroying an unsuccessful send cannot recursively send again.
+ let _ = tx.send(Value(None));
+ drops.fetch_add(1, Ordering::Relaxed);
+ })
+ };
+ for _ in 0..8192 {
+ assert!(tx.send(Value(Some(callback.clone()))).is_ok());
+ }
+ for _ in 0..17 {
+ drop(rx.try_recv().unwrap());
+ }
+ drop(rx);
+ assert_eq!(drops.load(Ordering::Relaxed), 8192);
+ });
+}
+
+#[cfg(panic = "unwind")]
+#[test]
+fn bounded_disconnect_finishes_cleanup_when_a_callback_panics() {
+ struct Value {
+ id: usize,
+ drops: Arc<[AtomicUsize; 5]>,
+ panic_on_drop: bool,
+ _sender: Option<mpsc::BoundedSender<Value>>,
+ }
+ impl Drop for Value {
+ fn drop(&mut self) {
+ self.drops[self.id].fetch_add(1, Ordering::Relaxed);
+ assert!(!self.panic_on_drop, "payload destructor panicked");
+ }
+ }
+ struct Notify {
+ woken: AtomicBool,
+ panic_on_wake: bool,
+ }
+ impl Wake for Notify {
+ fn wake(self: Arc<Self>) {
+ self.woken.store(true, Ordering::Relaxed);
+ assert!(!self.panic_on_wake, "wake callback panicked");
+ }
+ }
+
+ for panic_on_wake in [false, true] {
+ let (tx, rx) = mpsc::bounded(3);
+ let drops = Arc::new(std::array::from_fn(|_| AtomicUsize::new(0)));
+ for id in 0..3 {
+ assert!(
+ tx.try_send(Value {
+ id,
+ drops: drops.clone(),
+ panic_on_drop: id == 0 && !panic_on_wake,
+ _sender: Some(tx.clone()),
+ })
+ .is_ok()
+ );
+ }
+ let notify = [false, true].map(|second| {
+ Arc::new(Notify {
+ woken: AtomicBool::new(false),
+ panic_on_wake: !second && panic_on_wake,
+ })
+ });
+ let wakers = notify.each_ref().map(|notify| Waker::from(notify.clone()));
+ let mut sends = (3..5)
+ .map(|id| {
+ Box::pin(tx.send(Value {
+ id,
+ drops: drops.clone(),
+ panic_on_drop: false,
+ _sender: None,
+ }))
+ })
+ .collect::<Vec<_>>();
+ for (send, waker) in sends.iter_mut().zip(&wakers) {
+ assert!(poll_with(send.as_mut(), waker).is_pending());
+ }
+ assert!(std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(rx))).is_err());
+ assert!(
+ notify
+ .iter()
+ .all(|notify| notify.woken.load(Ordering::Relaxed))
+ );
+ for count in &drops[..3] {
+ assert_eq!(count.load(Ordering::Relaxed), 1);
+ }
+ for (id, send) in (3..5).zip(&mut sends) {
+ assert_eq!(drops[id].load(Ordering::Relaxed), 0);
+ let error = match expect_ready(poll_once(send.as_mut())) {
+ Err(error) => error,
+ Ok(()) => panic!("the receiver is disconnected"),
+ };
+ assert_eq!(error.into_inner().id, id);
+ assert_eq!(drops[id].load(Ordering::Relaxed), 1);
+ }
+ }
+}
diff --git a/tests-integration/tests/mpsc_test/concurrency.rs b/tests-integration/tests/mpsc_test/concurrency.rs
new file mode 100644
index 0000000..ce097fc
--- /dev/null
+++ b/tests-integration/tests/mpsc_test/concurrency.rs
@@ -0,0 +1,289 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::future::Future;
+use std::pin::Pin;
+use std::sync::Barrier;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+use std::task::Context;
+use std::task::Poll;
+use std::thread;
+
+use asyncband::mpsc;
+use asyncband::mpsc::RecvError;
+use asyncband::mpsc::TryRecvError;
+use asyncband::mpsc::TrySendError;
+use tests_integration::poll_once;
+use tests_integration::test_runtime;
+use tokio_test::assert_ok;
+
+use super::support::WakeCounter;
+use super::support::poll_with;
+
+#[test]
+fn bounded_receive_racing_with_send_registration_cannot_lose_wakeup() {
+ for _ in 0..128 {
+ let (tx, mut rx) = mpsc::bounded(1);
+ tx.try_send(1).unwrap();
+ let start = Barrier::new(2);
+ let (waker, notified) = WakeCounter::new();
+ let mut send = Box::pin(tx.send(2));
+ let poll = thread::scope(|scope| {
+ let receive = scope.spawn(|| {
+ start.wait();
+ assert_eq!(rx.try_recv(), Ok(1));
+ });
+ start.wait();
+ let poll = send.as_mut().poll(&mut Context::from_waker(&waker));
+ receive.join().unwrap();
+ poll
+ });
+ if poll.is_pending() {
+ assert!(notified.count() > 0);
+ assert_eq!(poll_once(send.as_mut()), Poll::Ready(Ok(())));
+ } else {
+ assert_eq!(poll, Poll::Ready(Ok(())));
+ }
+ assert_eq!(rx.try_recv(), Ok(2));
+ }
+}
+
+#[test]
+fn bounded_try_recv_does_not_report_empty_after_completed_sends() {
+ const PRODUCERS: usize = 4;
+ const MESSAGES_PER_PRODUCER: usize = if cfg!(miri) { 64 } else { 16_384 };
+ let (tx, mut rx) = mpsc::bounded(64);
+ let completed = AtomicUsize::new(0);
+ let mut premature_empty = 0;
+ let mut next = [0; PRODUCERS];
+ let mut out_of_order = 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((producer, sequence)) => {
+ out_of_order += usize::from(sequence != next[producer]);
+ next[producer] += 1;
+ 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);
+ assert_eq!(out_of_order, 0);
+ assert_eq!(next, [MESSAGES_PER_PRODUCER; PRODUCERS]);
+}
+
+fn receive_racing_with(
+ mut receive: Pin<&mut impl Future<Output = Result<usize, RecvError>>>,
+ publish: impl FnOnce() + Send,
+ expected: Result<usize, RecvError>,
+) {
+ let (waker, counter) = WakeCounter::new();
+ let start = Barrier::new(2);
+ let result = thread::scope(|scope| {
+ let producer = scope.spawn(|| {
+ start.wait();
+ publish();
+ });
+ start.wait();
+ let result = poll_with(receive.as_mut(), &waker);
+ producer.join().unwrap();
+ result
+ });
+ if result.is_pending() {
+ assert!(counter.count() > 0, "a pending receiver must be notified");
+ assert_eq!(poll_with(receive, &waker), Poll::Ready(expected));
+ } else {
+ assert_eq!(result, Poll::Ready(expected));
+ }
+}
+
+#[test]
+fn publication_racing_with_receiver_registration_cannot_lose_wakeup() {
+ for capacity in [1, 3] {
+ let (tx, mut rx) = mpsc::bounded(capacity);
+ for value in 0..128 {
+ receive_racing_with(
+ Box::pin(rx.recv()).as_mut(),
+ || tx.try_send(value).unwrap(),
+ Ok(value),
+ );
+ }
+ }
+ let (tx, mut rx) = mpsc::unbounded();
+ for value in 0..128 {
+ receive_racing_with(
+ Box::pin(rx.recv()).as_mut(),
+ || tx.send(value).unwrap(),
+ Ok(value),
+ );
+ }
+}
+
+#[test]
+fn last_sender_drop_racing_with_receiver_registration_cannot_lose_wakeup() {
+ for _ in 0..128 {
+ let (tx, mut rx) = mpsc::bounded::<usize>(1);
+ receive_racing_with(
+ Box::pin(rx.recv()).as_mut(),
+ move || drop(tx),
+ Err(RecvError::Disconnected),
+ );
+ let (tx, mut rx) = mpsc::unbounded::<usize>();
+ receive_racing_with(
+ Box::pin(rx.recv()).as_mut(),
+ move || drop(tx),
+ Err(RecvError::Disconnected),
+ );
+ }
+}
+
+#[test]
+#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")]
+fn unbounded_collects_from_multiple_producers() {
+ let (tx, mut rx) = mpsc::unbounded();
+
+ test_runtime().block_on(async move {
+ for i in 0..8 {
+ let tx = tx.clone();
+ tokio::spawn(async move {
+ tx.send(i).unwrap();
+ });
+ }
+ drop(tx);
+
+ let mut sum = 0;
+ while let Ok(i) = rx.recv().await {
+ sum += i;
+ }
+ assert_eq!(sum, 28);
+ });
+}
+
+#[tokio::test]
+#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")]
+async fn bounded_backpressure_progresses_on_an_executor() {
+ let (tx, mut rx) = mpsc::bounded(1);
+
+ tx.send(1).await.unwrap();
+ // This will block until the receiver is ready to receive.
+ tokio::spawn(async move {
+ tx.send(2).await.unwrap();
+ });
+
+ assert_eq!(Ok(1), rx.recv().await);
+ assert_eq!(Ok(2), rx.recv().await);
+ assert_eq!(Err(RecvError::Disconnected), rx.recv().await);
+}
+
+#[tokio::test]
+#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")]
+async fn selection_preserves_messages_across_cancelled_receives() {
+ let (tx1, mut rx1) = mpsc::unbounded::<i32>();
+ let (tx2, mut rx2) = mpsc::unbounded::<i32>();
+ let (tx3, mut rx3) = mpsc::bounded(1);
+ let (tx4, mut rx4) = mpsc::bounded(1);
+
+ tokio::spawn(async move {
+ assert_ok!(tx2.send(1));
+ tokio::task::yield_now().await;
+
+ assert_ok!(tx1.send(2));
+ tokio::task::yield_now().await;
+
+ assert_ok!(tx2.send(3));
+ tokio::task::yield_now().await;
+
+ assert_ok!(tx3.send(4).await);
+ tokio::task::yield_now().await;
+
+ assert_ok!(tx4.send(5).await);
+ tokio::task::yield_now().await;
+
+ assert_ok!(tx3.send(6).await);
+ tokio::task::yield_now().await;
+
+ drop((tx1, tx2));
+ });
+
+ let mut msgs = vec![];
+ let mut rx1_disconnected = false;
+ let mut rx2_disconnected = false;
+ let mut rx3_disconnected = false;
+ let mut rx4_disconnected = false;
+
+ while !(rx1_disconnected && rx2_disconnected && rx3_disconnected && rx4_disconnected) {
+ tokio::select! {
+ result = rx1.recv(), if !rx1_disconnected => {
+ match result {
+ Ok(x) => msgs.push(x),
+ Err(RecvError::Disconnected) => rx1_disconnected = true,
+ }
+ }
+ result = rx2.recv(), if !rx2_disconnected => {
+ match result {
+ Ok(y) => msgs.push(y),
+ Err(RecvError::Disconnected) => rx2_disconnected = true,
+ }
+ }
+ result = rx3.recv(), if !rx3_disconnected => {
+ match result {
+ Ok(z) => msgs.push(z),
+ Err(RecvError::Disconnected) => rx3_disconnected = true,
+ }
+ }
+ result = rx4.recv(), if !rx4_disconnected => {
+ match result {
+ Ok(w) => msgs.push(w),
+ Err(RecvError::Disconnected) => rx4_disconnected = true,
+ }
+ }
+ }
+ }
+
+ msgs.sort_unstable();
+ assert_eq!(&msgs[..], &[1, 2, 3, 4, 5, 6]);
+}
diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs
new file mode 100644
index 0000000..2490021
--- /dev/null
+++ b/tests-integration/tests/mpsc_test/main.rs
@@ -0,0 +1,175 @@
+// 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::task::Poll;
+
+use asyncband::mpsc;
+use asyncband::mpsc::RecvError;
+use asyncband::mpsc::TryRecvError;
+use asyncband::mpsc::TrySendError;
+use tests_integration::poll_once;
+
+use self::support::WakeCounter;
+use self::support::expect_ready;
+use self::support::poll_with;
+
+// Public channel contracts. The other suites cover backpressure, callbacks, and concurrency.
+mod backpressure;
+mod callbacks;
+mod concurrency;
+mod support;
+
+#[test]
+fn unbounded_try_recv_preserves_order_and_reports_state() {
+ let (tx, mut rx) = mpsc::unbounded();
+
+ for i in 0..4 {
+ tx.send(i).unwrap();
+ }
+
+ for i in 0..4 {
+ assert_eq!(rx.try_recv(), Ok(i));
+ }
+ assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
+ drop(tx);
+ assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
+}
+
+#[test]
+fn bounded_try_send_respects_capacity_and_order() {
+ for capacity in [1, 4, 16] {
+ let (tx, mut rx) = mpsc::bounded(capacity);
+
+ for i in 0..capacity {
+ tx.try_send(i).unwrap();
+ }
+
+ assert_eq!(tx.try_send(capacity), Err(TrySendError::Full(capacity)));
+
+ for i in 0..capacity {
+ assert_eq!(rx.try_recv(), Ok(i));
+ }
+
+ assert_eq!(rx.try_recv(), Err(TryRecvError::Empty));
+ drop(tx);
+ assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected));
+ }
+}
+
+#[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));
+}
+
+#[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 disconnected_sends_return_the_unsent_value() {
+ let (tx, rx) = mpsc::bounded(1);
+ tx.try_send(String::from("queued")).unwrap();
+ drop(rx);
+ assert_eq!(
+ tx.try_send(String::from("try")),
+ Err(TrySendError::Disconnected(String::from("try")))
+ );
+ let error =
+ expect_ready(poll_once(Box::pin(tx.send(String::from("async"))).as_mut())).unwrap_err();
+ assert_eq!(error.into_inner(), "async");
+
+ let (tx, rx) = mpsc::unbounded::<String>();
+ drop(rx);
+ assert_eq!(
+ tx.send(String::from("unbounded")).unwrap_err().into_inner(),
+ "unbounded"
+ );
+}
+
+#[test]
+fn receives_wake_for_messages_and_the_last_sender_drop() {
+ let (waker, counter) = WakeCounter::new();
+ let (tx, mut rx) = mpsc::bounded(1);
+ let other = tx.clone();
+ let mut receive = Box::pin(rx.recv());
+ assert!(poll_with(receive.as_mut(), &waker).is_pending());
+ tx.try_send(7).unwrap();
+ assert_eq!(counter.count(), 1);
+ assert_eq!(poll_with(receive.as_mut(), &waker), Poll::Ready(Ok(7)));
+ drop(receive);
+
+ let mut receive = Box::pin(rx.recv());
+ assert!(poll_with(receive.as_mut(), &waker).is_pending());
+ drop(tx);
+ assert_eq!(counter.count(), 1);
+ drop(other);
+ assert_eq!(counter.count(), 2);
+ assert_eq!(
+ poll_with(receive.as_mut(), &waker),
+ Poll::Ready(Err(RecvError::Disconnected))
+ );
+
+ let (waker, counter) = WakeCounter::new();
+ let (tx, mut rx) = mpsc::unbounded();
+ let other = tx.clone();
+ let mut receive = Box::pin(rx.recv());
+ assert!(poll_with(receive.as_mut(), &waker).is_pending());
+ tx.send(7).unwrap();
+ assert_eq!(counter.count(), 1);
+ assert_eq!(poll_with(receive.as_mut(), &waker), Poll::Ready(Ok(7)));
+ drop(receive);
+
+ let mut receive = Box::pin(rx.recv());
+ assert!(poll_with(receive.as_mut(), &waker).is_pending());
+ drop(tx);
+ assert_eq!(counter.count(), 1);
+ drop(other);
+ assert_eq!(counter.count(), 2);
+ assert_eq!(
+ poll_with(receive.as_mut(), &waker),
+ Poll::Ready(Err(RecvError::Disconnected))
+ );
+}
diff --git a/tests-integration/tests/mpsc_test/support.rs b/tests-integration/tests/mpsc_test/support.rs
new file mode 100644
index 0000000..360cb0a
--- /dev/null
+++ b/tests-integration/tests/mpsc_test/support.rs
@@ -0,0 +1,106 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::future::Future;
+use std::pin::Pin;
+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::RawWaker;
+use std::task::RawWakerVTable;
+use std::task::Wake;
+use std::task::Waker;
+use std::thread;
+
+pub fn expect_ready<T>(poll: Poll<T>) -> T {
+ match poll {
+ Poll::Ready(value) => value,
+ Poll::Pending => panic!("future should be ready"),
+ }
+}
+
+#[derive(Default)]
+pub struct WakeCounter(AtomicUsize);
+
+impl WakeCounter {
+ pub fn new() -> (Waker, Arc<Self>) {
+ let counter = Arc::new(Self::default());
+ (Waker::from(counter.clone()), counter)
+ }
+
+ pub fn count(&self) -> usize {
+ self.0.load(Ordering::Relaxed)
+ }
+}
+
+impl Wake for WakeCounter {
+ fn wake(self: Arc<Self>) {
+ self.0.fetch_add(1, Ordering::Relaxed);
+ }
+
+ fn wake_by_ref(self: &Arc<Self>) {
+ self.0.fetch_add(1, Ordering::Relaxed);
+ }
+}
+
+pub fn poll_with<F: Future>(future: Pin<&mut F>, waker: &Waker) -> Poll<F::Output> {
+ future.poll(&mut Context::from_waker(waker))
+}
+
+// RawWaker is needed only to exercise clone callbacks, which the safe Wake trait cannot override.
+pub fn waker_on_clone(callback: impl Fn() + Send + Sync + 'static) -> Waker {
+ struct OnClone(Box<dyn Fn() + Send + Sync>);
+
+ unsafe fn clone(data: *const ()) -> RawWaker {
+ let pointer = data.cast::<OnClone>();
+ // SAFETY: The input waker owns a live Arc; the returned waker gains its own reference.
+ unsafe {
+ ((*pointer).0)();
+ Arc::increment_strong_count(pointer);
+ }
+ RawWaker::new(data, &VTABLE)
+ }
+
+ unsafe fn release(data: *const ()) {
+ // SAFETY: Consumes the one Arc reference owned by this waker.
+ drop(unsafe { Arc::from_raw(data.cast::<OnClone>()) });
+ }
+
+ static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, release, |_| {}, release);
+ let pointer = Arc::into_raw(Arc::new(OnClone(Box::new(callback)))).cast();
+ // SAFETY: Each waker owns one Arc; its callback is Send + Sync and all vtable operations
+ // preserve that ownership. wake_by_ref borrows the reference without changing it.
+ unsafe { Waker::from_raw(RawWaker::new(pointer, &VTABLE)) }
+}
+
+pub 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();
+ });
+ #[cfg(not(miri))]
+ finished_rx
+ .recv_timeout(std::time::Duration::from_secs(10))
+ .expect("waker callback did not finish");
+ // Miri detects deadlock itself; its interpretation time must not determine test success.
+ #[cfg(miri)]
+ finished_rx.recv().expect("waker callback did not finish");
+ worker.join().unwrap();
+}
diff --git a/xtask/src/main.rs b/xtask/src/main.rs
index f84e1b6..6a16125 100644
--- a/xtask/src/main.rs
+++ b/xtask/src/main.rs
@@ -117,6 +117,7 @@ impl CommandMiri {
"tests-integration",
&["--test", "unsafe_paths_test"],
));
+ run_command(make_miri_cmd("tests-integration", &["--test", "mpsc_test"]));
}
}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@asyncband.apache.org
For additional commands, e-mail: commits-help@asyncband.apache.org