diff --git a/CHANGELOG.md b/CHANGELOG.md index 43ffd8e..7113dc6 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; pending sends and reservations receive capacity in wait-queue order, and 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. @@ -13,7 +17,6 @@ All notable changes to this project will be documented in this file. ### Improvements * Finish releasing buffered bounded MPSC messages even if one message destructor panics. -* Reduce bounded MPSC contention when senders or the receiver are not waiting, improving throughput without changing capacity or cancellation semantics. * Improve unbounded MPSC throughput with batched receiving and incremental storage reclamation; empty-buffer retention is bounded independently of previous peak occupancy. ## v0.7.2 diff --git a/LICENSE b/LICENSE index 942ddce..3bd8900 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 a6dcd20..0000000 --- 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/cache_padded.rs b/asyncband/src/internal/cache_padded.rs deleted file mode 100644 index 7268030..0000000 --- a/asyncband/src/internal/cache_padded.rs +++ /dev/null @@ -1,55 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -// Use 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 045a1c2..0252aa8 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", @@ -72,9 +67,6 @@ 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; diff --git a/asyncband/src/mpsc/bounded/mod.rs b/asyncband/src/mpsc/bounded/mod.rs index 455e709..a9ae677 100644 --- a/asyncband/src/mpsc/bounded/mod.rs +++ b/asyncband/src/mpsc/bounded/mod.rs @@ -18,300 +18,108 @@ //! A bounded multi-producer, single-consumer queue for sending values between asynchronous //! tasks with backpressure control. -use std::fmt; -use std::future::poll_fn; +use std::collections::VecDeque; 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 std::task::Waker; -use self::ring::Ring; -use self::waiters::SendWaiters; -use super::RecvError; -use super::SendError; -use super::TryRecvError; -use super::TrySendError; +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaitList; +use crate::mpsc::TryRecvError; +use crate::mpsc::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. -mod ring; -mod waiters; +mod receiver; +mod sender; + +pub use self::receiver::BoundedReceiver; +pub use self::sender::BoundedSender; +pub use self::sender::Permit; /// 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 -/// one slot for a waiting sender. +/// one slot for a waiting sender. Capacity is granted in the order that pending sends and +/// reservations enter the wait queue; new senders cannot take an already granted slot. +/// +/// Message storage is preallocated for `buffer` values. Queued messages and outstanding +/// reservations together occupy at most `buffer` capacity units. +/// +/// Operations briefly acquire an internal mutex; no lock is held across an await point or while +/// invoking waker callbacks or message destructors. The `try_*` methods do not wait for capacity +/// or messages, but may wait to acquire this mutex. /// /// # Panics /// -/// Panics if `buffer` is zero. +/// Panics if `buffer` is zero or the preallocated message storage exceeds the allocation size +/// limit. #[track_caller] pub fn bounded(buffer: usize) -> (BoundedSender, BoundedReceiver) { - assert!(buffer > 0, "mpsc bounded channel requires buffer > 0"); - let state = Arc::new(Shared { - buffer: Ring::new(buffer), - senders: AtomicUsize::new(1), - send_waiters: SendWaiters::new(), - }); - let sender = BoundedSender { - state: state.clone(), - }; - let receiver = BoundedReceiver { state }; - (sender, receiver) -} - -struct Shared { - buffer: Ring, - senders: AtomicUsize, - send_waiters: SendWaiters, + assert!( + buffer > 0, + "mpsc bounded channel capacity {buffer} must be nonzero", + ); + let shared = Arc::new(Mutex::new(State { + queue: VecDeque::with_capacity(buffer), + available: buffer, + senders: 1, + receiver: true, + recv_waker: None, + send_waiters: WaitList::new(), + })); + ( + BoundedSender::new(shared.clone()), + BoundedReceiver::new(shared), + ) } -/// The sending endpoint of a bounded mpsc channel. -/// -/// Instances are created by the [`bounded`] function. -pub struct BoundedSender { - state: Arc>, +// While open, capacity belongs to available, a queued message, a Permit, or a granted waiter. +// All transitions hold one mutex. Waker callbacks and message destruction run after unlocking. +struct State { + queue: VecDeque, + available: usize, + senders: usize, + // True while the receiving endpoint is alive. + receiver: bool, + recv_waker: Option, + send_waiters: WaitList, } -impl Clone for BoundedSender { - fn clone(&self) -> Self { - self.state.senders.fetch_add(1, Ordering::Release); - BoundedSender { - state: self.state.clone(), +impl State { + fn acquire(&mut self) -> Result<(), TrySendError<()>> { + if !self.receiver { + Err(TrySendError::Disconnected(())) + } else if self.available == 0 { + Err(TrySendError::Full(())) + } else { + self.available -= 1; + Ok(()) } } -} -impl fmt::Debug for BoundedSender { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("BoundedSender").finish_non_exhaustive() - } -} - -impl Drop for BoundedSender { - fn drop(&mut self) { - if self.state.senders.fetch_sub(1, Ordering::AcqRel) == 1 { - self.state.buffer.wake_receiver(); + fn release(&mut self) -> Option { + if !self.receiver { + return None; } - } -} - -impl BoundedSender { - /// Sends a message, waiting until the channel has capacity when necessary. - /// - /// If the receiver has been dropped, the returned error contains `value`. - /// - /// # Cancel safety - /// - /// 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. - 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); - poll_fn(|cx| { - let message = value.take().expect("send polled after completion"); - let message = match self.try_send(message) { - Ok(()) => { - waiter.finish(); - return Poll::Ready(Ok(())); - } - Err(TrySendError::Disconnected(message)) => { - waiter.finish(); - return Poll::Ready(Err(SendError::new(message))); - } - Err(TrySendError::Full(message)) => message, - }; - waiter.register(cx.waker()); - match self.try_send(message) { - Ok(()) => { - waiter.finish(); - Poll::Ready(Ok(())) - } - Err(TrySendError::Disconnected(message)) => { - waiter.finish(); - Poll::Ready(Err(SendError::new(message))) - } - Err(TrySendError::Full(message)) => { - value = Some(message); - Poll::Pending - } - } - }) - .await - } - - /// Attempts to send a message without waiting for capacity. - /// - /// A full buffer returns [`TrySendError::Full`], while a dropped receiver returns - /// [`TrySendError::Disconnected`]. Both errors return ownership of the unsent value. - /// - /// # Examples - /// - /// ``` - /// use asyncband::mpsc::TrySendError; - /// use asyncband::mpsc::bounded; - /// - /// let (tx, mut rx) = bounded(1); - /// tx.try_send(10).unwrap(); - /// assert_eq!(tx.try_send(20), Err(TrySendError::Full(20))); - /// - /// assert_eq!(rx.try_recv(), Ok(10)); - /// tx.try_send(20).unwrap(); - /// drop(rx); - /// 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) - } -} - -/// The receiving endpoint of a bounded mpsc channel. -/// -/// Instances are created by the [`bounded`] function. -pub struct BoundedReceiver { - state: Arc>, -} - -impl fmt::Debug for BoundedReceiver { - fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { - f.debug_struct("BoundedReceiver").finish_non_exhaustive() - } -} - -impl Drop for BoundedReceiver { - fn drop(&mut self) { - struct DrainOnDrop<'a, T>(&'a Ring); - impl Drop for DrainOnDrop<'_, T> { - fn drop(&mut self) { - // SAFETY: This guard lives only within the exclusive receiver's drop, after close. - unsafe { self.0.drain() }; - } - } - - self.state.buffer.close(); - let drain = DrainOnDrop(&self.state.buffer); - // A registered waker may own a sender; release it to break that ownership cycle. - let receiver_waker = self.state.buffer.take_receiver_waker(); - // Complete notifications before dropping messages. Either kind of callback may panic; - // the drain guard still releases buffered values if a wake or waker drop unwinds. - self.state.send_waiters.notify_all(); - drop(receiver_waker); - drop(drain); - } -} - -impl BoundedReceiver { - /// 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 - /// been dropped and all queued values have been consumed. - /// - /// # Examples - /// - /// ``` - /// use asyncband::mpsc::TryRecvError; - /// use asyncband::mpsc::bounded; - /// - /// let (tx, mut rx) = bounded(2); - /// tx.try_send("first").unwrap(); - /// tx.try_send("second").unwrap(); - /// - /// assert_eq!(rx.try_recv(), Ok("first")); - /// assert_eq!(rx.try_recv(), Ok("second")); - /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); - /// drop(tx); - /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); - /// ``` - pub fn try_recv(&mut self) -> Result { - loop { - if let Poll::Ready(result) = self.try_recv_once() { - return result; - } - std::thread::yield_now(); + if let Some((_, waiter)) = self.send_waiters.unlink_first_waiter(|_| true) { + // The detached node owns capacity until its future claims or cancels the grant. + waiter.grant = true; + return waiter.waker.take(); } + self.available += 1; + None } - 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 + fn pop(&mut self) -> Result<(T, Option), TryRecvError> { + if let Some(value) = self.queue.pop_front() { + Ok((value, self.release())) + } else if self.senders == 0 { + Err(TryRecvError::Disconnected) } else { - return Poll::Ready(Err(TryRecvError::Empty)); - }; - self.state.send_waiters.notify_one(); - 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 - /// dropped. It returns [`RecvError::Disconnected`] only after all senders are gone and the - /// buffer has been drained. - /// - /// # Cancel safety - /// - /// Dropping a pending `recv` does not remove a message from the channel. A later receive - /// operation can still observe the next queued value, so `recv` may safely be raced with other - /// futures in a selection construct. - /// - /// # Examples - /// - /// ``` - /// # #[tokio::main] - /// # async fn main() { - /// use asyncband::mpsc; - /// let (tx, mut rx) = mpsc::bounded(2); - /// - /// tx.send("first").await.unwrap(); - /// tx.send("second").await.unwrap(); - /// drop(tx); - /// - /// assert_eq!(rx.recv().await, Ok("first")); - /// assert_eq!(rx.recv().await, Ok("second")); - /// assert_eq!(rx.recv().await, Err(mpsc::RecvError::Disconnected)); - /// # } - /// ``` - pub async fn recv(&mut self) -> Result { - poll_fn(|cx| self.poll_recv(cx)).await - } - - fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { - 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.buffer.register_receiver(cx.waker()); - - 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, - } - } + Err(TryRecvError::Empty) } } } + +struct Waiter { + grant: bool, + waker: Option, +} diff --git a/asyncband/src/mpsc/bounded/receiver.rs b/asyncband/src/mpsc/bounded/receiver.rs new file mode 100644 index 0000000..8b46e05 --- /dev/null +++ b/asyncband/src/mpsc/bounded/receiver.rs @@ -0,0 +1,161 @@ +// 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::fmt; +use std::future::poll_fn; +use std::mem; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use super::State; +use crate::internal::mutex::Mutex; +use crate::internal::wake_all; +use crate::internal::waker_batch::WakerBatch; +use crate::mpsc::RecvError; +use crate::mpsc::TryRecvError; + +/// The receiving endpoint of a bounded mpsc channel. +/// +/// Instances are created by the [`bounded`](crate::mpsc::bounded) function. Dropping the receiver +/// discards queued values and disconnects pending sends and reservations. +pub struct BoundedReceiver { + shared: Arc>>, +} + +impl fmt::Debug for BoundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedReceiver").finish_non_exhaustive() + } +} + +impl Drop for BoundedReceiver { + fn drop(&mut self) { + let (queue, recv_waker, wakers) = { + let mut state = self.shared.lock(); + state.receiver = false; + let queue = mem::take(&mut state.queue); + let recv_waker = state.recv_waker.take(); + let mut wakers = WakerBatch::new(); + while let Some((_, waiter)) = state.send_waiters.unlink_first_waiter(|_| true) { + if let Some(waker) = waiter.waker.take() { + wakers.push(waker); + } + } + (queue, recv_waker, wakers) + }; + // Local ownership also drains the queue if a wake or waker destructor unwinds. + wake_all(wakers.into_iter()); + drop(recv_waker); + drop(queue); + } +} + +impl BoundedReceiver { + pub(super) fn new(shared: Arc>>) -> Self { + Self { shared } + } + + /// Attempts to receive the next queued value without waiting for a new message. + /// + /// 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 + /// been dropped and all queued values have been consumed. + /// + /// # Examples + /// + /// ``` + /// use asyncband::mpsc::TryRecvError; + /// use asyncband::mpsc::bounded; + /// + /// let (tx, mut rx) = bounded(2); + /// tx.try_send("first").unwrap(); + /// tx.try_send("second").unwrap(); + /// + /// assert_eq!(rx.try_recv(), Ok("first")); + /// assert_eq!(rx.try_recv(), Ok("second")); + /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); + /// drop(tx); + /// assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); + /// ``` + pub fn try_recv(&mut self) -> Result { + let (value, wake) = self.shared.lock().pop()?; + if let Some(waker) = wake { + waker.wake(); + } + 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 + /// dropped. It returns [`RecvError::Disconnected`] only after all senders are gone and the + /// buffer has been drained. + /// + /// # Cancel safety + /// + /// Dropping a pending `recv` does not remove a message from the channel. A later `recv` call + /// can still observe the next queued value, so `recv` may safely be raced with other futures + /// in a selection construct. + /// + /// # Examples + /// + /// ``` + /// # #[tokio::main] + /// # async fn main() { + /// use asyncband::mpsc; + /// let (tx, mut rx) = mpsc::bounded(2); + /// + /// tx.send("first").await.unwrap(); + /// tx.send("second").await.unwrap(); + /// drop(tx); + /// + /// assert_eq!(rx.recv().await, Ok("first")); + /// assert_eq!(rx.recv().await, Ok("second")); + /// assert_eq!(rx.recv().await, Err(mpsc::RecvError::Disconnected)); + /// # } + /// ``` + pub async fn recv(&mut self) -> Result { + poll_fn(|cx| self.poll_recv(cx)).await + } + + fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll> { + let waker = cx.waker().clone(); + let mut state = self.shared.lock(); + match state.pop() { + Ok((value, wake)) => { + drop(state); + if let Some(waker) = wake { + waker.wake(); + } + Poll::Ready(Ok(value)) + } + Err(TryRecvError::Disconnected) => { + let old = state.recv_waker.take(); + drop(state); + drop(old); + Poll::Ready(Err(RecvError::Disconnected)) + } + Err(TryRecvError::Empty) => { + let old = state.recv_waker.replace(waker); + drop(state); + drop(old); + Poll::Pending + } + } + } +} diff --git a/asyncband/src/mpsc/bounded/ring.rs b/asyncband/src/mpsc/bounded/ring.rs deleted file mode 100644 index ace5aa7..0000000 --- a/asyncband/src/mpsc/bounded/ring.rs +++ /dev/null @@ -1,301 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::cell::UnsafeCell; -use std::hint::spin_loop; -use std::mem::MaybeUninit; -use std::sync::atomic::AtomicBool; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::sync::atomic::fence; -use std::task::Poll; -use std::task::Waker; - -use crate::internal::cache_padded::CachePadded; -use crate::internal::mutex::Mutex; -use crate::mpsc::TrySendError; - -pub struct Ring { - slots: Box<[Slot]>, - head: CachePadded, - tail: CachePadded, - // This flag is usually stable while the consumer advances head. Sharing head - // for notifications would make every producer track a constantly invalidated cache line. - receiver_waiting: CachePadded, - receiver: Mutex>, - capacity: usize, - one_lap: usize, - mark_bit: usize, -} - -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)), - receiver_waiting: CachePadded::new(AtomicBool::new(false)), - receiver: Mutex::new(None), - 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); - // 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); - } - } - - /// 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 `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); - } - } - - /// Registers the exclusive receiver, which must retry `pop` before returning Pending. - pub fn register_receiver(&self, waker: &Waker) { - let mut receiver = self.receiver.lock(); - let old_waker = if receiver.as_ref().is_some_and(|old| old.will_wake(waker)) { - None - } else { - // Only the receiver registers. Producers can take the old waker while we clone, - // but cannot install a replacement. Clone/drop callbacks may send into this channel. - drop(receiver); - let waker = waker.clone(); - receiver = self.receiver.lock(); - receiver.replace(waker) - }; - self.receiver_waiting.store(true, Ordering::Relaxed); - // Paired with the publisher's fence: either it sees this flag, or the receiver's - // subsequent pop sees its stamp. Taking a waker clears the flag under the same lock, - // so it cannot erase a newer registration without also taking responsibility for it. - fence(Ordering::SeqCst); - drop(receiver); - drop(old_waker); - } - - pub fn take_receiver_waker(&self) -> Option { - let mut receiver = self.receiver.lock(); - self.receiver_waiting.store(false, Ordering::Relaxed); - receiver.take() - } - - pub fn wake_receiver(&self) { - if let Some(waker) = self.take_receiver_waker() { - waker.wake(); - } - } - - /// Prevents subsequent sends from reserving slots. Already reserved slots still publish. - pub fn close(&self) { - self.tail.fetch_or(self.mark_bit, Ordering::SeqCst); - } - - /// Drops all values after closing, including values whose publication is still in progress. - /// - /// # Safety - /// - /// The queue must be closed. The caller must serialize all calls to `pop` and `drain`. - pub unsafe fn drain(&self) { - struct DrainRemaining<'a, T> { - ring: &'a Ring, - tail: usize, - } - impl Drop for DrainRemaining<'_, T> { - fn drop(&mut self) { - // SAFETY: The guard is scoped to the exclusive consumer's drain of a closed ring. - unsafe { self.ring.discard_until(self.tail) }; - } - } - - let tail = self.tail.load(Ordering::Relaxed); - debug_assert_ne!(tail & self.mark_bit, 0); - let remaining = DrainRemaining { - ring: self, - tail: tail & !self.mark_bit, - }; - // SAFETY: The caller guarantees exclusive consumer access. The guard finishes draining if - // a value's destructor panics, so messages that own senders cannot retain the closed ring. - unsafe { self.discard_until(remaining.tail) }; - } - - fn advance(&self, position: usize) -> usize { - 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) { - self.close(); - // SAFETY: The queue is closed and its exclusive borrow rules out concurrent access. - unsafe { self.drain() }; - } -} - -#[cfg(test)] -#[path = "ring_tests.rs"] -mod tests; diff --git a/asyncband/src/mpsc/bounded/ring_tests.rs b/asyncband/src/mpsc/bounded/ring_tests.rs deleted file mode 100644 index d6d5043..0000000 --- a/asyncband/src/mpsc/bounded/ring_tests.rs +++ /dev/null @@ -1,200 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::Arc; -use std::sync::atomic::AtomicUsize; -use std::sync::atomic::Ordering; -use std::task::Poll; -use std::thread; - -use super::Ring; -use super::TrySendError; - -#[test] -fn bounded_queue_preserves_capacity_and_fifo_order() { - let queue = Ring::new(3); - for value in 0..3 { - assert!(queue.try_push(value).is_ok()); - } - assert!(matches!(queue.try_push(3), Err(TrySendError::Full(3)))); - for value in 0..3 { - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); - } - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(None)); - - for value in 3..12 { - assert!(queue.try_push(value).is_ok()); - // SAFETY: This thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(value))); - } -} - -#[test] -fn bounded_queue_does_not_report_empty_behind_an_unpublished_head() { - let queue = Ring::new(2); - - // Pause a synthetic producer after reserving and initializing slot 0, before publishing - // its stamp. Another producer can finish sending into slot 1 in the meantime. - queue.tail.store(1, Ordering::SeqCst); - let slot = &queue.slots[0]; - // SAFETY: advancing the tail reserved this initially empty slot for the synthetic producer. - unsafe { (*slot.value.get()).write(1) }; - let later_send = queue.try_push(2); - // SAFETY: This thread is the only consumer, even while a producer is unpublished. - let receive = unsafe { queue.pop() }; - - // Finish publication before asserting so even a failed assertion can safely drop the queue. - slot.stamp.store(1, Ordering::Release); - assert!(later_send.is_ok()); - assert_eq!(receive, Poll::Pending); - // SAFETY: This thread is the only consumer. - unsafe { - assert_eq!(queue.pop(), Poll::Ready(Some(1))); - assert_eq!(queue.pop(), Poll::Ready(Some(2))); - assert_eq!(queue.pop(), Poll::Ready(None)); - } -} - -#[test] -fn unpublished_reservations_count_toward_capacity() { - let queue = Arc::new(Ring::new(1)); - queue.tail.store(queue.one_lap, Ordering::SeqCst); - let (done, completed) = std::sync::mpsc::channel(); - let producer = { - let queue = queue.clone(); - thread::spawn(move || done.send(queue.try_push(2)).unwrap()) - }; - #[cfg(not(miri))] - let result = completed.recv_timeout(std::time::Duration::from_secs(10)); - // Miri reports a deadlock directly instead of relying on an interpretation-time deadline. - #[cfg(miri)] - let result = completed.recv(); - // Finish the synthetic reservation even if the other producer stalled. This lets the - // worker and the ring's destructor finish before the failure is reported. - let slot = &queue.slots[0]; - // SAFETY: Advancing the tail above exclusively reserved the initially empty slot. - unsafe { (*slot.value.get()).write(1) }; - slot.stamp.store(1, Ordering::Release); - producer.join().unwrap(); - assert!(matches!(result, Ok(Err(TrySendError::Full(2))))); - // SAFETY: Both producers have finished and this thread is the only consumer. - assert_eq!(unsafe { queue.pop() }, Poll::Ready(Some(1))); -} - -#[test] -fn bounded_queue_coordinates_multiple_producers() { - let queue = Arc::new(Ring::new(4)); - let producers: Vec<_> = (0..2) - .map(|producer| { - let queue = queue.clone(); - thread::spawn(move || { - for offset in 0..32 { - let mut value = producer * 32 + offset; - loop { - match queue.try_push(value) { - Ok(()) => break, - Err(TrySendError::Full(returned)) => { - value = returned; - thread::yield_now(); - } - Err(TrySendError::Disconnected(_)) => panic!("queue disconnected"), - } - } - } - }) - }) - .collect(); - - let mut values = Vec::new(); - while values.len() < 64 { - // SAFETY: Worker threads only push; this thread is the only consumer. - if let Poll::Ready(Some(value)) = unsafe { queue.pop() } { - values.push(value); - } else { - thread::yield_now(); - } - } - for producer in producers { - producer.join().unwrap(); - } - values.sort_unstable(); - assert_eq!(values, (0..64).collect::>()); -} - -#[test] -fn bounded_queue_discards_wrapped_values_once_after_receiver_disconnect() { - // This has no owning fields, so a buggy second drop remains observable as count == 2 - // instead of invalidating the tracker first. - struct DropSpy<'a>(&'a AtomicUsize); - - impl<'a> Drop for DropSpy<'a> { - fn drop(&mut self) { - self.0.fetch_add(1, Ordering::Relaxed); - } - } - - // Declare this before `queue` so the counters outlive values held by the queue. - let drops = [ - AtomicUsize::new(0), - AtomicUsize::new(0), - AtomicUsize::new(0), - AtomicUsize::new(0), - ]; - let queue = Ring::new(3); - - // Positions: 0, 1, 2 (then tail wraps to 8). - for counter in &drops[..3] { - assert!(queue.try_push(DropSpy(counter)).is_ok()); - } - - // Free slot 0, then reuse it on the next lap at position 8. - // SAFETY: This thread is the only consumer. - let popped = unsafe { queue.pop() }; - assert!(matches!(popped, Poll::Ready(Some(_)))); - drop(popped); - assert_eq!(drops[0].load(Ordering::Relaxed), 1); - assert!(queue.try_push(DropSpy(&drops[3])).is_ok()); - - // The pending range is positions 1 -> 2 -> 8 -> 9, not a contiguous integer range. - assert_eq!(queue.head.load(Ordering::Relaxed), 1); - assert_eq!(queue.tail.load(Ordering::Relaxed), queue.one_lap + 1); - - queue.close(); - // SAFETY: The queue is closed and this thread is the only consumer. - unsafe { queue.drain() }; - - // `discard_until` must dispose every value exactly once, including position 8. - for (value, counter) in drops.iter().enumerate() { - assert_eq!( - counter.load(Ordering::Relaxed), - 1, - "value {value} was dropped an unexpected number of times" - ); - } - - // Queue Drop calls discard_until again; it must see head == tail and not redrop. - drop(queue); - for (value, counter) in drops.iter().enumerate() { - assert_eq!( - counter.load(Ordering::Relaxed), - 1, - "value {value} was dropped more than once" - ); - } -} diff --git a/asyncband/src/mpsc/bounded/sender.rs b/asyncband/src/mpsc/bounded/sender.rs new file mode 100644 index 0000000..392fa15 --- /dev/null +++ b/asyncband/src/mpsc/bounded/sender.rs @@ -0,0 +1,293 @@ +// 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::fmt; +use std::future::poll_fn; +use std::mem; +use std::sync::Arc; +use std::task::Context; +use std::task::Poll; + +use super::State; +use super::Waiter; +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaiterId; +use crate::mpsc::SendError; +use crate::mpsc::TrySendError; + +/// The sending endpoint of a bounded mpsc channel. +/// +/// Instances are created by the [`bounded`](crate::mpsc::bounded) function. +pub struct BoundedSender { + shared: Arc>>, +} + +impl Clone for BoundedSender { + fn clone(&self) -> Self { + self.shared.lock().senders += 1; + BoundedSender { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for BoundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BoundedSender").finish_non_exhaustive() + } +} + +impl Drop for BoundedSender { + fn drop(&mut self) { + let wake = { + let mut state = self.shared.lock(); + state.senders -= 1; + if state.senders == 0 { + state.recv_waker.take() + } else { + None + } + }; + if let Some(waker) = wake { + waker.wake(); + } + } +} + +impl BoundedSender { + pub(super) fn new(shared: Arc>>) -> Self { + Self { shared } + } + + /// Sends a message, waiting until the channel has capacity when necessary. + /// + /// If the receiver has been dropped, the returned error contains `value`. + /// + /// # Cancel safety + /// + /// 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 [`try_send`](Self::try_send) when + /// the caller must retain ownership if capacity is unavailable, or [`reserve`](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, + }; + 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 loses its place in the wait queue. If capacity has already + /// been granted, it is released to the next waiter or made available to a new sender. + /// + /// # 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<()>> { + let mut reserve = Reserve { + shared: &self.shared, + waiter: None, + }; + poll_fn(|cx| reserve.poll(cx)).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.shared.lock().acquire()?; + Ok(Permit { + shared: &self.shared, + }) + } + + /// Attempts to send a message without waiting for capacity. + /// + /// A full buffer returns [`TrySendError::Full`], while a dropped receiver returns + /// [`TrySendError::Disconnected`]. Both errors return ownership of the unsent value. + /// + /// # Examples + /// + /// ``` + /// use asyncband::mpsc::TrySendError; + /// use asyncband::mpsc::bounded; + /// + /// let (tx, mut rx) = bounded(1); + /// tx.try_send(10).unwrap(); + /// assert_eq!(tx.try_send(20), Err(TrySendError::Full(20))); + /// + /// assert_eq!(rx.try_recv(), Ok(10)); + /// tx.try_send(20).unwrap(); + /// drop(rx); + /// assert_eq!(tx.try_send(30), Err(TrySendError::Disconnected(30))); + /// ``` + pub fn try_send(&self, value: T) -> Result<(), TrySendError> { + let mut state = self.shared.lock(); + match state.acquire() { + Ok(()) => { + state.queue.push_back(value); + let wake = state.recv_waker.take(); + drop(state); + if let Some(waker) = wake { + waker.wake(); + } + Ok(()) + } + 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> { + shared: &'a Mutex>, +} + +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 permit, 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> { + let mut state = self.shared.lock(); + if !state.receiver { + return Err(SendError::new(value)); + } + state.queue.push_back(value); + // The queued message now owns capacity, even if the wake callback panics. + mem::forget(self); + let wake = state.recv_waker.take(); + drop(state); + if let Some(waker) = wake { + waker.wake(); + } + Ok(()) + } +} + +impl Drop for Permit<'_, T> { + fn drop(&mut self) { + let wake = self.shared.lock().release(); + if let Some(waker) = wake { + waker.wake(); + } + } +} + +struct Reserve<'a, T> { + shared: &'a Mutex>, + waiter: Option, +} + +impl<'a, T> Reserve<'a, T> { + fn poll(&mut self, cx: &mut Context<'_>) -> Poll, SendError<()>>> { + let waker = cx.waker().clone(); + let mut state = self.shared.lock(); + if !state.receiver { + return Poll::Ready(Err(SendError::new(()))); + } + if let Some(index) = self.waiter { + let waiter = state.send_waiters.waiter_mut(index); + if waiter.grant { + let waiter = state.send_waiters.remove_unlinked_waiter(index); + self.waiter = None; + let permit = Permit { + shared: self.shared, + }; + drop(state); + drop(waiter); + drop(waker); + return Poll::Ready(Ok(permit)); + } + let old = waiter.waker.replace(waker); + drop(state); + drop(old); + return Poll::Pending; + } + if state.available != 0 { + state.available -= 1; + let permit = Permit { + shared: self.shared, + }; + drop(state); + drop(waker); + return Poll::Ready(Ok(permit)); + } + self.waiter = Some(state.send_waiters.push_back(Waiter { + grant: false, + waker: Some(waker), + })); + Poll::Pending + } +} + +impl Drop for Reserve<'_, T> { + fn drop(&mut self) { + let Some(index) = self.waiter else { return }; + let (waiter, wake) = { + let mut state = self.shared.lock(); + state.send_waiters.unlink_waiter(index, |_| true); + let waiter = state.send_waiters.remove_unlinked_waiter(index); + let wake = if waiter.grant { state.release() } else { None }; + (waiter, wake) + }; + if let Some(waker) = wake { + waker.wake(); + } + drop(waiter); + } +} diff --git a/asyncband/src/mpsc/bounded/waiters.rs b/asyncband/src/mpsc/bounded/waiters.rs deleted file mode 100644 index 3a34e4f..0000000 --- a/asyncband/src/mpsc/bounded/waiters.rs +++ /dev/null @@ -1,158 +0,0 @@ -// Licensed to the Apache Software Foundation (ASF) under one -// or more contributor license agreements. See the NOTICE file -// distributed with this work for additional information -// regarding copyright ownership. The ASF licenses this file -// to you under the Apache License, Version 2.0 (the -// "License"); you may not use this file except in compliance -// with the License. You may obtain a copy of the License at -// -// http://www.apache.org/licenses/LICENSE-2.0 -// -// Unless required by applicable law or agreed to in writing, -// software distributed under the License is distributed on an -// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY -// KIND, either express or implied. See the License for the -// specific language governing permissions and limitations -// under the License. - -use std::sync::atomic::AtomicBool; -use std::sync::atomic::Ordering; -use std::task::Waker; - -use crate::internal::mutex::Mutex; -use crate::internal::waitlist::WaitList; -use crate::internal::waitlist::WaiterId; -use crate::internal::wake_all; -use crate::internal::waker_batch::WakerBatch; - -// A wake grants a retry, not a capacity permit. Keeping notified nodes until their future -// consumes the notification lets cancellation pass an unused retry to the next sender. -pub struct SendWaiters { - waiting: AtomicBool, - queue: Mutex>>, -} - -impl SendWaiters { - pub fn new() -> Self { - Self { - waiting: AtomicBool::new(false), - queue: Mutex::new(WaitList::new()), - } - } - - pub fn waiter(&self) -> SendWaiter<'_> { - SendWaiter { - waiters: self, - index: None, - } - } - - pub fn notify_one(&self) { - // The receiver publishes its head with SeqCst before checking this flag. Registration - // publishes the flag with SeqCst before rechecking capacity (which uses a SeqCst fence). - // Either the receiver sees the registration or the sender sees the released capacity. - if !self.waiting.load(Ordering::SeqCst) { - return; - } - let waker = { - let mut queue = self.queue.lock(); - let waker = queue - .unlink_first_waiter(|_| true) - .and_then(|(_, waker)| waker.take()); - self.waiting.store(!queue.is_empty(), Ordering::SeqCst); - waker - }; - if let Some(waker) = waker { - waker.wake(); - } - } - - pub fn notify_all(&self) { - let mut wakers = WakerBatch::new(); - { - let mut queue = self.queue.lock(); - while let Some((_, waker)) = queue.unlink_first_waiter(|_| true) { - if let Some(waker) = waker.take() { - wakers.push(waker); - } - } - self.waiting.store(false, Ordering::SeqCst); - } - wake_all(wakers.into_iter()); - } -} - -pub struct SendWaiter<'a> { - waiters: &'a SendWaiters, - index: Option, -} - -impl SendWaiter<'_> { - // The caller must retry sending after registration, before returning Pending. - pub fn register(&mut self, waker: &Waker) { - let mut new_waker = None; - loop { - let mut queue = self.waiters.queue.lock(); - if let Some(index) = self.index { - if queue - .waiter_mut(index) - .as_ref() - .is_some_and(|current| current.will_wake(waker)) - { - return; - } - } - let Some(waker) = new_waker.take() else { - // Waker callbacks may reenter the channel, including clone and drop callbacks. - drop(queue); - new_waker = Some(waker.clone()); - continue; - }; - let old_waker = if let Some(index) = self.index { - let node = queue.waiter_mut(index); - if node.is_some() { - node.replace(waker) - } else { - queue.remove_unlinked_waiter(index); - self.index = Some(queue.push_back(Some(waker))); - None - } - } else { - self.index = Some(queue.push_back(Some(waker))); - None - }; - self.waiters.waiting.store(true, Ordering::SeqCst); - drop(queue); - drop(old_waker); - return; - } - } - - pub fn finish(&mut self) { - if let Some(index) = self.index.take() { - drop(self.remove(index)); - } - } - - fn remove(&self, index: WaiterId) -> Option { - let mut queue = self.waiters.queue.lock(); - queue.unlink_waiter(index, |_| true); - let waker = queue.remove_unlinked_waiter(index); - self.waiters - .waiting - .store(!queue.is_empty(), Ordering::SeqCst); - waker - } -} - -impl Drop for SendWaiter<'_> { - fn drop(&mut self) { - if let Some(index) = self.index.take() { - let waker = self.remove(index); - if waker.is_none() { - self.waiters.notify_one(); - } - drop(waker); - } - } -} diff --git a/asyncband/src/mpsc/error.rs b/asyncband/src/mpsc/error.rs index 586e4e9..69e545e 100644 --- a/asyncband/src/mpsc/error.rs +++ b/asyncband/src/mpsc/error.rs @@ -18,29 +18,29 @@ use std::any::type_name; use std::fmt; -/// A send failed because the receiving endpoint has been dropped. +/// A send or capacity reservation failed because the receiver has been dropped. /// -/// Returned from [`UnboundedSender::send`] or [`BoundedSender::send`] if the -/// corresponding [`UnboundedReceiver`] or [`BoundedReceiver`] has already been -/// dropped. +/// Returned by [`UnboundedSender::send`], [`BoundedSender::send`], [`Permit::send`], and +/// [`reserve`]. /// -/// The rejected message remains available through [`SendError::as_inner`] or -/// [`SendError::into_inner`]. +/// A failed send retains the unsent message. A failed reservation carries `()` because no +/// message has been provided yet. Access the value with [`as_inner`](Self::as_inner) or +/// [`into_inner`](Self::into_inner). /// /// [`UnboundedSender::send`]: crate::mpsc::UnboundedSender::send /// [`BoundedSender::send`]: crate::mpsc::BoundedSender::send -/// [`UnboundedReceiver`]: crate::mpsc::UnboundedReceiver -/// [`BoundedReceiver`]: crate::mpsc::BoundedReceiver +/// [`reserve`]: crate::mpsc::BoundedSender::reserve +/// [`Permit::send`]: crate::mpsc::Permit::send #[derive(Clone, PartialEq, Eq)] pub struct SendError(T); impl SendError { - /// Get a reference to the message that failed to be sent. + /// Gets a reference to the unsent message, or `()` for a failed reservation. pub fn as_inner(&self) -> &T { &self.0 } - /// Consumes the error and returns the message that failed to be sent. + /// Consumes the error and returns the unsent message, or `()` for a failed reservation. pub fn into_inner(self) -> T { self.0 } @@ -65,24 +65,28 @@ impl fmt::Debug for SendError { impl std::error::Error for SendError {} -/// A non-blocking send could not accept its message. +/// An attempt to send or reserve capacity failed. +/// +/// Returned by [`try_send`](crate::mpsc::BoundedSender::try_send) and +/// [`try_reserve`](crate::mpsc::BoundedSender::try_reserve). A failed send retains the unsent +/// message; a failed reservation carries `()` because no message has been provided yet. #[derive(Clone, PartialEq, Eq)] pub enum TrySendError { - /// The channel is full, so the message cannot be sent without waiting for capacity. + /// No capacity is available for sending or reserving a message. Full(T), - /// The receiver has been dropped, so the message can never be received. + /// The receiver has been dropped. Disconnected(T), } impl TrySendError { - /// Gets a reference to the message that failed to be sent. + /// Gets a reference to the unsent message, or `()` for a failed reservation. pub fn as_inner(&self) -> &T { match self { TrySendError::Full(msg) | TrySendError::Disconnected(msg) => msg, } } - /// Consumes the error and returns the message that failed to be sent. + /// Consumes the error and returns the unsent message, or `()` for a failed reservation. pub fn into_inner(self) -> T { match self { TrySendError::Full(msg) | TrySendError::Disconnected(msg) => msg, diff --git a/asyncband/src/mpsc/mod.rs b/asyncband/src/mpsc/mod.rs index 040c98b..bad23f0 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/asyncband/src/mpsc/unbounded/buffer.rs b/asyncband/src/mpsc/unbounded/buffer.rs index 7c11449..8df2640 100644 --- a/asyncband/src/mpsc/unbounded/buffer.rs +++ b/asyncband/src/mpsc/unbounded/buffer.rs @@ -38,10 +38,10 @@ impl Buffer { } fn segment_capacity() -> usize { - if mem::size_of::() == 0 { + if size_of::() == 0 { return usize::MAX; } - let limit = (SEGMENT_BYTES / mem::size_of::()).max(1); + let limit = (SEGMENT_BYTES / size_of::()).max(1); // Power-of-two limits let VecDeque grow naturally without exceeding the segment budget. 1 << (usize::BITS - 1 - limit.leading_zeros()) } @@ -65,9 +65,7 @@ impl Buffer { // Keep one empty segment for the next producer rollover. Every other consumed // segment is released, so retained payload storage does not track peak occupancy. self.spare = mem::replace(batch, sealed); - if self.sealed.is_empty() - && self.sealed.capacity() * mem::size_of::>() > 1024 - { + if self.sealed.is_empty() && self.sealed.capacity() * size_of::>() > 1024 { self.sealed = VecDeque::new(); } } else if !self.writable.is_empty() { @@ -78,7 +76,7 @@ impl Buffer { } pub fn pop_batch(batch: &mut VecDeque) -> T { - if batch.len() == 1 && batch.capacity().saturating_mul(mem::size_of::()) > SEGMENT_BYTES { + if batch.len() == 1 && batch.capacity().saturating_mul(size_of::()) > SEGMENT_BYTES { // Retire the allocation on the last value, outside the inbox lock. Keep this as a tail // expression to avoid intermediate storage for large inline values. mem::take(batch).pop_front() diff --git a/benchmarks/ecosystem/mpsc/adapters.rs b/benchmarks/ecosystem/mpsc/adapters.rs index 506f167..f9fe81b 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 d388cc6..44d2cd0 100644 --- a/benchmarks/ecosystem/mpsc/bounded.rs +++ b/benchmarks/ecosystem/mpsc/bounded.rs @@ -33,7 +33,11 @@ use super::support::RepeatedBatch; use super::support::RepeatedTasks; use crate::support::bench_context; -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +// Nanosecond-scale benches pin `sample_size`: divan's auto-tuning starts at 1 +// iteration per sample and stops once a sample exceeds 100x timer precision, +// so a cold first call (lazy initialization, cache misses) can end tuning +// immediately and quantize every sample to one timer tick (41 ns on macOS). +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume], sample_size = 512)] fn try_round_trip(bencher: Bencher) { let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); @@ -43,7 +47,7 @@ fn try_round_trip(bencher: Bencher) { }); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume], sample_size = 256)] fn ready_round_trip(bencher: Bencher) { let mut context = bench_context(); let (sender, mut receiver) = C::channel(BOUNDED_CAPACITY); @@ -80,7 +84,7 @@ fn sustained(bencher: Bencher, producer_count: usize) { bencher.bench_local(|| batch.run()); } -#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume])] +#[divan::bench(types = [Asyncband, Tokio, AsyncChannel, Flume], sample_size = 2048)] fn clone_drop_sender(bencher: Bencher) { let (sender, _receiver) = C::channel(BOUNDED_CAPACITY); bencher.bench_local(|| drop(black_box(sender.clone()))); @@ -106,7 +110,7 @@ fn sustained_capacity( #[divan::bench( types = [Asyncband, Tokio, AsyncChannel, Flume], consts = [1, 64, 4096], - args = [(1, 0), (4, 0), (1, 4), (4, 4), (8, 4)], + args = [(1, 0), (8, 0), (1, 4), (8, 4)], sample_count = 50, sample_size = 1, counter = ItemsCount::new(BATCH_MESSAGES), @@ -119,3 +123,52 @@ fn scheduled( batch.run(); bencher.bench_local(|| batch.run()); } + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + consts = [1, 64, 4096], + args = [(1, 0), (8, 0), (1, 4), (8, 4)], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn scheduled_inline, const CAPACITY: usize>( + bencher: Bencher, + (producers, workers): (usize, usize), +) { + let mut batch = RepeatedTasks::>::new(producers, workers); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + consts = [1, 64, 4096], + args = [1, 8], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn external_receiver(bencher: Bencher, producers: usize) { + let mut batch = RepeatedTasks::>::external_receiver(producers, 4); + batch.run(); + bencher.bench_local(|| batch.run()); +} + +#[divan::bench( + types = [Asyncband, Tokio, AsyncChannel, Flume], + consts = [1, 64, 4096], + args = [1, 8], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn external_receiver_inline, const CAPACITY: usize>( + bencher: Bencher, + producers: usize, +) { + let mut batch = + RepeatedTasks::>::external_receiver(producers, 4); + batch.run(); + bencher.bench_local(|| batch.run()); +} diff --git a/benchmarks/ecosystem/mpsc/mod.rs b/benchmarks/ecosystem/mpsc/mod.rs index dd09282..ed522cc 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 0000000..32defce --- /dev/null +++ b/benchmarks/ecosystem/mpsc/reservation.rs @@ -0,0 +1,123 @@ +// 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); + } +} + +// `sample_size` is pinned for the reason documented in `bounded.rs`. +#[divan::bench(types = [Asyncband, Tokio], sample_size = 512)] +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], sample_size = 1024)] +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 Message = usize; + 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 = [1, 64, 4096], + args = [(1, 0), (8, 0), (1, 4), (8, 4)], + sample_count = 50, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn scheduled( + bencher: Bencher, + (producers, workers): (usize, usize), +) { + let mut batch = RepeatedTasks::>::new(producers, workers); + batch.run(); + bencher.bench_local(|| batch.run()); +} diff --git a/benchmarks/ecosystem/mpsc/support.rs b/benchmarks/ecosystem/mpsc/support.rs index 3e777e0..1f1957e 100644 --- a/benchmarks/ecosystem/mpsc/support.rs +++ b/benchmarks/ecosystem/mpsc/support.rs @@ -33,20 +33,55 @@ pub const BOUNDED_CAPACITY: usize = 64; pub const BATCH_MESSAGES: usize = 16_384; pub const PRODUCER_COUNTS: &[usize] = &[1, 2, 4, 8]; +pub trait Message: Send + 'static { + fn new(sequence: usize) -> Self; + fn sequence(self) -> usize; +} + +impl Message for usize { + fn new(sequence: usize) -> Self { + sequence + } + + fn sequence(self) -> usize { + self + } +} + +impl Message for [u8; 1024] { + fn new(sequence: usize) -> Self { + let mut value = [1; 1024]; + value[..size_of::()].copy_from_slice(&sequence.to_le_bytes()); + value + } + + fn sequence(self) -> usize { + let value = black_box(self); + assert_eq!(value[1023], 1); + usize::from_le_bytes(value[..size_of::()].try_into().unwrap()) + } +} + pub trait ConcurrentMpsc: Send + Sync + 'static { + type Message: Message; type Sender: Clone + Send + Sync + 'static; type Receiver: Send + 'static; fn channel() -> (Self::Sender, Self::Receiver); - fn send(sender: &Self::Sender, value: usize); - fn recv(receiver: &mut Self::Receiver) -> usize; - fn send_async(sender: &Self::Sender, value: usize) -> impl Future + Send; - fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; + fn send(sender: &Self::Sender, value: Self::Message); + fn recv(receiver: &mut Self::Receiver) -> Self::Message; + fn send_async(sender: &Self::Sender, value: Self::Message) -> impl Future + Send; + fn recv_async(receiver: &mut Self::Receiver) -> impl Future + Send; } -pub struct Bounded(PhantomData); +pub struct Bounded( + PhantomData (C, T)>, +); -impl ConcurrentMpsc for Bounded { +impl, const CAPACITY: usize, T: Message> ConcurrentMpsc + for Bounded +{ + type Message = T; type Receiver = C::Receiver; type Sender = C::Sender; @@ -54,19 +89,19 @@ impl ConcurrentMpsc for Bounded usize { + async fn recv_async(receiver: &mut Self::Receiver) -> T { C::recv_async(receiver).await } - fn recv(receiver: &mut Self::Receiver) -> usize { + fn recv(receiver: &mut Self::Receiver) -> T { C::recv_blocking(receiver) } } @@ -74,6 +109,7 @@ impl ConcurrentMpsc for Bounded(PhantomData); impl ConcurrentMpsc for Unbounded { + type Message = usize; type Receiver = C::Receiver; type Sender = C::Sender; @@ -98,13 +134,13 @@ impl ConcurrentMpsc for Unbounded { } } -pub struct ConcurrentBatch { +pub struct ConcurrentBatch> { receiver: C::Receiver, start: Arc, workers: Vec>, } -impl ConcurrentBatch { +impl> ConcurrentBatch { pub fn new(producer_count: usize) -> Self { assert_eq!(BATCH_MESSAGES % producer_count, 0); @@ -143,7 +179,7 @@ impl ConcurrentBatch { } } -impl Drop for ConcurrentBatch { +impl> Drop for ConcurrentBatch { fn drop(&mut self) { let panicking = thread::panicking(); for worker in self.workers.drain(..) { @@ -156,14 +192,14 @@ impl Drop for ConcurrentBatch { } // Reuse worker threads and channel storage so steady-state samples exclude thread creation. -pub struct RepeatedBatch { +pub struct RepeatedBatch> { receiver: C::Receiver, start: Arc, stop: Arc, workers: Vec>, } -impl RepeatedBatch { +impl> RepeatedBatch { pub fn new(producer_count: usize) -> Self { assert_eq!(BATCH_MESSAGES % producer_count, 0); let (sender, receiver) = C::channel(); @@ -208,7 +244,7 @@ impl RepeatedBatch { } } -impl Drop for RepeatedBatch { +impl> Drop for RepeatedBatch { fn drop(&mut self) { self.stop.store(true, Ordering::Release); self.start.wait(); @@ -218,11 +254,21 @@ impl Drop for RepeatedBatch { } } -// Exercise executor wakeups as well as channel traffic. Reuse tasks, threads, and channel storage -// across samples; a current-thread runtime also exposes polling that monopolizes the executor. +// A spawned receiver shares the executor's scheduling with producers. Keeping the receiver in +// block_on instead measures worker-to-caller thread handoffs, which is a separate workload. +enum Receiver { + Task { + start: Arc, + completed: tokio::sync::mpsc::UnboundedReceiver, + }, + External(C::Receiver), +} + +// Reuse every task and the channel. The small control exchange happens once per 16,384-message +// batch; it never forwards measured messages. Both payload sizes use this same start protocol. pub struct RepeatedTasks { runtime: tokio::runtime::Runtime, - receiver: C::Receiver, + receiver: Receiver, start: Arc, stop: Arc, workers: Vec>, @@ -230,6 +276,18 @@ pub struct RepeatedTasks { impl RepeatedTasks { pub fn new(producer_count: usize, worker_threads: usize) -> Self { + Self::with_receiver(producer_count, worker_threads, false) + } + + pub fn external_receiver(producer_count: usize, worker_threads: usize) -> Self { + Self::with_receiver(producer_count, worker_threads, true) + } + + fn with_receiver( + producer_count: usize, + worker_threads: usize, + external_receiver: bool, + ) -> Self { assert_eq!(BATCH_MESSAGES % producer_count, 0); let runtime = if worker_threads == 0 { tokio::runtime::Builder::new_current_thread() @@ -241,11 +299,11 @@ impl RepeatedTasks { .build() .unwrap() }; - let (sender, receiver) = C::channel(); + let (sender, mut receiver) = C::channel(); let start = Arc::new(tokio::sync::Barrier::new(producer_count + 1)); let stop = Arc::new(AtomicBool::new(false)); let messages_per_producer = BATCH_MESSAGES / producer_count; - let workers = (0..producer_count) + let mut workers: Vec<_> = (0..producer_count) .map(|producer| { let sender = sender.clone(); let start = start.clone(); @@ -258,13 +316,38 @@ impl RepeatedTasks { } let first = producer * messages_per_producer; for offset in 0..messages_per_producer { - C::send_async(&sender, black_box(first + offset)).await; + C::send_async(&sender, black_box(C::Message::new(first + offset))) + .await; } } }) }) .collect(); drop(sender); + let receiver = if external_receiver { + Receiver::External(receiver) + } else { + let request = Arc::new(tokio::sync::Notify::new()); + let (completed_tx, completed) = tokio::sync::mpsc::unbounded_channel(); + let request_rx = request.clone(); + let start = start.clone(); + let stop = stop.clone(); + workers.push(runtime.spawn(async move { + loop { + request_rx.notified().await; + start.wait().await; + if stop.load(Ordering::Acquire) { + break; + } + let checksum = receive_batch::(&mut receiver).await; + completed_tx.send(checksum).unwrap(); + } + })); + Receiver::Task { + start: request, + completed, + } + }; Self { runtime, receiver, @@ -276,13 +359,16 @@ impl RepeatedTasks { pub fn run(&mut self) -> usize { self.runtime.block_on(async { - self.start.wait().await; - let mut checksum = 0usize; - for _ in 0..BATCH_MESSAGES { - checksum = checksum.wrapping_add(C::recv_async(&mut self.receiver).await); + match &mut self.receiver { + Receiver::Task { start, completed } => { + start.notify_one(); + completed.recv().await.expect("benchmark receiver panicked") + } + Receiver::External(receiver) => { + self.start.wait().await; + receive_batch::(receiver).await + } } - assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); - black_box(checksum) }) } } @@ -291,10 +377,24 @@ impl Drop for RepeatedTasks { fn drop(&mut self) { self.stop.store(true, Ordering::Release); self.runtime.block_on(async { - self.start.wait().await; + match &self.receiver { + Receiver::Task { start, .. } => start.notify_one(), + Receiver::External(_) => { + self.start.wait().await; + } + } for worker in self.workers.drain(..) { worker.await.expect("benchmark producer panicked"); } }); } } + +async fn receive_batch(receiver: &mut C::Receiver) -> usize { + let mut checksum = 0usize; + for _ in 0..BATCH_MESSAGES { + checksum = checksum.wrapping_add(C::recv_async(receiver).await.sequence()); + } + assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); + black_box(checksum) +} diff --git a/tests-integration/tests/mpsc_test/backpressure.rs b/tests-integration/tests/mpsc_test/backpressure.rs index ba21cfc..2db9919 100644 --- a/tests-integration/tests/mpsc_test/backpressure.rs +++ b/tests-integration/tests/mpsc_test/backpressure.rs @@ -67,6 +67,10 @@ fn cancelling_a_sender_preserves_capacity_and_notifies_the_next_waiter() { assert_eq!(second_wakes.count(), 0); } drop(first); + // Even a granted send must not publish its value until it is polled to completion. + if cancel_after_notification { + assert_eq!(rx.try_recv(), Err(mpsc::TryRecvError::Empty)); + } if !cancel_after_notification { assert_eq!(first_wakes.count(), 0); assert_eq!(second_wakes.count(), 0); diff --git a/tests-integration/tests/mpsc_test/callbacks.rs b/tests-integration/tests/mpsc_test/callbacks.rs index 2f654fc..a069f00 100644 --- a/tests-integration/tests/mpsc_test/callbacks.rs +++ b/tests-integration/tests/mpsc_test/callbacks.rs @@ -35,6 +35,7 @@ use super::support::assert_completes_without_deadlock; use super::support::expect_ready; use super::support::poll_with; use super::support::waker_on_clone; +use super::support::waker_on_drop; struct HoldSender { _sender: S, @@ -103,6 +104,65 @@ fn bounded_send_rechecks_capacity_freed_by_waker_clone() { }); } +#[cfg(panic = "unwind")] +#[test] +fn bounded_send_returns_capacity_when_an_unused_waker_panics_on_drop() { + use std::task::RawWaker; + use std::task::RawWakerVTable; + + struct Callbacks { + receiver: Mutex>, + drop_panics: AtomicBool, + } + + unsafe fn clone(data: *const ()) -> RawWaker { + let pointer = data.cast::(); + // SAFETY: The input waker owns a live Arc. The returned clone gains its own reference. + unsafe { + assert_eq!((*pointer).receiver.lock().unwrap().try_recv(), Ok(1)); + Arc::increment_strong_count(pointer); + } + RawWaker::new(data, &VTABLE) + } + + unsafe fn release(data: *const ()) { + // SAFETY: Consumes this waker's Arc reference, including if the callback unwinds. + let callbacks = unsafe { Arc::from_raw(data.cast::()) }; + assert!( + !callbacks.drop_panics.swap(false, Ordering::Relaxed), + "unused cloned waker panicked on drop" + ); + } + + // A raw vtable is needed to run callbacks for cloning and dropping each waker reference. + static VTABLE: RawWakerVTable = RawWakerVTable::new(clone, release, |_| {}, release); + + let (tx, rx) = mpsc::bounded(1); + tx.try_send(1).unwrap(); + let callbacks = Arc::new(Callbacks { + receiver: Mutex::new(rx), + drop_panics: AtomicBool::new(true), + }); + let data = Arc::into_raw(callbacks.clone()).cast(); + // SAFETY: Every waker owns an Arc reference. All callbacks preserve ownership and use only + // synchronized state; wake_by_ref does not touch the reference count. + let waker = unsafe { Waker::from_raw(RawWaker::new(data, &VTABLE)) }; + + let mut send = Box::pin(tx.send(2)); + assert!( + std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + poll_with(send.as_mut(), &waker) + })) + .is_err() + ); + drop(send); + let mut receiver = callbacks.receiver.lock().unwrap(); + assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); + tx.try_send(3) + .expect("unwinding must return acquired capacity"); + assert_eq!(receiver.try_recv(), Ok(3)); +} + #[test] fn receive_rechecks_messages_sent_by_waker_clone() { assert_completes_without_deadlock(|| { @@ -155,35 +215,68 @@ fn wake_callbacks_can_send_into_the_same_channel() { } #[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) {} - } +fn bounded_waiter_waker_replacement_and_cancellation_can_reenter() { + assert_completes_without_deadlock(|| { + for replace in [false, true] { + let (tx, mut rx) = mpsc::bounded(1); + tx.try_send(0).unwrap(); + let drops = Arc::new(AtomicUsize::new(0)); + let waker = waker_on_drop({ + let tx = tx.clone(); + let drops = drops.clone(); + move || { + assert_eq!(tx.try_send(9), Err(mpsc::TrySendError::Full(9))); + drops.fetch_add(1, Ordering::Relaxed); + } + }); + let mut send = Box::pin(tx.send(1)); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + drop(waker); + if replace { + assert!(poll_once(send.as_mut()).is_pending()); + assert_eq!(drops.load(Ordering::Relaxed), 1); + } + drop(send); + assert_eq!(drops.load(Ordering::Relaxed), 1); + assert_eq!(rx.try_recv(), Ok(0)); + tx.try_send(2).unwrap(); + assert_eq!(rx.try_recv(), Ok(2)); + } + }); +} - 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); +#[test] +fn bounded_receiver_waker_replacement_can_send() { + assert_completes_without_deadlock(|| { + let (tx, mut rx) = mpsc::bounded(1); + let waker = waker_on_drop(move || tx.try_send(7).unwrap()); + assert!(poll_with(Box::pin(rx.recv()).as_mut(), &waker).is_pending()); + drop(waker); + let (waker, wakes) = WakeCounter::new(); + let poll = poll_with(Box::pin(rx.recv()).as_mut(), &waker); + if poll.is_pending() { + assert!(wakes.count() > 0); + assert_eq!(rx.try_recv(), Ok(7)); + } else { + assert_eq!(poll, Poll::Ready(Ok(7))); } - } + assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); + }); +} +#[test] +fn unbounded_replaced_and_disconnected_wakers_can_send() { 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(), - })); + let waker = waker_on_drop({ + let drops = drops.clone(); + move || { + assert_eq!(tx.send(7).is_err(), disconnected); + drops.fetch_add(1, Ordering::Relaxed); + } + }); assert!( Box::pin(rx.recv()) .as_mut() diff --git a/tests-integration/tests/mpsc_test/concurrency.rs b/tests-integration/tests/mpsc_test/concurrency.rs index ce097fc..788fbcb 100644 --- a/tests-integration/tests/mpsc_test/concurrency.rs +++ b/tests-integration/tests/mpsc_test/concurrency.rs @@ -17,6 +17,7 @@ use std::future::Future; use std::pin::Pin; +use std::sync::Arc; use std::sync::Barrier; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; @@ -35,6 +36,54 @@ use tokio_test::assert_ok; use super::support::WakeCounter; use super::support::poll_with; +#[test] +fn publication_racing_with_close_drops_every_payload_once() { + #[derive(Debug)] + #[repr(align(128))] + struct Payload { + bytes: [u8; 1024], + drops: Arc<[AtomicUsize; 3]>, + // Queued messages must not retain the channel through a sender cycle. + _sender: mpsc::BoundedSender, + } + + impl Drop for Payload { + fn drop(&mut self) { + self.drops[self.bytes[0] as usize].fetch_add(1, Ordering::Relaxed); + } + } + + for _ in 0..if cfg!(miri) { 8 } else { 128 } { + let (tx, rx) = mpsc::bounded(3); + let drops = Arc::new(std::array::from_fn(|_| AtomicUsize::new(0))); + let start = Barrier::new(4); + thread::scope(|scope| { + for byte in 0..3 { + let permit = tx.try_reserve().unwrap(); + let value = Payload { + bytes: [byte; 1024], + drops: drops.clone(), + _sender: tx.clone(), + }; + let start = &start; + scope.spawn(move || { + start.wait(); + if let Err(error) = permit.send(value) { + let value = error.into_inner(); + assert_eq!(value.bytes, [byte; 1024]); + drop(value); + } + }); + } + start.wait(); + drop(rx); + }); + for count in drops.iter() { + assert_eq!(count.load(Ordering::Relaxed), 1); + } + } +} + #[test] fn bounded_receive_racing_with_send_registration_cannot_lose_wakeup() { for _ in 0..128 { @@ -63,6 +112,47 @@ fn bounded_receive_racing_with_send_registration_cannot_lose_wakeup() { } } +#[test] +fn bounded_competing_reservation_cannot_strand_a_waiting_send() { + for _ in 0..if cfg!(miri) { 32 } else { 512 } { + let (tx, mut rx) = mpsc::bounded(1); + tx.try_send(1).unwrap(); + let start = Barrier::new(3); + let received = Barrier::new(2); + let (waker, notified) = WakeCounter::new(); + let mut send = Box::pin(tx.send(2)); + + let poll = thread::scope(|scope| { + let receive = scope.spawn(|| { + start.wait(); + assert_eq!(rx.try_recv(), Ok(1)); + received.wait(); + }); + let competitor = scope.spawn(|| { + start.wait(); + received.wait(); + // Try to consume the returned capacity while the other sender registers. + tx.try_reserve().ok() + }); + start.wait(); + let poll = poll_with(send.as_mut(), &waker); + receive.join().unwrap(); + // Keep any competing reservation until registration has completed. Its release + // must notify a pending sender even if it won capacity during that registration. + drop(competitor.join().unwrap()); + poll + }); + + if poll.is_pending() { + assert!(notified.count() > 0, "available capacity stranded a sender"); + assert_eq!(poll_once(send.as_mut()), Poll::Ready(Ok(()))); + } else { + assert_eq!(poll, Poll::Ready(Ok(()))); + } + assert_eq!(rx.try_recv(), Ok(2)); + } +} + #[test] fn bounded_try_recv_does_not_report_empty_after_completed_sends() { const PRODUCERS: usize = 4; diff --git a/tests-integration/tests/mpsc_test/main.rs b/tests-integration/tests/mpsc_test/main.rs index 2490021..f2ab74f 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 { @@ -173,3 +174,19 @@ fn receives_wake_for_messages_and_the_last_sender_drop() { Poll::Ready(Err(RecvError::Disconnected)) ); } + +#[test] +#[should_panic(expected = "must be nonzero")] +fn bounded_rejects_zero_capacity() { + let _ = mpsc::bounded::(0); +} + +#[test] +fn bounded_supports_full_usize_capacity_for_zero_sized_messages() { + let (tx, mut rx) = mpsc::bounded::<()>(usize::MAX); + // Returning capacity at this boundary must not overflow the counter. + drop(tx.try_reserve().unwrap()); + tx.try_reserve().unwrap().send(()).unwrap(); + assert_eq!(rx.try_recv(), Ok(())); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); +} diff --git a/tests-integration/tests/mpsc_test/reservation.rs b/tests-integration/tests/mpsc_test/reservation.rs new file mode 100644 index 0000000..7d2a074 --- /dev/null +++ b/tests-integration/tests/mpsc_test/reservation.rs @@ -0,0 +1,226 @@ +// 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::mem; +use std::panic::AssertUnwindSafe; +use std::panic::catch_unwind; +use std::sync::Arc; +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(); + // The held permit stays usable while other messages repeatedly reuse the buffer. + for lap in 0..8 { + for offset in 1..capacity { + tx.try_send(lap * capacity + offset).unwrap(); + } + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + assert_eq!(tx.try_send(0), Err(TrySendError::Full(0))); + for offset in 1..capacity { + assert_eq!(rx.try_recv(), Ok(lap * capacity + offset)); + } + 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 released_capacity_is_granted_to_the_oldest_waiter() { + 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); + // The waiting future owns the released slot even before the executor polls it again. + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + assert_eq!(tx.try_send(9), Err(TrySendError::Full(9))); + 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 cancelling_a_granted_reservation_passes_capacity_to_a_waiting_send() { + let (tx, mut rx) = mpsc::bounded(1); + let held = tx.try_reserve().unwrap(); + let mut reservation = Box::pin(tx.reserve()); + let mut send = Box::pin(tx.send(7)); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_once(reservation.as_mut()).is_pending()); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + drop(held); + assert_eq!(wakes.count(), 0); + drop(reservation); + assert_eq!(wakes.count(), 1); + assert_eq!(tx.try_send(9), Err(TrySendError::Full(9))); + assert_eq!(expect_ready(poll_once(send.as_mut())), Ok(())); + assert_eq!(rx.try_recv(), Ok(7)); + let held = tx.try_reserve().unwrap(); + assert!(matches!(tx.try_reserve(), Err(TrySendError::Full(())))); + drop(held); +} + +#[test] +fn closing_after_a_grant_returns_the_unsent_message() { + let (tx, mut rx) = mpsc::bounded(1); + tx.try_send(String::from("queued")).unwrap(); + let mut send = Box::pin(tx.send(String::from("unsent"))); + let mut reservation = Box::pin(tx.reserve()); + assert!(poll_once(send.as_mut()).is_pending()); + assert!(poll_once(reservation.as_mut()).is_pending()); + assert_eq!(rx.try_recv().unwrap(), "queued"); + drop(rx); + let error = expect_ready(poll_once(send.as_mut())).unwrap_err(); + assert_eq!(error.into_inner(), "unsent"); + assert!(expect_ready(poll_once(reservation.as_mut())).is_err()); +} + +#[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(); + 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 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 channel. + 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); +} diff --git a/tests-integration/tests/mpsc_test/support.rs b/tests-integration/tests/mpsc_test/support.rs index 360cb0a..fdeb836 100644 --- a/tests-integration/tests/mpsc_test/support.rs +++ b/tests-integration/tests/mpsc_test/support.rs @@ -63,6 +63,24 @@ pub fn poll_with(future: Pin<&mut F>, waker: &Waker) -> Poll Waker { + struct OnDrop(Box); + + // Only destruction runs the callback; waking consumes the reference as usual. + #[allow(clippy::manual_noop_waker)] + impl Wake for OnDrop { + fn wake(self: Arc) {} + } + + impl Drop for OnDrop { + fn drop(&mut self) { + (self.0)(); + } + } + + Waker::from(Arc::new(OnDrop(Box::new(callback)))) +} + // RawWaker is needed only to exercise clone callbacks, which the safe Wake trait cannot override. pub fn waker_on_clone(callback: impl Fn() + Send + Sync + 'static) -> Waker { struct OnClone(Box);