diff --git a/CHANGELOG.md b/CHANGELOG.md index 0543a046..35e870e6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,14 @@ All notable changes to this project will be documented in this file. ## Unreleased +### Bug fixes + +* Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. + +### Improvements + +* Reduce unbounded MPSC synchronization overhead by transferring messages in batches and coordinating receiver notifications with queued messages; release large empty batch allocations while retaining small buffers for reuse. + ## v0.7.2 ### Improvements diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs index 3616691a..a6dcd20f 100644 --- a/asyncband/src/internal/atomic_waker.rs +++ b/asyncband/src/internal/atomic_waker.rs @@ -210,8 +210,10 @@ impl AtomicWaker { } } + /// Removes the registered waker if this call acquires the slot. A concurrent registration or + /// wake may instead take responsibility for notifying it. #[inline] - fn take(&self) -> Option { + pub fn take(&self) -> Option { // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the // previous owner. Release publishes the condition update that the caller performed before // calling wake, including when a registering thread already owns the slot. diff --git a/asyncband/src/internal/cache_padded.rs b/asyncband/src/internal/cache_padded.rs new file mode 100644 index 00000000..7268030d --- /dev/null +++ b/asyncband/src/internal/cache_padded.rs @@ -0,0 +1,55 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +// Use conservative architecture estimates, not a guarantee about every CPU's cache line. +// Keep 128 bytes for large ARM/PowerPC lines and adjacent-line prefetching on x86-64, +// 256 bytes for s390x, and at least 64 bytes elsewhere. +#[cfg_attr(target_arch = "s390x", repr(align(256)))] +#[cfg_attr( + any( + target_arch = "aarch64", + target_arch = "arm64ec", + target_arch = "powerpc64", + target_arch = "x86_64", + ), + repr(align(128)) +)] +#[cfg_attr( + not(any( + target_arch = "s390x", + target_arch = "aarch64", + target_arch = "arm64ec", + target_arch = "powerpc64", + target_arch = "x86_64", + )), + repr(align(64)) +)] +pub struct CachePadded(T); + +impl CachePadded { + pub const fn new(value: T) -> Self { + Self(value) + } +} + +impl std::ops::Deref for CachePadded { + type Target = T; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index c6363791..84309fc2 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -70,6 +70,9 @@ pub(crate) mod atomic_waker; #[allow(dead_code)] pub(crate) mod arena; +#[cfg(feature = "mpsc")] +pub(crate) mod cache_padded; + #[cfg(any(feature = "latch", feature = "once"))] pub(crate) mod countdown; @@ -92,6 +95,8 @@ pub(crate) mod value_cell; feature = "waitgroup", feature = "watch", ))] +// Some primitives use only shared access, leaving `Mutex::get_mut` unused. +#[allow(dead_code)] pub(crate) mod mutex; #[cfg(any( diff --git a/asyncband/src/internal/mutex.rs b/asyncband/src/internal/mutex.rs index ce5229f8..93baf8f2 100644 --- a/asyncband/src/internal/mutex.rs +++ b/asyncband/src/internal/mutex.rs @@ -34,4 +34,8 @@ impl Mutex { pub fn lock(&self) -> std::sync::MutexGuard<'_, T> { self.0.lock().unwrap_or_else(PoisonError::into_inner) } + + pub fn get_mut(&mut self) -> &mut T { + self.0.get_mut().unwrap_or_else(PoisonError::into_inner) + } } diff --git a/asyncband/src/mpsc/bounded.rs b/asyncband/src/mpsc/bounded.rs index 6818a47c..550bf165 100644 --- a/asyncband/src/mpsc/bounded.rs +++ b/asyncband/src/mpsc/bounded.rs @@ -19,15 +19,15 @@ //! tasks with backpressure control. use std::fmt; -use std::future::Future; use std::future::poll_fn; -use std::pin::pin; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use std::task::ready; +use self::ring::Ring; use super::RecvError; use super::SendError; use super::TryRecvError; @@ -36,6 +36,8 @@ use crate::internal::atomic_waker::AtomicWaker; use crate::internal::semaphore::Acquire; use crate::internal::semaphore::Semaphore; +mod ring; + /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// /// [`BoundedSender::send`] waits for capacity when the buffer is full. Receiving a message releases @@ -48,25 +50,23 @@ use crate::internal::semaphore::Semaphore; pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); let state = Arc::new(BoundedState { + buffer: Ring::new(buffer), senders: AtomicUsize::new(1), - tx_permits: Semaphore::new(0), + send_waiters: Semaphore::new(0), rx_waker: AtomicWaker::new(), }); - let (sender, receiver) = std::sync::mpsc::sync_channel(buffer); let sender = BoundedSender { state: state.clone(), - sender: Some(sender), - }; - let receiver = BoundedReceiver { - state: state.clone(), - receiver: Some(receiver), }; + let receiver = BoundedReceiver { state }; (sender, receiver) } -struct BoundedState { +struct BoundedState { + buffer: Ring, senders: AtomicUsize, - tx_permits: Semaphore, + // Notifications grant retries; only the ring determines whether buffer capacity is available. + send_waiters: Semaphore, rx_waker: AtomicWaker, } @@ -74,8 +74,7 @@ struct BoundedState { /// /// Instances are created by the [`bounded`] function. pub struct BoundedSender { - state: Arc, - sender: Option>, + state: Arc>, } impl Clone for BoundedSender { @@ -83,7 +82,6 @@ impl Clone for BoundedSender { self.state.senders.fetch_add(1, Ordering::Release); BoundedSender { state: self.state.clone(), - sender: self.sender.clone(), } } } @@ -96,9 +94,6 @@ impl fmt::Debug for BoundedSender { impl Drop for BoundedSender { fn drop(&mut self) { - // Dropping the final underlying sender disconnects the channel. - drop(self.sender.take()); - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { 1 => { // Wake the receiver so it can observe the channel's disconnected state. @@ -142,7 +137,7 @@ impl BoundedSender { }; loop { - let poll = pin!(&mut self.acquire).poll(cx); + let poll = self.acquire.poll_once(cx.waker()); value = match self.sender.try_send(value) { Ok(()) => return Poll::Ready(Ok(())), @@ -153,7 +148,7 @@ impl BoundedSender { }; if poll.is_ready() { - self.acquire = self.sender.state.tx_permits.poll_acquire(1); + self.acquire = self.sender.state.send_waiters.poll_acquire(1); } else { self.value = Some(value); return Poll::Pending; @@ -162,7 +157,7 @@ impl BoundedSender { } } - let acquire = self.state.tx_permits.poll_acquire(1); + let acquire = self.state.send_waiters.poll_acquire(1); let mut send = SendState { sender: self, value: Some(value), @@ -192,19 +187,9 @@ impl BoundedSender { /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - // INVARIANT: A shared borrow of the endpoint cannot overlap its destructor. - let sender = self.sender.as_ref().unwrap(); - match sender.try_send(value) { - Ok(()) => { - self.state.rx_waker.wake(); - - Ok(()) - } - Err(std::sync::mpsc::TrySendError::Full(value)) => Err(TrySendError::Full(value)), - Err(std::sync::mpsc::TrySendError::Disconnected(value)) => { - Err(TrySendError::Disconnected(value)) - } - } + self.state.buffer.try_push(value)?; + self.state.rx_waker.wake(); + Ok(()) } } @@ -212,14 +197,9 @@ impl BoundedSender { /// /// Instances are created by the [`bounded`] function. pub struct BoundedReceiver { - state: Arc, - receiver: Option>, + state: Arc>, } -/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`. -/// That is, `BoundedReceiver` can only be accessed by one thread at a time. -unsafe impl Sync for BoundedReceiver {} - impl fmt::Debug for BoundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("BoundedReceiver").finish_non_exhaustive() @@ -228,13 +208,20 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - drop(self.receiver.take()); - self.state.tx_permits.notify_all(); + // 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() }; + self.state.send_waiters.notify_all(); + drop(receiver_waker); } } impl BoundedReceiver { - /// Attempts to receive the next queued value without waiting. + /// Attempts to receive the next queued value without waiting for a new message. + /// + /// A producer already publishing a queued message may delay this call until publication + /// finishes. Use [`Self::recv`] to yield asynchronously while publication is in progress. /// /// Receiving a value frees one buffer slot. An empty channel returns [`TryRecvError::Empty`] /// while at least one sender remains, or [`TryRecvError::Disconnected`] after every sender has @@ -257,18 +244,33 @@ impl BoundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - // INVARIANT: A mutable borrow of the endpoint cannot overlap its destructor. - let receiver = self.receiver.as_ref().unwrap(); - match receiver.try_recv() { - Ok(v) => { - self.state.tx_permits.release_if_nonempty(1); - Ok(v) + loop { + if let Poll::Ready(result) = self.try_recv_once() { + return result; } - Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected), - Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty), + std::thread::yield_now(); } } + fn try_recv_once(&mut self) -> Poll> { + // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows. + let value = if let Some(value) = ready!(unsafe { self.state.buffer.pop() }) { + value + } else if self.state.senders.load(Ordering::Acquire) == 0 { + // The final sender can enqueue between the first empty observation and decrementing + // the sender count, so check the queue again before reporting disconnection. + // SAFETY: The exclusive receiver borrow still guarantees a single consumer. + let Some(value) = ready!(unsafe { self.state.buffer.pop() }) else { + return Poll::Ready(Err(TryRecvError::Disconnected)); + }; + value + } else { + return Poll::Ready(Err(TryRecvError::Empty)); + }; + self.state.send_waiters.release_if_nonempty(1); + Poll::Ready(Ok(value)) + } + /// Waits for and receives the next value, freeing one buffer slot. /// /// If no value is queued, this method waits until a sender adds one or the last sender is @@ -303,16 +305,20 @@ impl BoundedReceiver { } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => { + match self.try_recv_once() { + Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)), + Poll::Ready(Err(TryRecvError::Disconnected)) => { + Poll::Ready(Err(RecvError::Disconnected)) + } + Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => { self.state.rx_waker.register(cx.waker()); - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => Poll::Pending, + match self.try_recv_once() { + Poll::Ready(Ok(v)) => Poll::Ready(Ok(v)), + Poll::Ready(Err(TryRecvError::Disconnected)) => { + Poll::Ready(Err(RecvError::Disconnected)) + } + Poll::Pending | Poll::Ready(Err(TryRecvError::Empty)) => Poll::Pending, } } } diff --git a/asyncband/src/mpsc/bounded/ring.rs b/asyncband/src/mpsc/bounded/ring.rs new file mode 100644 index 00000000..4eb95085 --- /dev/null +++ b/asyncband/src/mpsc/bounded/ring.rs @@ -0,0 +1,372 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cell::UnsafeCell; +use std::hint::spin_loop; +use std::mem::MaybeUninit; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::sync::atomic::fence; +use std::task::Poll; + +use crate::internal::cache_padded::CachePadded; +use crate::mpsc::TrySendError; + +pub struct Ring { + slots: Box<[Slot]>, + head: CachePadded, + tail: CachePadded, + capacity: usize, + one_lap: usize, + mark_bit: usize, +} + +struct Slot { + stamp: AtomicUsize, + value: UnsafeCell>, +} + +// SAFETY: A successful tail CAS gives one producer exclusive access to a slot. That producer +// initializes the value before publishing the next stamp with Release ordering. The single +// consumer reads only after acquiring that stamp and publishes the following lap before reuse. +unsafe impl Sync for Slot {} + +// The ownership transition finishes before user code can unwind, and no stored-value reference is +// exposed. +impl std::panic::UnwindSafe for Slot {} +impl std::panic::RefUnwindSafe for Slot {} + +impl Ring { + pub fn new(capacity: usize) -> Self { + assert!(capacity <= usize::MAX / 4, "mpsc capacity is too large"); + let mark_bit = (capacity + 1).next_power_of_two(); + let one_lap = mark_bit * 2; + let slots = (0..capacity) + .map(|index| Slot { + stamp: AtomicUsize::new(index), + value: UnsafeCell::new(MaybeUninit::uninit()), + }) + .collect(); + Self { + slots, + head: CachePadded::new(AtomicUsize::new(0)), + tail: CachePadded::new(AtomicUsize::new(0)), + capacity, + one_lap, + mark_bit, + } + } + + pub fn try_push(&self, value: T) -> Result<(), TrySendError> { + let mut tail = self.tail.load(Ordering::Relaxed); + let mut backoff = 0; + loop { + if tail & self.mark_bit != 0 { + return Err(TrySendError::Disconnected(value)); + } + + let index = tail & (self.mark_bit - 1); + let slot = &self.slots[index]; + let stamp = slot.stamp.load(Ordering::Acquire); + if stamp == tail { + let next_tail = self.advance(tail); + match self.tail.compare_exchange_weak( + tail, + next_tail, + Ordering::SeqCst, + Ordering::Relaxed, + ) { + Ok(_) => { + // SAFETY: The successful CAS reserved this slot exclusively, and its + // matching stamp proves the consumer completed its previous lap. + unsafe { (*slot.value.get()).write(value) }; + slot.stamp.store(tail.wrapping_add(1), Ordering::Release); + return Ok(()); + } + Err(actual) => tail = actual, + } + } else if stamp.wrapping_add(self.one_lap) == tail.wrapping_add(1) { + fence(Ordering::SeqCst); + if self.head.load(Ordering::Relaxed).wrapping_add(self.one_lap) == tail { + return Err(TrySendError::Full(value)); + } + tail = self.tail.load(Ordering::Relaxed); + } else { + tail = self.tail.load(Ordering::Relaxed); + } + Self::spin(&mut backoff); + } + } + + /// Pending means the head slot is reserved but not published. It is distinct from an empty + /// queue: later producers may already have completed their sends. + /// + /// # Safety + /// + /// The caller must serialize all calls to `pop` and `disconnect_receiver` for this queue. + pub unsafe fn pop(&self) -> Poll> { + let mut head = self.head.load(Ordering::Relaxed); + let mut backoff = 0; + loop { + let index = head & (self.mark_bit - 1); + let slot = &self.slots[index]; + let stamp = slot.stamp.load(Ordering::Acquire); + if stamp == head.wrapping_add(1) { + let next_head = self.advance(head); + // SAFETY: Acquiring the matching stamp observes initialization by the producer. + // There is one consumer, so the value is read exactly once. + let value = unsafe { (*slot.value.get()).assume_init_read() }; + slot.stamp + .store(head.wrapping_add(self.one_lap), Ordering::Release); + self.head.store(next_head, Ordering::SeqCst); + return Poll::Ready(Some(value)); + } + + if stamp == head { + fence(Ordering::SeqCst); + if self.tail.load(Ordering::Relaxed) & !self.mark_bit == head { + return Poll::Ready(None); + } + } + if backoff == 8 { + return Poll::Pending; + } + Self::spin(&mut backoff); + head = self.head.load(Ordering::Relaxed); + } + } + + /// Closes the queue and drops all remaining values. + /// + /// # Safety + /// + /// The caller must serialize all calls to `pop` and `disconnect_receiver` for this queue. + pub unsafe fn disconnect_receiver(&self) { + let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; + // SAFETY: The caller guarantees exclusive consumer access. + unsafe { self.discard_until(tail) }; + } + + fn advance(&self, position: usize) -> usize { + let index = position & (self.mark_bit - 1); + if index + 1 < self.capacity { + position + 1 + } else { + let lap = position & !(self.one_lap - 1); + lap.wrapping_add(self.one_lap) + } + } + + // The caller must close the queue and have exclusive consumer access before discarding values. + unsafe fn discard_until(&self, tail: usize) { + let mut head = self.head.load(Ordering::Relaxed); + let mut backoff = 0; + while head != tail { + let index = head & (self.mark_bit - 1); + let slot = &self.slots[index]; + if slot.stamp.load(Ordering::Acquire) == head.wrapping_add(1) { + let next_head = self.advance(head); + // Move the head before dropping the value so unwinding cannot drop it twice. + slot.stamp + .store(head.wrapping_add(self.one_lap), Ordering::Release); + self.head.store(next_head, Ordering::SeqCst); + // SAFETY: The acquired matching stamp proves the slot contains an initialized + // value, and advancing the single-consumer head claims it exactly once. + unsafe { (*slot.value.get()).assume_init_drop() }; + head = next_head; + backoff = 0; + } else { + Self::spin(&mut backoff); + } + } + } + + fn spin(step: &mut u32) { + for _ in 0..(*step).min(6).pow(2) { + spin_loop(); + } + *step = (*step).saturating_add(1); + } +} + +impl Drop for Ring { + fn drop(&mut self) { + let tail = self.tail.fetch_or(self.mark_bit, Ordering::SeqCst) & !self.mark_bit; + // SAFETY: The queue is closed and its exclusive borrow rules out concurrent access. + unsafe { self.discard_until(tail) }; + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + use std::sync::atomic::AtomicUsize; + use std::sync::atomic::Ordering; + use std::task::Poll; + use std::thread; + + use super::Ring; + use super::TrySendError; + + #[test] + fn bounded_queue_preserves_capacity_and_fifo_order() { + let queue = Ring::new(3); + for value in 0..3 { + assert!(queue.try_push(value).is_ok()); + } + assert!(matches!(queue.try_push(3), Err(TrySendError::Full(3)))); + for value in 0..3 { + // SAFETY: This thread is the only consumer. + assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); + } + // SAFETY: This thread is the only consumer. + assert_eq!(unsafe { queue.pop() }, Poll::Ready(None)); + + for value in 3..12 { + assert!(queue.try_push(value).is_ok()); + // SAFETY: This thread is the only consumer. + assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); + } + } + + #[test] + fn bounded_queue_does_not_report_empty_behind_an_unpublished_head() { + let queue = Ring::new(2); + + // Pause a synthetic producer after reserving and initializing slot 0, before publishing + // its stamp. Another producer can finish sending into slot 1 in the meantime. + queue.tail.store(1, Ordering::SeqCst); + let slot = &queue.slots[0]; + // SAFETY: advancing the tail reserved this initially empty slot for the synthetic producer. + unsafe { (*slot.value.get()).write(1) }; + let later_send = queue.try_push(2); + // SAFETY: This thread is the only consumer, even while a producer is unpublished. + let receive = unsafe { queue.pop() }; + + // Finish publication before asserting so even a failed assertion can safely drop the queue. + slot.stamp.store(1, Ordering::Release); + assert!(later_send.is_ok()); + assert_eq!(receive, Poll::Pending); + // SAFETY: This thread is the only consumer. + unsafe { + assert_eq!(queue.pop(), Poll::Ready(Some(1))); + assert_eq!(queue.pop(), Poll::Ready(Some(2))); + assert_eq!(queue.pop(), Poll::Ready(None)); + } + } + + #[test] + fn bounded_queue_coordinates_multiple_producers() { + let queue = Arc::new(Ring::new(4)); + let producers: Vec<_> = (0..2) + .map(|producer| { + let queue = queue.clone(); + thread::spawn(move || { + for offset in 0..32 { + let mut value = producer * 32 + offset; + loop { + match queue.try_push(value) { + Ok(()) => break, + Err(TrySendError::Full(returned)) => { + value = returned; + thread::yield_now(); + } + Err(TrySendError::Disconnected(_)) => panic!("queue disconnected"), + } + } + } + }) + }) + .collect(); + + let mut values = Vec::new(); + while values.len() < 64 { + // SAFETY: Worker threads only push; this thread is the only consumer. + if let Poll::Ready(Some(value)) = unsafe { queue.pop() } { + values.push(value); + } else { + thread::yield_now(); + } + } + for producer in producers { + producer.join().unwrap(); + } + values.sort_unstable(); + assert_eq!(values, (0..64).collect::>()); + } + + #[test] + fn bounded_queue_discards_wrapped_values_once_after_receiver_disconnect() { + // This has no owning fields, so a buggy second drop remains observable as count == 2 + // instead of invalidating the tracker first. + struct DropSpy<'a>(&'a AtomicUsize); + + impl<'a> Drop for DropSpy<'a> { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::Relaxed); + } + } + + // Declare this before `queue` so the counters outlive values held by the queue. + let drops = [ + AtomicUsize::new(0), + AtomicUsize::new(0), + AtomicUsize::new(0), + AtomicUsize::new(0), + ]; + let queue = Ring::new(3); + + // Positions: 0, 1, 2 (then tail wraps to 8). + for counter in &drops[..3] { + assert!(queue.try_push(DropSpy(counter)).is_ok()); + } + + // Free slot 0, then reuse it on the next lap at position 8. + // SAFETY: This thread is the only consumer. + let popped = unsafe { queue.pop() }; + assert!(matches!(popped, Poll::Ready(Some(_)))); + drop(popped); + assert_eq!(drops[0].load(Ordering::Relaxed), 1); + assert!(queue.try_push(DropSpy(&drops[3])).is_ok()); + + // The pending range is positions 1 -> 2 -> 8 -> 9, not a contiguous integer range. + assert_eq!(queue.head.load(Ordering::Relaxed), 1); + assert_eq!(queue.tail.load(Ordering::Relaxed), queue.one_lap + 1); + + // SAFETY: This thread is the only consumer and no pop is in progress. + unsafe { queue.disconnect_receiver() }; + + // `discard_until` must dispose every value exactly once, including position 8. + for (value, counter) in drops.iter().enumerate() { + assert_eq!( + counter.load(Ordering::Relaxed), + 1, + "value {value} was dropped an unexpected number of times" + ); + } + + // Queue Drop calls discard_until again; it must see head == tail and not redrop. + drop(queue); + for (value, counter) in drops.iter().enumerate() { + assert_eq!( + counter.load(Ordering::Relaxed), + 1, + "value {value} was dropped more than once" + ); + } + } +} diff --git a/asyncband/src/mpsc/unbounded.rs b/asyncband/src/mpsc/unbounded.rs index cb5ca5e0..92be6990 100644 --- a/asyncband/src/mpsc/unbounded.rs +++ b/asyncband/src/mpsc/unbounded.rs @@ -18,60 +18,75 @@ //! An unbounded multi-producer, single-consumer queue for sending values between asynchronous //! tasks. +use std::collections::VecDeque; use std::fmt; use std::future::poll_fn; +use std::mem; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; use std::task::Poll; +use std::task::Waker; use super::RecvError; use super::SendError; use super::TryRecvError; -use crate::internal::atomic_waker::AtomicWaker; +use crate::internal::mutex::Mutex; /// Creates an unbounded mpsc channel whose send operation never waits for capacity. /// /// While the receiver is alive, each send appends its value immediately. Pending messages can /// therefore grow with producer demand and are limited only by successful memory allocation. Use a /// bounded channel or external admission control when producers may outpace the receiver. +/// +/// After all messages have been received, large backing allocations are released; small buffers +/// may be retained for reuse. Partially consumed batches can retain their original allocation. pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let state = Arc::new(UnboundedState { senders: AtomicUsize::new(1), - rx_waker: AtomicWaker::new(), + inbox: Mutex::new(Inbox { + messages: VecDeque::new(), + receiver_alive: true, + rx_waker: None, + }), }); - let (sender, receiver) = std::sync::mpsc::channel(); let sender = UnboundedSender { state: state.clone(), - sender: Some(sender), }; let receiver = UnboundedReceiver { - state: state.clone(), - receiver, + state, + batch: Mutex::new(VecDeque::new()), }; (sender, receiver) } -struct UnboundedState { +struct UnboundedState { + // Endpoint cloning and ordinary drops do not contend with message traffic. senders: AtomicUsize, - rx_waker: AtomicWaker, + inbox: Mutex>, +} + +// Queue contents, receiver liveness, and its wake registration share one lock. Registering a +// wait and checking its condition cannot race with sending or receiver disconnection. +struct Inbox { + messages: VecDeque, + receiver_alive: bool, + rx_waker: Option, } /// The sending endpoint of an unbounded mpsc channel. /// /// Instances are created by the [`unbounded`] function. pub struct UnboundedSender { - state: Arc, - sender: Option>, + state: Arc>, } impl Clone for UnboundedSender { fn clone(&self) -> Self { self.state.senders.fetch_add(1, Ordering::Release); - UnboundedSender { + Self { state: self.state.clone(), - sender: self.sender.clone(), } } } @@ -84,16 +99,10 @@ impl fmt::Debug for UnboundedSender { impl Drop for UnboundedSender { fn drop(&mut self) { - // Dropping the final underlying sender disconnects the channel. - drop(self.sender.take()); - - match self.state.senders.fetch_sub(1, Ordering::AcqRel) { - 1 => { - // Wake the receiver so it can observe the channel's disconnected state. - self.state.rx_waker.wake(); - } - _ => { - // there are still other senders left, do nothing + if self.state.senders.fetch_sub(1, Ordering::AcqRel) == 1 { + let waker = self.state.inbox.lock().rx_waker.take(); + if let Some(waker) = waker { + waker.wake(); } } } @@ -105,12 +114,17 @@ impl UnboundedSender { /// This operation is synchronous because the channel has no capacity limit. If the receiver has /// been dropped, the returned error contains `value`. pub fn send(&self, value: T) -> Result<(), SendError> { - // INVARIANT: A shared borrow of the endpoint cannot overlap its destructor. - let sender = self.sender.as_ref().unwrap(); - sender.send(value).map_err(|err| SendError::new(err.0))?; - - self.state.rx_waker.wake(); - + let waker = { + let mut state = self.state.inbox.lock(); + if !state.receiver_alive { + return Err(SendError::new(value)); + } + state.messages.push_back(value); + state.rx_waker.take() + }; + if let Some(waker) = waker { + waker.wake(); + } Ok(()) } } @@ -119,20 +133,30 @@ impl UnboundedSender { /// /// Instances are created by the [`unbounded`] function. pub struct UnboundedReceiver { - state: Arc, - receiver: std::sync::mpsc::Receiver, + state: Arc>, + // Only accessed through `get_mut`; the mutex preserves Sync for Send-only payloads. + batch: Mutex>, } -/// The only `!Sync` field `receiver` is protected by `&mut self` in `recv` and `try_recv`. -/// That is, `UnboundedReceiver` can only be accessed by one thread at a time. -unsafe impl Sync for UnboundedReceiver {} - impl fmt::Debug for UnboundedReceiver { fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { f.debug_struct("UnboundedReceiver").finish_non_exhaustive() } } +impl Drop for UnboundedReceiver { + fn drop(&mut self) { + let batch = mem::take(self.batch.get_mut()); + let (shared, waker) = { + let mut state = self.state.inbox.lock(); + state.receiver_alive = false; + (mem::take(&mut state.messages), state.rx_waker.take()) + }; + // Destructors may send again. A waker may also own a sender and form an ownership cycle. + drop((batch, shared, waker)); + } +} + impl UnboundedReceiver { /// Attempts to receive the next queued value without waiting. /// @@ -157,11 +181,21 @@ impl UnboundedReceiver { /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); /// ``` pub fn try_recv(&mut self) -> Result { - match self.receiver.try_recv() { - Ok(v) => Ok(v), - Err(std::sync::mpsc::TryRecvError::Disconnected) => Err(TryRecvError::Disconnected), - Err(std::sync::mpsc::TryRecvError::Empty) => Err(TryRecvError::Empty), + let batch = self.batch.get_mut(); + if batch.is_empty() { + let mut state = self.state.inbox.lock(); + if state.messages.is_empty() { + // Holding the inbox lock excludes a final send between the empty observation and + // the sender-count check; disconnection needs no second queue read. + return Err(if self.state.senders.load(Ordering::Acquire) == 0 { + TryRecvError::Disconnected + } else { + TryRecvError::Empty + }); + } + mem::swap(batch, &mut state.messages); } + Ok(pop_batch(batch)) } /// Waits for and receives the next value. @@ -198,18 +232,124 @@ impl UnboundedReceiver { } fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => { - self.state.rx_waker.register(cx.waker()); - - match self.try_recv() { - Ok(v) => Poll::Ready(Ok(v)), - Err(TryRecvError::Disconnected) => Poll::Ready(Err(RecvError::Disconnected)), - Err(TryRecvError::Empty) => Poll::Pending, - } + let batch = self.batch.get_mut(); + if !batch.is_empty() { + return Poll::Ready(Ok(pop_batch(batch))); + } + // Waker clone/drop callbacks can reenter this channel. Clone outside the lock, then + // recheck the condition before registering; keep replaced wakers outside the lock too. + let mut new_waker = None; + loop { + let mut state = self.state.inbox.lock(); + if !state.messages.is_empty() { + mem::swap(batch, &mut state.messages); + drop(state); + return Poll::Ready(Ok(pop_batch(batch))); } + if self.state.senders.load(Ordering::Acquire) == 0 { + return Poll::Ready(Err(RecvError::Disconnected)); + } + if state + .rx_waker + .as_ref() + .is_some_and(|waker| waker.will_wake(cx.waker())) + { + return Poll::Pending; + } + if let Some(waker) = new_waker.take() { + let old_waker = state.rx_waker.replace(waker); + drop(state); + drop(old_waker); + return Poll::Pending; + } + drop(state); + new_waker = Some(cx.waker().clone()); + } + } +} + +// No operation relies on a pinned location for the receiver batch or its values. +impl Unpin for UnboundedReceiver {} + +// Retain small buffers for reuse, measuring inline storage rather than separately boxed payloads. +const BATCH_CACHE_BYTES: usize = 64 * 1024; + +fn pop_batch(batch: &mut VecDeque) -> T { + if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::()) > BATCH_CACHE_BYTES + { + // Retire the allocation on the last value, outside the inbox lock. Keep this as a tail + // expression to avoid intermediate storage for large inline values. + mem::take(batch).pop_front() + } else { + batch.pop_front() + } + .expect("receiver batch must not be empty") +} + +#[cfg(test)] +mod tests { + use super::unbounded; + use crate::mpsc::TryRecvError; + + #[test] + fn batches_preserve_order_across_refills() { + let (tx, mut rx) = unbounded(); + tx.send(1).unwrap(); + tx.send(2).unwrap(); + assert_eq!(rx.try_recv(), Ok(1)); + tx.send(3).unwrap(); + assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv(), Ok(3)); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } + + #[test] + fn releases_large_batches_after_the_last_value() { + let (tx, mut rx) = unbounded(); + for value in 0..128u8 { + tx.send([value; 1024]).unwrap(); + } + for value in 0..64u8 { + assert_eq!(rx.try_recv(), Ok([value; 1024])); + } + // A partially consumed batch survives while producers start filling the next batch. + tx.send([128; 1024]).unwrap(); + for value in 64..128u8 { + assert_eq!(rx.try_recv(), Ok([value; 1024])); + } + // Reclaim on the final successful receive, without requiring an extra empty poll. + assert_eq!(rx.batch.get_mut().capacity(), 0); + assert_eq!(rx.try_recv(), Ok([128; 1024])); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + } + + #[test] + fn reuses_small_batches_on_refill() { + let (tx, mut rx) = unbounded(); + for value in 0..32usize { + tx.send(value).unwrap(); + } + assert_eq!(rx.try_recv(), Ok(0)); + let capacity = rx.batch.get_mut().capacity(); + for value in 1..32 { + assert_eq!(rx.try_recv(), Ok(value)); + } + assert_eq!(rx.batch.get_mut().capacity(), capacity); + + tx.send(32).unwrap(); + assert_eq!(rx.try_recv(), Ok(32)); + assert_eq!(rx.state.inbox.lock().messages.capacity(), capacity); + } + + #[test] + fn drains_zero_sized_values() { + let (tx, mut rx) = unbounded(); + for _ in 0..32 { + tx.send(()).unwrap(); + } + for _ in 0..32 { + assert_eq!(rx.try_recv(), Ok(())); } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } } diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs index 7e78c217..1fd4dbb0 100644 --- a/benchmarks/ecosystem/mpsc/adapters.rs +++ b/benchmarks/ecosystem/mpsc/adapters.rs @@ -15,6 +15,7 @@ // specific language governing permissions and limitations // under the License. +use std::fmt::Debug; use std::task::Context; use crate::support::poll_ready; @@ -37,15 +38,15 @@ pub trait BoundedMpsc: Send + Sync + 'static { fn recv_blocking(receiver: &mut Self::Receiver) -> usize; } -pub trait UnboundedMpsc: Send + Sync + 'static { +pub trait UnboundedMpsc: Send + Sync + 'static { type Sender: Clone + Send + 'static; type Receiver: Send + 'static; fn channel() -> (Self::Sender, Self::Receiver); - fn send(sender: &Self::Sender, value: usize); - fn try_recv(receiver: &mut Self::Receiver) -> usize; - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize; - fn recv_blocking(receiver: &mut Self::Receiver) -> usize; + fn send(sender: &Self::Sender, value: T); + fn try_recv(receiver: &mut Self::Receiver) -> T; + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T; + fn recv_blocking(receiver: &mut Self::Receiver) -> T; } impl BoundedMpsc for Asyncband { @@ -180,102 +181,102 @@ impl BoundedMpsc for Flume { } } -impl UnboundedMpsc for Asyncband { - type Receiver = asyncband::mpsc::UnboundedReceiver; - type Sender = asyncband::mpsc::UnboundedSender; +impl UnboundedMpsc for Asyncband { + type Receiver = asyncband::mpsc::UnboundedReceiver; + type Sender = asyncband::mpsc::UnboundedSender; fn channel() -> (Self::Sender, Self::Receiver) { asyncband::mpsc::unbounded() } - fn send(sender: &Self::Sender, value: usize) { + fn send(sender: &Self::Sender, value: T) { sender.send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl UnboundedMpsc for Tokio { - type Receiver = tokio::sync::mpsc::UnboundedReceiver; - type Sender = tokio::sync::mpsc::UnboundedSender; +impl UnboundedMpsc for Tokio { + type Receiver = tokio::sync::mpsc::UnboundedReceiver; + type Sender = tokio::sync::mpsc::UnboundedSender; fn channel() -> (Self::Sender, Self::Receiver) { tokio::sync::mpsc::unbounded_channel() } - fn send(sender: &Self::Sender, value: usize) { + fn send(sender: &Self::Sender, value: T) { sender.send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl UnboundedMpsc for AsyncChannel { - type Receiver = async_channel::Receiver; - type Sender = async_channel::Sender; +impl UnboundedMpsc for AsyncChannel { + type Receiver = async_channel::Receiver; + type Sender = async_channel::Sender; fn channel() -> (Self::Sender, Self::Receiver) { async_channel::unbounded() } - fn send(sender: &Self::Sender, value: usize) { + fn send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl UnboundedMpsc for Flume { - type Receiver = flume::Receiver; - type Sender = flume::Sender; +impl UnboundedMpsc for Flume { + type Receiver = flume::Receiver; + type Sender = flume::Sender; fn channel() -> (Self::Sender, Self::Receiver) { flume::unbounded() } - fn send(sender: &Self::Sender, value: usize) { + fn send(sender: &Self::Sender, value: T) { sender.send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv_async(), context).unwrap() } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv_async()).unwrap() } } diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index bda71079..80495e98 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -29,6 +29,7 @@ use super::support::BOUNDED_CAPACITY; use super::support::Bounded; use super::support::ConcurrentBatch; use super::support::PRODUCER_COUNTS; +use super::support::RepeatedBatch; use crate::support::bench_context; #[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] @@ -64,3 +65,22 @@ fn concurrent(bencher: Bencher, producer_count: usize) { .with_inputs(|| ConcurrentBatch::>::new(producer_count)) .bench_local_refs(|batch| batch.run()); } + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = PRODUCER_COUNTS, + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn sustained(bencher: Bencher, producer_count: usize) { + let mut batch = RepeatedBatch::>::new(producer_count); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +fn clone_drop_sender(bencher: Bencher) { + let (sender, _receiver) = C::channel(BOUNDED_CAPACITY); + bencher.bench_local(|| drop(black_box(sender.clone()))); +} diff --git a/benchmarks/ecosystem/mpsc/support.rs b/benchmarks/ecosystem/mpsc/support.rs index aea28858..11edb58f 100644 --- a/benchmarks/ecosystem/mpsc/support.rs +++ b/benchmarks/ecosystem/mpsc/support.rs @@ -18,6 +18,8 @@ use std::marker::PhantomData; use std::sync::Arc; use std::sync::Barrier; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; use std::thread; use std::thread::JoinHandle; @@ -133,3 +135,66 @@ impl Drop for ConcurrentBatch { } } } + +// Reuse worker threads and channel storage so steady-state samples exclude thread creation. +pub struct RepeatedBatch { + receiver: C::Receiver, + start: Arc, + stop: Arc, + workers: Vec>, +} + +impl RepeatedBatch { + pub fn new(producer_count: usize) -> Self { + assert_eq!(BATCH_MESSAGES % producer_count, 0); + let (sender, receiver) = C::channel(); + let start = Arc::new(Barrier::new(producer_count + 1)); + let stop = Arc::new(AtomicBool::new(false)); + let messages_per_producer = BATCH_MESSAGES / producer_count; + let workers = (0..producer_count) + .map(|producer| { + let sender = sender.clone(); + let start = start.clone(); + let stop = stop.clone(); + thread::spawn(move || { + loop { + start.wait(); + if stop.load(Ordering::Acquire) { + break; + } + let first = producer * messages_per_producer; + for offset in 0..messages_per_producer { + C::send(&sender, black_box(first + offset)); + } + } + }) + }) + .collect(); + drop(sender); + Self { + receiver, + start, + stop, + workers, + } + } + + pub fn run(&mut self) -> usize { + self.start.wait(); + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(C::recv(&mut self.receiver)); + } + black_box(checksum) + } +} + +impl Drop for RepeatedBatch { + fn drop(&mut self) { + self.stop.store(true, Ordering::Release); + self.start.wait(); + for worker in self.workers.drain(..) { + worker.join().expect("benchmark producer panicked"); + } + } +} diff --git a/benchmarks/ecosystem/mpsc/unbounded.rs b/benchmarks/ecosystem/mpsc/unbounded.rs index 6c8ac549..b08ed700 100644 --- a/benchmarks/ecosystem/mpsc/unbounded.rs +++ b/benchmarks/ecosystem/mpsc/unbounded.rs @@ -27,6 +27,7 @@ use super::adapters::UnboundedMpsc; use super::support::BATCH_MESSAGES; use super::support::ConcurrentBatch; use super::support::PRODUCER_COUNTS; +use super::support::RepeatedBatch; use super::support::Unbounded; use crate::support::bench_context; @@ -63,3 +64,92 @@ fn concurrent(bencher: Bencher, producer_count: usize) { .with_inputs(|| ConcurrentBatch::>::new(producer_count)) .bench_local_refs(|batch| batch.run()); } + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = [32, 1024, 65_536], + sample_count = 20, + sample_size = 1, +)] +fn burst_drain(bencher: Bencher, messages: usize) { + repeated_bursts::(bencher, messages, 0, || usize::MAX); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + consts = [64, 1024], + args = [32, 1024, 65_536], + sample_count = 20, + sample_size = 1, +)] +fn burst_drain_inline, const SIZE: usize>( + bencher: Bencher, + messages: usize, +) { + repeated_bursts::(bencher, messages, 0, || [1; SIZE]); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = [32, 1024, 65_536], + sample_count = 20, + sample_size = 1, +)] +fn burst_drain_boxed>>(bencher: Bencher, messages: usize) { + // Include payload allocation and destruction to compare the complete boxed-message lifecycle. + repeated_bursts::(bencher, messages, 0, || Box::new([1; 1024])); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = [1024, 65_536], + sample_count = 20, + sample_size = 1, +)] +fn burst_with_backlog(bencher: Bencher, messages: usize) { + repeated_bursts::(bencher, messages, messages / 2, || usize::MAX); +} + +fn repeated_bursts, T, F: Fn() -> T>( + bencher: Bencher, + messages: usize, + backlog: usize, + make_value: F, +) { + // Keep one channel alive across samples so allocation reuse and reclamation are measured. + let (sender, mut receiver) = C::channel(); + for _ in 0..backlog { + C::send(&sender, make_value()); + } + let mut run = || { + for _ in 0..messages { + C::send(&sender, black_box(make_value())); + } + for _ in 0..messages { + black_box(C::try_recv(&mut receiver)); + } + }; + // Measure recurring bursts after the initial allocation, including a deliberately retained + // backlog where requested. Do not require an extra empty receive to trigger reclamation. + run(); + bencher.counter(ItemsCount::new(messages)).bench_local(run); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + args = PRODUCER_COUNTS, + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn sustained(bencher: Bencher, producer_count: usize) { + let mut batch = RepeatedBatch::>::new(producer_count); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +fn clone_drop_sender(bencher: Bencher) { + let (sender, _receiver) = C::channel(); + bencher.bench_local(|| drop(black_box(sender.clone()))); +} diff --git a/tests-integration/tests/mpsc_test.rs b/tests-integration/tests/mpsc_test.rs index 50d5a108..7b84b933 100644 --- a/tests-integration/tests/mpsc_test.rs +++ b/tests-integration/tests/mpsc_test.rs @@ -15,7 +15,18 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Context; use std::task::Poll; +use std::task::RawWaker; +use std::task::RawWakerVTable; +use std::task::Wake; +use std::task::Waker; +use std::thread; +use std::time::Duration; use asyncband::mpsc; use asyncband::mpsc::RecvError; @@ -32,6 +43,177 @@ fn expect_ready(poll: Poll) -> T { } } +struct HoldSender { + _sender: S, +} + +// This waker must own the sender so its final drop can break the tested reference cycle. +#[allow(clippy::manual_noop_waker)] +impl Wake for HoldSender { + fn wake(self: Arc) {} +} + +#[test] +fn bounded_receiver_drop_releases_registered_waker() { + let (tx, mut rx) = mpsc::bounded::<()>(1); + let holder = Arc::new(HoldSender { _sender: tx }); + let retained = Arc::downgrade(&holder); + let waker = Waker::from(holder); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + drop(rx); + assert!(retained.upgrade().is_none()); +} + +#[test] +fn unbounded_receiver_drop_releases_registered_waker() { + let (tx, mut rx) = mpsc::unbounded::<()>(); + let holder = Arc::new(HoldSender { _sender: tx }); + let retained = Arc::downgrade(&holder); + let waker = Waker::from(holder); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + drop(rx); + assert!(retained.upgrade().is_none()); +} + +fn assert_completes_without_deadlock(test: impl FnOnce() + Send + 'static) { + let (finished_tx, finished_rx) = std::sync::mpsc::channel(); + let worker = thread::spawn(move || { + test(); + finished_tx.send(()).unwrap(); + }); + finished_rx + .recv_timeout(Duration::from_secs(10)) + .expect("waker callback did not finish"); + worker.join().unwrap(); +} + +#[test] +fn unbounded_wake_callback_can_send() { + struct SendOnWake(mpsc::UnboundedSender); + + impl Wake for SendOnWake { + fn wake(self: Arc) { + self.0.send(2).unwrap(); + } + } + + assert_completes_without_deadlock(|| { + let (tx, mut rx) = mpsc::unbounded(); + let waker = Waker::from(Arc::new(SendOnWake(tx.clone()))); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + tx.send(1).unwrap(); + assert_eq!(rx.try_recv(), Ok(1)); + assert_eq!(rx.try_recv(), Ok(2)); + }); +} + +#[test] +fn unbounded_replaced_and_disconnected_wakers_can_send() { + struct SendOnDrop { + sender: mpsc::UnboundedSender, + disconnected: bool, + drops: Arc, + } + + // The final waker drop must run a callback, even though waking itself does nothing. + #[allow(clippy::manual_noop_waker)] + impl Wake for SendOnDrop { + fn wake(self: Arc) {} + } + + impl Drop for SendOnDrop { + fn drop(&mut self) { + assert_eq!(self.sender.send(7).is_err(), self.disconnected); + self.drops.fetch_add(1, Ordering::Relaxed); + } + } + + assert_completes_without_deadlock(|| { + for disconnected in [false, true] { + let (tx, mut rx) = mpsc::unbounded(); + let drops = Arc::new(AtomicUsize::new(0)); + let waker = Waker::from(Arc::new(SendOnDrop { + sender: tx, + disconnected, + drops: drops.clone(), + })); + assert!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)) + .is_pending() + ); + drop(waker); + if disconnected { + drop(rx); + } else { + // Replacing the waker can enqueue a message during this poll. Either immediate + // completion or a notified Pending is valid, but the message must not be lost. + let poll = poll_once(Box::pin(rx.recv()).as_mut()); + if poll.is_pending() { + assert_eq!(rx.try_recv(), Ok(7)); + } else { + assert_eq!(poll, Poll::Ready(Ok(7))); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); + } + assert_eq!(drops.load(Ordering::Relaxed), 1); + } + }); +} + +#[test] +fn unbounded_waker_clone_rechecks_messages_sent_during_registration() { + unsafe fn clone_sender(data: *const ()) -> RawWaker { + let sender = data.cast::>(); + // SAFETY: Each raw waker owns an Arc to this sender; cloning borrows the live sender and + // then adds the strong reference owned by the returned waker. + unsafe { + (*sender).send(7).unwrap(); + Arc::increment_strong_count(sender); + } + RawWaker::new(data, &VTABLE) + } + + unsafe fn drop_sender(data: *const ()) { + // SAFETY: Consumes exactly the Arc reference owned by this raw waker. + drop(unsafe { Arc::from_raw(data.cast::>()) }); + } + + static VTABLE: RawWakerVTable = + RawWakerVTable::new(clone_sender, drop_sender, |_| {}, drop_sender); + + assert_completes_without_deadlock(|| { + let (tx, mut rx) = mpsc::unbounded::(); + let data = Arc::into_raw(Arc::new(tx)).cast(); + // SAFETY: The vtable manages one Arc reference per waker and the sender is Send + Sync. + let waker = unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }; + assert_eq!( + Box::pin(rx.recv()) + .as_mut() + .poll(&mut Context::from_waker(&waker)), + Poll::Ready(Ok(7)) + ); + }); +} + #[test] fn unbounded_collects_from_multiple_producers() { let (tx, mut rx) = mpsc::unbounded(); @@ -176,6 +358,44 @@ fn unbounded_try_recv_preserves_order_and_reports_state() { assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); } +#[test] +fn cancelled_receive_does_not_consume_a_later_message() { + let (unbounded_tx, mut unbounded_rx) = mpsc::unbounded(); + { + let mut receive = Box::pin(unbounded_rx.recv()); + assert!(poll_once(receive.as_mut()).is_pending()); + } + unbounded_tx.send(1).unwrap(); + assert_eq!(unbounded_rx.try_recv(), Ok(1)); + + let (bounded_tx, mut bounded_rx) = mpsc::bounded(1); + { + let mut receive = Box::pin(bounded_rx.recv()); + assert!(poll_once(receive.as_mut()).is_pending()); + } + bounded_tx.try_send(2).unwrap(); + assert_eq!(bounded_rx.try_recv(), Ok(2)); +} + +#[test] +fn buffered_messages_are_drained_before_disconnection() { + let (unbounded_tx, mut unbounded_rx) = mpsc::unbounded(); + unbounded_tx.send(1).unwrap(); + unbounded_tx.send(2).unwrap(); + drop(unbounded_tx); + assert_eq!(unbounded_rx.try_recv(), Ok(1)); + assert_eq!(unbounded_rx.try_recv(), Ok(2)); + assert_eq!(unbounded_rx.try_recv(), Err(TryRecvError::Disconnected)); + + let (bounded_tx, mut bounded_rx) = mpsc::bounded(2); + bounded_tx.try_send(3).unwrap(); + bounded_tx.try_send(4).unwrap(); + drop(bounded_tx); + assert_eq!(bounded_rx.try_recv(), Ok(3)); + assert_eq!(bounded_rx.try_recv(), Ok(4)); + assert_eq!(bounded_rx.try_recv(), Err(TryRecvError::Disconnected)); +} + #[tokio::test] async fn send_recv_bounded() { let (tx, mut rx) = mpsc::bounded(1); @@ -223,6 +443,51 @@ fn bounded_try_send_respects_capacity_and_order() { } } +#[test] +fn bounded_try_recv_does_not_report_empty_after_completed_sends() { + const PRODUCERS: usize = 4; + const MESSAGES_PER_PRODUCER: usize = 16_384; + let (tx, mut rx) = mpsc::bounded(64); + let completed = AtomicUsize::new(0); + let mut premature_empty = 0; + + thread::scope(|scope| { + for producer in 0..PRODUCERS { + let tx = tx.clone(); + let completed = &completed; + scope.spawn(move || { + for sequence in 0..MESSAGES_PER_PRODUCER { + loop { + match tx.try_send((producer, sequence)) { + Ok(()) => break, + Err(TrySendError::Full(_)) => thread::yield_now(), + Err(TrySendError::Disconnected(_)) => panic!("receiver is still alive"), + } + } + completed.fetch_add(1, Ordering::Release); + } + }); + } + + let mut received = 0; + while received < PRODUCERS * MESSAGES_PER_PRODUCER { + // Once more sends have completed than messages received, Empty cannot be correct. + let has_completed_send = completed.load(Ordering::Acquire) > received; + match rx.try_recv() { + Ok(_) => received += 1, + Err(TryRecvError::Empty) => { + premature_empty += usize::from(has_completed_send); + thread::yield_now(); + } + Err(TryRecvError::Disconnected) => panic!("original sender is still alive"), + } + } + }); + + // Drain and join before asserting so a failure cannot strand a producer on a full channel. + assert_eq!(premature_empty, 0); +} + #[tokio::test] async fn try_send_after_disconnection_bounded() { let (tx, rx) = mpsc::bounded(1); diff --git a/tests-integration/tests/traits_test.rs b/tests-integration/tests/traits_test.rs index bf10a0e7..dfbefb8c 100644 --- a/tests-integration/tests/traits_test.rs +++ b/tests-integration/tests/traits_test.rs @@ -16,6 +16,9 @@ // under the License. use std::cell::Cell; +use std::marker::PhantomPinned; +use std::panic::RefUnwindSafe; +use std::panic::UnwindSafe; use asyncband::barrier::Barrier; use asyncband::broadcast; @@ -184,6 +187,40 @@ fn public_types_are_unpin() { assert_unpin::(); } +#[test] +fn mpsc_endpoints_keep_legacy_traits_regardless_of_payload() { + fn assert_send() {} + fn assert_sync() {} + fn assert_unpin() {} + fn assert_unwind_safe() {} + fn assert_ref_unwind_safe() {} + + macro_rules! assert_endpoint_traits { + ($endpoint:ident, $payload:ty) => { + assert_send::>(); + assert_sync::>(); + assert_unpin::>(); + assert_unwind_safe::>(); + assert_ref_unwind_safe::>(); + }; + } + + macro_rules! assert_payload_traits { + ($endpoint:ident) => { + assert_endpoint_traits!($endpoint, i32); + assert_endpoint_traits!($endpoint, Cell); + assert_endpoint_traits!($endpoint, &'static mut i32); + assert_endpoint_traits!($endpoint, PhantomPinned); + }; + } + + // Four endpoint types × four payloads × five traits = 80 compile-time assertions. + assert_payload_traits!(BoundedSender); + assert_payload_traits!(BoundedReceiver); + assert_payload_traits!(UnboundedSender); + assert_payload_traits!(UnboundedReceiver); +} + #[test] fn unbounded_manual_manager_traits_do_not_depend_on_the_object() { fn assert_copy() {}