From badf60647ba062d9cad0db5899b25a24cf367095 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 13:06:49 +0800 Subject: [PATCH 1/3] feat(mpsc): add bounded permits with sequenced publication --- CHANGELOG.md | 4 + .../mpsc/bounded/{waiters.rs => capacity.rs} | 132 +++++++++- asyncband/src/mpsc/bounded/mod.rs | 184 ++++++++++--- asyncband/src/mpsc/bounded/ring.rs | 247 +++++++----------- asyncband/src/mpsc/bounded/ring_tests.rs | 212 ++++----------- asyncband/src/mpsc/mod.rs | 1 + tests-integration/tests/mpsc_test/main.rs | 3 +- .../tests/mpsc_test/reservation.rs | 210 +++++++++++++++ 8 files changed, 623 insertions(+), 370 deletions(-) rename asyncband/src/mpsc/bounded/{waiters.rs => capacity.rs} (51%) create mode 100644 tests-integration/tests/mpsc_test/reservation.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index 43ffd8e0..c7a38815 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,10 @@ All notable changes to this project will be documented in this file. ## Unreleased +### New features + +* Add bounded MPSC `reserve` and `try_reserve` methods returning a `Permit`, allowing callers to wait for capacity before constructing a message; unused permits release capacity without claiming message order. + ### Bug fixes * Release MPSC receiver wakers when the receiver is dropped, avoiding retained tasks and ownership cycles when a waker holds a sender. diff --git a/asyncband/src/mpsc/bounded/waiters.rs b/asyncband/src/mpsc/bounded/capacity.rs similarity index 51% rename from asyncband/src/mpsc/bounded/waiters.rs rename to asyncband/src/mpsc/bounded/capacity.rs index 3a34e4f5..1e53dfdc 100644 --- a/asyncband/src/mpsc/bounded/waiters.rs +++ b/asyncband/src/mpsc/bounded/capacity.rs @@ -16,40 +16,116 @@ // under the License. use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Waker; +use super::SEQUENCE_STEP; +use crate::internal::cache_padded::CachePadded; use crate::internal::mutex::Mutex; use crate::internal::waitlist::WaitList; use crate::internal::waitlist::WaiterId; use crate::internal::wake_all; use crate::internal::waker_batch::WakerBatch; +use crate::mpsc::TrySendError; + +const CLOSED: usize = 1; // A wake grants a retry, not a capacity permit. Keeping notified nodes until their future // consumes the notification lets cancellation pass an unused retry to the next sender. -pub struct SendWaiters { +pub struct Capacity { + claims: CachePadded, + consumed: CachePadded, + cancelled: CachePadded, + capacity: usize, waiting: AtomicBool, queue: Mutex>>, } -impl SendWaiters { - pub fn new() -> Self { +struct Claims { + next: AtomicUsize, + returned: AtomicUsize, +} + +impl Capacity { + pub fn new(capacity: usize) -> Self { Self { + claims: CachePadded::new(Claims { + next: AtomicUsize::new(0), + returned: AtomicUsize::new(0), + }), + consumed: CachePadded::new(AtomicUsize::new(0)), + cancelled: CachePadded::new(AtomicUsize::new(0)), + capacity: capacity * SEQUENCE_STEP, waiting: AtomicBool::new(false), queue: Mutex::new(WaitList::new()), } } - pub fn waiter(&self) -> SendWaiter<'_> { - SendWaiter { + pub fn try_acquire(&self) -> Result<(), TrySendError<()>> { + let mut claimed = self.claims.next.load(Ordering::Relaxed); + let mut returned = self.claims.returned.load(Ordering::Acquire); + loop { + if claimed & CLOSED != 0 { + return Err(TrySendError::Disconnected(())); + } + if claimed.wrapping_sub(returned) >= self.capacity { + returned = self + .consumed + .load(Ordering::SeqCst) + .wrapping_add(self.cancelled.load(Ordering::SeqCst)); + if claimed.wrapping_sub(returned) >= self.capacity { + // A newer return can overtake our claim snapshot. Refresh the snapshot + // before reporting Full, including a concurrent close. + let current = self.claims.next.load(Ordering::Relaxed); + if current != claimed { + claimed = current; + continue; + } + return Err(TrySendError::Full(())); + } + // Carry the acquired consumption edge with the cached progress. Producers only + // read the consumer's changing cache line when this capacity window runs out. + self.claims.returned.store(returned, Ordering::Release); + } + match self.claims.next.compare_exchange_weak( + claimed, + claimed.wrapping_add(SEQUENCE_STEP), + Ordering::AcqRel, + Ordering::Relaxed, + ) { + Ok(_) => return Ok(()), + Err(actual) => claimed = actual, + } + } + } + + pub fn consume(&self, head: usize) { + // Only the consumer advances this cursor. Producers never modify its cache line. + self.consumed.store(head, Ordering::SeqCst); + self.notify_one(); + } + + pub fn cancel(&self) { + self.cancelled.fetch_add(SEQUENCE_STEP, Ordering::SeqCst); + self.notify_one(); + } + + pub fn close(&self) { + self.claims.next.fetch_or(CLOSED, Ordering::SeqCst); + self.notify_all(); + } + + pub fn waiter(&self) -> ReserveWaiter<'_> { + ReserveWaiter { waiters: self, index: None, } } - pub fn notify_one(&self) { - // The receiver publishes its head with SeqCst before checking this flag. Registration - // publishes the flag with SeqCst before rechecking capacity (which uses a SeqCst fence). + fn notify_one(&self) { + // Releasing capacity precedes this SeqCst flag check. Registration publishes the flag + // with SeqCst before rechecking the SeqCst consumed and cancelled cursors. // Either the receiver sees the registration or the sender sees the released capacity. if !self.waiting.load(Ordering::SeqCst) { return; @@ -67,7 +143,7 @@ impl SendWaiters { } } - pub fn notify_all(&self) { + fn notify_all(&self) { let mut wakers = WakerBatch::new(); { let mut queue = self.queue.lock(); @@ -82,12 +158,12 @@ impl SendWaiters { } } -pub struct SendWaiter<'a> { - waiters: &'a SendWaiters, +pub struct ReserveWaiter<'a> { + waiters: &'a Capacity, index: Option, } -impl SendWaiter<'_> { +impl ReserveWaiter<'_> { // The caller must retry sending after registration, before returning Pending. pub fn register(&mut self, waker: &Waker) { let mut new_waker = None; @@ -145,7 +221,7 @@ impl SendWaiter<'_> { } } -impl Drop for SendWaiter<'_> { +impl Drop for ReserveWaiter<'_> { fn drop(&mut self) { if let Some(index) = self.index.take() { let waker = self.remove(index); @@ -156,3 +232,33 @@ impl Drop for SendWaiter<'_> { } } } + +#[cfg(test)] +mod tests { + use super::Capacity; + use super::Ordering; + use super::SEQUENCE_STEP; + use super::TrySendError; + + #[test] + fn consumption_and_cancellation_restore_capacity_across_counter_overflow() { + let capacity = Capacity::new(3); + // Simulate prior cancellations on the final lap without allocating billions of permits. + let position = usize::MAX - 5; + capacity.claims.next.store(position, Ordering::Relaxed); + capacity.cancelled.store(position, Ordering::Relaxed); + let mut consumed = 0; + for _ in 0..3 { + for _ in 0..3 { + capacity.try_acquire().unwrap(); + } + assert_eq!(capacity.try_acquire(), Err(TrySendError::Full(()))); + consumed += SEQUENCE_STEP; + capacity.consume(consumed); + capacity.cancel(); + capacity.cancel(); + } + capacity.close(); + assert_eq!(capacity.try_acquire(), Err(TrySendError::Disconnected(()))); + } +} diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 455e7095..f91e5354 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -20,6 +20,7 @@ use std::fmt; use std::future::poll_fn; +use std::mem; use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -27,17 +28,20 @@ use std::task::Context; use std::task::Poll; use std::task::ready; +use self::capacity::Capacity; use self::ring::Ring; -use self::waiters::SendWaiters; use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; -// Ring owns capacity, publication, and waiting for the head slot. SendWaiters only schedules -// retries after receiving frees capacity; a notification does not reserve a slot. +// Capacity accounts for permits and queued messages. Ring owns FIFO publication; the receiver +// alone advances its read cursor. A public reservation does not claim a position in the ring. +mod capacity; mod ring; -mod waiters; + +// The low bit marks closure; reservation and publication cursors advance in matching units. +const SEQUENCE_STEP: usize = 2; /// Creates a bounded mpsc channel with room for `buffer` queued messages. /// @@ -53,19 +57,19 @@ pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { let state = Arc::new(Shared { buffer: Ring::new(buffer), senders: AtomicUsize::new(1), - send_waiters: SendWaiters::new(), + capacity: Capacity::new(buffer), }); let sender = BoundedSender { state: state.clone(), }; - let receiver = BoundedReceiver { state }; + let receiver = BoundedReceiver { state, head: 0 }; (sender, receiver) } struct Shared { buffer: Ring, senders: AtomicUsize, - send_waiters: SendWaiters, + capacity: Capacity, } /// The sending endpoint of a bounded mpsc channel. @@ -107,47 +111,89 @@ impl BoundedSender { /// /// Dropping a pending `send` loses its place waiting for capacity and drops `value`; a call /// that has returned `Pending` has not sent the message. Use [`Self::try_send`] when the - /// caller must retain ownership if capacity is unavailable. + /// caller must retain ownership if capacity is unavailable, or [`Self::reserve`] to wait for + /// capacity before constructing the message. pub async fn send(&self, value: T) -> Result<(), SendError> { - let value = match self.try_send(value) { - Ok(()) => return Ok(()), - Err(TrySendError::Disconnected(value)) => return Err(SendError::new(value)), - Err(TrySendError::Full(value)) => value, - }; - let mut waiter = self.state.send_waiters.waiter(); - let mut value = Some(value); + match self.reserve().await { + Ok(permit) => permit.send(value), + Err(_) => Err(SendError::new(value)), + } + } + + /// Reserves capacity for one message before constructing it. + /// + /// A successful reservation returns a [`Permit`]. Dropping the permit releases capacity; + /// [`Permit::send`] publishes a value without waiting for space. Reservations do not establish + /// message order: other producers may send while a permit is held. + /// + /// Returns `SendError(())` if the receiver has been dropped. A permit obtained earlier does + /// not keep the receiver alive; sending with it can still return the unsent value on + /// disconnect. + /// + /// # Cancel safety + /// + /// Dropping a pending reservation removes its wait registration without consuming capacity. + /// Notifications grant a retry, so a new sender may acquire capacity before a woken waiter. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, mut rx) = asyncband::mpsc::bounded(1); + /// let permit = tx.reserve().await.unwrap(); + /// let message = String::from("constructed after capacity became available"); + /// permit.send(message).unwrap(); + /// assert_eq!( + /// rx.recv().await.unwrap(), + /// "constructed after capacity became available" + /// ); + /// # } + /// ``` + pub async fn reserve(&self) -> Result, SendError<()>> { + match self.try_reserve() { + Ok(permit) => return Ok(permit), + Err(TrySendError::Disconnected(())) => return Err(SendError::new(())), + Err(TrySendError::Full(())) => {} + } + let mut waiter = self.state.capacity.waiter(); poll_fn(|cx| { - let message = value.take().expect("send polled after completion"); - let message = match self.try_send(message) { - Ok(()) => { + match self.try_reserve() { + Ok(permit) => { waiter.finish(); - return Poll::Ready(Ok(())); + return Poll::Ready(Ok(permit)); } - Err(TrySendError::Disconnected(message)) => { + Err(TrySendError::Disconnected(())) => { waiter.finish(); - return Poll::Ready(Err(SendError::new(message))); + return Poll::Ready(Err(SendError::new(()))); } - Err(TrySendError::Full(message)) => message, - }; + Err(TrySendError::Full(())) => {} + } waiter.register(cx.waker()); - match self.try_send(message) { - Ok(()) => { + match self.try_reserve() { + Ok(permit) => { waiter.finish(); - Poll::Ready(Ok(())) + Poll::Ready(Ok(permit)) } - Err(TrySendError::Disconnected(message)) => { + Err(TrySendError::Disconnected(())) => { waiter.finish(); - Poll::Ready(Err(SendError::new(message))) - } - Err(TrySendError::Full(message)) => { - value = Some(message); - Poll::Pending + Poll::Ready(Err(SendError::new(()))) } + Err(TrySendError::Full(())) => Poll::Pending, } }) .await } + /// Reserves capacity for one message without waiting. + /// + /// Returns [`TrySendError::Full`] if queued messages and outstanding permits occupy the + /// buffer, or [`TrySendError::Disconnected`] if the receiver has been dropped. + pub fn try_reserve(&self) -> Result, TrySendError<()>> { + self.state.capacity.try_acquire()?; + Ok(Permit { sender: self }) + } + /// Attempts to send a message without waiting for capacity. /// /// A full buffer returns [`TrySendError::Full`], while a dropped receiver returns @@ -169,7 +215,54 @@ impl BoundedSender { /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { - self.state.buffer.try_push(value) + match self.try_reserve() { + Ok(permit) => permit + .send(value) + .map_err(|error| TrySendError::Disconnected(error.into_inner())), + Err(TrySendError::Full(())) => Err(TrySendError::Full(value)), + Err(TrySendError::Disconnected(())) => Err(TrySendError::Disconnected(value)), + } + } +} + +/// Capacity reserved for one message on a bounded channel. +/// +/// Created by [`BoundedSender::reserve`] or [`BoundedSender::try_reserve`]. Holding a permit +/// reduces available capacity but does not prevent other messages from being received. Dropping +/// it without sending releases capacity and notifies a waiting sender. +#[must_use = "dropping the permit releases its reserved capacity"] +pub struct Permit<'a, T> { + sender: &'a BoundedSender, +} + +impl fmt::Debug for Permit<'_, T> { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("Permit").finish_non_exhaustive() + } +} + +impl Permit<'_, T> { + /// Publishes a message using this reservation, without waiting for capacity. + /// + /// If the receiver has been dropped, the returned error contains the unsent value. + pub fn send(self, value: T) -> Result<(), SendError> { + // SAFETY: This permit owns one unit of capacity. No user code runs between claiming the + // position and publishing its value; the consumer returns the capacity after reading it. + let claim = match unsafe { self.sender.state.buffer.claim() } { + Ok(claim) => claim, + Err(()) => return Err(SendError::new(value)), + }; + // Publication can wake user code that panics. Transfer capacity ownership first so + // unwinding cannot return a permit for a message that is already in the ring. + mem::forget(self); + claim.publish(value); + Ok(()) + } +} + +impl Drop for Permit<'_, T> { + fn drop(&mut self) { + self.sender.state.capacity.cancel(); } } @@ -178,6 +271,7 @@ impl BoundedSender { /// Instances are created by the [`bounded`] function. pub struct BoundedReceiver { state: Arc>, + head: usize, } impl fmt::Debug for BoundedReceiver { @@ -188,21 +282,29 @@ impl fmt::Debug for BoundedReceiver { impl Drop for BoundedReceiver { fn drop(&mut self) { - struct DrainOnDrop<'a, T>(&'a Ring); + struct DrainOnDrop<'a, T> { + ring: &'a Ring, + head: &'a mut usize, + tail: usize, + } impl Drop for DrainOnDrop<'_, T> { fn drop(&mut self) { // SAFETY: This guard lives only within the exclusive receiver's drop, after close. - unsafe { self.0.drain() }; + unsafe { self.ring.drain(self.head, self.tail) }; } } - self.state.buffer.close(); - let drain = DrainOnDrop(&self.state.buffer); + let tail = self.state.buffer.close(); + let drain = DrainOnDrop { + ring: &self.state.buffer, + head: &mut self.head, + tail, + }; // A registered waker may own a sender; release it to break that ownership cycle. let receiver_waker = self.state.buffer.take_receiver_waker(); // Complete notifications before dropping messages. Either kind of callback may panic; // the drain guard still releases buffered values if a wake or waker drop unwinds. - self.state.send_waiters.notify_all(); + self.state.capacity.close(); drop(receiver_waker); drop(drain); } @@ -245,20 +347,20 @@ impl BoundedReceiver { fn try_recv_once(&mut self) -> Poll> { // SAFETY: Only this non-cloneable receiver consumes the queue, through exclusive borrows. - let value = if let Some(value) = ready!(unsafe { self.state.buffer.pop() }) { + let value = if let Some(value) = ready!(unsafe { self.state.buffer.pop(&mut self.head) }) { value } else if self.state.senders.load(Ordering::Acquire) == 0 { // The final sender can enqueue between the first empty observation and decrementing // the sender count, so check the queue again before reporting disconnection. // SAFETY: The exclusive receiver borrow still guarantees a single consumer. - let Some(value) = ready!(unsafe { self.state.buffer.pop() }) else { + let Some(value) = ready!(unsafe { self.state.buffer.pop(&mut self.head) }) else { return Poll::Ready(Err(TryRecvError::Disconnected)); }; value } else { return Poll::Ready(Err(TryRecvError::Empty)); }; - self.state.send_waiters.notify_one(); + self.state.capacity.consume(self.head); Poll::Ready(Ok(value)) } diff --git a/asyncband/src/mpsc/bounded/ring.rs b/asyncband/src/mpsc/bounded/ring.rs index ace5aa76..350a2eda 100644 --- a/asyncband/src/mpsc/bounded/ring.rs +++ b/asyncband/src/mpsc/bounded/ring.rs @@ -25,21 +25,19 @@ use std::sync::atomic::fence; use std::task::Poll; use std::task::Waker; +use super::SEQUENCE_STEP; use crate::internal::cache_padded::CachePadded; use crate::internal::mutex::Mutex; -use crate::mpsc::TrySendError; + +const CLOSED: usize = 1; pub struct Ring { slots: Box<[Slot]>, - head: CachePadded, tail: CachePadded, - // This flag is usually stable while the consumer advances head. Sharing head - // for notifications would make every producer track a constantly invalidated cache line. + // Publication and receiver registration synchronize independently of capacity release. receiver_waiting: CachePadded, receiver: Mutex>, - capacity: usize, - one_lap: usize, - mark_bit: usize, + mask: usize, } struct Slot { @@ -47,9 +45,9 @@ struct Slot { value: UnsafeCell>, } -// SAFETY: A successful tail CAS gives one producer exclusive access to a slot. That producer +// SAFETY: A successful tail increment gives one producer exclusive access to a slot. That producer // initializes the value before publishing the next stamp with Release ordering. The single -// consumer reads only after acquiring that stamp and publishes the following lap before reuse. +// consumer reads only after acquiring that stamp and returns capacity before reuse. unsafe impl Sync for Slot {} // The ownership transition finishes before user code can unwind, and no stored-value reference is @@ -60,122 +58,65 @@ impl std::panic::RefUnwindSafe for Slot {} impl Ring { pub fn new(capacity: usize) -> Self { assert!(capacity <= usize::MAX / 4, "mpsc capacity is too large"); - let mark_bit = (capacity + 1).next_power_of_two(); - let one_lap = mark_bit * 2; - let slots = (0..capacity) - .map(|index| Slot { - stamp: AtomicUsize::new(index), + // Physical storage is rounded up, while Capacity enforces the exact requested limit. + // A power-of-two ring keeps indexing cheap and continuous across sequence overflow. + let storage = capacity.next_power_of_two(); + let slots = (0..storage) + .map(|_| Slot { + stamp: AtomicUsize::new(CLOSED), value: UnsafeCell::new(MaybeUninit::uninit()), }) .collect(); Self { slots, - head: CachePadded::new(AtomicUsize::new(0)), tail: CachePadded::new(AtomicUsize::new(0)), receiver_waiting: CachePadded::new(AtomicBool::new(false)), receiver: Mutex::new(None), - capacity, - one_lap, - mark_bit, + mask: storage - 1, } } - pub fn try_push(&self, value: T) -> Result<(), TrySendError> { - let mut tail = self.tail.load(Ordering::Relaxed); - let mut backoff = 0; - loop { - if tail & self.mark_bit != 0 { - return Err(TrySendError::Disconnected(value)); - } - - let index = tail & (self.mark_bit - 1); - let slot = &self.slots[index]; - let stamp = slot.stamp.load(Ordering::Acquire); - if stamp == tail { - let next_tail = self.advance(tail); - match self.tail.compare_exchange_weak( - tail, - next_tail, - Ordering::SeqCst, - Ordering::Relaxed, - ) { - Ok(_) => { - // SAFETY: The successful CAS reserved this slot exclusively, and its - // matching stamp proves the consumer completed its previous lap. - unsafe { (*slot.value.get()).write(value) }; - slot.stamp.store(tail.wrapping_add(1), Ordering::Release); - // Publication precedes the wait check; registration pairs this fence - // with a second pop before the receiver is allowed to return Pending. - fence(Ordering::SeqCst); - if self.receiver_waiting.load(Ordering::Relaxed) - && self.receiver_waiting.swap(false, Ordering::Relaxed) - { - // Claim the notification before locking so concurrent publishers - // do not all queue behind the same receiver registration. - self.wake_receiver(); - } - return Ok(()); - } - Err(actual) => tail = actual, - } - } else if stamp.wrapping_add(self.one_lap) == tail.wrapping_add(1) { - fence(Ordering::SeqCst); - if self.head.load(Ordering::Relaxed).wrapping_add(self.one_lap) == tail { - return Err(TrySendError::Full(value)); - } - tail = self.tail.load(Ordering::Relaxed); - } else { - let actual = self.tail.load(Ordering::Relaxed); - if actual == tail { - // Reserved but unpublished messages also occupy capacity. In particular, a - // capacity-one queue must report Full without waiting for its producer to - // publish the slot's stamp. - fence(Ordering::SeqCst); - if self.head.load(Ordering::Relaxed).wrapping_add(self.one_lap) == tail { - return Err(TrySendError::Full(value)); - } - } - tail = actual; - } - Self::spin(&mut backoff); + /// Claims the next FIFO position. No payload is touched until `Claim::publish`. + /// + /// # Safety + /// + /// The caller must already own one capacity permit for this ring, and transfer that permit + /// to the consumer on publication. The claim must be published without invoking user code. + pub unsafe fn claim(&self) -> Result, ()> { + // A permit may predate the previous use of this slot. Carry earlier claimants' + // capacity acquires through the sequence so even an old permit observes that read. + let position = self.tail.fetch_add(SEQUENCE_STEP, Ordering::AcqRel); + if position & CLOSED != 0 { + Err(()) + } else { + Ok(Claim { + ring: self, + position, + }) } } - /// Pending means the head slot is reserved but not published. It is distinct from an empty - /// queue: later producers may already have completed their sends. + /// Pending means the head slot is claimed but not yet published. /// /// # Safety /// - /// The caller must serialize all calls to `pop` and `drain` for this queue. - pub unsafe fn pop(&self) -> Poll> { - let mut head = self.head.load(Ordering::Relaxed); - let mut backoff = 0; - loop { - let index = head & (self.mark_bit - 1); - let slot = &self.slots[index]; - let stamp = slot.stamp.load(Ordering::Acquire); - if stamp == head.wrapping_add(1) { - let next_head = self.advance(head); - // SAFETY: Acquiring the matching stamp observes initialization by the producer. - // There is one consumer, so the value is read exactly once. - let value = unsafe { (*slot.value.get()).assume_init_read() }; - slot.stamp - .store(head.wrapping_add(self.one_lap), Ordering::Release); - self.head.store(next_head, Ordering::SeqCst); - return Poll::Ready(Some(value)); - } - - if stamp == head { - fence(Ordering::SeqCst); - if self.tail.load(Ordering::Relaxed) & !self.mark_bit == head { - return Poll::Ready(None); - } - } - if backoff == 8 { - return Poll::Pending; - } - Self::spin(&mut backoff); - head = self.head.load(Ordering::Relaxed); + /// Only the exclusive consumer may call `pop` or `drain`, with its persistent head cursor. + /// Return one capacity permit after each successful pop, after the value has been read. + pub unsafe fn pop(&self, head: &mut usize) -> Poll> { + let index = (*head / SEQUENCE_STEP) & self.mask; + let slot = &self.slots[index]; + if slot.stamp.load(Ordering::Acquire) == *head { + // SAFETY: Acquiring the published stamp observes initialization. The consumer owns + // this cursor exclusively, and capacity is not returned until after reading the value. + let value = unsafe { (*slot.value.get()).assume_init_read() }; + *head = head.wrapping_add(SEQUENCE_STEP); + return Poll::Ready(Some(value)); + } + fence(Ordering::SeqCst); + if self.tail.load(Ordering::Relaxed) & !CLOSED == *head { + Poll::Ready(None) + } else { + Poll::Pending } } @@ -214,65 +155,53 @@ impl Ring { } /// Prevents subsequent sends from reserving slots. Already reserved slots still publish. - pub fn close(&self) { - self.tail.fetch_or(self.mark_bit, Ordering::SeqCst); + pub fn close(&self) -> usize { + // Failed claims may still advance tail after close. Freeze the drain boundary at the + // close operation itself; it includes every successful claim and no rejected claims. + self.tail.fetch_or(CLOSED, Ordering::SeqCst) } - /// Drops all values after closing, including values whose publication is still in progress. + /// Drops all values after closing, including short-lived claims still being published. /// /// # Safety /// - /// The queue must be closed. The caller must serialize all calls to `pop` and `drain`. - pub unsafe fn drain(&self) { + /// `tail` must be the value returned by the first close, and `head` the consumer's cursor. + pub unsafe fn drain(&self, head: &mut usize, tail: usize) { struct DrainRemaining<'a, T> { ring: &'a Ring, + head: &'a mut usize, tail: usize, } impl Drop for DrainRemaining<'_, T> { fn drop(&mut self) { - // SAFETY: The guard is scoped to the exclusive consumer's drain of a closed ring. - unsafe { self.ring.discard_until(self.tail) }; + // SAFETY: The guard owns the consumer cursor until this closed ring is drained. + unsafe { self.ring.discard_until(self.head, self.tail) }; } } - let tail = self.tail.load(Ordering::Relaxed); - debug_assert_ne!(tail & self.mark_bit, 0); + debug_assert_eq!(tail & CLOSED, 0); let remaining = DrainRemaining { ring: self, - tail: tail & !self.mark_bit, + head, + tail, }; - // SAFETY: The caller guarantees exclusive consumer access. The guard finishes draining if - // a value's destructor panics, so messages that own senders cannot retain the closed ring. - unsafe { self.discard_until(remaining.tail) }; - } - - fn advance(&self, position: usize) -> usize { - let index = position & (self.mark_bit - 1); - if index + 1 < self.capacity { - position + 1 - } else { - let lap = position & !(self.one_lap - 1); - lap.wrapping_add(self.one_lap) - } + // SAFETY: The caller guarantees exclusive consumer access. The guard finishes draining + // if a destructor panics, including messages that own senders and would retain the ring. + unsafe { self.discard_until(remaining.head, remaining.tail) }; } - // The caller must close the queue and have exclusive consumer access before discarding values. - unsafe fn discard_until(&self, tail: usize) { - let mut head = self.head.load(Ordering::Relaxed); + // The caller owns the cursor and has prevented new claims by closing the ring. + unsafe fn discard_until(&self, head: &mut usize, tail: usize) { let mut backoff = 0; - while head != tail { - let index = head & (self.mark_bit - 1); + while *head != tail { + let index = (*head / SEQUENCE_STEP) & self.mask; let slot = &self.slots[index]; - if slot.stamp.load(Ordering::Acquire) == head.wrapping_add(1) { - let next_head = self.advance(head); - // Move the head before dropping the value so unwinding cannot drop it twice. - slot.stamp - .store(head.wrapping_add(self.one_lap), Ordering::Release); - self.head.store(next_head, Ordering::SeqCst); - // SAFETY: The acquired matching stamp proves the slot contains an initialized - // value, and advancing the single-consumer head claims it exactly once. + if slot.stamp.load(Ordering::Acquire) == *head { + // Advance before running the destructor so unwinding cannot drop a value twice. + *head = head.wrapping_add(SEQUENCE_STEP); + // SAFETY: The acquired stamp proves initialization; the cursor claims the value + // exactly once, and close prevents any producer from reusing its slot. unsafe { (*slot.value.get()).assume_init_drop() }; - head = next_head; backoff = 0; } else { Self::spin(&mut backoff); @@ -288,11 +217,29 @@ impl Ring { } } -impl Drop for Ring { - fn drop(&mut self) { - self.close(); - // SAFETY: The queue is closed and its exclusive borrow rules out concurrent access. - unsafe { self.drain() }; +// Claims never escape a synchronous send. A public Permit owns capacity without claiming a +// position, so holding or forgetting one cannot leave an unpublished hole in the ring. +pub struct Claim<'a, T> { + ring: &'a Ring, + position: usize, +} + +impl Claim<'_, T> { + pub fn publish(self, value: T) { + let slot = &self.ring.slots[(self.position / SEQUENCE_STEP) & self.ring.mask]; + // SAFETY: Claiming required a capacity permit. Its acquire observes the consumer's + // completed read on the previous lap; the tail increment grants this producer exclusive + // access. + unsafe { (*slot.value.get()).write(value) }; + slot.stamp.store(self.position, Ordering::Release); + // Paired with receiver registration: either the publisher sees the wait flag or the + // receiver's second pop observes this publication before it can return Pending. + fence(Ordering::SeqCst); + if self.ring.receiver_waiting.load(Ordering::Relaxed) + && self.ring.receiver_waiting.swap(false, Ordering::Relaxed) + { + self.ring.wake_receiver(); + } } } diff --git a/asyncband/src/mpsc/bounded/ring_tests.rs b/asyncband/src/mpsc/bounded/ring_tests.rs index d6d50433..fb68e69e 100644 --- a/asyncband/src/mpsc/bounded/ring_tests.rs +++ b/asyncband/src/mpsc/bounded/ring_tests.rs @@ -15,186 +15,68 @@ // specific language governing permissions and limitations // under the License. -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Poll; -use std::thread; use super::Ring; -use super::TrySendError; #[test] -fn bounded_queue_preserves_capacity_and_fifo_order() { - let queue = Ring::new(3); - for value in 0..3 { - assert!(queue.try_push(value).is_ok()); - } - assert!(matches!(queue.try_push(3), Err(TrySendError::Full(3)))); - for value in 0..3 { - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); - } - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(None)); - - for value in 3..12 { - assert!(queue.try_push(value).is_ok()); - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); - } -} - -#[test] -fn bounded_queue_does_not_report_empty_behind_an_unpublished_head() { +fn receive_waits_for_the_first_claim_even_if_a_later_claim_is_published() { let queue = Ring::new(2); - - // Pause a synthetic producer after reserving and initializing slot 0, before publishing - // its stamp. Another producer can finish sending into slot 1 in the meantime. - queue.tail.store(1, Ordering::SeqCst); - let slot = &queue.slots[0]; - // SAFETY: advancing the tail reserved this initially empty slot for the synthetic producer. - unsafe { (*slot.value.get()).write(1) }; - let later_send = queue.try_push(2); - // SAFETY: This thread is the only consumer, even while a producer is unpublished. - let receive = unsafe { queue.pop() }; - - // Finish publication before asserting so even a failed assertion can safely drop the queue. - slot.stamp.store(1, Ordering::Release); - assert!(later_send.is_ok()); - assert_eq!(receive, Poll::Pending); - // SAFETY: This thread is the only consumer. + let mut head = 0; + // SAFETY: The two claims fit in the initially empty ring. Both publish before teardown. + let (first, second) = unsafe { (queue.claim().unwrap(), queue.claim().unwrap()) }; + second.publish(2); + // SAFETY: This test owns the only consumer cursor. + let pending = unsafe { queue.pop(&mut head) }; + first.publish(1); + assert_eq!(pending, Poll::Pending); + // SAFETY: No other consumer exists, and no slot will be reused by another producer. unsafe { - assert_eq!(queue.pop(), Poll::Ready(Some(1))); - assert_eq!(queue.pop(), Poll::Ready(Some(2))); - assert_eq!(queue.pop(), Poll::Ready(None)); + assert_eq!(queue.pop(&mut head), Poll::Ready(Some(1))); + assert_eq!(queue.pop(&mut head), Poll::Ready(Some(2))); + assert_eq!(queue.pop(&mut head), Poll::Ready(None)); } } #[test] -fn unpublished_reservations_count_toward_capacity() { - let queue = Arc::new(Ring::new(1)); - queue.tail.store(queue.one_lap, Ordering::SeqCst); - let (done, completed) = std::sync::mpsc::channel(); - let producer = { - let queue = queue.clone(); - thread::spawn(move || done.send(queue.try_push(2)).unwrap()) - }; - #[cfg(not(miri))] - let result = completed.recv_timeout(std::time::Duration::from_secs(10)); - // Miri reports a deadlock directly instead of relying on an interpretation-time deadline. - #[cfg(miri)] - let result = completed.recv(); - // Finish the synthetic reservation even if the other producer stalled. This lets the - // worker and the ring's destructor finish before the failure is reported. - let slot = &queue.slots[0]; - // SAFETY: Advancing the tail above exclusively reserved the initially empty slot. - unsafe { (*slot.value.get()).write(1) }; - slot.stamp.store(1, Ordering::Release); - producer.join().unwrap(); - assert!(matches!(result, Ok(Err(TrySendError::Full(2))))); - // SAFETY: Both producers have finished and this thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(1))); -} - -#[test] -fn bounded_queue_coordinates_multiple_producers() { - let queue = Arc::new(Ring::new(4)); - let producers: Vec<_> = (0..2) - .map(|producer| { - let queue = queue.clone(); - thread::spawn(move || { - for offset in 0..32 { - let mut value = producer * 32 + offset; - loop { - match queue.try_push(value) { - Ok(()) => break, - Err(TrySendError::Full(returned)) => { - value = returned; - thread::yield_now(); - } - Err(TrySendError::Disconnected(_)) => panic!("queue disconnected"), - } - } - } - }) - }) - .collect(); - - let mut values = Vec::new(); - while values.len() < 64 { - // SAFETY: Worker threads only push; this thread is the only consumer. - if let Poll::Ready(Some(value)) = unsafe { queue.pop() } { - values.push(value); - } else { - thread::yield_now(); - } - } - for producer in producers { - producer.join().unwrap(); - } - values.sort_unstable(); - assert_eq!(values, (0..64).collect::>()); +fn closed_ring_finishes_claimed_publications_and_rejects_new_claims() { + let queue = Ring::new(2); + let mut head = 0; + // SAFETY: The initially empty ring has two available slots. + let first = unsafe { queue.claim().unwrap() }; + let tail = queue.close(); + // SAFETY: The second capacity unit has not been claimed, even though close rejects it. + let rejected = unsafe { queue.claim().is_err() }; + first.publish(String::from("claimed before close")); + assert!(rejected); + // SAFETY: This test owns the only consumer cursor, and the queue is closed. + unsafe { queue.drain(&mut head, tail) }; + // SAFETY: Repeated draining must not revisit a consumed value. + unsafe { queue.drain(&mut head, tail) }; } #[test] -fn bounded_queue_discards_wrapped_values_once_after_receiver_disconnect() { - // This has no owning fields, so a buggy second drop remains observable as count == 2 - // instead of invalidating the tracker first. - struct DropSpy<'a>(&'a AtomicUsize); - - impl<'a> Drop for DropSpy<'a> { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::Relaxed); +fn publication_stamps_survive_cursor_overflow_and_non_power_of_two_capacity() { + for capacity in [1, 3, 4] { + let queue = Ring::new(capacity); + // Start on the final lap before usize overflow; the index and close bit are both zero. + let mut head = usize::MAX - (2 * queue.slots.len() - 1); + queue.tail.store(head, Ordering::Relaxed); + for lap in 0..3 { + for index in 0..capacity { + // SAFETY: The previous lap was completely drained, so these claims fit. + unsafe { queue.claim().unwrap() }.publish((lap, index)); + } + for index in 0..capacity { + // SAFETY: The only consumer owns this cursor; all reads finish before reuse. + assert_eq!( + unsafe { queue.pop(&mut head) }, + Poll::Ready(Some((lap, index))) + ); + } + // SAFETY: The only consumer owns this cursor. + assert_eq!(unsafe { queue.pop(&mut head) }, Poll::Ready(None)); } } - - // Declare this before `queue` so the counters outlive values held by the queue. - let drops = [ - AtomicUsize::new(0), - AtomicUsize::new(0), - AtomicUsize::new(0), - AtomicUsize::new(0), - ]; - let queue = Ring::new(3); - - // Positions: 0, 1, 2 (then tail wraps to 8). - for counter in &drops[..3] { - assert!(queue.try_push(DropSpy(counter)).is_ok()); - } - - // Free slot 0, then reuse it on the next lap at position 8. - // SAFETY: This thread is the only consumer. - let popped = unsafe { queue.pop() }; - assert!(matches!(popped, Poll::Ready(Some(_)))); - drop(popped); - assert_eq!(drops[0].load(Ordering::Relaxed), 1); - assert!(queue.try_push(DropSpy(&drops[3])).is_ok()); - - // The pending range is positions 1 -> 2 -> 8 -> 9, not a contiguous integer range. - assert_eq!(queue.head.load(Ordering::Relaxed), 1); - assert_eq!(queue.tail.load(Ordering::Relaxed), queue.one_lap + 1); - - queue.close(); - // SAFETY: The queue is closed and this thread is the only consumer. - unsafe { queue.drain() }; - - // `discard_until` must dispose every value exactly once, including position 8. - for (value, counter) in drops.iter().enumerate() { - assert_eq!( - counter.load(Ordering::Relaxed), - 1, - "value {value} was dropped an unexpected number of times" - ); - } - - // Queue Drop calls discard_until again; it must see head == tail and not redrop. - drop(queue); - for (value, counter) in drops.iter().enumerate() { - assert_eq!( - counter.load(Ordering::Relaxed), - 1, - "value {value} was dropped more than once" - ); - } } diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs index 040c98b2..bad23f06 100644 --- a/asyncband/src/mpsc/mod.rs +++ b/asyncband/src/mpsc/mod.rs @@ -27,6 +27,7 @@ mod unbounded; pub use self::bounded::BoundedReceiver; pub use self::bounded::BoundedSender; +pub use self::bounded::Permit; pub use self::bounded::bounded; pub use self::error::RecvError; pub use self::error::SendError; diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index 24900214..6b1268e3 100644 --- a/tests-integration/tests/mpsc_test/main.rs +++ b/tests-integration/tests/mpsc_test/main.rs @@ -31,6 +31,7 @@ use self::support::poll_with; mod backpressure; mod callbacks; mod concurrency; +mod reservation; mod support; #[test] @@ -51,7 +52,7 @@ fn unbounded_try_recv_preserves_order_and_reports_state() { #[test] fn bounded_try_send_respects_capacity_and_order() { - for capacity in [1, 4, 16] { + for capacity in [1, 3, 4, 16] { let (tx, mut rx) = mpsc::bounded(capacity); for i in 0..capacity { diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs new file mode 100644 index 00000000..df414652 --- /dev/null +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -0,0 +1,210 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::cell::Cell; +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; +use std::task::Wake; +use std::task::Waker; + +use asyncband::mpsc; +use asyncband::mpsc::TryRecvError; +use asyncband::mpsc::TrySendError; +use tests_integration::poll_once; + +use super::support::WakeCounter; +use super::support::expect_ready; +use super::support::poll_with; + +#[test] +fn held_permits_consume_capacity_without_claiming_message_order() { + for capacity in [1, 3, 64] { + let (tx, mut rx) = mpsc::bounded(capacity); + let permit = tx.try_reserve().unwrap(); + for value in 1..capacity { + tx.try_send(value).unwrap(); + } + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + assert_eq!(tx.try_send(0), Err(TrySendError::Full(0))); + for value in 1..capacity { + assert_eq!(rx.try_recv(), Ok(value)); + } + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + permit.send(0).unwrap(); + assert_eq!(rx.try_recv(), Ok(0)); + // Repeated reservation and cancellation must restore the exact original capacity. + for _ in 0..3 { + let permits: Vec<_> = (0..capacity).map(|_| tx.try_reserve().unwrap()).collect(); + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + drop(permits); + } + } +} + +#[test] +fn dropping_a_permit_wakes_a_pending_reservation() { + let (tx, mut rx) = mpsc::bounded(1); + let held = tx.try_reserve().unwrap(); + let mut waiting = Box::pin(tx.reserve()); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(waiting.as_mut(), &waker).is_pending()); + drop(held); + assert_eq!(wakes.count(), 1); + let permit = expect_ready(poll_with(waiting.as_mut(), &waker)).unwrap(); + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + permit.send(7).unwrap(); + assert_eq!(rx.try_recv(), Ok(7)); +} + +#[test] +fn receiver_drop_does_not_wait_for_held_or_forgotten_permits() { + let (tx, mut rx) = mpsc::bounded(3); + let held = tx.try_reserve().unwrap(); + std::mem::forget(tx.try_reserve().unwrap()); + tx.try_send(String::from("ready")).unwrap(); + assert_eq!(rx.try_recv().unwrap(), "ready"); + tx.try_send(String::from("discarded on close")).unwrap(); + let mut waiting = Box::pin(tx.reserve()); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(waiting.as_mut(), &waker).is_pending()); + drop(rx); + assert_eq!(wakes.count(), 1); + assert!(expect_ready(poll_with(waiting.as_mut(), &waker)).is_err()); + assert!(matches!( + tx.try_reserve(), + Err(TrySendError::Disconnected(())) + )); + assert_eq!( + held.send(String::from("unsent")).unwrap_err().into_inner(), + "unsent" + ); +} + +#[test] +fn a_permit_can_publish_send_only_payloads_from_another_thread() { + let (tx, mut rx) = mpsc::bounded(1); + let permit = tx.try_reserve().unwrap(); + std::thread::scope(|scope| { + scope + .spawn(move || permit.send(Cell::new(42)).unwrap()) + .join() + .unwrap(); + }); + assert_eq!(rx.try_recv().unwrap().get(), 42); +} + +#[test] +fn a_panicking_publication_wake_cannot_return_capacity_twice() { + struct PanicOnWake; + impl Wake for PanicOnWake { + fn wake(self: Arc) { + panic!("publication wake"); + } + } + let (tx, mut rx) = mpsc::bounded(1); + let permit = tx.try_reserve().unwrap(); + let waker = Waker::from(Arc::new(PanicOnWake)); + let mut receive = Box::pin(rx.recv()); + assert!(poll_with(receive.as_mut(), &waker).is_pending()); + assert!(catch_unwind(AssertUnwindSafe(|| permit.send(1))).is_err()); + assert_eq!(tx.try_send(2), Err(TrySendError::Full(2))); + assert_eq!(expect_ready(poll_once(receive.as_mut())), Ok(1)); + drop(receive); + tx.try_send(2).unwrap(); + assert_eq!(rx.try_recv(), Ok(2)); +} + +#[test] +fn an_old_permit_observes_consumption_before_reusing_a_slot() { + let (tx, mut rx) = mpsc::bounded(2); + let old = tx.try_reserve().unwrap(); + let recycled = AtomicBool::new(false); + std::thread::scope(|scope| { + let recycled = &recycled; + let producer = scope.spawn(move || { + // Coordinate the schedule without supplying the happens-before edge that the + // channel itself must provide between the previous read and this slot's reuse. + while !recycled.load(Ordering::Relaxed) { + std::thread::yield_now(); + } + old.send(String::from("reused")).unwrap(); + }); + for value in ["first", "second"] { + tx.try_send(String::from(value)).unwrap(); + assert_eq!(rx.try_recv().unwrap(), value); + } + recycled.store(true, Ordering::Relaxed); + producer.join().unwrap(); + }); + assert_eq!(rx.try_recv().unwrap(), "reused"); +} + +#[test] +fn concurrent_cancellation_preserves_capacity_and_message_order() { + const PRODUCERS: usize = 3; + const MESSAGES: usize = if cfg!(miri) { 8 } else { 256 }; + let (tx, mut rx) = mpsc::bounded(3); + let mut out_of_order = 0; + std::thread::scope(|scope| { + let mut workers = Vec::new(); + for producer in 0..PRODUCERS { + let tx = tx.clone(); + workers.push(scope.spawn(move || { + for sequence in 0..MESSAGES { + for cancel in [true, false] { + let permit = loop { + match tx.try_reserve() { + Ok(permit) => break permit, + Err(TrySendError::Full(())) => std::thread::yield_now(), + Err(TrySendError::Disconnected(())) => panic!("receiver is alive"), + } + }; + if cancel { + drop(permit); + } else { + permit.send((producer, sequence)).unwrap(); + } + } + } + })); + } + let mut next = [0; PRODUCERS]; + let mut count = 0; + while count < PRODUCERS * MESSAGES { + match rx.try_recv() { + Ok((producer, sequence)) => { + out_of_order += usize::from(next[producer] != sequence); + next[producer] += 1; + count += 1; + } + Err(TryRecvError::Empty) => std::thread::yield_now(), + Err(TryRecvError::Disconnected) => panic!("senders are alive"), + } + } + for worker in workers { + worker.join().unwrap(); + } + }); + // Drain and join before asserting so a regression cannot strand producers on a full ring. + assert_eq!(out_of_order, 0); + let permits: Vec<_> = (0..3).map(|_| tx.try_reserve().unwrap()).collect(); + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + drop(permits); +} From 22f6c6e380f85f4dcbc5faffcfb285c6bd69531c Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 13:07:00 +0800 Subject: [PATCH 2/3] test(mpsc): benchmark reservations and inline bounded messages --- benchmarks/ecosystem/mpsc/adapters.rs | 106 ++++++++++---------- benchmarks/ecosystem/mpsc/bounded.rs | 102 +++++++++++++++++++ benchmarks/ecosystem/mpsc/mod.rs | 1 + benchmarks/ecosystem/mpsc/reservation.rs | 121 +++++++++++++++++++++++ 4 files changed, 277 insertions(+), 53 deletions(-) create mode 100644 benchmarks/ecosystem/mpsc/reservation.rs diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs index 506f1678..f9fe81b6 100644 --- a/benchmarks/ecosystem/mpsc/adapters.rs +++ b/benchmarks/ecosystem/mpsc/adapters.rs @@ -26,19 +26,19 @@ pub struct Tokio; pub struct AsyncChannel; pub struct Flume; -pub trait BoundedMpsc: Send + Sync + 'static { +pub trait BoundedMpsc: Send + Sync + 'static { type Sender: Clone + Send + Sync + 'static; type Receiver: Send + 'static; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver); - fn try_send(sender: &Self::Sender, value: usize); - fn try_recv(receiver: &mut Self::Receiver) -> usize; - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>); - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize; - fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; - fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; - fn send_blocking(sender: &Self::Sender, value: usize); - fn recv_blocking(receiver: &mut Self::Receiver) -> usize; + fn try_send(sender: &Self::Sender, value: T); + fn try_recv(receiver: &mut Self::Receiver) -> T; + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>); + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T; + fn send_async(sender: &Self::Sender, value: T) -> impl Future + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; + fn send_blocking(sender: &Self::Sender, value: T); + fn recv_blocking(receiver: &mut Self::Receiver) -> T; } pub trait UnboundedMpsc: Send + Sync + 'static { @@ -53,166 +53,166 @@ pub trait UnboundedMpsc: Send + Sync + 'static { fn recv_blocking(receiver: &mut Self::Receiver) -> T; } -impl BoundedMpsc for Asyncband { - type Receiver = asyncband::mpsc::BoundedReceiver; - type Sender = asyncband::mpsc::BoundedSender; +impl BoundedMpsc for Asyncband { + type Receiver = asyncband::mpsc::BoundedReceiver; + type Sender = asyncband::mpsc::BoundedSender; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { asyncband::mpsc::bounded(capacity) } - fn try_send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { poll_ready(sender.send(value), context).unwrap(); } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - async fn send_async(sender: &Self::Sender, value: usize) { + async fn send_async(sender: &Self::Sender, value: T) { sender.send(value).await.unwrap(); } - async fn recv_async(receiver: &mut Self::Receiver) -> usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv().await.unwrap() } - fn send_blocking(sender: &Self::Sender, value: usize) { + fn send_blocking(sender: &Self::Sender, value: T) { pollster::block_on(sender.send(value)).unwrap(); } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl BoundedMpsc for Tokio { - type Receiver = tokio::sync::mpsc::Receiver; - type Sender = tokio::sync::mpsc::Sender; +impl BoundedMpsc for Tokio { + type Receiver = tokio::sync::mpsc::Receiver; + type Sender = tokio::sync::mpsc::Sender; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { tokio::sync::mpsc::channel(capacity) } - fn try_send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { poll_ready(sender.send(value), context).unwrap(); } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - async fn send_async(sender: &Self::Sender, value: usize) { + async fn send_async(sender: &Self::Sender, value: T) { sender.send(value).await.unwrap(); } - async fn recv_async(receiver: &mut Self::Receiver) -> usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv().await.unwrap() } - fn send_blocking(sender: &Self::Sender, value: usize) { + fn send_blocking(sender: &Self::Sender, value: T) { pollster::block_on(sender.send(value)).unwrap(); } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl BoundedMpsc for AsyncChannel { - type Receiver = async_channel::Receiver; - type Sender = async_channel::Sender; +impl BoundedMpsc for AsyncChannel { + type Receiver = async_channel::Receiver; + type Sender = async_channel::Sender; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { async_channel::bounded(capacity) } - fn try_send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { poll_ready(sender.send(value), context).unwrap(); } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv(), context).unwrap() } - async fn send_async(sender: &Self::Sender, value: usize) { + async fn send_async(sender: &Self::Sender, value: T) { sender.send(value).await.unwrap(); } - async fn recv_async(receiver: &mut Self::Receiver) -> usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv().await.unwrap() } - fn send_blocking(sender: &Self::Sender, value: usize) { + fn send_blocking(sender: &Self::Sender, value: T) { pollster::block_on(sender.send(value)).unwrap(); } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv()).unwrap() } } -impl BoundedMpsc for Flume { - type Receiver = flume::Receiver; - type Sender = flume::Sender; +impl BoundedMpsc for Flume { + type Receiver = flume::Receiver; + type Sender = flume::Sender; fn channel(capacity: usize) -> (Self::Sender, Self::Receiver) { flume::bounded(capacity) } - fn try_send(sender: &Self::Sender, value: usize) { + fn try_send(sender: &Self::Sender, value: T) { sender.try_send(value).unwrap(); } - fn try_recv(receiver: &mut Self::Receiver) -> usize { + fn try_recv(receiver: &mut Self::Receiver) -> T { receiver.try_recv().unwrap() } - fn send_ready(sender: &Self::Sender, value: usize, context: &mut Context<'_>) { + fn send_ready(sender: &Self::Sender, value: T, context: &mut Context<'_>) { poll_ready(sender.send_async(value), context).unwrap(); } - fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { + fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> T { poll_ready(receiver.recv_async(), context).unwrap() } - async fn send_async(sender: &Self::Sender, value: usize) { + async fn send_async(sender: &Self::Sender, value: T) { sender.send_async(value).await.unwrap(); } - async fn recv_async(receiver: &mut Self::Receiver) -> usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { receiver.recv_async().await.unwrap() } - fn send_blocking(sender: &Self::Sender, value: usize) { + fn send_blocking(sender: &Self::Sender, value: T) { pollster::block_on(sender.send_async(value)).unwrap(); } - fn recv_blocking(receiver: &mut Self::Receiver) -> usize { + fn recv_blocking(receiver: &mut Self::Receiver) -> T { pollster::block_on(receiver.recv_async()).unwrap() } } diff --git a/benchmarks/ecosystem/mpsc/bounded.rs b/benchmarks/ecosystem/mpsc/bounded.rs index d388cc67..42d875fc 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -15,6 +15,13 @@ // specific language governing permissions and limitations // under the License. +use std::future::Future; +use std::future::poll_fn; +use std::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::Ordering; + use divan::Bencher; use divan::black_box; use divan::counter::ItemsCount; @@ -119,3 +126,98 @@ fn scheduled( batch.run(); bencher.bench_local(|| batch.run()); } + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + consts = [64, 4096], + args = [1, 8], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn scheduled_inline, const CAPACITY: usize>( + bencher: Bencher, + producers: usize, +) { + let runtime = tokio::runtime::Builder::new_multi_thread() + .worker_threads(4) + .build() + .unwrap(); + let (sender, mut receiver) = C::channel(CAPACITY); + let start: Vec<_> = (0..producers) + .map(|_| Arc::new(tokio::sync::Notify::new())) + .collect(); + let stop = Arc::new(AtomicBool::new(false)); + let workers: Vec<_> = start + .iter() + .enumerate() + .map(|(producer, start)| { + let sender = sender.clone(); + let start = start.clone(); + let stop = stop.clone(); + runtime.spawn(async move { + let mut sequence = 0u64; + loop { + start.notified().await; + if stop.load(Ordering::Acquire) { + break; + } + for _ in 0..BATCH_MESSAGES / producers { + let mut value = [1; 1024]; + value[..8].copy_from_slice(&(producer as u64).to_le_bytes()); + value[8..16].copy_from_slice(&sequence.to_le_bytes()); + C::send_async(&sender, black_box(value)).await; + sequence += 1; + } + } + }) + }) + .collect(); + drop(sender); + let mut expected = vec![0u64; producers]; + let mut run = || { + runtime.block_on(async { + let first = { + let mut receive = pin!(tokio::task::unconstrained(C::recv_async(&mut receiver))); + let mut released = false; + poll_fn(|cx| { + let result = receive.as_mut().poll(cx); + if !released { + assert!(result.is_pending(), "each sample starts with an empty wait"); + released = true; + for producer in &start { + producer.notify_one(); + } + } + result + }) + .await + }; + let mut value = first; + for received in 0..BATCH_MESSAGES { + let producer = u64::from_le_bytes(value[..8].try_into().unwrap()) as usize; + let sequence = u64::from_le_bytes(value[8..16].try_into().unwrap()); + assert_eq!(sequence, expected[producer]); + expected[producer] += 1; + assert_eq!(black_box(value)[1023], 1); + if received + 1 < BATCH_MESSAGES { + value = C::recv_async(&mut receiver).await; + } + } + assert!(expected.iter().all(|count| *count == expected[0])); + }) + }; + // Reuse tasks and the channel. Include the initial empty wait, backpressure, and payload + // movement; verify per-producer order and payload integrity on every measured sample. + run(); + bencher.bench_local(run); + stop.store(true, Ordering::Release); + for producer in &start { + producer.notify_one(); + } + runtime.block_on(async { + for worker in workers { + worker.await.unwrap(); + } + }); +} diff --git a/benchmarks/ecosystem/mpsc/mod.rs b/benchmarks/ecosystem/mpsc/mod.rs index dd09282b..ed522cc3 100644 --- a/benchmarks/ecosystem/mpsc/mod.rs +++ b/benchmarks/ecosystem/mpsc/mod.rs @@ -17,5 +17,6 @@ mod adapters; mod bounded; +mod reservation; mod support; mod unbounded; diff --git a/benchmarks/ecosystem/mpsc/reservation.rs b/benchmarks/ecosystem/mpsc/reservation.rs new file mode 100644 index 00000000..b8f3a275 --- /dev/null +++ b/benchmarks/ecosystem/mpsc/reservation.rs @@ -0,0 +1,121 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::future::Future; +use std::marker::PhantomData; + +use divan::Bencher; +use divan::black_box; +use divan::counter::ItemsCount; + +use super::adapters::Asyncband; +use super::adapters::BoundedMpsc; +use super::adapters::Tokio; +use super::support::BATCH_MESSAGES; +use super::support::ConcurrentMpsc; +use super::support::RepeatedTasks; +use crate::support::bench_context; +use crate::support::poll_ready; + +trait Reservable: BoundedMpsc { + type Permit<'a>: Send; + fn try_reserve(sender: &Self::Sender) -> Self::Permit<'_>; + fn reserve(sender: &Self::Sender) -> impl Future> + Send; + fn publish(permit: Self::Permit<'_>, value: usize); +} + +impl Reservable for Asyncband { + type Permit<'a> = asyncband::mpsc::Permit<'a, usize>; + fn try_reserve(sender: &Self::Sender) -> Self::Permit<'_> { + sender.try_reserve().unwrap() + } + async fn reserve(sender: &Self::Sender) -> Self::Permit<'_> { + sender.reserve().await.unwrap() + } + fn publish(permit: Self::Permit<'_>, value: usize) { + permit.send(value).unwrap(); + } +} + +impl Reservable for Tokio { + type Permit<'a> = tokio::sync::mpsc::Permit<'a, usize>; + fn try_reserve(sender: &Self::Sender) -> Self::Permit<'_> { + sender.try_reserve().unwrap() + } + async fn reserve(sender: &Self::Sender) -> Self::Permit<'_> { + sender.reserve().await.unwrap() + } + fn publish(permit: Self::Permit<'_>, value: usize) { + permit.send(value); + } +} + +#[divan::bench(types = [Asyncband, Tokio])] +fn reserve_publish_receive(bencher: Bencher) { + let (sender, mut receiver) = C::channel(64); + let mut context = bench_context(); + bencher.bench_local(|| { + let permit = poll_ready(C::reserve(&sender), &mut context); + C::publish(permit, black_box(usize::MAX)); + black_box(C::try_recv(&mut receiver)) + }); +} + +#[divan::bench(types = [Asyncband, Tokio])] +fn cancel_reserved_capacity(bencher: Bencher) { + let (sender, _receiver) = C::channel(64); + bencher.bench_local(|| drop(black_box(C::try_reserve(&sender)))); +} + +struct Reserved(PhantomData); + +impl ConcurrentMpsc for Reserved { + type Sender = C::Sender; + type Receiver = C::Receiver; + fn channel() -> (Self::Sender, Self::Receiver) { + C::channel(CAPACITY) + } + fn send(sender: &Self::Sender, value: usize) { + C::publish(pollster::block_on(C::reserve(sender)), value); + } + fn recv(receiver: &mut Self::Receiver) -> usize { + C::recv_blocking(receiver) + } + async fn send_async(sender: &Self::Sender, value: usize) { + C::publish(C::reserve(sender).await, value); + } + async fn recv_async(receiver: &mut Self::Receiver) -> usize { + C::recv_async(receiver).await + } +} + +#[divan::bench( + types = [Asyncband, Tokio], + consts = [64, 4096], + args = [(1, 0), (8, 4)], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn scheduled( + bencher: Bencher, + (producers, workers): (usize, usize), +) { + let mut batch = RepeatedTasks::>::new(producers, workers); + batch.run(); + bencher.bench_local(|| batch.run()); +} From 82d6992028b2b25fd42b50ead4e41e08b7440491 Mon Sep 17 00:00:00 2001 From: tison Date: Tue, 8 Sep 2026 13:07:00 +0800 Subject: [PATCH 3/3] chore: remove the unused internal atomic waker --- LICENSE | 12 - asyncband/src/internal/atomic_waker.rs | 464 ------------------------- asyncband/src/internal/mod.rs | 5 - 3 files changed, 481 deletions(-) delete mode 100644 asyncband/src/internal/atomic_waker.rs diff --git a/LICENSE b/LICENSE index 942ddcee..3bd8900a 100644 --- a/LICENSE +++ b/LICENSE @@ -377,18 +377,6 @@ the Apache-2.0 option for the incorporated portions. Asyncband does not provide the upstream crate's synchronized receive operations and simplifies the incorporated implementation accordingly. -Portions of asyncband/src/internal/atomic_waker.rs are derived from futures-rs -0.3.34 at the following exact revision and source path: - - https://github.com/rust-lang/futures-rs/blob/705e6b5c0f06535b1aac1cb1989a172b3d45be8c/futures-core/src/task/__internal/atomic_waker.rs - -futures-rs is licensed under Apache-2.0 or MIT. Apache Asyncband uses the -Apache-2.0 option for the incorporated portions. The upstream source carries -the following copyright notices: - - Copyright (c) 2016 Alex Crichton - Copyright (c) 2017 The Tokio Authors - The polling loop in asyncband/src/blocking/executor.rs is adapted from Pollster 1.0.1 at the following exact revision and source path: diff --git a/asyncband/src/internal/atomic_waker.rs b/asyncband/src/internal/atomic_waker.rs deleted file mode 100644 index a6dcd20f..00000000 --- a/asyncband/src/internal/atomic_waker.rs +++ /dev/null @@ -1,464 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -// This file contains a state machine derived from futures-rs 0.3.34 and panic-recovery behavior -// informed by Tokio 1.53.1. -// Asyncband uses the Apache-2.0 license option for code incorporated from futures-rs. -// The incorporated code has been modified for use in Apache Asyncband. -// Upstream sources: -// https://github.com/rust-lang/futures-rs/blob/705e6b5c0f06535b1aac1cb1989a172b3d45be8c/futures-core/src/task/__internal/atomic_waker.rs -// https://github.com/tokio-rs/tokio/blob/75fef53d0a8590c2d1dbb63672aa7b7d1ef51155/tokio/src/sync/task/atomic_waker.rs - -use std::cell::UnsafeCell; -use std::panic::AssertUnwindSafe; -use std::panic::RefUnwindSafe; -use std::panic::UnwindSafe; -use std::panic::catch_unwind; -use std::panic::resume_unwind; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Waker; - -const WAITING: usize = 0; -const REGISTERING: usize = 0b01; -const WAKING: usize = 0b10; - -/// A single-registerer, multi-notifier cell for task wake-up. -/// -/// The atomic state both grants exclusive access to `waker` and records one coalesced wake request. -/// The operation that moves the state out of `WAITING` remains the only slot owner until it returns -/// the state to `WAITING`. -/// -/// * `WAITING`: the slot is unlocked and may contain a registered waker. -/// * `REGISTERING`: `register` exclusively owns the slot and no concurrent wake is pending. -/// * `WAKING`: `wake` exclusively owns the slot. A racing `register` self-wakes without touching -/// the slot. -/// * `REGISTERING | WAKING`: `register` still owns the slot and must complete a concurrent wake -/// before returning to `WAITING`. -/// -/// Valid state transitions are: -/// -/// ```text -/// register: WAITING ----------------Acquire CAS---------------> REGISTERING -/// REGISTERING ------------AcqRel CAS----------------> WAITING -/// -/// wake: WAITING ----------------AcqRel fetch_or-----------> WAKING -/// WAKING -----------------Release swap--------------> WAITING -/// -/// race: REGISTERING ------------AcqRel fetch_or-----------> REGISTERING | WAKING -/// REGISTERING | WAKING ---AcqRel swap---------------> WAITING -/// ``` -/// -/// Additional calls to `wake` while `WAKING` is set are coalesced. A wake completed before a -/// registration starts is not remembered, so callers must register before rechecking the condition -/// that determines whether to return `Pending`. -/// -/// Every transition that acquires slot ownership has an Acquire operation paired with the previous -/// owner's Release transition to `WAITING`. The Release half of `wake` also publishes the caller's -/// preceding condition update; a racing `register` acquires that publication before it returns. -pub struct AtomicWaker { - state: AtomicUsize, - waker: UnsafeCell>, -} - -// SAFETY: `state` grants exclusive access to `waker`, and losing concurrent registrations do not -// touch the slot. `Waker` itself is `Send + Sync`. -unsafe impl Sync for AtomicWaker {} - -// `Waker` callbacks may unwind, but no panic leaves a state bit owned by the unwinding operation. A -// failed clone leaves the old slot intact and completes any raced wake, while wake and drop -// callbacks run after that operation's critical section has been released. -impl RefUnwindSafe for AtomicWaker {} -impl UnwindSafe for AtomicWaker {} - -impl AtomicWaker { - #[inline] - pub const fn new() -> Self { - Self { - state: AtomicUsize::new(WAITING), - waker: UnsafeCell::new(None), - } - } - - /// Registers `waker`, replacing a previously registered task if it differs. - /// - /// Calls to this method must not overlap. It may run concurrently with any number of calls to - /// [`wake`](Self::wake). - #[inline] - pub fn register(&self, waker: &Waker) { - // ORDERING: On success, Acquire pairs with the Release operation that last returned the - // state to WAITING and transfers exclusive ownership of the waker slot to this thread. On - // failure, Acquire matters when this reads WAKING from a notifier's AcqRel fetch_or: it - // receives the condition update that preceded that wake before this method returns. - match self - .state - .compare_exchange(WAITING, REGISTERING, Ordering::Acquire, Ordering::Acquire) - .unwrap_or_else(|state| state) - { - WAITING => { - // SAFETY: changing WAITING to REGISTERING grants this thread exclusive access to - // the waker slot until the state is returned to WAITING. - unsafe { self.register_locked(waker) } - } - WAKING => { - // A concurrent wake owns the slot. Self-waking ensures that this registration is - // not lost even though it cannot replace the slot right now. - waker.wake_by_ref(); - } - state => { - // Concurrent registration violates this type's contract. Ignoring the losing - // registration preserves memory safety and lets the winner provide notification. - debug_assert!(state == REGISTERING || state == REGISTERING | WAKING); - } - } - } - - /// Registers a waker after this thread has acquired the REGISTERING state. - /// - /// # Safety - /// - /// The caller must have changed `state` from WAITING to REGISTERING and must be the only - /// thread accessing `waker`. - #[inline] - unsafe fn register_locked(&self, waker: &Waker) { - // Avoid both cloning and dropping the common case where an executor polls the receiver - // repeatedly with the same task waker. - let needs_replacement = match unsafe { &*self.waker.get() } { - Some(current) => !current.will_wake(waker), - None => true, - }; - - let mut clone_panic = None; - let old_waker = if needs_replacement { - match catch_unwind(AssertUnwindSafe(|| waker.clone())) { - Ok(new_waker) => unsafe { (*self.waker.get()).replace(new_waker) }, - Err(payload) => { - clone_panic = Some(payload); - None - } - } - } else { - None - }; - - // ORDERING: Release publishes a newly registered waker when the CAS succeeds. If it fails, - // Acquire receives the concurrent notifier's Release publication before the wake is - // completed below. AcqRel is the weakest success ordering that permits an Acquire failure - // ordering, although its Acquire half is not otherwise relied upon on the success path. - let concurrent_wake = match self.state.compare_exchange( - REGISTERING, - WAITING, - Ordering::AcqRel, - Ordering::Acquire, - ) { - Ok(_) => None, - Err(state) => { - debug_assert_eq!(state, REGISTERING | WAKING); - - // SAFETY: REGISTERING remains set, so this thread still owns the waker slot. - let registered = unsafe { (*self.waker.get()).take() }; - - // ORDERING: Acquire receives all coalesced wake publications. Release publishes - // the empty slot and makes it available to the next register or wake operation. - self.state.swap(WAITING, Ordering::AcqRel); - registered - } - }; - - if let Some(payload) = clone_panic { - // Preserve the original clone panic while still completing a wake that raced with it. - if let Some(waker) = concurrent_wake { - let _ = catch_unwind(AssertUnwindSafe(|| waker.wake())); - } - resume_unwind(payload); - } - - // User waker code runs only after the state machine is back in WAITING, so a panic cannot - // leave the cell locked. If the wake raced with a replacement, notify both tasks: the - // concurrent call may have targeted the old registration, while future progress relies on - // the new one. A panic from the superseded waker must not prevent the new task from waking. - if let Some(waker) = concurrent_wake { - if let Some(old_waker) = old_waker { - let _ = catch_unwind(AssertUnwindSafe(|| old_waker.wake())); - } - waker.wake(); - } else { - // Drop a replaced waker only after releasing the state lock. - drop(old_waker); - } - } - - /// Wakes and removes the most recently registered waker, if any. - #[inline] - pub fn wake(&self) { - if let Some(waker) = self.take() { - waker.wake(); - } - } - - /// Removes the registered waker if this call acquires the slot. A concurrent registration or - /// wake may instead take responsibility for notifying it. - #[inline] - pub fn take(&self) -> Option { - // ORDERING: When this reads WAITING, Acquire receives the registered waker published by the - // previous owner. Release publishes the condition update that the caller performed before - // calling wake, including when a registering thread already owns the slot. - match self.state.fetch_or(WAKING, Ordering::AcqRel) { - WAITING => { - // SAFETY: changing WAITING to WAKING grants this thread exclusive access to the - // waker slot until the state is returned to WAITING. - let waker = unsafe { (*self.waker.get()).take() }; - - // ORDERING: Release publishes the emptied slot before another operation acquires - // it. The fetch_or above already performed the required Acquire operation. - let old_state = self.state.swap(WAITING, Ordering::Release); - debug_assert_eq!(old_state, WAKING); - waker - } - state => { - // The thread registering a waker observes WAKING and completes this notification, - // or another waking thread has already taken responsibility for it. - debug_assert!( - state == REGISTERING || state == REGISTERING | WAKING || state == WAKING - ); - None - } - } - } -} - -#[cfg(test)] -mod tests { - use std::ptr; - use std::sync::Arc; - use std::sync::atomic::AtomicBool; - use std::sync::atomic::AtomicUsize; - use std::sync::atomic::Ordering; - use std::task::RawWaker; - use std::task::RawWakerVTable; - use std::task::Wake; - - use super::*; - - struct WakeCounter(AtomicUsize); - - impl Wake for WakeCounter { - fn wake(self: Arc) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } - - #[cfg(panic = "unwind")] - fn clone_panicking_waker() -> Waker { - static VTABLE: RawWakerVTable = RawWakerVTable::new( - |_| panic!("clone failed"), - |_| unreachable!(), - |_| unreachable!(), - |_| {}, - ); - - unsafe { Waker::from_raw(RawWaker::new(ptr::null(), &VTABLE)) } - } - - #[test] - fn wake_notifies_once() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.register(&waker); - atomic_waker.wake(); - atomic_waker.wake(); - - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn reregistering_same_task_does_not_clone_waker() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.register(&waker); - let registered_refs = Arc::strong_count(&counter); - atomic_waker.register(&waker); - - assert_eq!(Arc::strong_count(&counter), registered_refs); - } - - #[test] - fn wake_before_register_is_not_remembered() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - - atomic_waker.wake(); - atomic_waker.register(&waker); - - assert_eq!(counter.0.load(Ordering::Relaxed), 0); - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn wake_during_replacement_notifies_old_and_new_tasks() { - let old_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let old_waker = Waker::from(old_counter.clone()); - let new_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let new_waker = Waker::from(new_counter.clone()); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&old_waker); - - assert_eq!( - atomic_waker.state.compare_exchange( - WAITING, - REGISTERING, - Ordering::AcqRel, - Ordering::Acquire, - ), - Ok(WAITING) - ); - std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); - - // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching - // the slot. Calling the helper completes the interrupted registration. - unsafe { atomic_waker.register_locked(&new_waker) }; - - assert_eq!(old_counter.0.load(Ordering::Relaxed), 1); - assert_eq!(new_counter.0.load(Ordering::Relaxed), 1); - } - - #[test] - fn failed_wake_synchronizes_with_next_registration() { - for _ in 0..1_000 { - let did_publish = AtomicBool::new(false); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(Waker::noop()); - - std::thread::scope(|scope| { - let wake = scope.spawn(|| { - did_publish.store(true, Ordering::Relaxed); - atomic_waker.take() - }); - - let local_waker = atomic_waker.take(); - atomic_waker.register(Waker::noop()); - - let publication_is_visible = did_publish.load(Ordering::Relaxed); - let concurrent_thread_took_waker = wake.join().unwrap().is_some(); - assert!(publication_is_visible || concurrent_thread_took_waker); - drop(local_waker); - }); - } - } - - #[cfg(panic = "unwind")] - #[test] - fn clone_panic_does_not_poison_state() { - let atomic_waker = AtomicWaker::new(); - - assert!( - catch_unwind(|| { - atomic_waker.register(&clone_panicking_waker()); - }) - .is_err() - ); - - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - atomic_waker.register(&Waker::from(counter.clone())); - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } - - #[cfg(panic = "unwind")] - #[test] - fn clone_panic_completes_concurrent_wake() { - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&Waker::from(counter.clone())); - - assert_eq!( - atomic_waker.state.compare_exchange( - WAITING, - REGISTERING, - Ordering::Acquire, - Ordering::Acquire, - ), - Ok(WAITING) - ); - std::thread::scope(|scope| scope.spawn(|| atomic_waker.wake()).join().unwrap()); - - // SAFETY: this test acquired REGISTERING above and the waking thread has finished touching - // the state. Calling the helper completes the interrupted registration. - assert!( - catch_unwind(|| unsafe { - atomic_waker.register_locked(&clone_panicking_waker()); - }) - .is_err() - ); - - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - - let next_counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - atomic_waker.register(&Waker::from(next_counter.clone())); - atomic_waker.wake(); - assert_eq!(next_counter.0.load(Ordering::Relaxed), 1); - } - - #[cfg(panic = "unwind")] - #[test] - fn drop_panic_does_not_poison_state() { - unsafe fn clone_drop_panicker(data: *const ()) -> RawWaker { - RawWaker::new(data, &DROP_PANICKING_VTABLE) - } - - unsafe fn wake_drop_panicker(_: *const ()) {} - - unsafe fn drop_drop_panicker(data: *const ()) { - // SAFETY: the test keeps the pointed-to AtomicBool alive until every derived waker has - // been dropped. - let should_panic = unsafe { &*data.cast::() }; - if should_panic.swap(false, Ordering::Relaxed) { - panic!("drop failed"); - } - } - - static DROP_PANICKING_VTABLE: RawWakerVTable = RawWakerVTable::new( - clone_drop_panicker, - wake_drop_panicker, - wake_drop_panicker, - drop_drop_panicker, - ); - - let should_panic = AtomicBool::new(true); - let old_waker = unsafe { - Waker::from_raw(RawWaker::new( - ptr::from_ref(&should_panic).cast(), - &DROP_PANICKING_VTABLE, - )) - }; - let counter = Arc::new(WakeCounter(AtomicUsize::new(0))); - let new_waker = Waker::from(counter.clone()); - let atomic_waker = AtomicWaker::new(); - atomic_waker.register(&old_waker); - - assert!(catch_unwind(AssertUnwindSafe(|| atomic_waker.register(&new_waker))).is_err()); - - atomic_waker.wake(); - assert_eq!(counter.0.load(Ordering::Relaxed), 1); - } -} diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 045a1c2b..2bf8ac68 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -49,11 +49,6 @@ pub(crate) fn wake_all(mut wakers: impl Iterator) { } } -// MPSC owns its receiver wait protocol; the general-purpose waker currently has no production -// users. -#[cfg(test)] -pub(crate) mod atomic_waker; - #[cfg(any( feature = "barrier", feature = "broadcast",