You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@asyncband.apache.org by xy...@apache.org on 2026/09/29 11:22:07 UTC
(asyncband) branch main updated: feat(mpmc): add capacity reservation to bounded senders (#334)
This is an automated email from the ASF dual-hosted git repository.
BewareMyPower pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/asyncband.git
The following commit(s) were added to refs/heads/main by this push:
new b383e93 feat(mpmc): add capacity reservation to bounded senders (#334)
b383e93 is described below
commit b383e9351cda9df41b0316bd6dd9c4c2aaf6a920
Author: Gunther Xing <ji...@gmail.com>
AuthorDate: Tue Sep 29 19:22:03 2026 +0800
feat(mpmc): add capacity reservation to bounded senders (#334)
---
CHANGELOG.md | 4 +
asyncband/src/mpmc/bounded.rs | 53 ++++-
asyncband/src/mpmc/error.rs | 5 +-
asyncband/src/mpmc/mod.rs | 1 +
asyncband/src/mpmc/queue.rs | 269 ++++++++++++++++------
tests-integration/tests/mpmc_test/main.rs | 1 +
tests-integration/tests/mpmc_test/notification.rs | 53 -----
tests-integration/tests/mpmc_test/reservation.rs | 265 +++++++++++++++++++++
8 files changed, 517 insertions(+), 134 deletions(-)
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7fe6c26..8ab5099 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -23,6 +23,10 @@ All notable changes to this project will be documented in this file.
## Unreleased
+### New features
+
+* Add bounded MPMC `reserve` and `try_reserve` methods returning a borrowed `Permit`, so callers can wait for capacity before constructing a value; sends and reservations receive capacity in wait-queue order, and unused permits release it.
+
## v0.7.3 (2026-09-29)
### New features
diff --git a/asyncband/src/mpmc/bounded.rs b/asyncband/src/mpmc/bounded.rs
index 3122f83..cd93c39 100644
--- a/asyncband/src/mpmc/bounded.rs
+++ b/asyncband/src/mpmc/bounded.rs
@@ -22,12 +22,14 @@ use super::RecvError;
use super::SendError;
use super::TryRecvError;
use super::TrySendError;
+pub use super::queue::Permit;
use super::queue::Shared;
/// Creates a bounded multi-producer, multi-consumer queue.
///
-/// The queue stores at most `capacity` values. Sending waits for a receiver to free capacity when
-/// the queue is full.
+/// Queued values, held permits, and capacity granted to waiting senders occupy at most `capacity`
+/// slots. Pending sends and reservations receive capacity in wait-queue order. Sending waits for
+/// a receiver to free capacity when none is available.
///
/// The `try_*` methods do not wait for capacity or messages, but may briefly block on an internal
/// mutex.
@@ -82,16 +84,55 @@ impl<T> BoundedSender<T> {
///
/// # Cancel safety
///
- /// Dropping a pending `send` drops `value` without sending it or retaining capacity. Use
- /// [`try_send`](Self::try_send) when the caller must retain ownership if capacity is
- /// unavailable.
+ /// Dropping a pending `send` releases its waiting resources before dropping `value`, without
+ /// sending it or retaining capacity. Use [`reserve`](Self::reserve) to wait for capacity before
+ /// constructing a value.
pub async fn send(&self, value: T) -> Result<(), SendError<T>> {
self.shared.send(value).await
}
+ /// Reserves capacity for one value before constructing it.
+ ///
+ /// A successful reservation returns a [`Permit`]. Dropping it releases capacity. A permit
+ /// reserves space, not message order, and does not keep receivers alive.
+ ///
+ /// Returns `SendError(())` if all receivers have been dropped.
+ ///
+ /// # Cancel safety
+ ///
+ /// Dropping a pending reservation releases its place in the wait queue. If it was already
+ /// granted capacity, that capacity passes to the next waiter or becomes available again.
+ ///
+ /// # Examples
+ ///
+ /// ```
+ /// # #[tokio::main]
+ /// # async fn main() {
+ /// let (sender, receiver) = asyncband::mpmc::bounded(1);
+ /// let permit = sender.reserve().await.unwrap();
+ /// let message = String::from("constructed after capacity became available");
+ /// permit.send(message).unwrap();
+ /// assert_eq!(
+ /// receiver.recv().await.unwrap(),
+ /// "constructed after capacity became available"
+ /// );
+ /// # }
+ /// ```
+ pub async fn reserve(&self) -> Result<Permit<'_, T>, SendError<()>> {
+ self.shared.reserve().await
+ }
+
+ /// Reserves capacity for one value without waiting.
+ ///
+ /// Returns [`TrySendError::Full`] when all capacity belongs to queued values, held permits,
+ /// or granted waiters, and [`TrySendError::Disconnected`] when all receivers are gone.
+ pub fn try_reserve(&self) -> Result<Permit<'_, T>, TrySendError<()>> {
+ self.shared.try_reserve()
+ }
+
/// Attempts to send a value without waiting for capacity.
///
- /// Returns [`TrySendError::Full`] when the queue has reached its exact capacity and
+ /// Returns [`TrySendError::Full`] when no unassigned capacity remains and
/// [`TrySendError::Disconnected`] when all receivers have been dropped.
pub fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
self.shared.try_send(value)
diff --git a/asyncband/src/mpmc/error.rs b/asyncband/src/mpmc/error.rs
index 7bb7525..4cbe177 100644
--- a/asyncband/src/mpmc/error.rs
+++ b/asyncband/src/mpmc/error.rs
@@ -54,10 +54,11 @@ impl<T> fmt::Debug for SendError<T> {
impl<T> std::error::Error for SendError<T> {}
-/// Error returned by [`BoundedSender::try_send`](crate::mpmc::BoundedSender::try_send).
+/// Error returned by [`BoundedSender::try_send`](crate::mpmc::BoundedSender::try_send) or
+/// [`BoundedSender::try_reserve`](crate::mpmc::BoundedSender::try_reserve).
#[derive(Clone, PartialEq, Eq)]
pub enum TrySendError<T> {
- /// The queue is full, so the value cannot be sent without waiting for capacity.
+ /// No unassigned capacity remains, so the value cannot be sent without waiting.
Full(T),
/// All receivers have been dropped, so the value can never be received.
Disconnected(T),
diff --git a/asyncband/src/mpmc/mod.rs b/asyncband/src/mpmc/mod.rs
index 8a3f039..5ceda02 100644
--- a/asyncband/src/mpmc/mod.rs
+++ b/asyncband/src/mpmc/mod.rs
@@ -29,6 +29,7 @@ mod unbounded;
pub use self::bounded::BoundedReceiver;
pub use self::bounded::BoundedSender;
+pub use self::bounded::Permit;
pub use self::bounded::bounded;
pub use self::error::RecvError;
pub use self::error::SendError;
diff --git a/asyncband/src/mpmc/queue.rs b/asyncband/src/mpmc/queue.rs
index a279470..19b5e3a 100644
--- a/asyncband/src/mpmc/queue.rs
+++ b/asyncband/src/mpmc/queue.rs
@@ -16,6 +16,7 @@
// under the License.
use std::collections::VecDeque;
+use std::fmt;
use std::future::poll_fn;
use std::mem;
use std::task::Context;
@@ -40,17 +41,45 @@ pub struct Shared<T> {
/// outside the lock, because they may reenter the queue.
struct State<T> {
values: VecDeque<T>,
- capacity: Option<usize>,
+ // None is an unbounded queue. In a bounded queue, each slot belongs to available,
+ // a queued value, a Permit, or a detached waiter with a grant.
+ available: Option<usize>,
senders: usize,
receivers: usize,
- recv_waiters: WaitList<Waiter>,
- send_waiters: WaitList<Waiter>,
+ recv_waiters: WaitList<RecvWaiter>,
+ send_waiters: WaitList<SendWaiter>,
}
impl<T> State<T> {
fn has_capacity(&self) -> bool {
- self.capacity
- .is_none_or(|capacity| self.values.len() < capacity)
+ self.available.is_none_or(|available| available != 0)
+ }
+
+ fn acquire(&mut self) -> Result<(), TrySendError<()>> {
+ if self.receivers == 0 {
+ return Err(TrySendError::Disconnected(()));
+ }
+ if !self.has_capacity() {
+ return Err(TrySendError::Full(()));
+ }
+ if let Some(available) = &mut self.available {
+ *available -= 1;
+ }
+ Ok(())
+ }
+
+ fn release(&mut self) -> Option<Waker> {
+ if self.receivers == 0 {
+ return None;
+ }
+ let Some(available) = &mut self.available else {
+ return None;
+ };
+ if let Some(waker) = self.send_waiters.grant_one() {
+ return Some(waker);
+ }
+ *available += 1;
+ None
}
/// Queues a value and selects the receiver to wake.
@@ -59,11 +88,10 @@ impl<T> State<T> {
self.recv_waiters.notify_one()
}
- /// Takes the next value and selects the sender to wake.
+ /// Takes the next value and grants its capacity to the oldest waiting sender.
fn pop(&mut self) -> Result<(T, Option<Waker>), 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.release()))
} else if self.senders == 0 {
Err(TryRecvError::Disconnected)
} else {
@@ -72,37 +100,40 @@ impl<T> State<T> {
}
}
-/// 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 {
+/// A receiver notification makes the task runnable without reserving a value.
+enum RecvWaiter {
Waiting(Waker),
Notified,
}
-impl WaitList<Waiter> {
+/// A grant transfers capacity to a detached sender waiter until it claims or cancels it.
+enum SendWaiter {
+ Waiting(Waker),
+ Granted,
+}
+
+fn remove_waiter<W>(waiters: &mut WaitList<W>, id: WaiterId) -> W {
+ // Unlinking is idempotent, so detached waiters are removed the same way as linked ones.
+ waiters.unlink_waiter(id, |_| true);
+ waiters.remove_unlinked_waiter(id)
+}
+
+impl WaitList<RecvWaiter> {
fn notify_one(&mut self) -> Option<Waker> {
let (_, waiter) = self.unlink_first_waiter(|_| true)?;
- let Waiter::Waiting(waker) = mem::replace(waiter, Waiter::Notified) else {
+ let RecvWaiter::Waiting(waker) = mem::replace(waiter, RecvWaiter::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.
+ /// A notified receive that still found no value queues again at the back.
#[must_use = "drop the replaced waker after releasing the queue lock"]
fn register(&mut self, id: &mut Option<WaiterId>, current: &Waker) -> Option<Waker> {
if let Some(queued) = *id {
- if let Waiter::Waiting(waker) = self.waiter_mut(queued) {
+ if let RecvWaiter::Waiting(waker) = self.waiter_mut(queued) {
if waker.will_wake(current) {
return None;
}
@@ -112,13 +143,47 @@ impl WaitList<Waiter> {
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);
+ remove_waiter(self, notified);
}
- *id = Some(self.push_back(Waiter::Waiting(waker)));
+ *id = Some(self.push_back(RecvWaiter::Waiting(waker)));
None
}
}
+impl WaitList<SendWaiter> {
+ /// Queues a blocked sender or refreshes its waker without losing its place.
+ #[must_use = "drop the replaced waker after releasing the queue lock"]
+ fn register_waiter(&mut self, id: &mut Option<WaiterId>, current: &Waker) -> Option<Waker> {
+ if let Some(queued) = *id {
+ let SendWaiter::Waiting(waker) = self.waiter_mut(queued) else {
+ unreachable!("a granted waiter must be claimed before registration");
+ };
+ if waker.will_wake(current) {
+ return None;
+ }
+ return Some(mem::replace(waker, current.clone()));
+ }
+ *id = Some(self.push_back(SendWaiter::Waiting(current.clone())));
+ None
+ }
+
+ fn grant_one(&mut self) -> Option<Waker> {
+ let (_, waiter) = self.unlink_first_waiter(|_| true)?;
+ let SendWaiter::Waiting(waker) = mem::replace(waiter, SendWaiter::Granted) else {
+ unreachable!("only waiting operations remain linked");
+ };
+ Some(waker)
+ }
+
+ fn take_waiting_waker(&mut self) -> Option<Waker> {
+ let (id, _) = self.unlink_first_waiter(|_| true)?;
+ let SendWaiter::Waiting(waker) = self.remove_unlinked_waiter(id) else {
+ unreachable!("only waiting operations remain linked");
+ };
+ Some(waker)
+ }
+}
+
impl<T> Shared<T> {
pub fn bounded(capacity: usize) -> Self {
Self::new(Some(capacity))
@@ -132,7 +197,7 @@ impl<T> Shared<T> {
Self {
state: Mutex::new(State {
values: VecDeque::new(),
- capacity,
+ available: capacity,
senders: 1,
receivers: 1,
recv_waiters: WaitList::new(),
@@ -177,18 +242,19 @@ impl<T> Shared<T> {
};
// 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()));
+ wake_all(std::iter::from_fn(|| waiters.take_waiting_waker()));
drop(discarded);
}
pub fn try_send(&self, value: T) -> Result<(), TrySendError<T>> {
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));
+ match state.acquire() {
+ Ok(()) => {}
+ Err(TrySendError::Full(())) => return Err(TrySendError::Full(value)),
+ Err(TrySendError::Disconnected(())) => {
+ return Err(TrySendError::Disconnected(value));
+ }
}
state.push(value)
};
@@ -204,12 +270,23 @@ impl<T> Shared<T> {
Err(TrySendError::Disconnected(value)) => return Err(SendError::new(value)),
Err(TrySendError::Full(value)) => value,
};
- let mut send = Send {
+ match self.reserve().await {
+ Ok(permit) => permit.send(value),
+ Err(_) => Err(SendError::new(value)),
+ }
+ }
+
+ pub fn try_reserve(&self) -> Result<Permit<'_, T>, TrySendError<()>> {
+ self.state.lock().acquire()?;
+ Ok(Permit { shared: self })
+ }
+
+ pub async fn reserve(&self) -> Result<Permit<'_, T>, SendError<()>> {
+ let mut reserve = Reserve {
shared: self,
waiter: None,
- value: Some(value),
};
- poll_fn(|cx| send.poll(cx)).await
+ poll_fn(|cx| reserve.poll(cx)).await
}
pub fn try_recv(&self) -> Result<T, TryRecvError> {
@@ -234,51 +311,97 @@ impl<T> Shared<T> {
}
}
-struct Send<'a, T> {
+/// Capacity reserved for one value on a bounded MPMC queue.
+///
+/// A permit does not keep any receiver alive or claim message order. Dropping it without sending
+/// passes its capacity to the next waiting sender or makes it available again.
+#[must_use = "dropping the permit releases its reserved capacity"]
+pub struct Permit<'a, T> {
shared: &'a Shared<T>,
- waiter: Option<WaiterId>,
- // `Drop` passes an unconsumed notification on before this value is destroyed, because its
- // destructor may depend on another blocked sender making progress.
- value: Option<T>,
}
-impl<T> Send<'_, T> {
- fn take_value(&mut self) -> T {
- self.value.take().expect("pending send must own its value")
+impl<T> fmt::Debug for Permit<'_, T> {
+ fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+ f.debug_struct("Permit").finish_non_exhaustive()
}
+}
- fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<(), SendError<T>>> {
+impl<T> Permit<'_, T> {
+ /// Enqueues a value using the reserved capacity without waiting for space.
+ ///
+ /// If the last receiver has been dropped, the error returns ownership of the value.
+ pub fn send(self, value: T) -> Result<(), SendError<T>> {
+ let waker = {
+ let mut state = self.shared.state.lock();
+ if state.receivers == 0 {
+ return Err(SendError::new(value));
+ }
+ let waker = state.push(value);
+ // The queued value now owns this capacity, including if waking a receiver panics.
+ mem::forget(self);
+ waker
+ };
+ if let Some(waker) = waker {
+ waker.wake();
+ }
+ Ok(())
+ }
+}
+
+impl<T> Drop for Permit<'_, T> {
+ fn drop(&mut self) {
+ let waker = self.shared.state.lock().release();
+ if let Some(waker) = waker {
+ waker.wake();
+ }
+ }
+}
+
+struct Reserve<'a, T> {
+ shared: &'a Shared<T>,
+ waiter: Option<WaiterId>,
+}
+
+impl<'a, T> Reserve<'a, T> {
+ fn poll(&mut self, cx: &mut Context<'_>) -> Poll<Result<Permit<'a, T>, SendError<()>>> {
let mut state = self.shared.state.lock();
- let outcome = if state.receivers == 0 {
+ // The last receiver may have dropped while this reservation was waiting or granted.
+ if state.receivers == 0 {
self.waiter = None;
- Err(self.take_value())
+ return Poll::Ready(Err(SendError::new(())));
+ }
+ // A previous poll queued this reservation; a released slot may now belong to it.
+ if let Some(id) = self.waiter {
+ if matches!(state.send_waiters.waiter_mut(id), SendWaiter::Granted) {
+ let retired = state.send_waiters.remove_unlinked_waiter(id);
+ self.waiter = None;
+ drop(state);
+ drop(retired);
+ return Poll::Ready(Ok(Permit {
+ shared: self.shared,
+ }));
+ }
} else if state.has_capacity() {
- Ok(state.push(self.take_value()))
- } else {
- let retired = state.send_waiters.register(&mut self.waiter, cx.waker());
+ // This reservation has not queued yet and can claim unassigned capacity immediately.
+ if let Some(available) = &mut state.available {
+ *available -= 1;
+ }
drop(state);
- drop(retired);
- return Poll::Pending;
- };
- let retired = self
- .waiter
- .take()
- .map(|id| state.send_waiters.remove_waiter(id));
+ return Poll::Ready(Ok(Permit {
+ shared: self.shared,
+ }));
+ }
+ // Otherwise, wait for a slot or refresh the waker of the queued reservation.
+ let retired = state
+ .send_waiters
+ .register_waiter(&mut self.waiter, cx.waker());
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)
+ Poll::Pending
}
}
-impl<T> Drop for Send<'_, T> {
+impl<T> Drop for Reserve<'_, T> {
fn drop(&mut self) {
let Some(id) = self.waiter.take() else {
return;
@@ -288,10 +411,10 @@ impl<T> Drop for Send<'_, T> {
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()
+ let retired = remove_waiter(&mut state.send_waiters, id);
+ // A grant already owns a slot, so returning it does not require a capacity check.
+ let waker = if matches!(retired, SendWaiter::Granted) {
+ state.release()
} else {
None
};
@@ -329,7 +452,7 @@ impl<T> Recv<'_, T> {
let retired = self
.waiter
.take()
- .map(|id| state.recv_waiters.remove_waiter(id));
+ .map(|id| remove_waiter(&mut state.recv_waiters, id));
drop(state);
// Deliver the notification before running waker destructors, which may panic.
let result = outcome.map(|(value, waker)| {
@@ -353,9 +476,9 @@ impl<T> Drop for Recv<'_, T> {
if state.senders == 0 {
return;
}
- let retired = state.recv_waiters.remove_waiter(id);
+ let retired = remove_waiter(&mut state.recv_waiters, 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() {
+ let waker = if matches!(retired, RecvWaiter::Notified) && !state.values.is_empty() {
state.recv_waiters.notify_one()
} else {
None
diff --git a/tests-integration/tests/mpmc_test/main.rs b/tests-integration/tests/mpmc_test/main.rs
index 8795d7a..c4b6e2a 100644
--- a/tests-integration/tests/mpmc_test/main.rs
+++ b/tests-integration/tests/mpmc_test/main.rs
@@ -29,6 +29,7 @@ use tests_integration::poll_once;
mod callbacks;
mod concurrency;
mod notification;
+mod reservation;
/// Either receiver flavor, so one case can cover both queues.
trait Receiver<T>: Clone {
diff --git a/tests-integration/tests/mpmc_test/notification.rs b/tests-integration/tests/mpmc_test/notification.rs
index 1b2a944..5b8845c 100644
--- a/tests-integration/tests/mpmc_test/notification.rs
+++ b/tests-integration/tests/mpmc_test/notification.rs
@@ -124,59 +124,6 @@ fn notified_receiver_that_loses_the_value_queues_behind_waiting_receivers() {
);
}
-#[test]
-fn notified_sender_that_loses_capacity_queues_behind_waiting_senders() {
- let (sender, receiver) = mpmc::bounded(1);
- sender.try_send(0).unwrap();
- let competing = sender.clone();
- let mut first = Box::pin(sender.send(1));
- let mut second = Box::pin(competing.send(2));
- let (first_waker, first_wakes) = WakeCounter::new();
- let (second_waker, second_wakes) = WakeCounter::new();
-
- assert!(poll_with(first.as_mut(), &first_waker).is_pending());
- assert!(poll_with(second.as_mut(), &second_waker).is_pending());
- assert_eq!(receiver.try_recv(), Ok(0));
- assert_eq!(first_wakes.count(), 1);
- sender.try_send(3).unwrap();
- assert!(poll_with(first.as_mut(), &first_waker).is_pending());
-
- assert_eq!(receiver.try_recv(), Ok(3));
- assert_eq!(first_wakes.count(), 1);
- assert_eq!(second_wakes.count(), 1);
- expect_ready(poll_with(second.as_mut(), &second_waker)).unwrap();
- assert_eq!(receiver.try_recv(), Ok(2));
- assert_eq!(first_wakes.count(), 2);
- expect_ready(poll_with(first.as_mut(), &first_waker)).unwrap();
- assert_eq!(receiver.try_recv(), Ok(1));
-}
-
-#[test]
-fn cancelling_a_sender_after_capacity_is_taken_does_not_wake_next() {
- let (sender, receiver) = mpmc::bounded(1);
- sender.try_send(0).unwrap();
- let competing = sender.clone();
- let mut cancelled = Box::pin(sender.send(1));
- let mut waiting = Box::pin(competing.send(2));
- let (cancelled_waker, cancelled_wakes) = WakeCounter::new();
- let (waiting_waker, waiting_wakes) = WakeCounter::new();
-
- assert!(poll_with(cancelled.as_mut(), &cancelled_waker).is_pending());
- assert!(poll_with(waiting.as_mut(), &waiting_waker).is_pending());
- assert_eq!(receiver.try_recv(), Ok(0));
- assert_eq!(cancelled_wakes.count(), 1);
- assert_eq!(waiting_wakes.count(), 0);
- sender.try_send(3).unwrap();
- drop(cancelled);
- assert_eq!(waiting_wakes.count(), 0);
-
- assert_eq!(receiver.try_recv(), Ok(3));
- assert_eq!(waiting_wakes.count(), 1);
- expect_ready(poll_with(waiting.as_mut(), &waiting_waker)).unwrap();
- assert_eq!(receiver.try_recv(), Ok(2));
- assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty));
-}
-
#[test]
fn bounded_cancelled_sender_notifies_next_sender_before_dropping_value() {
// A message destructor may depend on another blocked sender making progress.
diff --git a/tests-integration/tests/mpmc_test/reservation.rs b/tests-integration/tests/mpmc_test/reservation.rs
new file mode 100644
index 0000000..577feec
--- /dev/null
+++ b/tests-integration/tests/mpmc_test/reservation.rs
@@ -0,0 +1,265 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::sync::Arc;
+use std::sync::atomic::AtomicUsize;
+use std::sync::atomic::Ordering;
+
+use asyncband::mpmc;
+use asyncband::mpmc::TryRecvError;
+use asyncband::mpmc::TrySendError;
+use tests_integration::WakeCounter;
+use tests_integration::expect_ready;
+use tests_integration::poll_once;
+use tests_integration::poll_with;
+
+#[test]
+fn permits_and_messages_share_exact_capacity_without_reserving_message_order() {
+ let (sender, receiver) = mpmc::bounded(3);
+ let competing = receiver.clone();
+ let permit: mpmc::Permit<'_, _> = sender.try_reserve().unwrap();
+ sender.try_send(1).unwrap();
+ expect_ready(poll_once(Box::pin(sender.send(2)).as_mut())).unwrap();
+ assert_eq!(sender.try_send(3), Err(TrySendError::Full(3)));
+ assert!(matches!(sender.try_reserve(), Err(TrySendError::Full(()))));
+
+ assert_eq!(receiver.try_recv(), Ok(1));
+ let second_permit = sender.try_reserve().unwrap();
+ assert_eq!(sender.try_send(4), Err(TrySendError::Full(4)));
+ permit.send(3).unwrap();
+ assert_eq!(competing.try_recv(), Ok(2));
+ second_permit.send(4).unwrap();
+ assert_eq!(receiver.try_recv(), Ok(3));
+ assert_eq!(competing.try_recv(), Ok(4));
+ assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty));
+}
+
+#[test]
+fn mixed_sends_and_reservations_receive_grants_in_wait_queue_order() {
+ let (sender, receiver) = mpmc::bounded(1);
+ let competing = receiver.clone();
+ sender.try_send(0).unwrap();
+ let mut first = Box::pin(sender.reserve());
+ let mut second = Box::pin(sender.send(2));
+ let mut third = Box::pin(sender.reserve());
+ let (first_waker, first_wakes) = WakeCounter::new();
+ let (second_waker, second_wakes) = WakeCounter::new();
+ let (third_waker, third_wakes) = WakeCounter::new();
+ assert!(poll_with(first.as_mut(), &first_waker).is_pending());
+ assert!(poll_with(second.as_mut(), &second_waker).is_pending());
+ assert!(poll_with(third.as_mut(), &third_waker).is_pending());
+
+ assert_eq!(receiver.try_recv(), Ok(0));
+ assert_eq!(
+ (
+ first_wakes.count(),
+ second_wakes.count(),
+ third_wakes.count()
+ ),
+ (1, 0, 0)
+ );
+ assert_eq!(sender.try_send(9), Err(TrySendError::Full(9)));
+ assert!(matches!(sender.try_reserve(), Err(TrySendError::Full(()))));
+ let first_permit = expect_ready(poll_with(first.as_mut(), &first_waker)).unwrap();
+ first_permit.send(1).unwrap();
+ assert_eq!(competing.try_recv(), Ok(1));
+ assert_eq!((second_wakes.count(), third_wakes.count()), (1, 0));
+ expect_ready(poll_with(second.as_mut(), &second_waker)).unwrap();
+ assert_eq!(receiver.try_recv(), Ok(2));
+ assert_eq!(third_wakes.count(), 1);
+ let third_permit = expect_ready(poll_with(third.as_mut(), &third_waker)).unwrap();
+ third_permit.send(3).unwrap();
+ assert_eq!(competing.try_recv(), Ok(3));
+}
+
+#[test]
+fn dropping_an_unused_permit_transfers_its_one_slot_to_a_waiter() {
+ let (sender, receiver) = mpmc::bounded(1);
+ let held = sender.try_reserve().unwrap();
+ let mut waiting = Box::pin(sender.reserve());
+ let (waker, wakes) = WakeCounter::new();
+ assert!(poll_with(waiting.as_mut(), &waker).is_pending());
+ drop(held);
+ assert_eq!(wakes.count(), 1);
+ assert!(matches!(sender.try_reserve(), Err(TrySendError::Full(()))));
+ assert_eq!(sender.try_send(9), Err(TrySendError::Full(9)));
+ let mut barging = Box::pin(sender.send(10));
+ assert!(poll_once(barging.as_mut()).is_pending());
+ drop(barging);
+ let granted = expect_ready(poll_with(waiting.as_mut(), &waker)).unwrap();
+ drop(granted);
+ let available = sender.try_reserve().unwrap();
+ assert!(matches!(sender.try_reserve(), Err(TrySendError::Full(()))));
+ available.send(7).unwrap();
+ assert_eq!(receiver.try_recv(), Ok(7));
+}
+
+#[test]
+fn cancelling_a_reservation_before_its_grant_skips_it() {
+ let (sender, _receiver) = mpmc::bounded::<usize>(1);
+ let held = sender.try_reserve().unwrap();
+ let mut cancelled = Box::pin(sender.reserve());
+ let mut waiting = Box::pin(sender.reserve());
+ let (waker, wakes) = WakeCounter::new();
+ assert!(poll_once(cancelled.as_mut()).is_pending());
+ assert!(poll_with(waiting.as_mut(), &waker).is_pending());
+ drop(cancelled);
+ assert_eq!(wakes.count(), 0);
+ drop(held);
+ assert_eq!(wakes.count(), 1);
+ let granted = expect_ready(poll_with(waiting.as_mut(), &waker)).unwrap();
+ assert!(matches!(sender.try_reserve(), Err(TrySendError::Full(()))));
+ drop(granted);
+ drop(sender.try_reserve().unwrap());
+}
+
+#[test]
+fn cancelling_a_granted_reservation_hands_capacity_to_the_next_sender() {
+ let (sender, receiver) = mpmc::bounded(1);
+ let held = sender.try_reserve().unwrap();
+ let mut cancelled = Box::pin(sender.reserve());
+ let mut waiting = Box::pin(sender.send(7));
+ let (first_waker, first_wakes) = WakeCounter::new();
+ let (second_waker, second_wakes) = WakeCounter::new();
+ assert!(poll_with(cancelled.as_mut(), &first_waker).is_pending());
+ assert!(poll_with(waiting.as_mut(), &second_waker).is_pending());
+ drop(held);
+ assert_eq!((first_wakes.count(), second_wakes.count()), (1, 0));
+ assert_eq!(sender.try_send(9), Err(TrySendError::Full(9)));
+ drop(cancelled);
+ assert_eq!(second_wakes.count(), 1);
+ assert!(matches!(sender.try_reserve(), Err(TrySendError::Full(()))));
+ expect_ready(poll_with(waiting.as_mut(), &second_waker)).unwrap();
+ assert_eq!(receiver.try_recv(), Ok(7));
+}
+
+#[test]
+fn cancelling_a_granted_send_wakes_the_next_waiter_before_dropping_its_value() {
+ struct ObserveDrop {
+ wakes: Arc<WakeCounter>,
+ seen: Arc<AtomicUsize>,
+ }
+
+ impl Drop for ObserveDrop {
+ fn drop(&mut self) {
+ self.seen.store(self.wakes.count(), Ordering::Relaxed);
+ }
+ }
+
+ let (sender, receiver) = mpmc::bounded(1);
+ let held = sender.try_reserve().unwrap();
+ let (waiting_waker, waiting_wakes) = WakeCounter::new();
+ let wakes_at_drop = Arc::new(AtomicUsize::new(usize::MAX));
+ let mut cancelled = Box::pin(sender.send(Some(ObserveDrop {
+ wakes: waiting_wakes.clone(),
+ seen: wakes_at_drop.clone(),
+ })));
+ let mut waiting = Box::pin(sender.reserve());
+ let (cancelled_waker, cancelled_wakes) = WakeCounter::new();
+ assert!(poll_with(cancelled.as_mut(), &cancelled_waker).is_pending());
+ assert!(poll_with(waiting.as_mut(), &waiting_waker).is_pending());
+ drop(held);
+ assert_eq!((cancelled_wakes.count(), waiting_wakes.count()), (1, 0));
+ drop(cancelled);
+ assert_eq!(wakes_at_drop.load(Ordering::Relaxed), 1);
+ let permit = expect_ready(poll_with(waiting.as_mut(), &waiting_waker)).unwrap();
+ permit.send(None).unwrap();
+ assert!(receiver.try_recv().unwrap().is_none());
+}
+
+#[test]
+fn cancelling_an_ungranted_send_drops_its_value_without_taking_capacity() {
+ struct CountDrop(Arc<AtomicUsize>);
+
+ impl Drop for CountDrop {
+ fn drop(&mut self) {
+ self.0.fetch_add(1, Ordering::Relaxed);
+ }
+ }
+
+ let (sender, receiver) = mpmc::bounded(1);
+ let held = sender.try_reserve().unwrap();
+ let drops = Arc::new(AtomicUsize::new(0));
+ let mut cancelled = Box::pin(sender.send(Some(CountDrop(drops.clone()))));
+ let mut waiting = Box::pin(sender.reserve());
+ let (waker, wakes) = WakeCounter::new();
+ assert!(poll_once(cancelled.as_mut()).is_pending());
+ assert!(poll_with(waiting.as_mut(), &waker).is_pending());
+ drop(cancelled);
+ assert_eq!(drops.load(Ordering::Relaxed), 1);
+ assert_eq!(wakes.count(), 0);
+ drop(held);
+ assert_eq!(wakes.count(), 1);
+ expect_ready(poll_with(waiting.as_mut(), &waker))
+ .unwrap()
+ .send(None)
+ .unwrap();
+ assert!(receiver.try_recv().unwrap().is_none());
+ assert_eq!(drops.load(Ordering::Relaxed), 1);
+}
+
+#[test]
+fn last_receiver_drop_fails_pending_and_granted_reservations() {
+ let (sender, receiver) = mpmc::bounded::<usize>(1);
+ let competing = receiver.clone();
+ sender.try_send(0).unwrap();
+ let mut granted = Box::pin(sender.reserve());
+ let mut pending = Box::pin(sender.reserve());
+ let (first_waker, first_wakes) = WakeCounter::new();
+ let (second_waker, second_wakes) = WakeCounter::new();
+ assert!(poll_with(granted.as_mut(), &first_waker).is_pending());
+ assert!(poll_with(pending.as_mut(), &second_waker).is_pending());
+ drop(receiver);
+ assert_eq!((first_wakes.count(), second_wakes.count()), (0, 0));
+ assert_eq!(competing.try_recv(), Ok(0));
+ assert_eq!((first_wakes.count(), second_wakes.count()), (1, 0));
+ drop(competing);
+ assert!(first_wakes.count() >= 1);
+ assert!(second_wakes.count() >= 1);
+ assert!(expect_ready(poll_with(granted.as_mut(), &first_waker)).is_err());
+ assert!(expect_ready(poll_with(pending.as_mut(), &second_waker)).is_err());
+ assert!(matches!(
+ sender.try_reserve(),
+ Err(TrySendError::Disconnected(()))
+ ));
+}
+
+#[test]
+fn permit_does_not_keep_receivers_alive_and_returns_the_unsent_value() {
+ let (sender, receiver) = mpmc::bounded(2);
+ let competing = receiver.clone();
+ let permit = sender.try_reserve().unwrap();
+ drop(receiver);
+ sender.try_send(String::from("still connected")).unwrap();
+ drop(competing);
+ assert_eq!(
+ permit
+ .send(String::from("unsent"))
+ .unwrap_err()
+ .into_inner(),
+ "unsent"
+ );
+ assert!(matches!(
+ sender.try_reserve(),
+ Err(TrySendError::Disconnected(()))
+ ));
+ let error = match expect_ready(poll_once(Box::pin(sender.reserve()).as_mut())) {
+ Ok(_) => panic!("reservation succeeded after the last receiver dropped"),
+ Err(error) => error,
+ };
+ assert_eq!(error.into_inner(), ());
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@asyncband.apache.org
For additional commands, e-mail: commits-help@asyncband.apache.org