diff --git a/CHANGELOG.md b/CHANGELOG.md index 76f0b4b..e6430ea 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 ### 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 BoundedSender { /// /// # 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> { 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, 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, 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> { 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 fmt::Debug for SendError { impl std::error::Error for SendError {} -/// 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 { - /// 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 { /// outside the lock, because they may reenter the queue. struct State { values: VecDeque, - capacity: Option, + // 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, senders: usize, receivers: usize, - recv_waiters: WaitList, - send_waiters: WaitList, + recv_waiters: WaitList, + send_waiters: WaitList, } impl State { 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 { + 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 State { 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), 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 State { } } -/// 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 { +/// A grant transfers capacity to a detached sender waiter until it claims or cancels it. +enum SendWaiter { + Waiting(Waker), + Granted, +} + +fn remove_waiter(waiters: &mut WaitList, 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 { fn notify_one(&mut self) -> Option { 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, current: &Waker) -> Option { 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 { 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 { + /// 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, current: &Waker) -> Option { + 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 { + 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 { + 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 Shared { pub fn bounded(capacity: usize) -> Self { Self::new(Some(capacity)) @@ -132,7 +197,7 @@ impl Shared { Self { state: Mutex::new(State { values: VecDeque::new(), - capacity, + available: capacity, senders: 1, receivers: 1, recv_waiters: WaitList::new(), @@ -177,18 +242,19 @@ impl Shared { }; // 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> { 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 Shared { 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, TrySendError<()>> { + self.state.lock().acquire()?; + Ok(Permit { shared: self }) + } + + pub async fn reserve(&self) -> Result, 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 { @@ -234,51 +311,97 @@ impl Shared { } } -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, - 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") +impl 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>> { +impl 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> { + 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 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, + waiter: Option, +} + +impl<'a, T> Reserve<'a, T> { + fn poll(&mut self, cx: &mut Context<'_>) -> Poll, 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 Drop for Send<'_, T> { +impl Drop for Reserve<'_, T> { fn drop(&mut self) { let Some(id) = self.waiter.take() else { return; @@ -288,10 +411,10 @@ impl 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 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 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: 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::(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, + seen: Arc, + } + + 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); + + 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::(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(), ()); +}