diff --git a/CHANGELOG.md b/CHANGELOG.md index f5e01ff8..4b7d5899 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -28,6 +28,7 @@ All notable changes to this project will be documented in this file. * Add `event::AutoResetEvent`, a reusable signal that releases one waiter, retains at most one unassigned signal that can be cleared with `reset`, and transfers assigned signals when waits are cancelled. * Add `ManualResetEvent::try_wait` to check readiness without registering a waiter or consuming the set state. * Implement `broadcast::mpmc::bounded`, a lossless bounded broadcast channel that retains at most the requested capacity and makes producers wait for the slowest active receiver. +* Add opt-in bounded and unbounded `asyncband::spmc` queues with one non-cloneable sender requiring exclusive access, cloneable competing receivers, and cancellation-safe receive notification handoff. * Add opt-in bounded and unbounded `asyncband::mpmc` queues with cloneable producers and competing consumers, delivering each accepted value to exactly one receiver while a receiver remains. * Add an opt-in runtime-agnostic `Phaser` with shared observer handles, dynamic RAII participants registered individually or in batches through an owning iterator, `u64` phase numbers, split arrival/wait with cancellation-resilient retries, and a `close` operation that releases unfinished waits with `Closed`. * 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. diff --git a/README.md b/README.md index 2a1bcad4..4e94e213 100644 --- a/README.md +++ b/README.md @@ -103,6 +103,7 @@ Runnable examples live in the [`examples`](examples) workspace crate. They demon | | [`oneshot`](https://docs.rs/asyncband/*/asyncband/oneshot/) | `oneshot` | Send one value from one sender to one receiver. | | | [`mpmc`](https://docs.rs/asyncband/*/asyncband/mpmc/) | `mpmc` | Distribute each value to exactly one of multiple competing receivers. | | | [`mpsc`](https://docs.rs/asyncband/*/asyncband/mpsc/) | `mpsc` | Send each value from multiple producers to one receiver with bounded backpressure or an unbounded queue. | +| | [`spmc`](https://docs.rs/asyncband/*/asyncband/spmc/) | `spmc` | Distribute work from one exclusive sender to multiple competing receivers, with bounded or unbounded storage. | | | [`broadcast`](https://docs.rs/asyncband/*/asyncband/broadcast/) | `broadcast` | Deliver every value to active receivers with bounded backpressure or unbounded retention. | | | [`watch`](https://docs.rs/asyncband/*/asyncband/watch/) | `watch` | Publish cloneable latest state from one or more senders; receivers independently coalesce intermediate updates. | | Object reuse | [`pool`](https://docs.rs/asyncband/*/asyncband/pool/) | `pool` | Reuse objects through bounded or unbounded pool variants. | diff --git a/asyncband/Cargo.toml b/asyncband/Cargo.toml index d34a926b..819ef36b 100644 --- a/asyncband/Cargo.toml +++ b/asyncband/Cargo.toml @@ -65,6 +65,7 @@ rwlock = [] semaphore = [] shutdown = ["latch", "waitgroup"] singleflight = ["dep:hashbrown", "once-cell"] +spmc = [] waitgroup = [] watch = [] diff --git a/asyncband/src/internal/mod.rs b/asyncband/src/internal/mod.rs index 0b8211bc..0b831c48 100644 --- a/asyncband/src/internal/mod.rs +++ b/asyncband/src/internal/mod.rs @@ -72,6 +72,7 @@ pub(crate) fn wake_all(mut wakers: impl Iterator) { feature = "completion", feature = "latch", feature = "mpmc", + feature = "spmc", feature = "mpsc", feature = "mutex", feature = "phaser", @@ -101,6 +102,7 @@ pub(crate) mod value_cell; feature = "completion", feature = "latch", feature = "mpmc", + feature = "spmc", feature = "mpsc", feature = "mutex", feature = "phaser", @@ -129,6 +131,7 @@ pub(crate) mod semaphore; feature = "broadcast", feature = "event", feature = "mpmc", + feature = "spmc", feature = "mpsc", feature = "mutex", feature = "rwlock", diff --git a/asyncband/src/lib.rs b/asyncband/src/lib.rs index a0472fa7..375ebe9a 100644 --- a/asyncband/src/lib.rs +++ b/asyncband/src/lib.rs @@ -76,6 +76,7 @@ //! | | [`oneshot`] | `oneshot` | Send one value from one sender to one receiver. | //! | | [`mpmc`] | `mpmc` | Distribute each value to exactly one of multiple competing receivers. | //! | | [`mpsc`] | `mpsc` | Send each value from multiple producers to one receiver with bounded backpressure or an unbounded queue. | +//! | | [`spmc`] | `spmc` | Distribute work from one exclusive sender to multiple competing receivers, with bounded or unbounded storage. | //! | | [`broadcast`] | `broadcast` | Deliver every value to active receivers with bounded backpressure or unbounded retention. | //! | | [`watch`] | `watch` | Publish cloneable latest state from one or more senders; receivers independently coalesce intermediate updates. | //! | Object reuse | [`pool`] | `pool` | Reuse objects through bounded or unbounded pool variants. | @@ -162,6 +163,8 @@ pub mod semaphore; pub mod shutdown; #[cfg(feature = "singleflight")] pub mod singleflight; +#[cfg(feature = "spmc")] +pub mod spmc; #[cfg(feature = "waitgroup")] pub mod waitgroup; #[cfg(feature = "watch")] diff --git a/asyncband/src/spmc/bounded.rs b/asyncband/src/spmc/bounded.rs new file mode 100644 index 00000000..f164013b --- /dev/null +++ b/asyncband/src/spmc/bounded.rs @@ -0,0 +1,145 @@ +// 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::sync::Arc; + +use super::RecvError; +use super::SendError; +use super::TryRecvError; +use super::TrySendError; +use super::queue::Shared; + +/// Creates a bounded single-producer, multi-consumer queue. +/// +/// The queue stores at most `capacity` values. Sending waits for a receiver to free capacity when +/// the queue is full. +/// +/// Operations briefly acquire internal mutexes. No lock is held across an await point, while +/// waking tasks, or while dropping messages. The `try_*` methods do not wait for capacity or +/// messages, but may wait to acquire a mutex. +/// +/// # Panics +/// +/// Panics if `capacity` is zero. +#[track_caller] +pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { + assert!(capacity > 0, "spmc bounded queue requires capacity > 0"); + let shared = Arc::new(Shared::bounded(capacity)); + ( + BoundedSender { + shared: shared.clone(), + }, + BoundedReceiver { shared }, + ) +} + +/// Sends values to the associated [`BoundedReceiver`] handles. +/// +/// Instances are created by [`bounded`] and cannot be cloned. Sending requires exclusive access to +/// this endpoint. +pub struct BoundedSender { + shared: Arc>, +} + +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) { + self.shared.drop_sender(); + } +} + +impl BoundedSender { + /// Sends a value, waiting until capacity is available if the queue is full. + /// + /// If all receivers have been dropped, the value is returned in [`SendError`]. + /// + /// # Cancel safety + /// + /// Dropping a pending `send` removes it from the wait queue and drops `value`; a call that has + /// returned `Pending` has not sent the value. Cancelling releases the exclusive sender borrow + /// and leaves available capacity usable by the next send. Use [`try_send`](Self::try_send) when + /// the caller must retain ownership if capacity is unavailable. + pub async fn send(&mut self, value: T) -> Result<(), SendError> { + self.shared.send(value).await + } + + /// Attempts to send a value without waiting for capacity. + /// + /// Returns [`TrySendError::Full`] when the queue has reached its exact capacity and + /// [`TrySendError::Disconnected`] when all receivers have been dropped. + pub fn try_send(&mut self, value: T) -> Result<(), TrySendError> { + self.shared.try_send(value) + } +} + +/// Receives values from the associated [`BoundedSender`] handles. +/// +/// Cloned receivers compete for values, and every accepted value is returned by exactly one +/// receiver while a receiver remains. Dropping the final receiver releases buffered values. +pub struct BoundedReceiver { + shared: Arc>, +} + +impl Clone for BoundedReceiver { + fn clone(&self) -> Self { + self.shared.clone_receiver(); + Self { + shared: self.shared.clone(), + } + } +} + +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) { + self.shared.drop_receiver(); + } +} + +impl BoundedReceiver { + /// Receives the next available value. + /// + /// Buffered values remain available after the sender is dropped. Once they are drained, + /// this method returns [`RecvError::Disconnected`]. + /// + /// # Cancel safety + /// + /// Dropping a pending `recv` does not consume a value. Any selected value notification is + /// passed to another waiting receiver, so cancellation does not prevent it from receiving. + pub async fn recv(&self) -> Result { + self.shared.recv().await + } + + /// Attempts to receive the next available value without waiting for a message. + /// + /// Returns [`TryRecvError::Empty`] while the queue is empty and a sender remains, or + /// [`TryRecvError::Disconnected`] once the queue is empty and the sender has been dropped. + pub fn try_recv(&self) -> Result { + self.shared.try_recv() + } +} diff --git a/asyncband/src/spmc/error.rs b/asyncband/src/spmc/error.rs new file mode 100644 index 00000000..18451db2 --- /dev/null +++ b/asyncband/src/spmc/error.rs @@ -0,0 +1,138 @@ +// 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::any::type_name; +use std::fmt; + +/// An error returned when trying to send on a disconnected queue. +/// +/// The value that could not be sent can be retrieved with [`SendError::into_inner`]. +#[derive(Clone, PartialEq, Eq)] +pub struct SendError(T); + +impl SendError { + /// Gets a reference to the value that failed to be sent. + pub fn as_inner(&self) -> &T { + &self.0 + } + + /// Consumes the error and returns the value that failed to be sent. + pub fn into_inner(self) -> T { + self.0 + } + + pub(super) fn new(value: T) -> Self { + Self(value) + } +} + +impl fmt::Display for SendError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("sending on a disconnected queue") + } +} + +impl fmt::Debug for SendError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "SendError<{}>(..)", type_name::()) + } +} + +impl std::error::Error for SendError {} + +/// Error returned by [`BoundedSender::try_send`](crate::spmc::BoundedSender::try_send). +#[derive(Clone, PartialEq, Eq)] +pub enum TrySendError { + /// The queue is full, so the value cannot be sent without waiting for capacity. + Full(T), + /// All receivers have been dropped, so the value can never be received. + Disconnected(T), +} + +impl TrySendError { + /// Gets a reference to the value that failed to be sent. + pub fn as_inner(&self) -> &T { + match self { + TrySendError::Full(value) | TrySendError::Disconnected(value) => value, + } + } + + /// Consumes the error and returns the value that failed to be sent. + pub fn into_inner(self) -> T { + match self { + TrySendError::Full(value) | TrySendError::Disconnected(value) => value, + } + } +} + +impl fmt::Display for TrySendError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + TrySendError::Full(_) => "sending on a full queue", + TrySendError::Disconnected(_) => "sending on a disconnected queue", + }) + } +} + +impl fmt::Debug for TrySendError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + let ty = type_name::(); + match self { + TrySendError::Full(_) => write!(f, "TrySendError<{ty}>::Full(..)"), + TrySendError::Disconnected(_) => { + write!(f, "TrySendError<{ty}>::Disconnected(..)") + } + } + } +} + +impl std::error::Error for TrySendError {} + +/// Error returned by a receive operation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum RecvError { + /// The sender has been dropped, and no buffered values remain. + Disconnected, +} + +impl fmt::Display for RecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str("receiving on a disconnected queue") + } +} + +impl std::error::Error for RecvError {} + +/// Error returned by a non-blocking receive operation. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum TryRecvError { + /// No value is currently available, but the sender remains. + Empty, + /// The sender has been dropped, and no buffered values remain. + Disconnected, +} + +impl fmt::Display for TryRecvError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.write_str(match self { + TryRecvError::Empty => "receiving on an empty queue", + TryRecvError::Disconnected => "receiving on a disconnected queue", + }) + } +} + +impl std::error::Error for TryRecvError {} diff --git a/asyncband/src/spmc/mod.rs b/asyncband/src/spmc/mod.rs new file mode 100644 index 00000000..613985ea --- /dev/null +++ b/asyncband/src/spmc/mod.rs @@ -0,0 +1,101 @@ +// 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. + +//! Single-producer, multi-consumer queues for distributing work between asynchronous tasks. +//! +//! Enable the `spmc` Cargo feature to use this module. [`bounded`] applies backpressure at its +//! exact capacity; [`unbounded`] sends synchronously and can grow until memory is exhausted. +//! Receivers are cloneable and compete for messages: each accepted message is delivered to one +//! receiver while receivers remain. Values leave the queue in FIFO order, but consumer completion +//! order and an equal distribution of work are not guaranteed. +//! +//! The sender cannot be cloned, and every send operation requires `&mut self`, including for the +//! lifetime of a bounded send future. It can move between tasks, but shared references cannot send. +//! Dropping the sender lets receivers drain buffered messages before observing disconnection. +//! Dropping the last receiver releases buffered messages and makes sending return the unsent value. +//! +//! # Example +//! +//! ``` +//! # #[tokio::main] +//! # async fn main() { +//! use asyncband::spmc; +//! +//! let (mut sender, receiver) = spmc::bounded(2); +//! let competing = receiver.clone(); +//! sender.send("first").await.unwrap(); +//! sender.send("second").await.unwrap(); +//! drop(sender); +//! +//! assert_eq!(receiver.recv().await, Ok("first")); +//! assert_eq!(competing.recv().await, Ok("second")); +//! assert_eq!(receiver.recv().await, Err(spmc::RecvError::Disconnected)); +//! # } +//! ``` +//! +//! # Single-producer capability +//! +//! Neither sender supports cloning: +//! +//! ```compile_fail,E0599 +//! let (sender, _receiver) = asyncband::spmc::bounded::(1); +//! let second_producer = sender.clone(); +//! ``` +//! +//! ```compile_fail,E0599 +//! let (sender, _receiver) = asyncband::spmc::unbounded::(); +//! let second_producer = sender.clone(); +//! ``` +//! +//! Sending through a shared reference is rejected: +//! +//! ```compile_fail,E0596 +//! fn send(sender: &asyncband::spmc::BoundedSender) { +//! let _ = sender.try_send(1); +//! } +//! ``` +//! +//! ```compile_fail,E0596 +//! fn send(sender: &asyncband::spmc::UnboundedSender) { +//! let _ = sender.send(1); +//! } +//! ``` +//! +//! A bounded send future retains the exclusive borrow until completion or cancellation: +//! +//! ```compile_fail,E0499 +//! let (mut sender, _receiver) = asyncband::spmc::bounded(1); +//! let pending = sender.send(1); +//! let _ = sender.try_send(2); +//! drop(pending); +//! ``` + +mod bounded; +mod error; +mod queue; +mod unbounded; + +pub use self::bounded::BoundedReceiver; +pub use self::bounded::BoundedSender; +pub use self::bounded::bounded; +pub use self::error::RecvError; +pub use self::error::SendError; +pub use self::error::TryRecvError; +pub use self::error::TrySendError; +pub use self::unbounded::UnboundedReceiver; +pub use self::unbounded::UnboundedSender; +pub use self::unbounded::unbounded; diff --git a/asyncband/src/spmc/queue.rs b/asyncband/src/spmc/queue.rs new file mode 100644 index 00000000..c8be6bc1 --- /dev/null +++ b/asyncband/src/spmc/queue.rs @@ -0,0 +1,363 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::VecDeque; +use std::future::poll_fn; +use std::mem; +use std::task::Context; +use std::task::Poll; +use std::task::Waker; + +use super::RecvError; +use super::SendError; +use super::TryRecvError; +use super::TrySendError; +use crate::internal::mutex::Mutex; +use crate::internal::waitlist::WaitList; +use crate::internal::waitlist::WaiterId; +use crate::internal::wake_all; + +pub struct Shared { + state: Mutex>, +} + +/// Values, endpoint state, and both waiter queues share one lock, so each transition and the +/// waiter it selects are decided together. Wake callbacks and waker or value destructors run +/// outside the lock, because they may reenter the queue. +struct State { + values: VecDeque, + capacity: Option, + sender_alive: bool, + receivers: usize, + recv_waiters: WaitList, + send_waiters: WaitList, +} + +impl State { + fn has_capacity(&self) -> bool { + self.capacity + .is_none_or(|capacity| self.values.len() < capacity) + } + + /// Queues a value and selects the receiver to wake. + fn push(&mut self, value: T) -> Option { + self.values.push_back(value); + self.recv_waiters.notify_one() + } + + /// Takes the next value and selects the sender to wake. + fn pop(&mut self) -> Result<(T, Option), TryRecvError> { + if let Some(value) = self.values.pop_front() { + // Unbounded queues never block senders, so their sender queue is always empty. + Ok((value, self.send_waiters.notify_one())) + } else if !self.sender_alive { + Err(TryRecvError::Disconnected) + } else { + Err(TryRecvError::Empty) + } + } +} + +/// A pending receive or bounded send. +/// +/// Notification makes a waiter runnable; it does not reserve a value or slot. The detached node +/// remains owned by its future until it retries or is dropped. +enum Waiter { + Waiting(Waker), + Notified, +} + +impl WaitList { + fn notify_one(&mut self) -> Option { + let (_, waiter) = self.unlink_first_waiter(|_| true)?; + let Waiter::Waiting(waker) = mem::replace(waiter, Waiter::Notified) else { + unreachable!("only waiting operations remain linked"); + }; + Some(waker) + } + + fn remove_waiter(&mut self, id: WaiterId) -> Waiter { + // Unlinking is idempotent, so notified waiters are removed the same way as linked ones. + self.unlink_waiter(id, |_| true); + self.remove_unlinked_waiter(id) + } + + /// Queues a blocked operation or refreshes the waker of a queued one. + /// + /// A notified operation that still found no value or slot queues again at the back. + #[must_use = "drop the replaced waker after releasing the queue lock"] + fn register(&mut self, id: &mut Option, current: &Waker) -> Option { + if let Some(queued) = *id { + if let Waiter::Waiting(waker) = self.waiter_mut(queued) { + if waker.will_wake(current) { + return None; + } + return Some(mem::replace(waker, current.clone())); + } + } + let waker = current.clone(); + if let Some(notified) = id.take() { + // The notification already took this node's waker, so nothing is retired. + self.remove_waiter(notified); + } + *id = Some(self.push_back(Waiter::Waiting(waker))); + None + } +} + +impl Shared { + pub fn bounded(capacity: usize) -> Self { + Self::new(Some(capacity)) + } + + pub fn unbounded() -> Self { + Self::new(None) + } + + fn new(capacity: Option) -> Self { + Self { + state: Mutex::new(State { + values: VecDeque::new(), + capacity, + sender_alive: true, + receivers: 1, + recv_waiters: WaitList::new(), + send_waiters: WaitList::new(), + }), + } + } + + pub fn drop_sender(&self) { + let mut waiters = { + let mut state = self.state.lock(); + state.sender_alive = false; + // Disconnection invalidates every receiver waiter ID. Move the storage out so both + // notification and reclamation happen without holding the queue lock. + mem::replace(&mut state.recv_waiters, WaitList::new()) + }; + wake_all(std::iter::from_fn(|| waiters.notify_one())); + } + + pub fn clone_receiver(&self) { + self.state.lock().receivers += 1; + } + + pub fn drop_receiver(&self) { + let (discarded, mut waiters) = { + let mut state = self.state.lock(); + state.receivers -= 1; + if state.receivers != 0 { + return; + } + ( + mem::take(&mut state.values), + mem::replace(&mut state.send_waiters, WaitList::new()), + ) + }; + // Release blocked senders before destroying buffered values. Local ownership still drops + // the values if a wake callback unwinds. + wake_all(std::iter::from_fn(|| waiters.notify_one())); + drop(discarded); + } + + pub fn try_send(&self, value: T) -> Result<(), TrySendError> { + let waker = { + let mut state = self.state.lock(); + if state.receivers == 0 { + return Err(TrySendError::Disconnected(value)); + } + if !state.has_capacity() { + return Err(TrySendError::Full(value)); + } + state.push(value) + }; + if let Some(waker) = waker { + waker.wake(); + } + Ok(()) + } + + 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 send = Send { + shared: self, + waiter: None, + value: Some(value), + }; + poll_fn(|cx| send.poll(cx)).await + } + + pub fn try_recv(&self) -> Result { + let (value, waker) = self.state.lock().pop()?; + if let Some(waker) = waker { + waker.wake(); + } + Ok(value) + } + + pub async fn recv(&self) -> Result { + match self.try_recv() { + Ok(value) => return Ok(value), + Err(TryRecvError::Disconnected) => return Err(RecvError::Disconnected), + Err(TryRecvError::Empty) => {} + } + let mut recv = Recv { + shared: self, + waiter: None, + }; + poll_fn(|cx| recv.poll(cx)).await + } +} + +struct Send<'a, T> { + shared: &'a Shared, + waiter: Option, + // `Drop` passes an unconsumed notification on before this value is destroyed, because its + // destructor may depend on another blocked sender making progress. + value: Option, +} + +impl Send<'_, T> { + fn take_value(&mut self) -> T { + self.value.take().expect("pending send must own its value") + } + + fn poll(&mut self, cx: &mut Context<'_>) -> Poll>> { + let mut state = self.shared.state.lock(); + let outcome = if state.receivers == 0 { + self.waiter = None; + Err(self.take_value()) + } else if state.has_capacity() { + Ok(state.push(self.take_value())) + } else { + let retired = state.send_waiters.register(&mut self.waiter, cx.waker()); + drop(state); + drop(retired); + return Poll::Pending; + }; + let retired = self + .waiter + .take() + .map(|id| state.send_waiters.remove_waiter(id)); + drop(state); + // Deliver the notification before running waker destructors, which may panic. + let result = outcome + .map(|waker| { + if let Some(waker) = waker { + waker.wake(); + } + }) + .map_err(SendError::new); + drop(retired); + Poll::Ready(result) + } +} + +impl Drop for Send<'_, T> { + fn drop(&mut self) { + let Some(id) = self.waiter.take() else { + return; + }; + let (retired, waker) = { + let mut state = self.shared.state.lock(); + if state.receivers == 0 { + return; + } + let retired = state.send_waiters.remove_waiter(id); + // Hand an unconsumed notification to the next sender while the slot is still free. + let waker = if matches!(retired, Waiter::Notified) && state.has_capacity() { + state.send_waiters.notify_one() + } else { + None + }; + (retired, waker) + }; + if let Some(waker) = waker { + waker.wake(); + } + drop(retired); + } +} + +struct Recv<'a, T> { + shared: &'a Shared, + waiter: Option, +} + +impl Recv<'_, T> { + fn poll(&mut self, cx: &mut Context<'_>) -> Poll> { + let mut state = self.shared.state.lock(); + if !state.sender_alive { + // Buffered values remain readable after the waiter storage has been detached. + self.waiter = None; + } + let outcome = match state.pop() { + Ok(popped) => Ok(popped), + Err(TryRecvError::Disconnected) => Err(RecvError::Disconnected), + Err(TryRecvError::Empty) => { + let retired = state.recv_waiters.register(&mut self.waiter, cx.waker()); + drop(state); + drop(retired); + return Poll::Pending; + } + }; + let retired = self + .waiter + .take() + .map(|id| state.recv_waiters.remove_waiter(id)); + drop(state); + // Deliver the notification before running waker destructors, which may panic. + let result = outcome.map(|(value, waker)| { + if let Some(waker) = waker { + waker.wake(); + } + value + }); + drop(retired); + Poll::Ready(result) + } +} + +impl Drop for Recv<'_, T> { + fn drop(&mut self) { + let Some(id) = self.waiter.take() else { + return; + }; + let (retired, waker) = { + let mut state = self.shared.state.lock(); + if !state.sender_alive { + return; + } + let retired = state.recv_waiters.remove_waiter(id); + // Hand an unconsumed notification to the next receiver while a value still waits. + let waker = if matches!(retired, Waiter::Notified) && !state.values.is_empty() { + state.recv_waiters.notify_one() + } else { + None + }; + (retired, waker) + }; + if let Some(waker) = waker { + waker.wake(); + } + drop(retired); + } +} diff --git a/asyncband/src/spmc/unbounded.rs b/asyncband/src/spmc/unbounded.rs new file mode 100644 index 00000000..60e89473 --- /dev/null +++ b/asyncband/src/spmc/unbounded.rs @@ -0,0 +1,127 @@ +// 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::sync::Arc; + +use super::RecvError; +use super::SendError; +use super::TryRecvError; +use super::TrySendError; +use super::queue::Shared; + +/// Creates an unbounded single-producer, multi-consumer queue. +/// +/// Sends are synchronous and values may be buffered until available memory is exhausted. +/// +/// Operations briefly acquire internal mutexes. No lock is held across an await point, while +/// waking tasks, or while dropping messages. Sending and trying to receive may wait to acquire +/// a mutex, but never wait for capacity or new messages. +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { + let shared = Arc::new(Shared::unbounded()); + ( + UnboundedSender { + shared: shared.clone(), + }, + UnboundedReceiver { shared }, + ) +} + +/// Sends values to the associated [`UnboundedReceiver`] handles. +/// +/// Instances are created by [`unbounded`] and cannot be cloned. Sending requires exclusive access +/// to this endpoint. +pub struct UnboundedSender { + shared: Arc>, +} + +impl fmt::Debug for UnboundedSender { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedSender").finish_non_exhaustive() + } +} + +impl Drop for UnboundedSender { + fn drop(&mut self) { + self.shared.drop_sender(); + } +} + +impl UnboundedSender { + /// Sends a value without waiting for capacity. + /// + /// If all receivers have been dropped, the value is returned in [`SendError`]. + pub fn send(&mut self, value: T) -> Result<(), SendError> { + match self.shared.try_send(value) { + Ok(()) => Ok(()), + Err(TrySendError::Disconnected(value)) => Err(SendError::new(value)), + Err(TrySendError::Full(_)) => unreachable!("unbounded queue cannot be full"), + } + } +} + +/// Receives values from the associated [`UnboundedSender`] handles. +/// +/// Cloned receivers compete for values, and every accepted value is returned by exactly one +/// receiver while a receiver remains. Dropping the final receiver releases buffered values. +pub struct UnboundedReceiver { + shared: Arc>, +} + +impl Clone for UnboundedReceiver { + fn clone(&self) -> Self { + self.shared.clone_receiver(); + Self { + shared: self.shared.clone(), + } + } +} + +impl fmt::Debug for UnboundedReceiver { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("UnboundedReceiver").finish_non_exhaustive() + } +} + +impl Drop for UnboundedReceiver { + fn drop(&mut self) { + self.shared.drop_receiver(); + } +} + +impl UnboundedReceiver { + /// Receives the next available value. + /// + /// Buffered values remain available after the sender is dropped. Once they are drained, + /// this method returns [`RecvError::Disconnected`]. + /// + /// # Cancel safety + /// + /// Dropping a pending `recv` does not consume a value. Any selected value notification is + /// passed to another waiting receiver, so cancellation does not prevent it from receiving. + pub async fn recv(&self) -> Result { + self.shared.recv().await + } + + /// Attempts to receive the next available value without waiting for a message. + /// + /// Returns [`TryRecvError::Empty`] while the queue is empty and a sender remains, or + /// [`TryRecvError::Disconnected`] once the queue is empty and the sender has been dropped. + pub fn try_recv(&self) -> Result { + self.shared.try_recv() + } +} diff --git a/benchmarks/Cargo.toml b/benchmarks/Cargo.toml index bcab8f24..b49b8eb1 100644 --- a/benchmarks/Cargo.toml +++ b/benchmarks/Cargo.toml @@ -46,6 +46,7 @@ asyncband = { workspace = true, features = [ "semaphore", "shutdown", "singleflight", + "spmc", "waitgroup", "watch", ] } diff --git a/benchmarks/asyncband/main.rs b/benchmarks/asyncband/main.rs index a389c0ae..b0d38ab2 100644 --- a/benchmarks/asyncband/main.rs +++ b/benchmarks/asyncband/main.rs @@ -23,10 +23,6 @@ mod condvar; mod event; mod latch; mod mpmc; - -#[allow(dead_code)] -#[path = "../mpmc/mod.rs"] -mod mpmc_support; mod mpsc; mod mutex; mod once; @@ -41,6 +37,10 @@ mod singleflight; mod support; mod waitgroup; +#[allow(dead_code)] +#[path = "../channels/mod.rs"] +mod channels; + fn main() { divan::main(); } diff --git a/benchmarks/asyncband/mpmc/bounded.rs b/benchmarks/asyncband/mpmc/bounded.rs index 1bf7253f..6cf7c044 100644 --- a/benchmarks/asyncband/mpmc/bounded.rs +++ b/benchmarks/asyncband/mpmc/bounded.rs @@ -23,15 +23,15 @@ use divan::black_box; use divan::counter::ItemsCount; use super::FAST_SAMPLE_SIZE; -use crate::mpmc_support::adapters::Asyncband; -use crate::mpmc_support::support::BATCH_MESSAGES; -use crate::mpmc_support::support::BOUNDED_CAPACITY; -use crate::mpmc_support::support::Bounded; -use crate::mpmc_support::support::TOPOLOGIES; -use crate::mpmc_support::support::TaskBatch; -use crate::mpmc_support::support::ThreadBatch; -use crate::mpmc_support::support::Topology; -use crate::mpmc_support::support::runtime; +use crate::channels::BATCH_MESSAGES; +use crate::channels::BOUNDED_CAPACITY; +use crate::channels::adapters::Bounded; +use crate::channels::adapters::Mpmc; +use crate::channels::mpmc::TOPOLOGIES; +use crate::channels::mpmc::TaskBatch; +use crate::channels::mpmc::ThreadBatch; +use crate::channels::mpmc::Topology; +use crate::channels::runtime; use crate::support::bench_context; use crate::support::poll_pending; use crate::support::poll_pinned_ready; @@ -45,7 +45,7 @@ use crate::support::poll_ready; )] fn blocking_threads(bencher: Bencher, topology: Topology) { bencher - .with_inputs(|| ThreadBatch::new_bounded::(BOUNDED_CAPACITY, topology)) + .with_inputs(|| ThreadBatch::new_bounded::(BOUNDED_CAPACITY, topology)) .bench_local_refs(|batch| batch.run()); } @@ -59,7 +59,7 @@ fn blocking_threads(bencher: Bencher, topology: Topology) { fn tokio_tasks(bencher: Bencher, topology: Topology) { let runtime = runtime(WORKERS); bencher - .with_inputs(|| TaskBatch::new::>(&runtime, topology)) + .with_inputs(|| TaskBatch::new::>(&runtime, topology)) .bench_local_refs(|batch| runtime.block_on(batch.run())); } diff --git a/benchmarks/asyncband/mpmc/unbounded.rs b/benchmarks/asyncband/mpmc/unbounded.rs index 34234c42..475d07cc 100644 --- a/benchmarks/asyncband/mpmc/unbounded.rs +++ b/benchmarks/asyncband/mpmc/unbounded.rs @@ -23,14 +23,14 @@ use divan::black_box; use divan::counter::ItemsCount; use super::FAST_SAMPLE_SIZE; -use crate::mpmc_support::adapters::Asyncband; -use crate::mpmc_support::support::BATCH_MESSAGES; -use crate::mpmc_support::support::TOPOLOGIES; -use crate::mpmc_support::support::TaskBatch; -use crate::mpmc_support::support::ThreadBatch; -use crate::mpmc_support::support::Topology; -use crate::mpmc_support::support::Unbounded; -use crate::mpmc_support::support::runtime; +use crate::channels::BATCH_MESSAGES; +use crate::channels::adapters::Mpmc; +use crate::channels::adapters::Unbounded; +use crate::channels::mpmc::TOPOLOGIES; +use crate::channels::mpmc::TaskBatch; +use crate::channels::mpmc::ThreadBatch; +use crate::channels::mpmc::Topology; +use crate::channels::runtime; use crate::support::bench_context; use crate::support::poll_pending; use crate::support::poll_pinned_ready; @@ -44,7 +44,7 @@ use crate::support::poll_ready; )] fn blocking_threads(bencher: Bencher, topology: Topology) { bencher - .with_inputs(|| ThreadBatch::new_unbounded::(topology)) + .with_inputs(|| ThreadBatch::new_unbounded::(topology)) .bench_local_refs(|batch| batch.run()); } @@ -58,7 +58,7 @@ fn blocking_threads(bencher: Bencher, topology: Topology) { fn tokio_tasks(bencher: Bencher, topology: Topology) { let runtime = runtime(WORKERS); bencher - .with_inputs(|| TaskBatch::new::>(&runtime, topology)) + .with_inputs(|| TaskBatch::new::>(&runtime, topology)) .bench_local_refs(|batch| runtime.block_on(batch.run())); } diff --git a/benchmarks/mpmc/adapters.rs b/benchmarks/channels/adapters.rs similarity index 66% rename from benchmarks/mpmc/adapters.rs rename to benchmarks/channels/adapters.rs index 209869dd..6c78dce6 100644 --- a/benchmarks/mpmc/adapters.rs +++ b/benchmarks/channels/adapters.rs @@ -16,13 +16,98 @@ // under the License. use std::future::Future; +use std::marker::PhantomData; use asyncband::blocking::FutureExt; -pub struct Asyncband; +use super::BOUNDED_CAPACITY; + +pub struct Mpmc; +pub struct Spmc; pub struct AsyncChannel; pub struct Flume; +pub struct Bounded(PhantomData); +pub struct Unbounded(PhantomData); + +// Each task owns its sender. Only workloads with multiple producers require Sender: Clone. +pub trait Channel: Send + Sync + 'static { + type Sender: Send + 'static; + type Receiver: Clone + Send + Sync + 'static; + + fn channel() -> (Self::Sender, Self::Receiver); + fn send(sender: &mut Self::Sender, value: usize) -> impl Future + Send; + fn recv(receiver: &Self::Receiver) -> impl Future> + Send; +} + +impl Channel for Bounded { + type Sender = asyncband::spmc::BoundedSender; + type Receiver = asyncband::spmc::BoundedReceiver; + + fn channel() -> (Self::Sender, Self::Receiver) { + asyncband::spmc::bounded(CAPACITY) + } + + async fn send(sender: &mut Self::Sender, value: usize) { + sender.send(value).await.unwrap(); + } + + async fn recv(receiver: &Self::Receiver) -> Option { + receiver.recv().await.ok() + } +} + +impl Channel for Unbounded { + type Sender = asyncband::spmc::UnboundedSender; + type Receiver = asyncband::spmc::UnboundedReceiver; + + fn channel() -> (Self::Sender, Self::Receiver) { + asyncband::spmc::unbounded() + } + + async fn send(sender: &mut Self::Sender, value: usize) { + sender.send(value).unwrap(); + } + + async fn recv(receiver: &Self::Receiver) -> Option { + receiver.recv().await.ok() + } +} + +impl Channel for Bounded { + type Sender = C::Sender; + type Receiver = C::Receiver; + + fn channel() -> (Self::Sender, Self::Receiver) { + C::channel(CAPACITY) + } + + async fn send(sender: &mut Self::Sender, value: usize) { + C::send_async(sender, value).await; + } + + async fn recv(receiver: &Self::Receiver) -> Option { + C::recv_async(receiver).await + } +} + +impl Channel for Unbounded { + type Sender = C::Sender; + type Receiver = C::Receiver; + + fn channel() -> (Self::Sender, Self::Receiver) { + C::channel() + } + + async fn send(sender: &mut Self::Sender, value: usize) { + C::send(sender, value); + } + + async fn recv(receiver: &Self::Receiver) -> Option { + C::recv_async(receiver).await + } +} + pub trait BoundedMpmc: Send + Sync + 'static { type Sender: Clone + Send + Sync + 'static; type Receiver: Clone + Send + Sync + 'static; @@ -57,7 +142,7 @@ pub trait UnboundedMpmc: Send + Sync + 'static { } } -impl BoundedMpmc for Asyncband { +impl BoundedMpmc for Mpmc { type Receiver = asyncband::mpmc::BoundedReceiver; type Sender = asyncband::mpmc::BoundedSender; @@ -77,7 +162,7 @@ impl BoundedMpmc for Asyncband { } } -impl UnboundedMpmc for Asyncband { +impl UnboundedMpmc for Mpmc { type Receiver = asyncband::mpmc::UnboundedReceiver; type Sender = asyncband::mpmc::UnboundedSender; diff --git a/benchmarks/mpmc/mod.rs b/benchmarks/channels/mod.rs similarity index 63% rename from benchmarks/mpmc/mod.rs rename to benchmarks/channels/mod.rs index 440df0ac..b70e390e 100644 --- a/benchmarks/mpmc/mod.rs +++ b/benchmarks/channels/mod.rs @@ -15,5 +15,24 @@ // specific language governing permissions and limitations // under the License. +use tokio::runtime::Runtime; + pub mod adapters; -pub mod support; +pub mod mpmc; +pub mod spmc; + +pub const BATCH_MESSAGES: usize = 16_384; +pub const BOUNDED_CAPACITY: usize = 64; + +pub fn runtime(worker_threads: usize) -> Runtime { + if worker_threads == 0 { + tokio::runtime::Builder::new_current_thread() + .build() + .unwrap() + } else { + tokio::runtime::Builder::new_multi_thread() + .worker_threads(worker_threads) + .build() + .unwrap() + } +} diff --git a/benchmarks/mpmc/support.rs b/benchmarks/channels/mpmc.rs similarity index 77% rename from benchmarks/mpmc/support.rs rename to benchmarks/channels/mpmc.rs index 0eb3d714..d8f04d80 100644 --- a/benchmarks/mpmc/support.rs +++ b/benchmarks/channels/mpmc.rs @@ -15,8 +15,6 @@ // specific language governing permissions and limitations // under the License. -use std::future::Future; -use std::marker::PhantomData; use std::sync::Arc; use std::sync::Barrier; use std::thread; @@ -26,12 +24,11 @@ use divan::black_box; use tokio::runtime::Runtime; use tokio::task::JoinSet; +use super::BATCH_MESSAGES; use super::adapters::BoundedMpmc; +use super::adapters::Channel; use super::adapters::UnboundedMpmc; -pub const BATCH_MESSAGES: usize = 16_384; -pub const BOUNDED_CAPACITY: usize = 64; - #[derive(Clone, Copy, Debug)] pub struct Topology { pub producers: usize, @@ -57,66 +54,6 @@ pub const TOPOLOGIES: &[Topology] = &[ }, ]; -pub trait ConcurrentMpmc: Send + Sync + 'static { - type Sender: Clone + Send + Sync + 'static; - type Receiver: Clone + Send + Sync + 'static; - - fn channel() -> (Self::Sender, Self::Receiver); - fn send(sender: &Self::Sender, value: usize) -> impl Future + Send; - fn recv(receiver: &Self::Receiver) -> impl Future> + Send; -} - -pub struct Bounded(PhantomData); - -impl ConcurrentMpmc for Bounded { - type Sender = C::Sender; - type Receiver = C::Receiver; - - fn channel() -> (Self::Sender, Self::Receiver) { - C::channel(CAPACITY) - } - - async fn send(sender: &Self::Sender, value: usize) { - C::send_async(sender, value).await; - } - - async fn recv(receiver: &Self::Receiver) -> Option { - C::recv_async(receiver).await - } -} - -pub struct Unbounded(PhantomData); - -impl ConcurrentMpmc for Unbounded { - type Sender = C::Sender; - type Receiver = C::Receiver; - - fn channel() -> (Self::Sender, Self::Receiver) { - C::channel() - } - - async fn send(sender: &Self::Sender, value: usize) { - C::send(sender, value); - } - - async fn recv(receiver: &Self::Receiver) -> Option { - C::recv_async(receiver).await - } -} - -pub fn runtime(worker_threads: usize) -> Runtime { - if worker_threads == 0 { - tokio::runtime::Builder::new_current_thread() - .build() - .unwrap() - } else { - tokio::runtime::Builder::new_multi_thread() - .worker_threads(worker_threads) - .build() - .unwrap() - } -} - // The caller only coordinates the batch. All measured sends and receives run in spawned tasks, // including on the current-thread runtime; no data is received by Runtime::block_on itself. pub struct TaskBatch { @@ -125,7 +62,10 @@ pub struct TaskBatch { } impl TaskBatch { - pub fn new(runtime: &Runtime, topology: Topology) -> Self { + pub fn new(runtime: &Runtime, topology: Topology) -> Self + where + C::Sender: Clone, + { assert_eq!(BATCH_MESSAGES % topology.producers, 0); let (sender, receiver) = C::channel(); let start = Arc::new(tokio::sync::Barrier::new( @@ -134,14 +74,14 @@ impl TaskBatch { let messages_per_producer = BATCH_MESSAGES / topology.producers; let mut workers = JoinSet::new(); for producer in 0..topology.producers { - let sender = sender.clone(); + let mut sender = sender.clone(); let start = start.clone(); workers.spawn_on( async move { start.wait().await; let first = producer * messages_per_producer; for value in first..first + messages_per_producer { - C::send(&sender, black_box(value)).await; + C::send(&mut sender, black_box(value)).await; } // Completion drops this sender so receivers can observe the end of input. (0, 0) diff --git a/benchmarks/channels/spmc.rs b/benchmarks/channels/spmc.rs new file mode 100644 index 00000000..2e99c5c1 --- /dev/null +++ b/benchmarks/channels/spmc.rs @@ -0,0 +1,86 @@ +// 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 divan::black_box; +use tokio::runtime::Runtime; +use tokio::task::JoinSet; + +use super::BATCH_MESSAGES; +use super::adapters::Channel; + +pub const CONSUMERS: &[usize] = &[1, 2, 4, 8]; + +pub struct TaskBatch { + start: Arc, + tasks: JoinSet<(usize, usize)>, +} + +impl TaskBatch { + pub fn new(runtime: &Runtime, consumers: usize) -> Self { + let (mut sender, receiver) = C::channel(); + let start = Arc::new(tokio::sync::Barrier::new(consumers + 2)); + let mut tasks = JoinSet::new(); + let producer_start = start.clone(); + tasks.spawn_on( + async move { + producer_start.wait().await; + for value in 0..BATCH_MESSAGES { + C::send(&mut sender, black_box(value)).await; + } + // The sender is moved once and dropped on completion so consumers can drain. + (0, 0) + }, + runtime.handle(), + ); + for _ in 0..consumers { + let receiver = receiver.clone(); + let start = start.clone(); + tasks.spawn_on( + async move { + start.wait().await; + let mut count = 0; + let mut checksum = 0usize; + // Consumers compete freely, with no fixed per-consumer quota. + while let Some(value) = C::recv(&receiver).await { + count += 1; + checksum = checksum.wrapping_add(value); + } + (count, checksum) + }, + runtime.handle(), + ); + } + drop(receiver); + Self { start, tasks } + } + + pub async fn run(&mut self) -> (usize, usize) { + self.start.wait().await; + let mut count = 0; + let mut checksum = 0usize; + while let Some(result) = self.tasks.join_next().await { + let (received, sum) = result.expect("benchmark task panicked"); + count += received; + checksum = checksum.wrapping_add(sum); + } + assert_eq!(count, BATCH_MESSAGES); + assert_eq!(checksum, BATCH_MESSAGES * (BATCH_MESSAGES - 1) / 2); + black_box((count, checksum)) + } +} diff --git a/benchmarks/ecosystem/main.rs b/benchmarks/ecosystem/main.rs index 6cc187c2..c24d5ef4 100644 --- a/benchmarks/ecosystem/main.rs +++ b/benchmarks/ecosystem/main.rs @@ -17,14 +17,15 @@ mod broadcast; mod mpmc; - -#[allow(dead_code)] -#[path = "../mpmc/mod.rs"] -mod mpmc_support; mod mpsc; +mod spmc; mod waitgroup; mod watch; +#[allow(dead_code)] +#[path = "../channels/mod.rs"] +mod channels; + #[allow(dead_code)] #[path = "../asyncband/support.rs"] mod support; diff --git a/benchmarks/ecosystem/mpmc/bounded.rs b/benchmarks/ecosystem/mpmc/bounded.rs index c908b416..87b61ce4 100644 --- a/benchmarks/ecosystem/mpmc/bounded.rs +++ b/benchmarks/ecosystem/mpmc/bounded.rs @@ -18,21 +18,21 @@ use divan::Bencher; use divan::counter::ItemsCount; -use crate::mpmc_support::adapters::AsyncChannel; -use crate::mpmc_support::adapters::Asyncband; -use crate::mpmc_support::adapters::BoundedMpmc; -use crate::mpmc_support::adapters::Flume; -use crate::mpmc_support::support::BATCH_MESSAGES; -use crate::mpmc_support::support::BOUNDED_CAPACITY; -use crate::mpmc_support::support::Bounded; -use crate::mpmc_support::support::TOPOLOGIES; -use crate::mpmc_support::support::TaskBatch; -use crate::mpmc_support::support::ThreadBatch; -use crate::mpmc_support::support::Topology; -use crate::mpmc_support::support::runtime; +use crate::channels::BATCH_MESSAGES; +use crate::channels::BOUNDED_CAPACITY; +use crate::channels::adapters::AsyncChannel; +use crate::channels::adapters::Bounded; +use crate::channels::adapters::BoundedMpmc; +use crate::channels::adapters::Flume; +use crate::channels::adapters::Mpmc; +use crate::channels::mpmc::TOPOLOGIES; +use crate::channels::mpmc::TaskBatch; +use crate::channels::mpmc::ThreadBatch; +use crate::channels::mpmc::Topology; +use crate::channels::runtime; #[divan::bench( - types = [Asyncband, AsyncChannel, Flume], + types = [Mpmc, AsyncChannel, Flume], args = TOPOLOGIES, sample_count = 20, sample_size = 1, @@ -45,7 +45,7 @@ fn blocking_threads(bencher: Bencher, topology: Topology) { } #[divan::bench( - types = [Asyncband, AsyncChannel, Flume], + types = [Mpmc, AsyncChannel, Flume], consts = [0, 4], args = TOPOLOGIES, sample_count = 20, diff --git a/benchmarks/ecosystem/mpmc/unbounded.rs b/benchmarks/ecosystem/mpmc/unbounded.rs index 6013ff55..3d244ef1 100644 --- a/benchmarks/ecosystem/mpmc/unbounded.rs +++ b/benchmarks/ecosystem/mpmc/unbounded.rs @@ -18,20 +18,20 @@ use divan::Bencher; use divan::counter::ItemsCount; -use crate::mpmc_support::adapters::AsyncChannel; -use crate::mpmc_support::adapters::Asyncband; -use crate::mpmc_support::adapters::Flume; -use crate::mpmc_support::adapters::UnboundedMpmc; -use crate::mpmc_support::support::BATCH_MESSAGES; -use crate::mpmc_support::support::TOPOLOGIES; -use crate::mpmc_support::support::TaskBatch; -use crate::mpmc_support::support::ThreadBatch; -use crate::mpmc_support::support::Topology; -use crate::mpmc_support::support::Unbounded; -use crate::mpmc_support::support::runtime; +use crate::channels::BATCH_MESSAGES; +use crate::channels::adapters::AsyncChannel; +use crate::channels::adapters::Flume; +use crate::channels::adapters::Mpmc; +use crate::channels::adapters::Unbounded; +use crate::channels::adapters::UnboundedMpmc; +use crate::channels::mpmc::TOPOLOGIES; +use crate::channels::mpmc::TaskBatch; +use crate::channels::mpmc::ThreadBatch; +use crate::channels::mpmc::Topology; +use crate::channels::runtime; #[divan::bench( - types = [Asyncband, AsyncChannel, Flume], + types = [Mpmc, AsyncChannel, Flume], args = TOPOLOGIES, sample_count = 20, sample_size = 1, @@ -44,7 +44,7 @@ fn blocking_threads(bencher: Bencher, topology: Topology) { } #[divan::bench( - types = [Asyncband, AsyncChannel, Flume], + types = [Mpmc, AsyncChannel, Flume], consts = [0, 4], args = TOPOLOGIES, sample_count = 20, diff --git a/benchmarks/ecosystem/spmc/bounded.rs b/benchmarks/ecosystem/spmc/bounded.rs new file mode 100644 index 00000000..8ee5ec82 --- /dev/null +++ b/benchmarks/ecosystem/spmc/bounded.rs @@ -0,0 +1,48 @@ +// 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 divan::Bencher; +use divan::counter::ItemsCount; + +use crate::channels::BATCH_MESSAGES; +use crate::channels::adapters::AsyncChannel; +use crate::channels::adapters::Bounded; +use crate::channels::adapters::Channel; +use crate::channels::adapters::Flume; +use crate::channels::adapters::Mpmc; +use crate::channels::adapters::Spmc; +use crate::channels::runtime; +use crate::channels::spmc::CONSUMERS; +use crate::channels::spmc::TaskBatch; + +#[divan::bench( + types = [Spmc, Mpmc, AsyncChannel, Flume], + consts = [0, 4], + args = CONSUMERS, + sample_count = 20, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn tokio_tasks(bencher: Bencher, consumers: usize) +where + Bounded: Channel, +{ + let runtime = runtime(WORKERS); + bencher + .with_inputs(|| TaskBatch::new::>(&runtime, consumers)) + .bench_local_refs(|batch| runtime.block_on(batch.run())); +} diff --git a/benchmarks/ecosystem/spmc/mod.rs b/benchmarks/ecosystem/spmc/mod.rs new file mode 100644 index 00000000..06a557d9 --- /dev/null +++ b/benchmarks/ecosystem/spmc/mod.rs @@ -0,0 +1,28 @@ +// 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. + +//! One producer moves its sender into a task; consumers compete until disconnection. Runtime, +//! channel, and task creation are excluded from timing; data operations and task completion are +//! included. Synchronous unbounded sends can finish before consumers run on a current-thread +//! executor, so the four-worker cases measure concurrent consumer contention. +//! +//! Compile with `cargo x bench --no-run`, then run the ecosystem executable with +//! `--bench --color never --sample-count 100 'spmc::'`. Record repeated measurements serially, +//! without concurrent builds or tests, together with the commit, toolchain, and machine details. + +mod bounded; +mod unbounded; diff --git a/benchmarks/ecosystem/spmc/unbounded.rs b/benchmarks/ecosystem/spmc/unbounded.rs new file mode 100644 index 00000000..a16c4d5d --- /dev/null +++ b/benchmarks/ecosystem/spmc/unbounded.rs @@ -0,0 +1,48 @@ +// 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 divan::Bencher; +use divan::counter::ItemsCount; + +use crate::channels::BATCH_MESSAGES; +use crate::channels::adapters::AsyncChannel; +use crate::channels::adapters::Channel; +use crate::channels::adapters::Flume; +use crate::channels::adapters::Mpmc; +use crate::channels::adapters::Spmc; +use crate::channels::adapters::Unbounded; +use crate::channels::runtime; +use crate::channels::spmc::CONSUMERS; +use crate::channels::spmc::TaskBatch; + +#[divan::bench( + types = [Spmc, Mpmc, AsyncChannel, Flume], + consts = [0, 4], + args = CONSUMERS, + sample_count = 20, + sample_size = 1, + counter = ItemsCount::new(BATCH_MESSAGES), +)] +fn tokio_tasks(bencher: Bencher, consumers: usize) +where + Unbounded: Channel, +{ + let runtime = runtime(WORKERS); + bencher + .with_inputs(|| TaskBatch::new::>(&runtime, consumers)) + .bench_local_refs(|batch| runtime.block_on(batch.run())); +} diff --git a/benchmarks/tests/mpmc_batch.rs b/benchmarks/tests/channel_batch.rs similarity index 55% rename from benchmarks/tests/mpmc_batch.rs rename to benchmarks/tests/channel_batch.rs index d4362790..edc14e68 100644 --- a/benchmarks/tests/mpmc_batch.rs +++ b/benchmarks/tests/channel_batch.rs @@ -20,11 +20,12 @@ use std::thread; use std::time::Duration; #[allow(dead_code)] -#[path = "../mpmc/mod.rs"] -mod mpmc_support; +#[path = "../channels/mod.rs"] +mod channels; -use mpmc_support::adapters; -use mpmc_support::support; +use channels::adapters; +use channels::mpmc; +use channels::spmc; // Deliberately drain only after the last producer drops its sender. This makes // retaining senders across the completion barrier deadlock deterministically. @@ -72,8 +73,8 @@ fn assert_completes(run: impl FnOnce() + Send + 'static) { #[test] fn thread_batch_closes_before_waiting_for_consumers() { assert_completes(|| { - for &topology in support::TOPOLOGIES { - let batch = support::ThreadBatch::new_unbounded::(topology); + for &topology in mpmc::TOPOLOGIES { + let batch = mpmc::ThreadBatch::new_unbounded::(topology); batch.run(); } }); @@ -83,7 +84,7 @@ fn thread_batch_closes_before_waiting_for_consumers() { fn flume_batches_complete_with_competing_consumers() { assert_completes(|| { for _ in 0..100 { - let batch = support::ThreadBatch::new_unbounded::(support::Topology { + let batch = mpmc::ThreadBatch::new_unbounded::(mpmc::Topology { producers: 1, consumers: 8, }); @@ -94,31 +95,55 @@ fn flume_batches_complete_with_competing_consumers() { #[test] fn tokio_batches_drain_all_messages_before_completion() { - fn check(runtime: &tokio::runtime::Runtime) { - for topology in support::TOPOLOGIES - .iter() - .copied() - .chain([support::Topology { - producers: 1, - consumers: 3, - }]) - { + fn check(runtime: &tokio::runtime::Runtime) + where + C::Sender: Clone, + { + for topology in mpmc::TOPOLOGIES.iter().copied().chain([mpmc::Topology { + producers: 1, + consumers: 3, + }]) { // Three consumers cannot receive equal quotas from a 16,384-message batch. - let mut batch = support::TaskBatch::new::(runtime, topology); + let mut batch = mpmc::TaskBatch::new::(runtime, topology); + runtime.block_on(batch.run()); + } + } + + assert_completes(|| { + for workers in [0, 4] { + let runtime = channels::runtime(workers); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + } + }); +} + +#[test] +fn spmc_batches_drain_all_messages_before_completion() { + fn check(runtime: &tokio::runtime::Runtime) { + for consumers in [1, 3, 8] { + let mut batch = spmc::TaskBatch::new::(runtime, consumers); runtime.block_on(batch.run()); } } assert_completes(|| { for workers in [0, 4] { - let runtime = support::runtime(workers); - check::>(&runtime); - check::>(&runtime); - check::>(&runtime); - check::>(&runtime); - check::>(&runtime); - check::>(&runtime); - check::>(&runtime); + let runtime = channels::runtime(workers); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); + check::>(&runtime); } }); } diff --git a/tests-integration/Cargo.toml b/tests-integration/Cargo.toml index 44942255..e93720d7 100644 --- a/tests-integration/Cargo.toml +++ b/tests-integration/Cargo.toml @@ -48,6 +48,7 @@ asyncband = { workspace = true, features = [ "semaphore", "shutdown", "singleflight", + "spmc", "waitgroup", "watch", ] } diff --git a/tests-integration/tests/spmc_test/concurrency.rs b/tests-integration/tests/spmc_test/concurrency.rs new file mode 100644 index 00000000..6ef9c839 --- /dev/null +++ b/tests-integration/tests/spmc_test/concurrency.rs @@ -0,0 +1,105 @@ +// 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::time::Duration; + +use asyncband::spmc; +use tokio::sync::Barrier; +use tokio::task::JoinHandle; + +const CONSUMERS: usize = 8; +const TOTAL: usize = 2_048; + +async fn assert_delivered_exactly_once( + producer: JoinHandle<()>, + consumers: Vec>>, +) { + tokio::time::timeout(Duration::from_secs(10), async { + producer.await.unwrap(); + let mut received = Vec::with_capacity(TOTAL); + for consumer in consumers { + let values = consumer.await.unwrap(); + assert!(values.windows(2).all(|pair| pair[0] < pair[1])); + received.extend(values); + } + received.sort_unstable(); + assert_eq!(received, (0..TOTAL).collect::>()); + }) + .await + .expect("SPMC sender and all consumers must make progress"); +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] +async fn bounded_values_are_delivered_exactly_once_to_eight_consumers() { + for capacity in [1, 64] { + let (mut sender, receiver) = spmc::bounded(capacity); + let start = Arc::new(Barrier::new(CONSUMERS + 1)); + let consumers = (0..CONSUMERS) + .map(|_| { + let receiver = receiver.clone(); + let start = start.clone(); + tokio::spawn(async move { + start.wait().await; + let mut values = Vec::new(); + while let Ok(value) = receiver.recv().await { + values.push(value); + } + values + }) + }) + .collect(); + drop(receiver); + let producer = tokio::spawn(async move { + start.wait().await; + for value in 0..TOTAL { + sender.send(value).await.unwrap(); + } + }); + assert_delivered_exactly_once(producer, consumers).await; + } +} + +#[tokio::test(flavor = "multi_thread", worker_threads = 4)] +#[cfg_attr(miri, ignore = "requires an OS-backed Tokio runtime")] +async fn unbounded_values_are_delivered_exactly_once_to_eight_consumers() { + let (mut sender, receiver) = spmc::unbounded(); + let start = Arc::new(Barrier::new(CONSUMERS + 1)); + let consumers = (0..CONSUMERS) + .map(|_| { + let receiver = receiver.clone(); + let start = start.clone(); + tokio::spawn(async move { + start.wait().await; + let mut values = Vec::new(); + while let Ok(value) = receiver.recv().await { + values.push(value); + } + values + }) + }) + .collect(); + drop(receiver); + let producer = tokio::spawn(async move { + start.wait().await; + for value in 0..TOTAL { + sender.send(value).unwrap(); + } + }); + assert_delivered_exactly_once(producer, consumers).await; +} diff --git a/tests-integration/tests/spmc_test/main.rs b/tests-integration/tests/spmc_test/main.rs new file mode 100644 index 00000000..684c3f35 --- /dev/null +++ b/tests-integration/tests/spmc_test/main.rs @@ -0,0 +1,180 @@ +// 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::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Poll; + +use asyncband::spmc; +use asyncband::spmc::RecvError; +use asyncband::spmc::TryRecvError; +use tests_integration::poll_once; + +// Public queue contracts. The other suites cover notifications, cancellation, and concurrency. +mod concurrency; +mod notification; + +#[derive(Debug)] +struct DropSpy(Arc); + +impl Drop for DropSpy { + fn drop(&mut self) { + self.0.fetch_add(1, Ordering::SeqCst); + } +} + +#[test] +fn bounded_receivers_compete_in_fifo_order_and_drain_after_sender_drop() { + let (mut sender, receiver) = spmc::bounded(2); + let competing = receiver.clone(); + sender.try_send(10).unwrap(); + sender.try_send(20).unwrap(); + assert_eq!(receiver.try_recv(), Ok(10)); + assert_eq!(competing.try_recv(), Ok(20)); + assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); + + sender.try_send(30).unwrap(); + sender.try_send(40).unwrap(); + drop(sender); + assert_eq!(poll_once(pin!(receiver.recv())), Poll::Ready(Ok(30))); + assert_eq!(poll_once(pin!(competing.recv())), Poll::Ready(Ok(40))); + assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected)); + assert_eq!( + poll_once(pin!(competing.recv())), + Poll::Ready(Err(RecvError::Disconnected)) + ); +} + +#[test] +fn bounded_only_last_receiver_disconnects_and_returns_unsent_value() { + let (mut sender, receiver) = spmc::bounded(2); + let competing = receiver.clone(); + drop(receiver); + sender.try_send(10).unwrap(); + assert_eq!(competing.try_recv(), Ok(10)); + drop(competing); + let error = sender.try_send(20).unwrap_err(); + assert_eq!(error.as_inner(), &20); + assert_eq!(error.into_inner(), 20); +} + +#[test] +fn bounded_buffered_received_and_rejected_values_are_each_dropped_once() { + let (mut sender, receiver) = spmc::bounded(2); + let competing = receiver.clone(); + let drops: Vec<_> = (0..4).map(|_| Arc::new(AtomicUsize::new(0))).collect(); + sender.try_send(DropSpy(drops[0].clone())).unwrap(); + sender.try_send(DropSpy(drops[1].clone())).unwrap(); + let received = receiver.try_recv().unwrap(); + sender.try_send(DropSpy(drops[2].clone())).unwrap(); + + drop(receiver); + assert!(drops.iter().all(|count| count.load(Ordering::SeqCst) == 0)); + drop(competing); + assert_eq!(drops[1].load(Ordering::SeqCst), 1); + assert_eq!(drops[2].load(Ordering::SeqCst), 1); + + let rejected = sender.try_send(DropSpy(drops[3].clone())).unwrap_err(); + assert_eq!(drops[3].load(Ordering::SeqCst), 0); + drop(rejected.into_inner()); + drop(received); + drop(sender); + assert!(drops.iter().all(|count| count.load(Ordering::SeqCst) == 1)); +} + +#[test] +fn unbounded_receivers_compete_in_fifo_order_and_drain_after_sender_drop() { + let (mut sender, receiver) = spmc::unbounded(); + let competing = receiver.clone(); + sender.send(10).unwrap(); + sender.send(20).unwrap(); + assert_eq!(receiver.try_recv(), Ok(10)); + assert_eq!(competing.try_recv(), Ok(20)); + assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); + + sender.send(30).unwrap(); + sender.send(40).unwrap(); + drop(sender); + assert_eq!(poll_once(pin!(receiver.recv())), Poll::Ready(Ok(30))); + assert_eq!(poll_once(pin!(competing.recv())), Poll::Ready(Ok(40))); + assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected)); + assert_eq!( + poll_once(pin!(competing.recv())), + Poll::Ready(Err(RecvError::Disconnected)) + ); +} + +#[test] +fn unbounded_only_last_receiver_disconnects_and_returns_unsent_value() { + let (mut sender, receiver) = spmc::unbounded(); + let competing = receiver.clone(); + drop(receiver); + sender.send(10).unwrap(); + assert_eq!(competing.try_recv(), Ok(10)); + drop(competing); + let error = sender.send(20).unwrap_err(); + assert_eq!(error.as_inner(), &20); + assert_eq!(error.into_inner(), 20); +} + +#[test] +fn unbounded_buffered_received_and_rejected_values_are_each_dropped_once() { + let (mut sender, receiver) = spmc::unbounded(); + let competing = receiver.clone(); + let drops: Vec<_> = (0..4).map(|_| Arc::new(AtomicUsize::new(0))).collect(); + sender.send(DropSpy(drops[0].clone())).unwrap(); + sender.send(DropSpy(drops[1].clone())).unwrap(); + let received = receiver.try_recv().unwrap(); + sender.send(DropSpy(drops[2].clone())).unwrap(); + + drop(receiver); + assert!(drops.iter().all(|count| count.load(Ordering::SeqCst) == 0)); + drop(competing); + assert_eq!(drops[1].load(Ordering::SeqCst), 1); + assert_eq!(drops[2].load(Ordering::SeqCst), 1); + + let rejected = sender.send(DropSpy(drops[3].clone())).unwrap_err(); + assert_eq!(drops[3].load(Ordering::SeqCst), 0); + drop(rejected.into_inner()); + drop(received); + drop(sender); + assert!(drops.iter().all(|count| count.load(Ordering::SeqCst) == 1)); +} + +#[test] +#[should_panic(expected = "spmc bounded queue requires capacity > 0")] +fn bounded_rejects_zero_capacity() { + let _ = spmc::bounded::<()>(0); +} + +#[test] +fn endpoint_and_future_traits_allow_send_but_not_sync_payloads() { + fn assert_traits() {} + fn assert_send(_: T) {} + assert_traits::>>(); + assert_traits::>>(); + assert_traits::>>(); + assert_traits::>>(); + let (mut sender, receiver) = spmc::bounded::>(1); + assert_send(sender.send(Cell::new(1))); + assert_send(receiver.recv()); + let (_sender, receiver) = spmc::unbounded::>(); + assert_send(receiver.recv()); +} diff --git a/tests-integration/tests/spmc_test/notification.rs b/tests-integration/tests/spmc_test/notification.rs new file mode 100644 index 00000000..51f0e0fb --- /dev/null +++ b/tests-integration/tests/spmc_test/notification.rs @@ -0,0 +1,254 @@ +// 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::pin::pin; +use std::sync::Arc; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; +use std::task::Poll; + +use asyncband::spmc; +use asyncband::spmc::RecvError; +use asyncband::spmc::TryRecvError; +use asyncband::spmc::TrySendError; +use tests_integration::WakeCounter; +use tests_integration::expect_ready; +use tests_integration::poll_once; +use tests_integration::poll_with; + +use super::DropSpy; + +#[test] +fn bounded_one_send_wakes_one_of_eight_receivers_and_cancellation_hands_off() { + let (mut sender, receiver) = spmc::bounded::(2); + let receivers: Vec<_> = (0..8).map(|_| receiver.clone()).collect(); + let mut pending: Vec<_> = receivers + .iter() + .map(|receiver| { + let (waker, wakes) = WakeCounter::new(); + let mut receive = Box::pin(receiver.recv()); + assert!(poll_with(receive.as_mut(), &waker).is_pending()); + (receive, wakes) + }) + .collect(); + + sender.try_send(7).unwrap(); + // Each cancellation passes the notification to exactly one remaining receiver. + for _ in 0..7 { + let (cancelled, wakes) = pending.remove(0); + assert_eq!(wakes.count(), 1); + assert!(pending.iter().all(|(_, wakes)| wakes.count() == 0)); + drop(cancelled); + } + let (mut last, wakes) = pending.pop().unwrap(); + assert_eq!(wakes.count(), 1); + assert_eq!(poll_once(last.as_mut()), Poll::Ready(Ok(7))); + assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); +} + +#[test] +fn bounded_cancelling_before_notification_removes_the_waiter() { + let (mut sender, receiver) = spmc::bounded::(2); + let competing = receiver.clone(); + let (cancelled_waker, cancelled_wakes) = WakeCounter::new(); + let (waiting_waker, waiting_wakes) = WakeCounter::new(); + let mut cancelled = Box::pin(receiver.recv()); + let mut waiting = Box::pin(competing.recv()); + assert!(poll_with(cancelled.as_mut(), &cancelled_waker).is_pending()); + assert!(poll_with(waiting.as_mut(), &waiting_waker).is_pending()); + + drop(cancelled); + sender.try_send(5).unwrap(); + assert_eq!(cancelled_wakes.count(), 0); + assert_eq!(waiting_wakes.count(), 1); + assert_eq!(poll_once(waiting.as_mut()), Poll::Ready(Ok(5))); +} + +#[test] +fn bounded_sender_disconnection_wakes_all_receivers() { + let (sender, receiver) = spmc::bounded::(2); + let receivers: Vec<_> = (0..8).map(|_| receiver.clone()).collect(); + let pending: Vec<_> = receivers + .iter() + .map(|receiver| { + let (waker, wakes) = WakeCounter::new(); + let mut receive = Box::pin(receiver.recv()); + assert!(poll_with(receive.as_mut(), &waker).is_pending()); + (receive, wakes) + }) + .collect(); + + drop(sender); + for (mut receive, wakes) in pending { + assert_eq!(wakes.count(), 1); + assert_eq!( + poll_once(receive.as_mut()), + Poll::Ready(Err(RecvError::Disconnected)) + ); + } +} + +#[test] +fn unbounded_one_send_wakes_one_of_eight_receivers_and_cancellation_hands_off() { + let (mut sender, receiver) = spmc::unbounded::(); + let receivers: Vec<_> = (0..8).map(|_| receiver.clone()).collect(); + let mut pending: Vec<_> = receivers + .iter() + .map(|receiver| { + let (waker, wakes) = WakeCounter::new(); + let mut receive = Box::pin(receiver.recv()); + assert!(poll_with(receive.as_mut(), &waker).is_pending()); + (receive, wakes) + }) + .collect(); + + sender.send(7).unwrap(); + // Each cancellation passes the notification to exactly one remaining receiver. + for _ in 0..7 { + let (cancelled, wakes) = pending.remove(0); + assert_eq!(wakes.count(), 1); + assert!(pending.iter().all(|(_, wakes)| wakes.count() == 0)); + drop(cancelled); + } + let (mut last, wakes) = pending.pop().unwrap(); + assert_eq!(wakes.count(), 1); + assert_eq!(poll_once(last.as_mut()), Poll::Ready(Ok(7))); + assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); +} + +#[test] +fn unbounded_cancelling_before_notification_removes_the_waiter() { + let (mut sender, receiver) = spmc::unbounded::(); + let competing = receiver.clone(); + let (cancelled_waker, cancelled_wakes) = WakeCounter::new(); + let (waiting_waker, waiting_wakes) = WakeCounter::new(); + let mut cancelled = Box::pin(receiver.recv()); + let mut waiting = Box::pin(competing.recv()); + assert!(poll_with(cancelled.as_mut(), &cancelled_waker).is_pending()); + assert!(poll_with(waiting.as_mut(), &waiting_waker).is_pending()); + + drop(cancelled); + sender.send(5).unwrap(); + assert_eq!(cancelled_wakes.count(), 0); + assert_eq!(waiting_wakes.count(), 1); + assert_eq!(poll_once(waiting.as_mut()), Poll::Ready(Ok(5))); +} + +#[test] +fn unbounded_sender_disconnection_wakes_all_receivers() { + let (sender, receiver) = spmc::unbounded::(); + let receivers: Vec<_> = (0..8).map(|_| receiver.clone()).collect(); + let pending: Vec<_> = receivers + .iter() + .map(|receiver| { + let (waker, wakes) = WakeCounter::new(); + let mut receive = Box::pin(receiver.recv()); + assert!(poll_with(receive.as_mut(), &waker).is_pending()); + (receive, wakes) + }) + .collect(); + + drop(sender); + for (mut receive, wakes) in pending { + assert_eq!(wakes.count(), 1); + assert_eq!( + poll_once(receive.as_mut()), + Poll::Ready(Err(RecvError::Disconnected)) + ); + } +} + +#[test] +fn bounded_capacity_and_pending_send_progress() { + for capacity in [1, 2, 3, 8] { + let (mut sender, receiver) = spmc::bounded(capacity); + for value in 0..capacity { + sender.try_send(value).unwrap(); + } + assert_eq!(sender.try_send(capacity), Err(TrySendError::Full(capacity))); + let mut waiting = Box::pin(sender.send(capacity)); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(waiting.as_mut(), &waker).is_pending()); + + assert_eq!(receiver.try_recv(), Ok(0)); + assert_eq!(wakes.count(), 1); + assert_eq!(poll_once(waiting.as_mut()), Poll::Ready(Ok(()))); + drop(waiting); + drop(sender); + for value in 1..=capacity { + assert_eq!(receiver.try_recv(), Ok(value)); + } + assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected)); + } +} + +#[test] +fn cancelling_an_unnotified_send_preserves_the_buffered_value() { + let (mut sender, receiver) = spmc::bounded(1); + let drops = Arc::new(AtomicUsize::new(0)); + sender.try_send(DropSpy(drops.clone())).unwrap(); + let mut cancelled = Box::pin(sender.send(DropSpy(drops.clone()))); + assert!(poll_once(cancelled.as_mut()).is_pending()); + + drop(cancelled); + assert_eq!(drops.load(Ordering::SeqCst), 1); + drop(receiver.try_recv().unwrap()); + assert_eq!(drops.load(Ordering::SeqCst), 2); + + expect_ready(poll_once(pin!(sender.send(DropSpy(drops.clone()))))).unwrap(); + drop(receiver); + drop(sender); + assert_eq!(drops.load(Ordering::SeqCst), 3); +} + +#[test] +fn cancelling_a_notified_send_leaves_capacity_for_the_next_send() { + let (mut sender, receiver) = spmc::bounded(1); + let drops = Arc::new(AtomicUsize::new(0)); + sender.try_send(DropSpy(drops.clone())).unwrap(); + let mut cancelled = Box::pin(sender.send(DropSpy(drops.clone()))); + let (waker, wakes) = WakeCounter::new(); + assert!(poll_with(cancelled.as_mut(), &waker).is_pending()); + + drop(receiver.try_recv().unwrap()); + assert_eq!(wakes.count(), 1); + assert_eq!(drops.load(Ordering::SeqCst), 1); + drop(cancelled); + assert_eq!(drops.load(Ordering::SeqCst), 2); + + expect_ready(poll_once(pin!(sender.send(DropSpy(drops.clone()))))).unwrap(); + drop(receiver); + drop(sender); + assert_eq!(drops.load(Ordering::SeqCst), 3); +} + +#[test] +fn last_receiver_wakes_pending_sender_and_returns_its_value() { + let (mut sender, receiver) = spmc::bounded(1); + let competing = receiver.clone(); + sender.try_send(0).unwrap(); + let (waker, wakes) = WakeCounter::new(); + let mut pending = Box::pin(sender.send(1)); + assert!(poll_with(pending.as_mut(), &waker).is_pending()); + + drop(receiver); + assert_eq!(wakes.count(), 0); + drop(competing); + assert_eq!(wakes.count(), 1); + let error = expect_ready(poll_once(pending.as_mut())).unwrap_err(); + assert_eq!(error.into_inner(), 1); +} diff --git a/xtask/src/main.rs b/xtask/src/main.rs index f2733663..6ae91be7 100644 --- a/xtask/src/main.rs +++ b/xtask/src/main.rs @@ -122,6 +122,7 @@ impl CommandMiri { )); run_command(make_miri_cmd("tests-integration", &["--test", "mpsc_test"])); run_command(make_miri_cmd("tests-integration", &["--test", "mpmc_test"])); + run_command(make_miri_cmd("tests-integration", &["--test", "spmc_test"])); run_command(make_miri_cmd( "tests-integration", &["--test", "phaser_test"],