From 8383b68984234ef980d6e0540473c885e31e3753 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 02:25:16 +0800 Subject: [PATCH 01/10] bench(mpsc): cover capacity and executor scheduling --- 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 | 14 ++++ 5 files changed, 208 insertions(+), 7 deletions(-) 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 + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; fn send_blocking(sender: &Self::Sender, value: usize); fn recv_blocking(receiver: &mut Self::Receiver) -> usize; } pub trait UnboundedMpsc: 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 + 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 UnboundedMpsc 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 UnboundedMpsc 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 UnboundedMpsc 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 UnboundedMpsc 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(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( + bencher: Bencher, + producer_count: usize, +) { + let mut batch = RepeatedBatch::>::new(producer_count); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + 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( + bencher: Bencher, + (producers, workers): (usize, usize), +) { + let mut batch = RepeatedTasks::>::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 + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; } -pub struct Bounded(PhantomData); +pub struct Bounded(PhantomData); -impl ConcurrentMpsc for Bounded { +impl ConcurrentMpsc for Bounded { 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 ConcurrentMpsc for Unbounded { 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 Drop for RepeatedBatch { } } } + +// Exercise executor wakeups as well as channel traffic. Reuse tasks, threads, and channel storage +// across samples; a current-thread runtime also exposes polling that monopolizes the executor. +pub struct RepeatedTasks { + runtime: tokio::runtime::Runtime, + receiver: C::Receiver, + start: Arc, + stop: Arc, + workers: Vec>, +} + +impl RepeatedTasks { + 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 Drop for RepeatedTasks { + 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..229bf2d 100644 --- a/benchmarks/ecosystem/mpsc/unbounded.rs +++ b/benchmarks/ecosystem/mpsc/unbounded.rs @@ -28,6 +28,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 +154,16 @@ fn clone_drop_sender(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(bencher: Bencher, (producers, workers): (usize, usize)) { + let mut batch = RepeatedTasks::>::new(producers, workers); + batch.run(); + bencher.bench_local(|| batch.run()); +} From 6c0d8bd049d0202cc3b8e0cd7116ceac277b73eb Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 02:44:07 +0800 Subject: [PATCH 02/10] fix(mpsc): complete bounded publication and disconnection transitions --- CHANGELOG.md | 2 + asyncband/src/mpsc/bounded.rs | 15 +++++- asyncband/src/mpsc/bounded/ring.rs | 79 +++++++++++++++++++++++----- tests-integration/tests/mpsc_test.rs | 72 +++++++++++++++++++++++++ 4 files changed, 154 insertions(+), 14 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 35e870e..397c9f1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ All notable changes to this project will be documented in this file. * Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. +* Complete bounded MPSC disconnection notifications and buffered-message cleanup even when a wake callback or message destructor panics. + ### 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. diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index 550bf16..de561ef 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -208,12 +208,23 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { + struct DrainOnDrop<'a, T>(&'a Ring); + impl Drop for DrainOnDrop<'_, T> { + fn drop(&mut self) { + // SAFETY: This guard lives only within the exclusive receiver's drop, after close. + unsafe { self.0.drain() }; + } + } + + 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() }; + // 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); } } diff --git a/asyncband/src/mpsc/bounded/ring.rs b/asyncband/src/mpsc/bounded/ring.rs index 4eb9508..c4e4e47 100644 --- a/asyncband/src/mpsc/bounded/ring.rs +++ b/asyncband/src/mpsc/bounded/ring.rs @@ -106,7 +106,17 @@ impl Ring { } 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 +127,7 @@ impl Ring { /// /// # 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> { let mut head = self.head.load(Ordering::Relaxed); let mut backoff = 0; @@ -150,15 +160,37 @@ impl Ring { } } - /// Closes the queue and drops all remaining values. + /// 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, + tail: usize, + } + impl Drop for DrainRemaining<'_, T> { + fn drop(&mut self) { + // SAFETY: The guard is scoped to the exclusive consumer's drain of a closed ring. + unsafe { self.ring.discard_until(self.tail) }; + } + } + + 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,9 +237,9 @@ impl Ring { impl Drop for Ring { 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() }; } } @@ -269,6 +301,28 @@ mod tests { } } + #[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()) + }; + let result = completed.recv_timeout(std::time::Duration::from_secs(10)); + // 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)); @@ -347,8 +401,9 @@ mod tests { 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() }; + 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() { diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 7b84b93..1d39c2d 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -17,6 +17,7 @@ use std::future::Future; use std::sync::Arc; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; @@ -572,3 +573,74 @@ fn bounded_receiver_drop_returns_values_to_all_blocked_senders() { assert_eq!(first_error.into_inner(), 1); assert_eq!(second_error.into_inner(), 2); } + +#[cfg(panic = "unwind")] +#[test] +fn bounded_disconnect_finishes_cleanup_when_a_callback_panics() { + struct Value { + id: usize, + drops: Arc<[AtomicUsize; 4]>, + panic_on_drop: bool, + _sender: Option>, + } + 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.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 = Arc::new(Notify { + woken: AtomicBool::new(false), + panic_on_wake, + }); + let waker = Waker::from(notify.clone()); + let mut send = Box::pin(tx.send(Value { + id: 3, + drops: drops.clone(), + panic_on_drop: false, + _sender: None, + })); + assert!( + send.as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + assert!(std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(rx))).is_err()); + assert!(notify.woken.load(Ordering::Relaxed)); + for count in &drops[..3] { + assert_eq!(count.load(Ordering::Relaxed), 1); + } + assert_eq!(drops[3].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, 3); + assert_eq!(drops[3].load(Ordering::Relaxed), 1); + } +} From d7abed350d86b10fffeed03134e91db895fe9ad0 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 02:46:43 +0800 Subject: [PATCH 03/10] refactor(mpsc): model bounded sender waits as retry notifications --- CHANGELOG.md | 2 + asyncband/src/internal/mod.rs | 12 +- asyncband/src/internal/semaphore.rs | 34 ------ asyncband/src/mpsc/bounded.rs | 81 ++++++------- asyncband/src/mpsc/bounded/waiters.rs | 158 ++++++++++++++++++++++++++ tests-integration/tests/mpsc_test.rs | 82 +++++++++++++ 6 files changed, 280 insertions(+), 89 deletions(-) create mode 100644 asyncband/src/mpsc/bounded/waiters.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 397c9f1..c1bdc02 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,8 @@ All notable changes to this project will be documented in this file. * Complete bounded MPSC disconnection notifications and buffered-message cleanup even when a wake callback or 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. diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 84309fc..8d4f1b7 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -99,15 +99,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.rs index de561ef..89267b6 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -28,15 +28,15 @@ 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; mod ring; +mod waiters; /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// @@ -52,7 +52,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { let state = Arc::new(BoundedState { buffer: Ring::new(buffer), senders: AtomicUsize::new(1), - send_waiters: Semaphore::new(0), + send_waiters: SendWaiters::new(), rx_waker: AtomicWaker::new(), }); let sender = BoundedSender { @@ -65,8 +65,7 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { struct BoundedState { buffer: Ring, senders: AtomicUsize, - // Notifications grant retries; only the ring determines whether buffer capacity is available. - send_waiters: Semaphore, + send_waiters: SendWaiters, rx_waker: AtomicWaker, } @@ -122,48 +121,38 @@ impl BoundedSender { Err(TrySendError::Disconnected(value)) => return Err(SendError::new(value)), Err(TrySendError::Full(value)) => value, }; - - struct SendState<'a, T> { - sender: &'a BoundedSender, - value: Option, - acquire: Acquire<'a>, - } - - impl SendState<'_, T> { - fn poll_send(&mut self, cx: &mut Context<'_>) -> Poll>> { - 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. @@ -278,7 +267,7 @@ impl BoundedReceiver { } 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)) } 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>>, +} + +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, +} + +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 { + 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/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 1d39c2d..1b469eb 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -17,6 +17,8 @@ use std::future::Future; use std::sync::Arc; +use std::sync::Barrier; +use std::sync::Mutex; use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -100,6 +102,86 @@ fn assert_completes_without_deadlock(test: impl FnOnce() + Send + 'static) { worker.join().unwrap(); } +#[test] +fn bounded_send_rechecks_capacity_freed_by_waker_clone() { + struct ReceiveOnClone { + receiver: Mutex>, + received: AtomicBool, + } + + unsafe fn clone(data: *const ()) -> RawWaker { + let pointer = data.cast::(); + // SAFETY: Each raw waker owns one Arc reference, and this callback borrows that reference. + let state = unsafe { &*pointer }; + if !state.received.swap(true, Ordering::Relaxed) { + assert_eq!(state.receiver.lock().unwrap().try_recv(), Ok(1)); + } + // SAFETY: The live reference owned by the input waker keeps the allocation alive. + unsafe { Arc::increment_strong_count(pointer) }; + RawWaker::new(data, &VTABLE) + } + unsafe fn release(data: *const ()) { + // SAFETY: Consumes exactly the Arc reference owned by this waker. + drop(unsafe { Arc::from_raw(data.cast::()) }); + } + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, release, |_| {}, release); + + assert_completes_without_deadlock(|| { + let (tx, rx) = mpsc::bounded(1); + tx.try_send(1).unwrap(); + let state = Arc::new(ReceiveOnClone { + receiver: Mutex::new(rx), + received: AtomicBool::new(false), + }); + let data = Arc::into_raw(state.clone()).cast(); + // SAFETY: The vtable owns one Arc per waker; the callback state is Send + Sync. + let waker = unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }; + assert_eq!( + Box::pin(tx.send(2)) + .as_mut() + .poll(&mut Context::from_waker(&waker)), + Poll::Ready(Ok(())) + ); + assert_eq!(state.receiver.lock().unwrap().try_recv(), Ok(2)); + }); +} + +#[test] +fn bounded_receive_racing_with_send_registration_cannot_lose_wakeup() { + struct Notified(AtomicBool); + impl Wake for Notified { + fn wake(self: Arc) { + self.0.store(true, Ordering::Relaxed); + } + } + + for _ in 0..128 { + let (tx, mut rx) = mpsc::bounded(1); + tx.try_send(1).unwrap(); + let start = Barrier::new(2); + let notified = Arc::new(Notified(AtomicBool::new(false))); + let waker = Waker::from(notified.clone()); + 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.0.load(Ordering::Relaxed)); + 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 unbounded_wake_callback_can_send() { struct SendOnWake(mpsc::UnboundedSender); From 313598fc02b3f1ea93be69383a4ae851e3661f28 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 02:47:20 +0800 Subject: [PATCH 04/10] perf(mpsc): skip unarmed receiver notification updates --- CHANGELOG.md | 1 + asyncband/src/internal/atomic_waker.rs | 44 ++++++++++++++++++++++++-- 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index c1bdc02..99e1c62 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -14,6 +14,7 @@ All notable changes to this project will be documented in this file. ### Improvements +* Reduce bounded MPSC contention when senders or the receiver are not waiting, improving throughput without changing capacity or cancellation semantics. * Reduce unbounded MPSC synchronization overhead by transferring messages in batches and coordinating receiver notifications with queued messages; release large empty batch allocations while retaining small buffers for reuse. ## v0.7.2 diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs index a6dcd20..809e6c8 100644 --- a/asyncband/src/internal/atomic_waker.rs +++ b/asyncband/src/internal/atomic_waker.rs @@ -29,8 +29,10 @@ use std::panic::RefUnwindSafe; use std::panic::UnwindSafe; use std::panic::catch_unwind; use std::panic::resume_unwind; +use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; +use std::sync::atomic::fence; use std::task::Waker; const WAITING: usize = 0; @@ -68,9 +70,14 @@ const WAKING: usize = 0b10; /// that determines whether to return `Pending`. /// /// Every transition that acquires slot ownership has an Acquire operation paired with the previous -/// owner's Release transition to `WAITING`. The Release half of `wake` also publishes the caller's -/// preceding condition update; a racing `register` acquires that publication before it returns. +/// owner's Release transition to `WAITING`. An additional `armed` hint lets notifications skip the +/// state machine when no registration needs waking. SeqCst fences order the notifier's condition +/// update and hint check against the registerer's hint update and subsequent condition check. +/// Either the notifier observes the registration, or the registerer observes the condition update. pub struct AtomicWaker { + // Only the state machine grants slot ownership. A false hint may also mean another notifier + // has claimed responsibility for a registration, but has not taken the slot yet. + armed: AtomicBool, state: AtomicUsize, waker: UnsafeCell>, } @@ -89,6 +96,7 @@ impl AtomicWaker { #[inline] pub const fn new() -> Self { Self { + armed: AtomicBool::new(false), state: AtomicUsize::new(WAITING), waker: UnsafeCell::new(None), } @@ -125,6 +133,11 @@ impl AtomicWaker { debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); } } + // Publish the hint after installing the waker. A notifier that skipped an unfinished + // registration must publish its condition before its hint check; the paired fences ensure + // the caller's post-registration condition check cannot miss that publication too. + self.armed.store(true, Ordering::SeqCst); + fence(Ordering::SeqCst); } /// Registers a waker after this thread has acquired the REGISTERING state. @@ -214,6 +227,13 @@ impl AtomicWaker { /// wake may instead take responsibility for notifying it. #[inline] pub fn take(&self) -> Option { + fence(Ordering::SeqCst); + if !self.armed.load(Ordering::SeqCst) { + return None; + } + // Clear the hint before taking the slot. A registration racing after this clear either + // gets taken below or arms the hint again; clearing after taking could erase a newer wait. + self.armed.store(false, Ordering::SeqCst); // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the // previous owner. Release publishes the condition update that the caller performed before // calling wake, including when a registering thread already owns the slot. @@ -314,6 +334,26 @@ mod tests { assert_eq!(counter.0.load(Ordering::Relaxed), 1); } + #[test] + fn first_registration_cannot_miss_a_skipped_notification() { + for _ in 0..128 { + let published = AtomicBool::new(false); + let atomic_waker = AtomicWaker::new(); + let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); + let waker = Waker::from(counter.clone()); + std::thread::scope(|scope| { + let sender = scope.spawn(|| { + published.store(true, Ordering::Relaxed); + atomic_waker.wake(); + }); + atomic_waker.register(&waker); + let publication_seen = published.load(Ordering::Relaxed); + sender.join().unwrap(); + assert!(publication_seen || counter.0.load(Ordering::Relaxed) > 0); + }); + } + } + #[test] fn wake_during_replacement_notifies_old_and_new_tasks() { let old_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); From 5c200ef919755be6efa93fdefe0c334a9f2a709d Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 02:49:09 +0800 Subject: [PATCH 05/10] perf(mpsc): reclaim unbounded storage in byte-limited segments --- CHANGELOG.md | 4 +- asyncband/src/mpsc/unbounded.rs | 135 +++++++++++++++++++++------ tests-integration/tests/mpsc_test.rs | 35 +++++++ 3 files changed, 145 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 99e1c62..db61a12 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,15 +7,13 @@ 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. - * Complete bounded MPSC disconnection notifications and buffered-message cleanup even when a wake callback or message destructor panics. - * Avoid deadlocks when a bounded MPSC sender's waker clone callback receives from the same channel. ### Improvements * Reduce bounded MPSC contention when senders or the receiver are not waiting, improving throughput without changing capacity or cancellation semantics. -* 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. +* 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/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index 92be699..790c8c3 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -40,13 +40,13 @@ use crate::internal::mutex::Mutex; /// 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() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { senders: AtomicUsize::new(1), inbox: Mutex::new(Inbox { - messages: VecDeque::new(), + buffer: Buffer::new(), receiver_alive: true, rx_waker: None, }), @@ -70,7 +70,7 @@ struct UnboundedState { // Queue contents, receiver liveness, and its wake registration share one lock. Registering a // wait and checking its condition cannot race with sending or receiver disconnection. struct Inbox { - messages: VecDeque, + buffer: Buffer, receiver_alive: bool, rx_waker: Option, } @@ -119,7 +119,7 @@ impl UnboundedSender { 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 { @@ -150,7 +150,10 @@ impl Drop for UnboundedReceiver { 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 +187,8 @@ impl UnboundedReceiver { 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 +197,6 @@ impl UnboundedReceiver { TryRecvError::Empty }); } - mem::swap(batch, &mut state.messages); } Ok(pop_batch(batch)) } @@ -241,8 +244,8 @@ impl UnboundedReceiver { 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))); } @@ -271,12 +274,67 @@ impl UnboundedReceiver { // No operation relies on a pinned location for the receiver batch or its values. impl Unpin for UnboundedReceiver {} -// Retain small buffers for reuse, measuring inline storage rather than separately boxed payloads. -const BATCH_CACHE_BYTES: usize = 64 * 1024; +// 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. +const SEGMENT_BYTES: usize = 32 * 1024; + +struct Buffer { + writable: VecDeque, + sealed: VecDeque>, + spare: VecDeque, +} + +impl Buffer { + fn new() -> Self { + Self { + writable: VecDeque::new(), + sealed: VecDeque::new(), + spare: VecDeque::new(), + } + } + + fn segment_capacity() -> usize { + if mem::size_of::() == 0 { + return usize::MAX; + } + let limit = (SEGMENT_BYTES / mem::size_of::()).max(1); + // Power-of-two limits let VecDeque grow naturally without exceeding the segment budget. + 1 << (usize::BITS - 1 - limit.leading_zeros()) + } + + 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); + } + + fn refill(&mut self, batch: &mut VecDeque) { + 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::>() > 1024 + { + self.sealed = VecDeque::new(); + } + } else if !self.writable.is_empty() { + self.spare = VecDeque::new(); + mem::swap(batch, &mut self.writable); + } + } +} fn pop_batch(batch: &mut VecDeque) -> T { - if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::()) > BATCH_CACHE_BYTES - { + if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::()) > SEGMENT_BYTES { // Retire the allocation on the last value, outside the inbox lock. Keep this as a tail // expression to avoid intermediate storage for large inline values. mem::take(batch).pop_front() @@ -288,6 +346,7 @@ fn pop_batch(batch: &mut VecDeque) -> T { #[cfg(test)] mod tests { + use super::SEGMENT_BYTES; use super::unbounded; use crate::mpsc::TryRecvError; @@ -304,22 +363,46 @@ mod tests { } #[test] - fn releases_large_batches_after_the_last_value() { + fn reclaims_storage_while_a_burst_is_partially_consumed() { + fn allocated_bytes(rx: &mut super::UnboundedReceiver) -> usize { + let batch = rx.batch.get_mut().capacity(); + let inbox = rx.state.inbox.lock(); + let buffer = &inbox.buffer; + let slots = batch + + buffer.writable.capacity() + + buffer.spare.capacity() + + buffer + .sealed + .iter() + .map(|batch| batch.capacity()) + .sum::(); + slots * size_of::() + } + let (tx, mut rx) = unbounded(); - for value in 0..128u8 { - tx.send([value; 1024]).unwrap(); + for value in 0..1024usize { + tx.send([value; 128]).unwrap(); } - for value in 0..64u8 { - assert_eq!(rx.try_recv(), Ok([value; 1024])); + let peak = allocated_bytes(&mut rx); + for value in 0..512 { + assert_eq!(rx.try_recv(), Ok([value; 128])); } - // 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])); + assert!(allocated_bytes(&mut rx) <= peak * 3 / 4); + // New sends must remain behind both the receiver's current segment and sealed segments. + tx.send([1024; 128]).unwrap(); + for value in 512..=1024 { + assert_eq!(rx.try_recv(), Ok([value; 128])); } - // Reclaim on the final successful receive, without requiring an extra empty poll. + assert!(allocated_bytes(&mut rx) <= 2 * SEGMENT_BYTES); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } + + #[test] + fn does_not_cache_an_oversized_inline_value() { + let (tx, mut rx) = unbounded(); + tx.send([7u8; SEGMENT_BYTES + 1]).unwrap(); + assert_eq!(rx.try_recv(), Ok([7u8; SEGMENT_BYTES + 1])); assert_eq!(rx.batch.get_mut().capacity(), 0); - assert_eq!(rx.try_recv(), Ok([128; 1024])); assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } @@ -338,7 +421,7 @@ mod tests { tx.send(32).unwrap(); assert_eq!(rx.try_recv(), Ok(32)); - assert_eq!(rx.state.inbox.lock().messages.capacity(), capacity); + assert_eq!(rx.state.inbox.lock().buffer.writable.capacity(), capacity); } #[test] diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 1b469eb..60a7390 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -182,6 +182,41 @@ fn bounded_receive_racing_with_send_registration_cannot_lose_wakeup() { } } +#[test] +fn unbounded_disconnect_drops_partial_and_queued_batches_outside_lock() { + struct Value(Option>); + 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 = { + 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); + }); +} + #[test] fn unbounded_wake_callback_can_send() { struct SendOnWake(mpsc::UnboundedSender); From 4a3615fdde0dbf0fd163e3a4a0830943e69526f3 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 02:49:12 +0800 Subject: [PATCH 06/10] test(mpsc): include channel lifecycle coverage in the Miri workflow --- tests-integration/tests/mpsc_test.rs | 15 ++++++++++++++- xtask/src/main.rs | 1 + 2 files changed, 15 insertions(+), 1 deletion(-) diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 60a7390..6a54ac0 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -29,6 +29,7 @@ use std::task::RawWakerVTable; use std::task::Wake; use std::task::Waker; use std::thread; +#[cfg(not(miri))] use std::time::Duration; use asyncband::mpsc; @@ -96,9 +97,13 @@ fn assert_completes_without_deadlock(test: impl FnOnce() + Send + 'static) { test(); finished_tx.send(()).unwrap(); }); + #[cfg(not(miri))] finished_rx .recv_timeout(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(); } @@ -333,6 +338,7 @@ fn unbounded_waker_clone_rechecks_messages_sent_during_registration() { } #[test] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] fn unbounded_collects_from_multiple_producers() { let (tx, mut rx) = mpsc::unbounded(); @@ -354,6 +360,7 @@ fn unbounded_collects_from_multiple_producers() { } #[tokio::test] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] async fn select_streams() { let (tx1, mut rx1) = mpsc::unbounded::(); let (tx2, mut rx2) = mpsc::unbounded::(); @@ -431,6 +438,7 @@ async fn select_streams() { } #[tokio::test] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] async fn send_recv_unbounded() { let (tx, mut rx) = mpsc::unbounded::(); @@ -447,6 +455,7 @@ async fn send_recv_unbounded() { } #[tokio::test] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] async fn async_send_recv_unbounded() { let (tx, mut rx) = mpsc::unbounded(); @@ -515,6 +524,7 @@ fn buffered_messages_are_drained_before_disconnection() { } #[tokio::test] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] async fn send_recv_bounded() { let (tx, mut rx) = mpsc::bounded(1); @@ -526,6 +536,7 @@ async fn send_recv_bounded() { } #[tokio::test] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] async fn async_send_recv_bounded() { let (tx, mut rx) = mpsc::bounded(1); @@ -564,7 +575,7 @@ fn bounded_try_send_respects_capacity_and_order() { #[test] fn bounded_try_recv_does_not_report_empty_after_completed_sends() { const PRODUCERS: usize = 4; - const MESSAGES_PER_PRODUCER: usize = 16_384; + 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; @@ -607,6 +618,7 @@ fn bounded_try_recv_does_not_report_empty_after_completed_sends() { } #[tokio::test] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] async fn try_send_after_disconnection_bounded() { let (tx, rx) = mpsc::bounded(1); @@ -617,6 +629,7 @@ async fn try_send_after_disconnection_bounded() { } #[tokio::test] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] async fn send_after_disconnection_bounded() { let (tx, mut rx) = mpsc::bounded(1); 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"])); } } From 95ef9eb7af2ebc461935cfb8da5eac40bc4e0020 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 10:40:11 +0800 Subject: [PATCH 07/10] bench(mpsc): cover concurrent inline payload bursts and wakeups --- benchmarks/ecosystem/mpsc/unbounded.rs | 119 +++++++++++++++++++++++++ 1 file changed, 119 insertions(+) diff --git a/benchmarks/ecosystem/mpsc/unbounded.rs b/benchmarks/ecosystem/mpsc/unbounded.rs index 229bf2d..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; @@ -167,3 +174,115 @@ fn scheduled(bencher: Bencher, (producers, workers): (usize, u 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>( + 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]) +} From 0adca285c45514b9a5ca2068b739634cc39cab0e Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 11:56:08 +0800 Subject: [PATCH 08/10] refactor(mpsc): keep bounded receiver waiting local to the ring --- asyncband/src/internal/atomic_waker.rs | 44 +--- asyncband/src/internal/mod.rs | 4 +- .../src/mpsc/{bounded.rs => bounded/mod.rs} | 31 +-- asyncband/src/mpsc/bounded/ring.rs | 236 ++++-------------- asyncband/src/mpsc/bounded/ring_tests.rs | 200 +++++++++++++++ 5 files changed, 271 insertions(+), 244 deletions(-) rename asyncband/src/mpsc/{bounded.rs => bounded/mod.rs} (93%) create mode 100644 asyncband/src/mpsc/bounded/ring_tests.rs diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs index 809e6c8..a6dcd20 100644 --- a/asyncband/src/internal/atomic_waker.rs +++ b/asyncband/src/internal/atomic_waker.rs @@ -29,10 +29,8 @@ use std::panic::RefUnwindSafe; use std::panic::UnwindSafe; use std::panic::catch_unwind; use std::panic::resume_unwind; -use std::sync::atomic::AtomicBool; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; -use std::sync::atomic::fence; use std::task::Waker; const WAITING: usize = 0; @@ -70,14 +68,9 @@ const WAKING: usize = 0b10; /// that determines whether to return `Pending`. /// /// Every transition that acquires slot ownership has an Acquire operation paired with the previous -/// owner's Release transition to `WAITING`. An additional `armed` hint lets notifications skip the -/// state machine when no registration needs waking. SeqCst fences order the notifier's condition -/// update and hint check against the registerer's hint update and subsequent condition check. -/// Either the notifier observes the registration, or the registerer observes the condition update. +/// owner's Release transition to `WAITING`. The Release half of `wake` also publishes the caller's +/// preceding condition update; a racing `register` acquires that publication before it returns. pub struct AtomicWaker { - // Only the state machine grants slot ownership. A false hint may also mean another notifier - // has claimed responsibility for a registration, but has not taken the slot yet. - armed: AtomicBool, state: AtomicUsize, waker: UnsafeCell>, } @@ -96,7 +89,6 @@ impl AtomicWaker { #[inline] pub const fn new() -> Self { Self { - armed: AtomicBool::new(false), state: AtomicUsize::new(WAITING), waker: UnsafeCell::new(None), } @@ -133,11 +125,6 @@ impl AtomicWaker { debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); } } - // Publish the hint after installing the waker. A notifier that skipped an unfinished - // registration must publish its condition before its hint check; the paired fences ensure - // the caller's post-registration condition check cannot miss that publication too. - self.armed.store(true, Ordering::SeqCst); - fence(Ordering::SeqCst); } /// Registers a waker after this thread has acquired the REGISTERING state. @@ -227,13 +214,6 @@ impl AtomicWaker { /// wake may instead take responsibility for notifying it. #[inline] pub fn take(&self) -> Option { - fence(Ordering::SeqCst); - if !self.armed.load(Ordering::SeqCst) { - return None; - } - // Clear the hint before taking the slot. A registration racing after this clear either - // gets taken below or arms the hint again; clearing after taking could erase a newer wait. - self.armed.store(false, Ordering::SeqCst); // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the // previous owner. Release publishes the condition update that the caller performed before // calling wake, including when a registering thread already owns the slot. @@ -334,26 +314,6 @@ mod tests { assert_eq!(counter.0.load(Ordering::Relaxed), 1); } - #[test] - fn first_registration_cannot_miss_a_skipped_notification() { - for _ in 0..128 { - let published = AtomicBool::new(false); - let atomic_waker = AtomicWaker::new(); - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - std::thread::scope(|scope| { - let sender = scope.spawn(|| { - published.store(true, Ordering::Relaxed); - atomic_waker.wake(); - }); - atomic_waker.register(&waker); - let publication_seen = published.load(Ordering::Relaxed); - sender.join().unwrap(); - assert!(publication_seen || counter.0.load(Ordering::Relaxed) > 0); - }); - } - } - #[test] fn wake_during_replacement_notifies_old_and_new_tasks() { let old_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 8d4f1b7..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) { } } -#[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( diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded/mod.rs similarity index 93% rename from asyncband/src/mpsc/bounded.rs rename to asyncband/src/mpsc/bounded/mod.rs index 89267b6..455e709 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -33,8 +33,9 @@ use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; -use crate::internal::atomic_waker::AtomicWaker; +// 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; @@ -49,11 +50,10 @@ mod waiters; #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { 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: SendWaiters::new(), - rx_waker: AtomicWaker::new(), }); let sender = BoundedSender { state: state.clone(), @@ -62,18 +62,17 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { (sender, receiver) } -struct BoundedState { +struct Shared { buffer: Ring, senders: AtomicUsize, send_waiters: SendWaiters, - rx_waker: AtomicWaker, } /// The sending endpoint of a bounded mpsc channel. /// /// Instances are created by the [`bounded`] function. pub struct BoundedSender { - state: Arc>, + state: Arc>, } impl Clone for BoundedSender { @@ -93,14 +92,8 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { 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(); } } } @@ -176,9 +169,7 @@ impl BoundedSender { /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - self.state.buffer.try_push(value)?; - self.state.rx_waker.wake(); - Ok(()) + self.state.buffer.try_push(value) } } @@ -186,7 +177,7 @@ impl BoundedSender { /// /// Instances are created by the [`bounded`] function. pub struct BoundedReceiver { - state: Arc>, + state: Arc>, } impl fmt::Debug for BoundedReceiver { @@ -208,7 +199,7 @@ impl Drop for BoundedReceiver { 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(); + 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(); @@ -311,7 +302,7 @@ impl BoundedReceiver { 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 c4e4e47..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 { slots: Box<[Slot]>, head: CachePadded, tail: CachePadded, + // This flag is usually stable while the consumer advances head. Sharing head + // for notifications would make every producer track a constantly invalidated cache line. + receiver_waiting: CachePadded, + receiver: Mutex>, capacity: usize, one_lap: usize, mark_bit: usize, @@ -65,6 +72,8 @@ impl Ring { 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 Ring { // 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, @@ -160,6 +179,40 @@ impl Ring { } } + /// Registers the exclusive receiver, which must retry `pop` before returning Pending. + pub fn register_receiver(&self, waker: &Waker) { + let mut receiver = self.receiver.lock(); + let old_waker = if receiver.as_ref().is_some_and(|old| old.will_wake(waker)) { + None + } else { + // Only the receiver registers. Producers can take the old waker while we clone, + // but cannot install a replacement. Clone/drop callbacks may send into this channel. + drop(receiver); + let waker = waker.clone(); + receiver = self.receiver.lock(); + receiver.replace(waker) + }; + self.receiver_waiting.store(true, Ordering::Relaxed); + // Paired with the publisher's fence: either it sees this flag, or the receiver's + // subsequent pop sees its stamp. Taking a waker clears the flag under the same lock, + // so it cannot erase a newer registration without also taking responsibility for it. + fence(Ordering::SeqCst); + drop(receiver); + drop(old_waker); + } + + pub fn take_receiver_waker(&self) -> Option { + let mut receiver = self.receiver.lock(); + self.receiver_waiting.store(false, Ordering::Relaxed); + receiver.take() + } + + pub fn wake_receiver(&self) { + if let Some(waker) = self.take_receiver_waker() { + waker.wake(); + } + } + /// Prevents subsequent sends from reserving slots. Already reserved slots still publish. pub fn close(&self) { self.tail.fetch_or(self.mark_bit, Ordering::SeqCst); @@ -244,184 +297,5 @@ impl Drop for Ring { } #[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 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()) - }; - let result = completed.recv_timeout(std::time::Duration::from_secs(10)); - // 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::>()); - } - - #[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" - ); - } - } -} +#[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::>()); +} + +#[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" + ); + } +} From 7be7d2f2add0c25cffc1c0102a00c50b70d1b516 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 11:56:08 +0800 Subject: [PATCH 09/10] refactor(mpsc): separate unbounded storage from channel lifecycle --- asyncband/src/mpsc/unbounded/buffer.rs | 93 +++++++++ asyncband/src/mpsc/unbounded/buffer_tests.rs | 114 ++++++++++++ .../mpsc/{unbounded.rs => unbounded/mod.rs} | 176 +----------------- 3 files changed, 216 insertions(+), 167 deletions(-) create mode 100644 asyncband/src/mpsc/unbounded/buffer.rs create mode 100644 asyncband/src/mpsc/unbounded/buffer_tests.rs rename asyncband/src/mpsc/{unbounded.rs => unbounded/mod.rs} (62%) 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 { + writable: VecDeque, + sealed: VecDeque>, + spare: VecDeque, +} + +impl Buffer { + pub fn new() -> Self { + Self { + writable: VecDeque::new(), + sealed: VecDeque::new(), + spare: VecDeque::new(), + } + } + + fn segment_capacity() -> usize { + if mem::size_of::() == 0 { + return usize::MAX; + } + let limit = (SEGMENT_BYTES / mem::size_of::()).max(1); + // Power-of-two limits let VecDeque grow naturally without exceeding the segment budget. + 1 << (usize::BITS - 1 - limit.leading_zeros()) + } + + pub fn push(&mut self, value: T) { + if self.writable.len() == Self::segment_capacity() { + let next = 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) { + 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::>() > 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(batch: &mut VecDeque) -> T { + if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::()) > SEGMENT_BYTES { + // Retire the allocation on the last value, outside the inbox lock. Keep this as a tail + // expression to avoid intermediate storage for large inline values. + mem::take(batch).pop_front() + } 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(buffer: &Buffer, batch: &VecDeque) -> usize { + let slots = batch.capacity() + + buffer.writable.capacity() + + buffer.spare.capacity() + + buffer.sealed.iter().map(VecDeque::capacity).sum::(); + slots * size_of::() +} + +fn receive(buffer: &mut Buffer, batch: &mut VecDeque) -> 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 62% rename from asyncband/src/mpsc/unbounded.rs rename to asyncband/src/mpsc/unbounded/mod.rs index 790c8c3..7c69ac4 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded/mod.rs @@ -29,11 +29,16 @@ 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 @@ -43,7 +48,7 @@ use crate::internal::mutex::Mutex; /// 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() -> (UnboundedSender, UnboundedReceiver) { - let state = Arc::new(UnboundedState { + let state = Arc::new(Shared { senders: AtomicUsize::new(1), inbox: Mutex::new(Inbox { buffer: Buffer::new(), @@ -61,7 +66,7 @@ pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { (sender, receiver) } -struct UnboundedState { +struct Shared { // Endpoint cloning and ordinary drops do not contend with message traffic. senders: AtomicUsize, inbox: Mutex>, @@ -79,7 +84,7 @@ struct Inbox { /// /// Instances are created by the [`unbounded`] function. pub struct UnboundedSender { - state: Arc>, + state: Arc>, } impl Clone for UnboundedSender { @@ -133,7 +138,7 @@ impl UnboundedSender { /// /// Instances are created by the [`unbounded`] function. pub struct UnboundedReceiver { - state: Arc>, + state: Arc>, // Only accessed through `get_mut`; the mutex preserves Sync for Send-only payloads. batch: Mutex>, } @@ -273,166 +278,3 @@ impl UnboundedReceiver { // No operation relies on a pinned location for the receiver batch or its values. impl Unpin for UnboundedReceiver {} - -// 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. -const SEGMENT_BYTES: usize = 32 * 1024; - -struct Buffer { - writable: VecDeque, - sealed: VecDeque>, - spare: VecDeque, -} - -impl Buffer { - fn new() -> Self { - Self { - writable: VecDeque::new(), - sealed: VecDeque::new(), - spare: VecDeque::new(), - } - } - - fn segment_capacity() -> usize { - if mem::size_of::() == 0 { - return usize::MAX; - } - let limit = (SEGMENT_BYTES / mem::size_of::()).max(1); - // Power-of-two limits let VecDeque grow naturally without exceeding the segment budget. - 1 << (usize::BITS - 1 - limit.leading_zeros()) - } - - 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); - } - - fn refill(&mut self, batch: &mut VecDeque) { - 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::>() > 1024 - { - self.sealed = VecDeque::new(); - } - } else if !self.writable.is_empty() { - self.spare = VecDeque::new(); - mem::swap(batch, &mut self.writable); - } - } -} - -fn pop_batch(batch: &mut VecDeque) -> T { - if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::()) > SEGMENT_BYTES { - // Retire the allocation on the last value, outside the inbox lock. Keep this as a tail - // expression to avoid intermediate storage for large inline values. - mem::take(batch).pop_front() - } else { - batch.pop_front() - } - .expect("receiver batch must not be empty") -} - -#[cfg(test)] -mod tests { - use super::SEGMENT_BYTES; - 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 reclaims_storage_while_a_burst_is_partially_consumed() { - fn allocated_bytes(rx: &mut super::UnboundedReceiver) -> usize { - let batch = rx.batch.get_mut().capacity(); - let inbox = rx.state.inbox.lock(); - let buffer = &inbox.buffer; - let slots = batch - + buffer.writable.capacity() - + buffer.spare.capacity() - + buffer - .sealed - .iter() - .map(|batch| batch.capacity()) - .sum::(); - slots * size_of::() - } - - let (tx, mut rx) = unbounded(); - for value in 0..1024usize { - tx.send([value; 128]).unwrap(); - } - let peak = allocated_bytes(&mut rx); - for value in 0..512 { - assert_eq!(rx.try_recv(), Ok([value; 128])); - } - assert!(allocated_bytes(&mut rx) <= peak * 3 / 4); - // New sends must remain behind both the receiver's current segment and sealed segments. - tx.send([1024; 128]).unwrap(); - for value in 512..=1024 { - assert_eq!(rx.try_recv(), Ok([value; 128])); - } - assert!(allocated_bytes(&mut rx) <= 2 * SEGMENT_BYTES); - assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - } - - #[test] - fn does_not_cache_an_oversized_inline_value() { - let (tx, mut rx) = unbounded(); - tx.send([7u8; SEGMENT_BYTES + 1]).unwrap(); - assert_eq!(rx.try_recv(), Ok([7u8; SEGMENT_BYTES + 1])); - assert_eq!(rx.batch.get_mut().capacity(), 0); - 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().buffer.writable.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)); - } -} From f2ce32294cf94c0d98bff61a11abdaa90113b0c5 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 11:56:08 +0800 Subject: [PATCH 10/10] test(mpsc): organize coverage around channel contracts --- CHANGELOG.md | 3 +- tests-integration/tests/mpsc_test.rs | 776 ------------------ .../tests/mpsc_test/backpressure.rs | 107 +++ .../tests/mpsc_test/callbacks.rs | 328 ++++++++ .../tests/mpsc_test/concurrency.rs | 289 +++++++ tests-integration/tests/mpsc_test/main.rs | 175 ++++ tests-integration/tests/mpsc_test/support.rs | 106 +++ 7 files changed, 1007 insertions(+), 777 deletions(-) delete mode 100644 tests-integration/tests/mpsc_test.rs create mode 100644 tests-integration/tests/mpsc_test/backpressure.rs create mode 100644 tests-integration/tests/mpsc_test/callbacks.rs create mode 100644 tests-integration/tests/mpsc_test/concurrency.rs create mode 100644 tests-integration/tests/mpsc_test/main.rs create mode 100644 tests-integration/tests/mpsc_test/support.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index db61a12..43ffd8e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,11 +7,12 @@ 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. -* Complete bounded MPSC disconnection notifications and buffered-message cleanup even when a wake callback or message destructor panics. +* 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 +* 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. diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs deleted file mode 100644 index 6a54ac0..0000000 --- a/tests-integration/tests/mpsc_test.rs +++ /dev/null @@ -1,776 +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::Barrier; -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::RawWaker; -use std::task::RawWakerVTable; -use std::task::Wake; -use std::task::Waker; -use std::thread; -#[cfg(not(miri))] -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(poll: Poll) -> T { - match poll { - Poll::Ready(value) => value, - Poll::Pending => panic!("future should be ready"), - } -} - -struct HoldSender { - _sender: S, -} - -// This waker must own the sender so its final drop can break the tested reference cycle. -#[allow(clippy::manual_noop_waker)] -impl Wake for HoldSender { - fn wake(self: Arc) {} -} - -#[test] -fn bounded_receiver_drop_releases_registered_waker() { - let (tx, mut rx) = mpsc::bounded::<()>(1); - let holder = Arc::new(HoldSender { _sender: tx }); - let retained = Arc::downgrade(&holder); - let waker = Waker::from(holder); - assert!( - Box::pin(rx.recv()) - .as_mut() - .poll(&mut Context::from_waker(&waker)) - .is_pending() - ); - drop(waker); - drop(rx); - assert!(retained.upgrade().is_none()); -} - -#[test] -fn unbounded_receiver_drop_releases_registered_waker() { - let (tx, mut rx) = mpsc::unbounded::<()>(); - let holder = Arc::new(HoldSender { _sender: tx }); - let retained = Arc::downgrade(&holder); - let waker = Waker::from(holder); - assert!( - Box::pin(rx.recv()) - .as_mut() - .poll(&mut Context::from_waker(&waker)) - .is_pending() - ); - drop(waker); - drop(rx); - assert!(retained.upgrade().is_none()); -} - -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(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(); -} - -#[test] -fn bounded_send_rechecks_capacity_freed_by_waker_clone() { - struct ReceiveOnClone { - receiver: Mutex>, - received: AtomicBool, - } - - unsafe fn clone(data: *const ()) -> RawWaker { - let pointer = data.cast::(); - // SAFETY: Each raw waker owns one Arc reference, and this callback borrows that reference. - let state = unsafe { &*pointer }; - if !state.received.swap(true, Ordering::Relaxed) { - assert_eq!(state.receiver.lock().unwrap().try_recv(), Ok(1)); - } - // SAFETY: The live reference owned by the input waker keeps the allocation alive. - unsafe { Arc::increment_strong_count(pointer) }; - RawWaker::new(data, &VTABLE) - } - unsafe fn release(data: *const ()) { - // SAFETY: Consumes exactly the Arc reference owned by this waker. - drop(unsafe { Arc::from_raw(data.cast::()) }); - } - static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, release, |_| {}, release); - - assert_completes_without_deadlock(|| { - let (tx, rx) = mpsc::bounded(1); - tx.try_send(1).unwrap(); - let state = Arc::new(ReceiveOnClone { - receiver: Mutex::new(rx), - received: AtomicBool::new(false), - }); - let data = Arc::into_raw(state.clone()).cast(); - // SAFETY: The vtable owns one Arc per waker; the callback state is Send + Sync. - let waker = unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }; - assert_eq!( - Box::pin(tx.send(2)) - .as_mut() - .poll(&mut Context::from_waker(&waker)), - Poll::Ready(Ok(())) - ); - assert_eq!(state.receiver.lock().unwrap().try_recv(), Ok(2)); - }); -} - -#[test] -fn bounded_receive_racing_with_send_registration_cannot_lose_wakeup() { - struct Notified(AtomicBool); - impl Wake for Notified { - fn wake(self: Arc) { - self.0.store(true, Ordering::Relaxed); - } - } - - for _ in 0..128 { - let (tx, mut rx) = mpsc::bounded(1); - tx.try_send(1).unwrap(); - let start = Barrier::new(2); - let notified = Arc::new(Notified(AtomicBool::new(false))); - let waker = Waker::from(notified.clone()); - 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.0.load(Ordering::Relaxed)); - 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 unbounded_disconnect_drops_partial_and_queued_batches_outside_lock() { - struct Value(Option>); - 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 = { - 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); - }); -} - -#[test] -fn unbounded_wake_callback_can_send() { - struct SendOnWake(mpsc::UnboundedSender); - - impl Wake for SendOnWake { - fn wake(self: Arc) { - self.0.send(2).unwrap(); - } - } - - assert_completes_without_deadlock(|| { - let (tx, mut rx) = mpsc::unbounded(); - let waker = Waker::from(Arc::new(SendOnWake(tx.clone()))); - assert!( - Box::pin(rx.recv()) - .as_mut() - .poll(&mut Context::from_waker(&waker)) - .is_pending() - ); - tx.send(1).unwrap(); - assert_eq!(rx.try_recv(), Ok(1)); - assert_eq!(rx.try_recv(), Ok(2)); - }); -} - -#[test] -fn unbounded_replaced_and_disconnected_wakers_can_send() { - struct SendOnDrop { - sender: mpsc::UnboundedSender, - disconnected: bool, - drops: Arc, - } - - // The final waker drop must run a callback, even though waking itself does nothing. - #[allow(clippy::manual_noop_waker)] - impl Wake for SendOnDrop { - fn wake(self: Arc) {} - } - - impl Drop for SendOnDrop { - fn drop(&mut self) { - assert_eq!(self.sender.send(7).is_err(), self.disconnected); - self.drops.fetch_add(1, Ordering::Relaxed); - } - } - - assert_completes_without_deadlock(|| { - for disconnected in [false, true] { - let (tx, mut rx) = mpsc::unbounded(); - let drops = Arc::new(AtomicUsize::new(0)); - let waker = Waker::from(Arc::new(SendOnDrop { - sender: tx, - disconnected, - drops: drops.clone(), - })); - assert!( - Box::pin(rx.recv()) - .as_mut() - .poll(&mut Context::from_waker(&waker)) - .is_pending() - ); - drop(waker); - if disconnected { - drop(rx); - } else { - // Replacing the waker can enqueue a message during this poll. Either immediate - // completion or a notified Pending is valid, but the message must not be lost. - let poll = poll_once(Box::pin(rx.recv()).as_mut()); - if poll.is_pending() { - assert_eq!(rx.try_recv(), Ok(7)); - } else { - assert_eq!(poll, Poll::Ready(Ok(7))); - } - assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); - } - assert_eq!(drops.load(Ordering::Relaxed), 1); - } - }); -} - -#[test] -fn unbounded_waker_clone_rechecks_messages_sent_during_registration() { - unsafe fn clone_sender(data: *const ()) -> RawWaker { - let sender = data.cast::>(); - // SAFETY: Each raw waker owns an Arc to this sender; cloning borrows the live sender and - // then adds the strong reference owned by the returned waker. - unsafe { - (*sender).send(7).unwrap(); - Arc::increment_strong_count(sender); - } - RawWaker::new(data, &VTABLE) - } - - unsafe fn drop_sender(data: *const ()) { - // SAFETY: Consumes exactly the Arc reference owned by this raw waker. - drop(unsafe { Arc::from_raw(data.cast::>()) }); - } - - static VTABLE: RawWakerVTable = - RawWakerVTable::new(clone_sender, drop_sender, |_| {}, drop_sender); - - assert_completes_without_deadlock(|| { - let (tx, mut rx) = mpsc::unbounded::(); - let data = Arc::into_raw(Arc::new(tx)).cast(); - // SAFETY: The vtable manages one Arc reference per waker and the sender is Send + Sync. - let waker = unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }; - assert_eq!( - Box::pin(rx.recv()) - .as_mut() - .poll(&mut Context::from_waker(&waker)), - Poll::Ready(Ok(7)) - ); - }); -} - -#[test] -#[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 select_streams() { - let (tx1, mut rx1) = mpsc::unbounded::(); - let (tx2, mut rx2) = mpsc::unbounded::(); - 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] -#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] -async fn send_recv_unbounded() { - let (tx, mut rx) = mpsc::unbounded::(); - - // 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] -#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] -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] -#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] -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] -#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] -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 = if cfg!(miri) { 64 } else { 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] -#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] -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] -#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] -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); -} - -#[cfg(panic = "unwind")] -#[test] -fn bounded_disconnect_finishes_cleanup_when_a_callback_panics() { - struct Value { - id: usize, - drops: Arc<[AtomicUsize; 4]>, - panic_on_drop: bool, - _sender: Option>, - } - 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.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 = Arc::new(Notify { - woken: AtomicBool::new(false), - panic_on_wake, - }); - let waker = Waker::from(notify.clone()); - let mut send = Box::pin(tx.send(Value { - id: 3, - drops: drops.clone(), - panic_on_drop: false, - _sender: None, - })); - assert!( - send.as_mut() - .poll(&mut Context::from_waker(&waker)) - .is_pending() - ); - assert!(std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| drop(rx))).is_err()); - assert!(notify.woken.load(Ordering::Relaxed)); - for count in &drops[..3] { - assert_eq!(count.load(Ordering::Relaxed), 1); - } - assert_eq!(drops[3].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, 3); - assert_eq!(drops[3].load(Ordering::Relaxed), 1); - } -} 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 { + _sender: S, +} + +// This waker must own the sender so its final drop can break the tested reference cycle. +#[allow(clippy::manual_noop_waker)] +impl Wake for HoldSender { + fn wake(self: Arc) {} +} + +#[test] +fn bounded_receiver_drop_releases_registered_waker() { + let (tx, mut rx) = mpsc::bounded::<()>(1); + let holder = Arc::new(HoldSender { _sender: tx }); + let retained = Arc::downgrade(&holder); + let waker = Waker::from(holder); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + drop(rx); + assert!(retained.upgrade().is_none()); +} + +#[test] +fn unbounded_receiver_drop_releases_registered_waker() { + let (tx, mut rx) = mpsc::unbounded::<()>(); + let holder = Arc::new(HoldSender { _sender: tx }); + let retained = Arc::downgrade(&holder); + let waker = Waker::from(holder); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + drop(rx); + assert!(retained.upgrade().is_none()); +} + +#[test] +fn 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); + impl Wake for SendOnWake { + fn wake(self: Arc) { + (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, + disconnected: bool, + drops: Arc, + } + + // The final waker drop must run a callback, even though waking itself does nothing. + #[allow(clippy::manual_noop_waker)] + impl Wake for SendOnDrop { + fn wake(self: Arc) {} + } + + impl Drop for SendOnDrop { + fn drop(&mut self) { + assert_eq!(self.sender.send(7).is_err(), self.disconnected); + self.drops.fetch_add(1, Ordering::Relaxed); + } + } + + assert_completes_without_deadlock(|| { + for disconnected in [false, true] { + let (tx, mut rx) = mpsc::unbounded(); + let drops = Arc::new(AtomicUsize::new(0)); + let waker = Waker::from(Arc::new(SendOnDrop { + sender: tx, + disconnected, + drops: drops.clone(), + })); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + if disconnected { + drop(rx); + } else { + // Replacing the waker can enqueue a message during this poll. Either immediate + // completion or a notified Pending is valid, but the message must not be lost. + let (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>); + 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 = { + 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>, + } + 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.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::>(); + 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>>, + publish: impl FnOnce() + Send, + expected: Result, +) { + 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::(1); + receive_racing_with( + Box::pin(rx.recv()).as_mut(), + move || drop(tx), + Err(RecvError::Disconnected), + ); + let (tx, mut rx) = mpsc::unbounded::(); + 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::(); + let (tx2, mut rx2) = mpsc::unbounded::(); + 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::(); + 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(poll: Poll) -> 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) { + 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.0.fetch_add(1, Ordering::Relaxed); + } + + fn wake_by_ref(self: &Arc) { + self.0.fetch_add(1, Ordering::Relaxed); + } +} + +pub fn poll_with(future: Pin<&mut F>, waker: &Waker) -> Poll { + 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); + + unsafe fn clone(data: *const ()) -> RawWaker { + let pointer = data.cast::(); + // 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::()) }); + } + + 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(); +}