From 026699efd2f7e394f717173bab6e98686840b10e Mon Sep 17 00:00:00 2001 From: tison Date: Sat, 3 Oct 2026 16:45:49 +0800 Subject: [PATCH] perf(spmc): specialize single-producer coordination --- CHANGELOG.md | 4 + asyncband/src/spmc/bounded.rs | 36 ++- asyncband/src/spmc/queue.rs | 231 +++++++++--------- asyncband/src/spmc/unbounded.rs | 28 +-- benchmarks/benches/primitives/main.rs | 1 + benchmarks/benches/primitives/spmc/bounded.rs | 93 +++++++ benchmarks/benches/primitives/spmc/mod.rs | 22 ++ .../benches/primitives/spmc/unbounded.rs | 68 ++++++ .../tests/spmc_test/callbacks.rs | 138 +++++++++++ .../tests/spmc_test/concurrency.rs | 31 +++ tests-integration/tests/spmc_test/main.rs | 1 + .../tests/spmc_test/notification.rs | 42 ++++ 12 files changed, 540 insertions(+), 155 deletions(-) create mode 100644 benchmarks/benches/primitives/spmc/bounded.rs create mode 100644 benchmarks/benches/primitives/spmc/mod.rs create mode 100644 benchmarks/benches/primitives/spmc/unbounded.rs create mode 100644 tests-integration/tests/spmc_test/callbacks.rs diff --git a/CHANGELOG.md b/CHANGELOG.md index a59b9983..f607d0a6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -31,6 +31,10 @@ All notable changes to this project will be documented in this file. * Executor waker operations, including cloning, waking, and dropping, are expected not to panic; recovery from panicking waker callbacks is no longer supported. +### Improvements + +* Reduce SPMC waiting and cancellation overhead. Bounded SPMC queues now preallocate storage for the requested capacity when created instead of growing it during sending. + ## v0.7.3 (2026-09-29) ### New features diff --git a/asyncband/src/spmc/bounded.rs b/asyncband/src/spmc/bounded.rs index b4d7e529..6139b49a 100644 --- a/asyncband/src/spmc/bounded.rs +++ b/asyncband/src/spmc/bounded.rs @@ -22,12 +22,14 @@ use super::RecvError; use super::SendError; use super::TryRecvError; use super::TrySendError; +use super::queue::Producer; use super::queue::Shared; +use super::queue::channel; /// 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. +/// The queue stores at most `capacity` values in preallocated storage. Sending waits for a +/// receiver to free capacity when the queue is full. /// /// Operations briefly acquire an internal mutex. No lock is held across an await point, while /// waking tasks, or while dropping messages. The `try_*` methods do not wait for capacity or @@ -35,15 +37,14 @@ use super::queue::Shared; /// /// # Panics /// -/// Panics if `capacity` is zero. +/// Panics if `capacity` is zero or the preallocated message storage exceeds the allocation size +/// limit. #[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)); + let (producer, shared) = channel(capacity); ( - BoundedSender { - shared: shared.clone(), - }, + BoundedSender { producer, capacity }, BoundedReceiver { shared }, ) } @@ -53,7 +54,9 @@ pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { /// Instances are created by [`bounded`] and cannot be cloned. Sending requires exclusive access to /// this endpoint. pub struct BoundedSender { - shared: Arc>, + producer: Producer, + // Only the producer can increase the queue length, so consumers need no capacity counter. + capacity: usize, } impl fmt::Debug for BoundedSender { @@ -62,12 +65,6 @@ impl fmt::Debug for BoundedSender { } } -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. /// @@ -75,13 +72,12 @@ impl BoundedSender { /// /// # Cancel safety /// - /// Dropping a pending `send` future removes its waiter and drops `value`. A future that has - /// returned `Pending` has not sent the value. Cancellation releases the exclusive sender borrow - /// and allows the next operation to use any available capacity. Use + /// Dropping a pending `send` future drops `value` without enqueueing it and releases the + /// exclusive sender borrow. Any available capacity remains usable by the next operation. 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 + self.producer.send(value, self.capacity).await } /// Attempts to send a value without waiting for capacity. @@ -89,11 +85,11 @@ impl BoundedSender { /// 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) + self.producer.try_send(value, self.capacity) } } -/// Receives values from the associated [`BoundedSender`] handles. +/// Receives values from the associated [`BoundedSender`]. /// /// 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. diff --git a/asyncband/src/spmc/queue.rs b/asyncband/src/spmc/queue.rs index ad47bddf..4d4a59a6 100644 --- a/asyncband/src/spmc/queue.rs +++ b/asyncband/src/spmc/queue.rs @@ -18,6 +18,7 @@ use std::collections::VecDeque; use std::future::poll_fn; use std::mem; +use std::sync::Arc; use std::task::Context; use std::task::Poll; use std::task::Waker; @@ -27,6 +28,7 @@ use super::SendError; use super::TryRecvError; use super::TrySendError; use crate::internal::mutex::Mutex; +use crate::internal::register_waker; use crate::internal::waitlist::WaitList; use crate::internal::waitlist::WaiterId; @@ -34,24 +36,41 @@ 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. +/// The only capability that can append values. Exclusive borrowing also limits the channel to +/// one pending `send` future, including inside this module. +pub struct Producer { + shared: Arc>, +} + +pub fn channel(initial_capacity: usize) -> (Producer, Arc>) { + let shared = Arc::new(Shared { + state: Mutex::new(State { + values: VecDeque::with_capacity(initial_capacity), + sender_alive: true, + receivers: 1, + recv_waiters: WaitList::new(), + send_waker: None, + }), + }); + let producer = Producer { + shared: shared.clone(), + }; + (producer, shared) +} + +/// Consumers serialize removal with publication and waiter registration. Wake callbacks and +/// waker or value destructors run outside this lock, because they may reenter the queue. struct State { values: VecDeque, - capacity: Option, sender_alive: bool, receivers: usize, recv_waiters: WaitList, - send_waiters: WaitList, + // Only the bounded producer can wait for capacity. Once a consumer frees a slot, no other + // producer can take it, so notification needs neither a queue nor a capacity grant. + send_waker: Option, } impl State { - fn has_capacity(&self) -> bool { - self.capacity - .is_none_or(|capacity| self.values.len() < capacity) - } - /// Queues a value and selects a waiting receiver's waker. fn push(&mut self, value: T) -> Option { self.values.push_back(value); @@ -61,8 +80,7 @@ impl State { /// Takes the next value and selects a waiting sender's waker. 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())) + Ok((value, self.send_waker.take())) } else if !self.sender_alive { Err(TryRecvError::Disconnected) } else { @@ -71,9 +89,9 @@ impl State { } } -/// Notification state for an operation waiting to send or receive a value. +/// Notification state for a pending `recv` future. /// -/// Notification wakes the waiting task; it does not reserve a value or queue slot. +/// Notification wakes the waiting task; it does not reserve a value. /// The future retains its waiter ID so it can reclaim the detached node when polled again or /// dropped. Disconnection clears the waiter storage instead. enum Waiter { @@ -98,7 +116,7 @@ impl WaitList { /// Registers a pending operation's waker or refreshes an existing registration. /// - /// If the future finds no value or slot after notification, its waiter rejoins the queue. + /// If the future finds no value after notification, its waiter rejoins the queue. #[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 { @@ -118,93 +136,91 @@ impl WaitList { } } -impl Shared { - pub fn bounded(capacity: usize) -> Self { - Self::new(Some(capacity)) +impl Producer { + pub fn try_send(&mut self, value: T, capacity: usize) -> Result<(), TrySendError> { + let waker = { + let mut state = self.shared.state.lock(); + if state.receivers == 0 { + return Err(TrySendError::Disconnected(value)); + } + if state.values.len() == capacity { + return Err(TrySendError::Full(value)); + } + state.push(value) + }; + if let Some(waker) = waker { + waker.wake(); + } + Ok(()) } - pub fn unbounded() -> Self { - Self::new(None) + pub fn send_unbounded(&mut self, value: T) -> Result<(), SendError> { + let waker = { + let mut state = self.shared.state.lock(); + if state.receivers == 0 { + return Err(SendError::new(value)); + } + state.push(value) + }; + if let Some(waker) = waker { + waker.wake(); + } + Ok(()) } - 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 async fn send(&mut self, value: T, capacity: usize) -> Result<(), SendError> { + let value = match self.try_send(value, capacity) { + Ok(()) => return Ok(()), + Err(TrySendError::Disconnected(value)) => return Err(SendError::new(value)), + Err(TrySendError::Full(value)) => value, + }; + let mut send = Send { + shared: &self.shared, + capacity, + registered: false, + value: Some(value), + }; + poll_fn(|cx| send.poll(cx)).await } +} - pub fn drop_sender(&self) { - let mut waiters = { - let mut state = self.state.lock(); +impl Drop for Producer { + fn drop(&mut self) { + let (mut waiters, retired) = { + let mut state = self.shared.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()) + let waiters = mem::replace(&mut state.recv_waiters, WaitList::new()); + // An explicitly forgotten `send` future may have left a registration behind. + (waiters, state.send_waker.take()) }; while let Some(waker) = waiters.notify_one() { waker.wake(); } + drop(retired); } +} +impl Shared { pub fn clone_receiver(&self) { self.state.lock().receivers += 1; } pub fn drop_receiver(&self) { - let (discarded, mut waiters) = { + let (discarded, waker) = { 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()), - ) - }; - // Notify blocked senders before destroying buffered values, whose destructors may panic. - while let Some(waker) = waiters.notify_one() { - waker.wake(); - } - 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) + (mem::take(&mut state.values), state.send_waker.take()) }; + // Notify the producer before destroying buffered values, whose destructors may panic. 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 + drop(discarded); } pub fn try_recv(&self) -> Result { @@ -216,11 +232,6 @@ impl Shared { } 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, @@ -230,33 +241,33 @@ impl Shared { } struct Send<'a, T> { + // `Producer::send` retains the exclusive producer borrow for this future's lifetime. shared: &'a Shared, - waiter: Option, + capacity: usize, + registered: bool, 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()); + if state.receivers != 0 && state.values.len() == self.capacity { + let retired = register_waker(&mut state.send_waker, cx.waker()); + self.registered = true; drop(state); drop(retired); return Poll::Pending; + } + let value = self.value.take().expect("pending send must own its value"); + let outcome = if state.receivers == 0 { + Err(value) + } else { + Ok(state.push(value)) }; - let retired = self - .waiter - .take() - .map(|id| state.send_waiters.remove_waiter(id)); + // A consumer that frees capacity, or the last receiver that disconnects, takes the + // registration. With no competing producer, readiness cannot be stolen before this poll. + debug_assert!(state.send_waker.is_none()); + self.registered = false; drop(state); let result = outcome .map(|waker| { @@ -265,24 +276,18 @@ impl Send<'_, T> { } }) .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 = { - let mut state = self.shared.state.lock(); - if state.receivers == 0 { - return; - } - state.send_waiters.remove_waiter(id) - }; - drop(retired); + if self.registered { + // Remove the registration before dropping `value`, without running either destructor + // under the lock. No other producer can have replaced this future's registration. + let retired = self.shared.state.lock().send_waker.take(); + drop(retired); + } } } @@ -294,10 +299,6 @@ struct Recv<'a, T> { 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), @@ -308,10 +309,12 @@ impl Recv<'_, T> { return Poll::Pending; } }; - let retired = self - .waiter - .take() - .map(|id| state.recv_waiters.remove_waiter(id)); + let retired = self.waiter.take().and_then(|id| { + // Disconnection detaches waiter storage, but buffered values remain readable. + state + .sender_alive + .then(|| state.recv_waiters.remove_waiter(id)) + }); drop(state); let result = outcome.map(|(value, waker)| { if let Some(waker) = waker { diff --git a/asyncband/src/spmc/unbounded.rs b/asyncband/src/spmc/unbounded.rs index b2b4851e..f7b08be8 100644 --- a/asyncband/src/spmc/unbounded.rs +++ b/asyncband/src/spmc/unbounded.rs @@ -21,8 +21,9 @@ use std::sync::Arc; use super::RecvError; use super::SendError; use super::TryRecvError; -use super::TrySendError; +use super::queue::Producer; use super::queue::Shared; +use super::queue::channel; /// Creates an unbounded single-producer, multi-consumer queue. /// @@ -32,13 +33,8 @@ use super::queue::Shared; /// 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 }, - ) + let (producer, shared) = channel(0); + (UnboundedSender { producer }, UnboundedReceiver { shared }) } /// Sends values to the associated [`UnboundedReceiver`] handles. @@ -46,7 +42,7 @@ pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { /// Instances are created by [`unbounded`] and cannot be cloned. Sending requires exclusive access /// to this endpoint. pub struct UnboundedSender { - shared: Arc>, + producer: Producer, } impl fmt::Debug for UnboundedSender { @@ -55,26 +51,16 @@ impl fmt::Debug for UnboundedSender { } } -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"), - } + self.producer.send_unbounded(value) } } -/// Receives values from the associated [`UnboundedSender`] handles. +/// Receives values from the associated [`UnboundedSender`]. /// /// 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. diff --git a/benchmarks/benches/primitives/main.rs b/benchmarks/benches/primitives/main.rs index 6d9f4c32..703bf0e5 100644 --- a/benchmarks/benches/primitives/main.rs +++ b/benchmarks/benches/primitives/main.rs @@ -36,6 +36,7 @@ mod rwlock; mod semaphore; mod shutdown; mod singleflight; +mod spmc; mod waitgroup; fn main() { diff --git a/benchmarks/benches/primitives/spmc/bounded.rs b/benchmarks/benches/primitives/spmc/bounded.rs new file mode 100644 index 00000000..2b780155 --- /dev/null +++ b/benchmarks/benches/primitives/spmc/bounded.rs @@ -0,0 +1,93 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::pin::pin; + +use asyncband::spmc; +use benchmarks::support::bench_context; +use benchmarks::support::poll_pending; +use benchmarks::support::poll_pinned_ready; +use benchmarks::support::poll_ready; +use divan::Bencher; +use divan::black_box; + +use super::FAST_SAMPLE_SIZE; + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn try_send_then_try_recv(bencher: Bencher) { + let (mut sender, receiver) = spmc::bounded(1); + bencher.bench_local(|| { + sender.try_send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn send_then_recv(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, receiver) = spmc::bounded(1); + bencher.bench_local(|| { + poll_ready(sender.send(black_box(1usize)), &mut context).unwrap(); + black_box(poll_ready(receiver.recv(), &mut context).unwrap()) + }); +} + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn wake_blocked_sender(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, receiver) = spmc::bounded(1); + sender.try_send(0usize).unwrap(); + bencher.bench_local(|| { + let mut send = pin!(sender.send(black_box(1))); + poll_pending(send.as_mut(), &mut context); + black_box(receiver.try_recv().unwrap()); + poll_pinned_ready(send.as_mut(), &mut context).unwrap(); + }); +} + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn cancel_pending_send(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, _receiver) = spmc::bounded(1); + sender.try_send(0usize).unwrap(); + bencher.bench_local(|| { + let mut send = pin!(sender.send(black_box(1))); + poll_pending(send.as_mut(), &mut context); + }); +} + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn repoll_pending_send(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, _receiver) = spmc::bounded(1); + sender.try_send(0usize).unwrap(); + let mut send = pin!(sender.send(1)); + poll_pending(send.as_mut(), &mut context); + bencher.bench_local(|| poll_pending(send.as_mut(), &mut context)); +} + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn wake_pending_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, receiver) = spmc::bounded(1); + bencher.bench_local(|| { + let mut recv = pin!(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + sender.try_send(black_box(usize::MAX)).unwrap(); + black_box(poll_pinned_ready(recv.as_mut(), &mut context).unwrap()) + }); +} diff --git a/benchmarks/benches/primitives/spmc/mod.rs b/benchmarks/benches/primitives/spmc/mod.rs new file mode 100644 index 00000000..a00276ab --- /dev/null +++ b/benchmarks/benches/primitives/spmc/mod.rs @@ -0,0 +1,22 @@ +// 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. + +mod bounded; +mod unbounded; + +// Fixed sample sizes keep one-time warm-up work from changing Divan's iteration granularity. +const FAST_SAMPLE_SIZE: u32 = 256; diff --git a/benchmarks/benches/primitives/spmc/unbounded.rs b/benchmarks/benches/primitives/spmc/unbounded.rs new file mode 100644 index 00000000..e426a4d8 --- /dev/null +++ b/benchmarks/benches/primitives/spmc/unbounded.rs @@ -0,0 +1,68 @@ +// 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 asyncband::spmc; +use benchmarks::support::bench_context; +use benchmarks::support::poll_pending; +use benchmarks::support::poll_pinned_ready; +use benchmarks::support::poll_ready; +use divan::Bencher; +use divan::black_box; + +use super::FAST_SAMPLE_SIZE; + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn send_then_try_recv(bencher: Bencher) { + let (mut sender, receiver) = spmc::unbounded(); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(receiver.try_recv().unwrap()) + }); +} + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn send_then_recv(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, receiver) = spmc::unbounded(); + bencher.bench_local(|| { + sender.send(black_box(1usize)).unwrap(); + black_box(poll_ready(receiver.recv(), &mut context).unwrap()) + }); +} + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn wake_pending_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (mut sender, receiver) = spmc::unbounded(); + bencher.bench_local(|| { + let mut recv = pin!(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + sender.send(black_box(usize::MAX)).unwrap(); + black_box(poll_pinned_ready(recv.as_mut(), &mut context).unwrap()) + }); +} + +#[divan::bench(sample_size = FAST_SAMPLE_SIZE)] +fn repoll_pending_receiver(bencher: Bencher) { + let mut context = bench_context(); + let (_sender, receiver) = spmc::unbounded::(); + let mut recv = pin!(receiver.recv()); + poll_pending(recv.as_mut(), &mut context); + bencher.bench_local(|| poll_pending(recv.as_mut(), &mut context)); +} diff --git a/tests-integration/tests/spmc_test/callbacks.rs b/tests-integration/tests/spmc_test/callbacks.rs new file mode 100644 index 00000000..d63d4d0b --- /dev/null +++ b/tests-integration/tests/spmc_test/callbacks.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::panic::AssertUnwindSafe; +use std::panic::catch_unwind; + +use asyncband::spmc; +use tests_integration::WakeCounter; +use tests_integration::assert_completes_without_deadlock; +use tests_integration::expect_ready; +use tests_integration::poll_once; +use tests_integration::poll_with; +use tests_integration::waker_on_drop; +use tests_integration::waker_on_wake; + +#[test] +fn replacing_a_send_waker_allows_its_destructor_to_receive() { + assert_completes_without_deadlock(|| { + let (mut sender, receiver) = spmc::bounded(1); + sender.try_send(0).unwrap(); + let reentrant = receiver.clone(); + let first = waker_on_drop(move || assert_eq!(reentrant.try_recv(), Ok(0))); + let (second, second_wakes) = WakeCounter::new(); + let mut send = Box::pin(sender.send(1)); + + assert!(poll_with(send.as_mut(), &first).is_pending()); + drop(first); + // The retired waker frees capacity and wakes the replacement registration before this + // poll returns. The next poll must still observe that capacity. + assert!(poll_with(send.as_mut(), &second).is_pending()); + assert_eq!(second_wakes.count(), 1); + expect_ready(poll_with(send.as_mut(), &second)).unwrap(); + assert_eq!(receiver.try_recv(), Ok(1)); + }); +} + +#[test] +fn cancelling_a_send_allows_its_waker_destructor_to_receive() { + assert_completes_without_deadlock(|| { + let (mut sender, receiver) = spmc::bounded(1); + sender.try_send(0).unwrap(); + let reentrant = receiver.clone(); + let waker = waker_on_drop(move || assert_eq!(reentrant.try_recv(), Ok(0))); + let mut send = Box::pin(sender.send(1)); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + drop(waker); + drop(send); + + sender.try_send(2).unwrap(); + assert_eq!(receiver.try_recv(), Ok(2)); + }); +} + +#[test] +fn waking_the_sender_allows_its_callback_to_receive() { + assert_completes_without_deadlock(|| { + let (mut sender, receiver) = spmc::bounded(2); + sender.try_send(0).unwrap(); + sender.try_send(1).unwrap(); + let reentrant = receiver.clone(); + let waker = waker_on_wake(move || assert_eq!(reentrant.try_recv(), Ok(1))); + let mut send = Box::pin(sender.send(2)); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + + assert_eq!(receiver.try_recv(), Ok(0)); + expect_ready(poll_once(send.as_mut())).unwrap(); + assert_eq!(receiver.try_recv(), Ok(2)); + }); +} + +#[derive(Debug)] +struct Payload { + panic_on_drop: bool, +} + +impl Drop for Payload { + fn drop(&mut self) { + assert!(!self.panic_on_drop, "payload destructor panicked"); + } +} + +#[test] +fn cancellation_unregisters_before_a_payload_destructor_panics() { + let (mut sender, receiver) = spmc::bounded(1); + sender + .try_send(Payload { + panic_on_drop: false, + }) + .unwrap(); + let (waker, wakes) = WakeCounter::new(); + let mut send = Box::pin(sender.send(Payload { + panic_on_drop: true, + })); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + + assert!(catch_unwind(AssertUnwindSafe(|| drop(send))).is_err()); + drop(receiver.try_recv().unwrap()); + assert_eq!(wakes.count(), 0); + sender + .try_send(Payload { + panic_on_drop: false, + }) + .unwrap(); +} + +#[test] +fn disconnect_notifies_the_sender_before_a_buffered_payload_destructor_panics() { + let (mut sender, receiver) = spmc::bounded(1); + sender + .try_send(Payload { + panic_on_drop: true, + }) + .unwrap(); + let (waker, wakes) = WakeCounter::new(); + let mut send = Box::pin(sender.send(Payload { + panic_on_drop: false, + })); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + + assert!(catch_unwind(AssertUnwindSafe(|| drop(receiver))).is_err()); + assert_eq!(wakes.count(), 1); + let rejected = expect_ready(poll_once(send.as_mut())).unwrap_err(); + assert!(!rejected.into_inner().panic_on_drop); +} diff --git a/tests-integration/tests/spmc_test/concurrency.rs b/tests-integration/tests/spmc_test/concurrency.rs index d44ba4d1..7a5c02fb 100644 --- a/tests-integration/tests/spmc_test/concurrency.rs +++ b/tests-integration/tests/spmc_test/concurrency.rs @@ -16,15 +16,46 @@ // under the License. use std::sync::Arc; +use std::task::Poll; use std::time::Duration; use asyncband::spmc; +use tests_integration::WakeCounter; +use tests_integration::poll_with; use tokio::sync::Barrier; use tokio::task::JoinHandle; const CONSUMERS: usize = 8; const TOTAL: usize = 2_048; +#[test] +fn freeing_capacity_during_sender_registration_cannot_lose_the_notification() { + let (mut sender, receiver) = spmc::bounded(1); + for value in 0..32 { + sender.try_send(value).unwrap(); + let start = std::sync::Barrier::new(2); + let (waker, wakes) = WakeCounter::new(); + let mut pending = Box::pin(sender.send(value + 1)); + let first_poll = std::thread::scope(|scope| { + scope.spawn(|| { + start.wait(); + assert_eq!(receiver.try_recv(), Ok(value)); + }); + start.wait(); + poll_with(pending.as_mut(), &waker) + }); + + if first_poll.is_pending() { + assert_eq!(wakes.count(), 1); + assert_eq!(poll_with(pending.as_mut(), &waker), Poll::Ready(Ok(()))); + } else { + assert_eq!(first_poll, Poll::Ready(Ok(()))); + } + drop(pending); + assert_eq!(receiver.try_recv(), Ok(value + 1)); + } +} + async fn assert_delivered_exactly_once( producer: JoinHandle<()>, consumers: Vec>>, diff --git a/tests-integration/tests/spmc_test/main.rs b/tests-integration/tests/spmc_test/main.rs index 684c3f35..5f053422 100644 --- a/tests-integration/tests/spmc_test/main.rs +++ b/tests-integration/tests/spmc_test/main.rs @@ -28,6 +28,7 @@ use asyncband::spmc::TryRecvError; use tests_integration::poll_once; // Public queue contracts. The other suites cover notifications, cancellation, and concurrency. +mod callbacks; mod concurrency; mod notification; diff --git a/tests-integration/tests/spmc_test/notification.rs b/tests-integration/tests/spmc_test/notification.rs index 51f0e0fb..05fa79f6 100644 --- a/tests-integration/tests/spmc_test/notification.rs +++ b/tests-integration/tests/spmc_test/notification.rs @@ -196,6 +196,48 @@ fn bounded_capacity_and_pending_send_progress() { } } +#[test] +fn freeing_multiple_slots_wakes_only_the_latest_sender_task_once() { + let (mut sender, receiver) = spmc::bounded(2); + let competing = receiver.clone(); + sender.try_send(0).unwrap(); + sender.try_send(1).unwrap(); + let (first, first_wakes) = WakeCounter::new(); + let (latest, latest_wakes) = WakeCounter::new(); + let mut pending = Box::pin(sender.send(2)); + assert!(poll_with(pending.as_mut(), &first).is_pending()); + assert!(poll_with(pending.as_mut(), &first).is_pending()); + assert!(poll_with(pending.as_mut(), &latest).is_pending()); + + assert_eq!(receiver.try_recv(), Ok(0)); + assert_eq!(competing.try_recv(), Ok(1)); + assert_eq!(first_wakes.count(), 0); + assert_eq!(latest_wakes.count(), 1); + expect_ready(poll_with(pending.as_mut(), &latest)).unwrap(); + drop(pending); + assert_eq!(receiver.try_recv(), Ok(2)); + assert_eq!(latest_wakes.count(), 1); +} + +#[test] +fn cancelling_a_pending_send_allows_a_new_task_to_wait_for_capacity() { + let (mut sender, receiver) = spmc::bounded(1); + sender.try_send(0).unwrap(); + let (cancelled, cancelled_wakes) = WakeCounter::new(); + let mut first = Box::pin(sender.send(1)); + assert!(poll_with(first.as_mut(), &cancelled).is_pending()); + drop(first); + + let (current, current_wakes) = WakeCounter::new(); + let mut second = Box::pin(sender.send(2)); + assert!(poll_with(second.as_mut(), ¤t).is_pending()); + assert_eq!(receiver.try_recv(), Ok(0)); + assert_eq!(cancelled_wakes.count(), 0); + assert_eq!(current_wakes.count(), 1); + expect_ready(poll_with(second.as_mut(), ¤t)).unwrap(); + assert_eq!(receiver.try_recv(), Ok(2)); +} + #[test] fn cancelling_an_unnotified_send_preserves_the_buffered_value() { let (mut sender, receiver) = spmc::bounded(1);