diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f226004..d1d97fbf 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -23,6 +23,10 @@ All notable changes to this project will be documented in this file. ## Unreleased +### Breaking changes + +* Return shared messages from `broadcast::mpmc` receivers: `recv` and `try_recv` now yield `Arc`, and `bounded` and `unbounded` no longer require `T: Clone`. Callers that need an owned value use the new `recv_cloned` and `try_recv_cloned`, which still move the payload out instead of cloning it when nothing else still refers to the message; callers that sent `Arc` to avoid cloning can send `T` directly. A retained `Arc` keeps its payload alive after the message leaves the backlog, so a bounded channel's capacity limits retained messages rather than memory. + ### 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. diff --git a/asyncband/src/broadcast/mpmc/bounded/mod.rs b/asyncband/src/broadcast/mpmc/bounded/mod.rs index 0bfbc8b5..3d3f6f07 100644 --- a/asyncband/src/broadcast/mpmc/bounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/bounded/mod.rs @@ -66,10 +66,10 @@ //! tx.send(10).await; //! tx.send(20).await; //! -//! assert_eq!(rx1.recv().await, Ok(10)); -//! assert_eq!(rx1.recv().await, Ok(20)); -//! assert_eq!(rx2.recv().await, Ok(10)); -//! assert_eq!(rx2.recv().await, Ok(20)); +//! assert_eq!(*rx1.recv().await.unwrap(), 10); +//! assert_eq!(*rx1.recv().await.unwrap(), 20); +//! assert_eq!(*rx2.recv().await.unwrap(), 10); +//! assert_eq!(*rx2.recv().await.unwrap(), 20); //! # } //! ``` //! @@ -87,7 +87,7 @@ //! assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); //! //! // `rx1` draining is not enough: `rx2` has read neither message, so both stay retained. -//! assert_eq!(rx1.try_recv(), Ok(1)); +//! assert_eq!(*rx1.try_recv().unwrap(), 1); //! assert_eq!(tx.retained_message_count(), 2); //! assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); //! @@ -112,6 +112,7 @@ use std::task::Poll; use super::common; use super::common::Backlog; use super::common::Inner; +use super::common::Received; use super::error::RecvError; use super::error::TryRecvError; use super::error::TrySendError; @@ -142,10 +143,10 @@ mod tests; /// /// let (tx, mut rx) = mpmc::bounded(1); /// tx.try_send(10).unwrap(); -/// assert_eq!(rx.try_recv(), Ok(10)); +/// assert_eq!(*rx.try_recv().unwrap(), 10); /// ``` #[track_caller] -pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { +pub fn bounded(capacity: usize) -> (BoundedSender, BoundedReceiver) { assert!( capacity > 0, "broadcast bounded channel requires capacity > 0" @@ -283,7 +284,7 @@ impl BoundedSender { /// # async fn main() { /// let (tx, mut rx) = mpmc::bounded(1); /// tx.send(10).await; - /// assert_eq!(rx.recv().await, Ok(10)); + /// assert_eq!(*rx.recv().await.unwrap(), 10); /// # } /// ``` pub async fn send(&self, value: T) { @@ -380,7 +381,7 @@ impl BoundedSender { /// tx.try_send(10).unwrap(); /// assert_eq!(tx.try_send(20), Err(TrySendError::Full(20))); /// - /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(*rx.try_recv().unwrap(), 10); /// tx.try_send(20).unwrap(); /// ``` pub fn try_send(&self, value: T) -> Result<(), TrySendError> { @@ -445,7 +446,7 @@ impl BoundedSender { /// tx.try_send(10).unwrap(); /// assert_eq!(tx.retained_message_count(), 1); /// - /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(*rx.try_recv().unwrap(), 10); /// assert_eq!(tx.retained_message_count(), 0); /// ``` pub fn retained_message_count(&self) -> usize { @@ -488,7 +489,7 @@ impl BoundedSender { /// let mut rx = tx.subscribe(); /// assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); /// tx.send(20).await; - /// assert_eq!(rx.recv().await, Ok(20)); + /// assert_eq!(*rx.recv().await.unwrap(), 20); /// # } /// ``` #[must_use = "the receiver is dropped immediately if it is not retained"] @@ -537,12 +538,16 @@ impl Drop for BoundedReceiver { } } -impl BoundedReceiver { - /// Receives the next value for this receiver. +impl BoundedReceiver { + /// Receives the next value for this receiver as the channel's shared `Arc`. + /// + /// Receiving never clones `T`. The `Arc` outlives the message's stay in the backlog: a + /// caller that keeps it keeps the payload alive but holds no capacity. Use + /// [`recv_cloned`](BoundedReceiver::recv_cloned) to take an owned value instead. /// /// # Returns /// - /// * `Ok(T)`: The next message. + /// * `Ok(Arc)`: The next message. /// * `Err(RecvError::Disconnected)`: All senders have been dropped and this receiver has no /// remaining messages. /// @@ -561,22 +566,23 @@ impl BoundedReceiver { /// # async fn main() { /// let (tx, mut rx) = mpmc::bounded(4); /// tx.send(10).await; - /// assert_eq!(rx.recv().await, Ok(10)); + /// assert_eq!(*rx.recv().await.unwrap(), 10); /// # } /// ``` - pub async fn recv(&mut self) -> Result { - Recv { - receiver: self, - token: None, - } - .await + pub async fn recv(&mut self) -> Result, RecvError> { + let (msg, reclaimed) = self.wait_next().await?; + Ok(common::share_msg(msg, reclaimed)) } - /// Attempts to receive the next value for this receiver without blocking. + /// Attempts to receive the next value for this receiver without blocking, as the channel's + /// shared `Arc`. + /// + /// See [`recv`](BoundedReceiver::recv) for what keeping the `Arc` means, and + /// [`try_recv_cloned`](BoundedReceiver::try_recv_cloned) to take an owned value instead. /// /// # Returns /// - /// * `Ok(T)`: The next message. + /// * `Ok(Arc)`: The next message. /// * `Err(TryRecvError::Empty)`: No message is currently available. /// * `Err(TryRecvError::Disconnected)`: All senders have been dropped and this receiver has no /// remaining messages. @@ -588,21 +594,26 @@ impl BoundedReceiver { /// /// let (tx, mut rx) = mpmc::bounded(4); /// tx.try_send(10).unwrap(); - /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(*rx.try_recv().unwrap(), 10); /// ``` - pub fn try_recv(&mut self) -> Result { + pub fn try_recv(&mut self) -> Result, TryRecvError> { + let (msg, reclaimed) = self.receive_released()?; + Ok(common::share_msg(msg, reclaimed)) + } + + /// Receives without waiting and hands the capacity this receive released to blocked producers + /// before the caller can touch the payload. + /// + /// `T::clone` and `T::drop` may panic, and if either ran first the reclaimed slots would never + /// be handed to a waiting producer, stalling it permanently. The `Recv` future orders its poll + /// the same way. + fn receive_released(&mut self) -> Result, TryRecvError> { let (msg, reclaimed) = common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; - - // Release before taking the payload: `take_msg` runs `T::clone` and `T::drop`, and if - // either panics the reclaimed slots would otherwise never be handed to a waiting producer, - // stalling it permanently. self.shared.release_reclaimed(reclaimed.len()); - Ok(common::take_msg(msg, reclaimed)) + Ok((msg, reclaimed)) } -} -impl BoundedReceiver { /// Re-subscribes to the channel, returning a new receiver that starts receiving messages from /// the *current* tail of the channel. /// @@ -622,7 +633,7 @@ impl BoundedReceiver { /// let mut rx2 = rx.resubscribe(); /// tx.try_send(3).unwrap(); /// - /// assert_eq!(rx2.try_recv(), Ok(3)); + /// assert_eq!(*rx2.try_recv().unwrap(), 3); /// ``` #[must_use = "the receiver is dropped immediately if it is not retained"] pub fn resubscribe(&self) -> Self { @@ -654,12 +665,83 @@ impl BoundedReceiver { /// tx.try_send(20).unwrap(); /// assert_eq!(rx.unread_message_count(), 2); /// - /// assert_eq!(rx.try_recv(), Ok(10)); + /// assert_eq!(*rx.try_recv().unwrap(), 10); /// assert_eq!(rx.unread_message_count(), 1); /// ``` pub fn unread_message_count(&self) -> usize { self.shared.inner.lock().log.unread(self.key) } + + fn wait_next(&mut self) -> Recv<'_, T> { + Recv { + receiver: self, + token: None, + } + } +} + +impl BoundedReceiver { + /// Receives the next value for this receiver as an owned `T`. + /// + /// The payload is moved out of the channel without cloning when this read reclaims the message + /// and nothing else still refers to it, such as an `Arc` kept from a shared receive or another + /// receiver's clone still in progress. Otherwise it is cloned with the channel unlocked. If + /// `T::clone` panics, this receiver has still advanced past the value. + /// + /// # Returns + /// + /// * `Ok(T)`: The next message. + /// * `Err(RecvError::Disconnected)`: All senders have been dropped and this receiver has no + /// remaining messages. + /// + /// # Cancel safety + /// + /// Dropping a pending `recv_cloned` future leaves this receiver's cursor unchanged. A + /// subsequent receive can still return the same next value, so these futures may safely be + /// raced with other futures in a selection construct. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (tx, mut rx) = mpmc::bounded(4); + /// tx.send(String::from("owned")).await; + /// assert_eq!(rx.recv_cloned().await, Ok(String::from("owned"))); + /// # } + /// ``` + pub async fn recv_cloned(&mut self) -> Result { + let (msg, reclaimed) = self.wait_next().await?; + Ok(common::take_msg(msg, reclaimed)) + } + + /// Attempts to receive the next value for this receiver without blocking, as an owned `T`. + /// + /// See [`recv_cloned`](BoundedReceiver::recv_cloned) for when the payload is moved out rather + /// than cloned. + /// + /// # Returns + /// + /// * `Ok(T)`: The next message. + /// * `Err(TryRecvError::Empty)`: No message is currently available. + /// * `Err(TryRecvError::Disconnected)`: All senders have been dropped and this receiver has no + /// remaining messages. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc; + /// + /// let (tx, mut rx) = mpmc::bounded(4); + /// tx.try_send(String::from("owned")).unwrap(); + /// assert_eq!(rx.try_recv_cloned(), Ok(String::from("owned"))); + /// ``` + pub fn try_recv_cloned(&mut self) -> Result { + let (msg, reclaimed) = self.receive_released()?; + Ok(common::take_msg(msg, reclaimed)) + } } struct Recv<'a, T> { @@ -683,8 +765,8 @@ impl Drop for Recv<'_, T> { } } -impl Future for Recv<'_, T> { - type Output = Result; +impl Future for Recv<'_, T> { + type Output = Result, RecvError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let Self { receiver, token } = self.get_mut(); @@ -701,9 +783,10 @@ impl Future for Recv<'_, T> { Poll::Ready(Ok(received)) => received, }; - // Release before taking the payload, for the same reason as `try_recv`: a panicking - // `T::clone` must not strand producers waiting for the reclaimed slots. + // Release before handing out the payload, for the same reason as `receive_released`: a + // panicking `T::clone` or `T::drop` must not strand producers waiting for the reclaimed + // slots. receiver.shared.release_reclaimed(reclaimed.len()); - Poll::Ready(Ok(common::take_msg(msg, reclaimed))) + Poll::Ready(Ok((msg, reclaimed))) } } diff --git a/asyncband/src/broadcast/mpmc/bounded/tests.rs b/asyncband/src/broadcast/mpmc/bounded/tests.rs index 07c16df1..e1992337 100644 --- a/asyncband/src/broadcast/mpmc/bounded/tests.rs +++ b/asyncband/src/broadcast/mpmc/bounded/tests.rs @@ -50,10 +50,10 @@ fn buffer_is_preallocated_and_never_shrinks() { tx.try_send(i).unwrap(); } for i in 0..capacity { - assert_eq!(rx.try_recv(), Ok(i)); + assert_eq!(rx.try_recv().as_deref(), Ok(&i)); } tx.try_send(0).unwrap(); - assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(rx.try_recv().as_deref(), Ok(&0)); assert_eq!(tx.retained_message_count(), 0); assert_eq!(tx.shared.inner.lock().log.buffer_capacity(), allocated); diff --git a/asyncband/src/broadcast/mpmc/common.rs b/asyncband/src/broadcast/mpmc/common.rs index 74a04a9a..fefc5f44 100644 --- a/asyncband/src/broadcast/mpmc/common.rs +++ b/asyncband/src/broadcast/mpmc/common.rs @@ -503,12 +503,22 @@ pub fn poll_receive( } } -/// Drops the reclaimed backlog, then yields the received message, both with the channel unlocked. +/// Releases the backlog's reference to the message, then hands out the stored `Arc`, both with the +/// channel unlocked. +/// +/// A receive reclaims at most the message itself, and `msg` keeps the payload alive, so a shared +/// receive runs neither `T::clone` nor `T::drop`. +pub fn share_msg(msg: Arc, reclaimed: Reclaimed) -> Arc { + drop(reclaimed); + msg +} + +/// Drops the reclaimed backlog, then yields the received message owned, both with the channel +/// unlocked. /// /// A non-empty backlog means receiving this message removed `msg` from the buffer. Once the /// reclaimed backlog is dropped, the caller may hold the only remaining reference and can then -/// move the payload out instead of cloning it. A channel with a single receiver never clones a -/// payload. +/// move the payload out instead of cloning it. /// /// Ownership is decided from that bookkeeping rather than by probing the reference count. An /// [`Arc::try_unwrap`] call for every received message would fail under fan-out, and its failed @@ -524,9 +534,9 @@ pub fn take_msg(msg: Arc, reclaimed: Reclaimed) -> T { return (*msg).clone(); } - // Another receiver can still hold an in-flight reference to the same message, so the clone - // remains the fallback. - Arc::try_unwrap(msg).unwrap_or_else(|msg| (*msg).clone()) + // Another receiver can still hold an in-flight reference to the same message, or an `Arc` it + // kept from a shared receive, so the clone remains the fallback. + Arc::unwrap_or_clone(msg) } #[cfg(test)] diff --git a/asyncband/src/broadcast/mpmc/mod.rs b/asyncband/src/broadcast/mpmc/mod.rs index b5eb4afd..8cdfb926 100644 --- a/asyncband/src/broadcast/mpmc/mod.rs +++ b/asyncband/src/broadcast/mpmc/mod.rs @@ -27,9 +27,13 @@ //! //! Receiving a value advances the subscription before returning that value. The channel tracks //! unread messages, not application work: retaining a received value or processing it -//! asynchronously does not hold backlog capacity. There is no acknowledgement or -//! processing-completion barrier. If cloning a received value panics, that subscription has still -//! advanced past the value. +//! asynchronously does not hold backlog capacity. `recv` and `try_recv` hand out the channel's own +//! `Arc`, so an `Arc` kept past the receive keeps its payload alive after the message has left +//! the backlog, and a bounded channel's capacity limits retained messages rather than memory. +//! Callers that wrapped values in an `Arc` only to avoid broadcast cloning can send `T` directly, +//! since the channel already shares one `Arc` with every receiver. There is no acknowledgement +//! or processing-completion barrier. If `T::clone` panics inside `recv_cloned` or +//! `try_recv_cloned`, that subscription has still advanced past the value. //! //! Sending with no subscriptions discards the value and succeeds. A later subscription starts //! with future publications; it does not replay discarded or previously retained values. diff --git a/asyncband/src/broadcast/mpmc/unbounded/mod.rs b/asyncband/src/broadcast/mpmc/unbounded/mod.rs index de2eab06..0339b9d4 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/mod.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/mod.rs @@ -46,9 +46,9 @@ //! let mut late = publisher.subscribe(); //! publisher.send("after subscription"); //! -//! assert_eq!(early.try_recv(), Ok("before subscription")); -//! assert_eq!(early.try_recv(), Ok("after subscription")); -//! assert_eq!(late.try_recv(), Ok("after subscription")); +//! assert_eq!(*early.try_recv().unwrap(), "before subscription"); +//! assert_eq!(*early.try_recv().unwrap(), "after subscription"); +//! assert_eq!(*late.try_recv().unwrap(), "after subscription"); //! assert_eq!(late.try_recv(), Err(TryRecvError::Empty)); //! ``` //! @@ -66,6 +66,7 @@ use std::task::Poll; use super::common; use super::common::Backlog; use super::common::Inner; +use super::common::Received; use super::error::RecvError; use super::error::TryRecvError; use crate::internal::arena::SlotId; @@ -87,9 +88,9 @@ mod tests; /// /// let (publisher, mut receiver) = unbounded(); /// publisher.send("ready"); -/// assert_eq!(receiver.try_recv(), Ok("ready")); +/// assert_eq!(*receiver.try_recv().unwrap(), "ready"); /// ``` -pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { +pub fn unbounded() -> (UnboundedSender, UnboundedReceiver) { let (inner, key) = Inner::with_first_subscription(Backlog::elastic()); let shared = Arc::new(Shared { inner, @@ -163,8 +164,8 @@ impl UnboundedSender { /// let mut second = publisher.subscribe(); /// publisher.send("update"); /// - /// assert_eq!(first.try_recv(), Ok("update")); - /// assert_eq!(second.try_recv(), Ok("update")); + /// assert_eq!(*first.try_recv().unwrap(), "update"); + /// assert_eq!(*second.try_recv().unwrap(), "update"); /// ``` pub fn send(&self, msg: T) { let msg = Arc::new(msg); @@ -197,9 +198,9 @@ impl UnboundedSender { /// publisher.send("update"); /// assert_eq!(publisher.retained_message_count(), 1); /// - /// assert_eq!(fast.try_recv(), Ok("update")); + /// assert_eq!(*fast.try_recv().unwrap(), "update"); /// assert_eq!(publisher.retained_message_count(), 1); - /// assert_eq!(slow.try_recv(), Ok("update")); + /// assert_eq!(*slow.try_recv().unwrap(), "update"); /// assert_eq!(publisher.retained_message_count(), 0); /// ``` pub fn retained_message_count(&self) -> usize { @@ -222,7 +223,7 @@ impl UnboundedSender { /// let mut receiver = publisher.subscribe(); /// assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); /// publisher.send("later"); - /// assert_eq!(receiver.try_recv(), Ok("later")); + /// assert_eq!(*receiver.try_recv().unwrap(), "later"); /// ``` #[must_use = "the receiver is dropped immediately if it is not retained"] pub fn subscribe(&self) -> UnboundedReceiver { @@ -259,8 +260,12 @@ impl Drop for UnboundedReceiver { } } -impl UnboundedReceiver { - /// Waits for this receiver's next value. +impl UnboundedReceiver { + /// Waits for this receiver's next value and returns the channel's shared `Arc`. + /// + /// Receiving never clones `T`. The `Arc` outlives the value's stay in the backlog, so a + /// caller that keeps it keeps the payload alive. Use + /// [`recv_cloned`](UnboundedReceiver::recv_cloned) to take an owned value instead. /// /// Values already published for this receiver are returned before disconnection. /// [`RecvError::Disconnected`] is returned only when no sender remains and this receiver's @@ -284,19 +289,20 @@ impl UnboundedReceiver { /// publisher.send("final update"); /// drop(publisher); /// - /// assert_eq!(receiver.recv().await, Ok("final update")); + /// assert_eq!(*receiver.recv().await.unwrap(), "final update"); /// assert_eq!(receiver.recv().await, Err(RecvError::Disconnected)); /// # } /// ``` - pub async fn recv(&mut self) -> Result { - Recv { - receiver: self, - token: None, - } - .await + pub async fn recv(&mut self) -> Result, RecvError> { + let (msg, reclaimed) = self.wait_next().await?; + Ok(common::share_msg(msg, reclaimed)) } - /// Attempts to take this receiver's next value without waiting. + /// Attempts to take this receiver's next value without waiting, as the channel's shared + /// `Arc`. + /// + /// See [`recv`](UnboundedReceiver::recv) for what keeping the `Arc` means, and + /// [`try_recv_cloned`](UnboundedReceiver::try_recv_cloned) to take an owned value instead. /// /// [`TryRecvError::Empty`] means this receiver is currently caught up while a sender remains. /// [`TryRecvError::Disconnected`] means no sender remains and this receiver has drained its @@ -312,18 +318,16 @@ impl UnboundedReceiver { /// assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty)); /// /// publisher.send("update"); - /// assert_eq!(receiver.try_recv(), Ok("update")); + /// assert_eq!(*receiver.try_recv().unwrap(), "update"); /// drop(publisher); /// assert_eq!(receiver.try_recv(), Err(TryRecvError::Disconnected)); /// ``` - pub fn try_recv(&mut self) -> Result { + pub fn try_recv(&mut self) -> Result, TryRecvError> { let (msg, reclaimed) = common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; - Ok(common::take_msg(msg, reclaimed)) + Ok(common::share_msg(msg, reclaimed)) } -} -impl UnboundedReceiver { /// Re-subscribes to the channel, returning a new receiver that starts receiving messages from /// the *current* tail of the channel. /// @@ -343,9 +347,9 @@ impl UnboundedReceiver { /// assert_eq!(fresh.try_recv(), Err(TryRecvError::Empty)); /// publisher.send("visible to both"); /// - /// assert_eq!(original.try_recv(), Ok("pending for original")); - /// assert_eq!(original.try_recv(), Ok("visible to both")); - /// assert_eq!(fresh.try_recv(), Ok("visible to both")); + /// assert_eq!(*original.try_recv().unwrap(), "pending for original"); + /// assert_eq!(*original.try_recv().unwrap(), "visible to both"); + /// assert_eq!(*fresh.try_recv().unwrap(), "visible to both"); /// ``` #[must_use = "the receiver is dropped immediately if it is not retained"] pub fn resubscribe(&self) -> Self { @@ -374,12 +378,80 @@ impl UnboundedReceiver { /// publisher.send("second"); /// assert_eq!(receiver.unread_message_count(), 2); /// - /// assert_eq!(receiver.try_recv(), Ok("first")); + /// assert_eq!(*receiver.try_recv().unwrap(), "first"); /// assert_eq!(receiver.unread_message_count(), 1); /// ``` pub fn unread_message_count(&self) -> usize { self.shared.inner.lock().log.unread(self.key) } + + fn wait_next(&mut self) -> Recv<'_, T> { + Recv { + receiver: self, + token: None, + } + } +} + +impl UnboundedReceiver { + /// Waits for this receiver's next value and returns it as an owned `T`. + /// + /// The payload is moved out of the channel without cloning when this read reclaims the message + /// and nothing else still refers to it, such as an `Arc` kept from a shared receive or another + /// receiver's clone still in progress. Otherwise it is cloned with the channel unlocked. If + /// `T::clone` panics, this receiver has still advanced past the value. + /// + /// Values already published for this receiver are returned before disconnection. + /// [`RecvError::Disconnected`] is returned only when no sender remains and this receiver's + /// backlog is empty. + /// + /// # Cancel safety + /// + /// Dropping a pending `recv_cloned` future leaves this receiver's cursor unchanged. A + /// subsequent receive can still return the same next value, so these futures may safely be + /// raced with other futures in a selection construct. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc::unbounded; + /// + /// # #[tokio::main] + /// # async fn main() { + /// let (publisher, mut receiver) = unbounded(); + /// publisher.send(String::from("owned")); + /// assert_eq!(receiver.recv_cloned().await, Ok(String::from("owned"))); + /// # } + /// ``` + pub async fn recv_cloned(&mut self) -> Result { + let (msg, reclaimed) = self.wait_next().await?; + Ok(common::take_msg(msg, reclaimed)) + } + + /// Attempts to take this receiver's next value without waiting and returns it as an owned + /// `T`. + /// + /// See [`recv_cloned`](UnboundedReceiver::recv_cloned) for when the payload is moved out + /// rather than cloned. + /// + /// [`TryRecvError::Empty`] means this receiver is currently caught up while a sender remains. + /// [`TryRecvError::Disconnected`] means no sender remains and this receiver has drained its + /// backlog. + /// + /// # Examples + /// + /// ``` + /// use asyncband::broadcast::mpmc::unbounded; + /// + /// let (publisher, mut receiver) = unbounded(); + /// publisher.send(String::from("owned")); + /// assert_eq!(receiver.try_recv_cloned(), Ok(String::from("owned"))); + /// ``` + pub fn try_recv_cloned(&mut self) -> Result { + let (msg, reclaimed) = + common::try_receive(&self.shared.inner, &self.shared.senders, self.key)?; + Ok(common::take_msg(msg, reclaimed)) + } } struct Recv<'a, T> { @@ -403,24 +475,18 @@ impl Drop for Recv<'_, T> { } } -impl Future for Recv<'_, T> { - type Output = Result; +impl Future for Recv<'_, T> { + type Output = Result, RecvError>; fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { let Self { receiver, token } = self.get_mut(); - let (msg, reclaimed) = match common::poll_receive( + common::poll_receive( &receiver.shared.inner, &receiver.shared.senders, receiver.key, token, cx, - ) { - Poll::Pending => return Poll::Pending, - Poll::Ready(Err(err)) => return Poll::Ready(Err(err)), - Poll::Ready(Ok(received)) => received, - }; - - Poll::Ready(Ok(common::take_msg(msg, reclaimed))) + ) } } diff --git a/asyncband/src/broadcast/mpmc/unbounded/tests.rs b/asyncband/src/broadcast/mpmc/unbounded/tests.rs index 4bf330aa..f91c8e70 100644 --- a/asyncband/src/broadcast/mpmc/unbounded/tests.rs +++ b/asyncband/src/broadcast/mpmc/unbounded/tests.rs @@ -38,7 +38,7 @@ fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { assert!(tx.shared.inner.lock().log.buffer_capacity() >= burst); for i in 0..burst { - assert_eq!(rx.try_recv(), Ok(i)); + assert_eq!(rx.try_recv().as_deref(), Ok(&i)); } // Draining evaluates the cycle that just peaked, so the burst allocation is still held. @@ -47,7 +47,7 @@ fn one_off_burst_allocation_is_returned_once_it_is_behind_us() { // The next cycle stays small, which is what releases the memory. tx.send(0); - assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(rx.try_recv().as_deref(), Ok(&0)); assert!(tx.shared.inner.lock().log.buffer_capacity() < burst); } @@ -61,7 +61,7 @@ fn repeated_bursts_keep_their_allocation() { tx.send(i); } for i in 0..burst { - assert_eq!(rx.try_recv(), Ok(i)); + assert_eq!(rx.try_recv().as_deref(), Ok(&i)); } } diff --git a/benchmarks/benches/ecosystem/broadcast/mpmc/adapters.rs b/benchmarks/benches/ecosystem/broadcast/mpmc/adapters.rs index 48e72e34..daf3d7af 100644 --- a/benchmarks/benches/ecosystem/broadcast/mpmc/adapters.rs +++ b/benchmarks/benches/ecosystem/broadcast/mpmc/adapters.rs @@ -21,6 +21,8 @@ use std::task::Context; use asyncband::blocking::FutureExt; use benchmarks::support::poll_ready; +/// The peers hand out owned values, so the asyncband adapters use the cloned receives to keep the +/// comparison fair. pub struct Asyncband; pub struct Tokio; pub struct AsyncBroadcast; @@ -54,7 +56,7 @@ impl BroadcastMpmc for Asyncband { } fn try_recv(receiver: &mut Self::Receiver) -> Option { - match receiver.try_recv() { + match receiver.try_recv_cloned() { Ok(value) => Some(value), Err(asyncband::broadcast::mpmc::TryRecvError::Empty) => None, Err(asyncband::broadcast::mpmc::TryRecvError::Disconnected) => { @@ -64,7 +66,7 @@ impl BroadcastMpmc for Asyncband { } fn recv_ready(receiver: &mut Self::Receiver, context: &mut Context<'_>) -> usize { - poll_ready(receiver.recv(), context).unwrap() + poll_ready(receiver.recv_cloned(), context).unwrap() } } @@ -183,11 +185,11 @@ impl BoundedBroadcastMpmc for Asyncband { } async fn recv_async(receiver: &mut Self::Receiver) -> usize { - receiver.recv().await.unwrap() + receiver.recv_cloned().await.unwrap() } fn try_recv(receiver: &mut Self::Receiver) -> Option { - match receiver.try_recv() { + match receiver.try_recv_cloned() { Ok(value) => Some(value), Err(asyncband::broadcast::mpmc::TryRecvError::Empty) => None, Err(asyncband::broadcast::mpmc::TryRecvError::Disconnected) => { diff --git a/benchmarks/benches/primitives/broadcast/mpmc/bounded.rs b/benchmarks/benches/primitives/broadcast/mpmc/bounded.rs index dd03bdf5..39f3d9e9 100644 --- a/benchmarks/benches/primitives/broadcast/mpmc/bounded.rs +++ b/benchmarks/benches/primitives/broadcast/mpmc/bounded.rs @@ -146,7 +146,7 @@ fn drop_lagging_receiver_wakes_senders(bencher: Bencher, backlog: usize) { let slow = sender.subscribe(); for value in 0..backlog { sender.try_send(value).unwrap(); - assert_eq!(fast.try_recv().unwrap(), value); + assert_eq!(*fast.try_recv().unwrap(), value); } let mut context = bench_context(); let mut sends = (0..backlog) diff --git a/benchmarks/benches/primitives/broadcast/mpmc/unbounded.rs b/benchmarks/benches/primitives/broadcast/mpmc/unbounded.rs index 7728c17b..d25bb4e0 100644 --- a/benchmarks/benches/primitives/broadcast/mpmc/unbounded.rs +++ b/benchmarks/benches/primitives/broadcast/mpmc/unbounded.rs @@ -68,8 +68,9 @@ fn send_and_try_recv_shared(bencher: Bencher) { }); } -// The `usize` benchmarks above hide the cost of receiving a payload that owns memory: cloning it -// requires an allocation rather than a register move. +// The `usize` benchmarks above hide the cost of receiving a payload that owns memory as an owned +// value: cloning it requires an allocation rather than a register move, and moving it out skips +// the clone only while no other `Arc` shares it. fn payload() -> String { "x".repeat(64) } @@ -79,7 +80,7 @@ fn send_and_try_recv_owned(bencher: Bencher) { let (sender, mut receiver) = mpmc::unbounded(); bencher.bench_local(|| { sender.send(black_box(payload())); - black_box(receiver.try_recv().unwrap()) + black_box(receiver.try_recv_cloned().unwrap()) }); } @@ -89,8 +90,8 @@ fn send_and_try_recv_owned_shared(bencher: Bencher) { let mut second = sender.subscribe(); bencher.bench_local(|| { sender.send(black_box(payload())); - black_box(first.try_recv().unwrap()); - black_box(second.try_recv().unwrap()) + black_box(first.try_recv_cloned().unwrap()); + black_box(second.try_recv_cloned().unwrap()) }); } diff --git a/tests-integration/tests/broadcast_test/bounded.rs b/tests-integration/tests/broadcast_test/bounded.rs index 5423783e..186b6eea 100644 --- a/tests-integration/tests/broadcast_test/bounded.rs +++ b/tests-integration/tests/broadcast_test/bounded.rs @@ -18,6 +18,8 @@ use std::future::Future; use std::sync::Arc; use std::sync::Barrier; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering; use std::task::Context; use std::task::Waker; use std::thread; @@ -26,7 +28,9 @@ use asyncband::blocking::FutureExt; use asyncband::broadcast::mpmc::*; 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; /// A payload whose destructor re-enters the channel it was sent through. struct Reentrant { @@ -105,10 +109,10 @@ async fn bounded_delivers_every_message_to_every_receiver() { tx.send(10).await; tx.send(20).await; - assert_eq!(rx1.recv().await, Ok(10)); - assert_eq!(rx1.recv().await, Ok(20)); - assert_eq!(rx2.recv().await, Ok(10)); - assert_eq!(rx2.recv().await, Ok(20)); + assert_eq!(rx1.recv().await.as_deref(), Ok(&10)); + assert_eq!(rx1.recv().await.as_deref(), Ok(&20)); + assert_eq!(rx2.recv().await.as_deref(), Ok(&10)); + assert_eq!(rx2.recv().await.as_deref(), Ok(&20)); } #[test] @@ -121,18 +125,18 @@ fn bounded_slow_receiver_keeps_every_message_under_backpressure() { // The fast subscription draining does not release anything, because the slow one has read // nothing — being bounded must not turn into dropping what the slow subscription still owes. - assert_eq!(fast.try_recv(), Ok(1)); - assert_eq!(fast.try_recv(), Ok(2)); + assert_eq!(fast.try_recv().as_deref(), Ok(&1)); + assert_eq!(fast.try_recv().as_deref(), Ok(&2)); assert_eq!(tx.try_send(3), Err(TrySendError::Full(3))); // One read by the slow subscription frees exactly one slot. - assert_eq!(slow.try_recv(), Ok(1)); + assert_eq!(slow.try_recv().as_deref(), Ok(&1)); tx.try_send(3).unwrap(); // Every value accepted while both were active reaches both, in order. - assert_eq!(slow.try_recv(), Ok(2)); - assert_eq!(slow.try_recv(), Ok(3)); - assert_eq!(fast.try_recv(), Ok(3)); + assert_eq!(slow.try_recv().as_deref(), Ok(&2)); + assert_eq!(slow.try_recv().as_deref(), Ok(&3)); + assert_eq!(fast.try_recv().as_deref(), Ok(&3)); assert_eq!(tx.retained_message_count(), 0); } @@ -145,7 +149,7 @@ fn bounded_subscribe_starts_at_the_committed_tail() { assert_eq!(late.try_recv(), Err(TryRecvError::Empty)); tx.try_send(2).unwrap(); - assert_eq!(late.try_recv(), Ok(2)); + assert_eq!(late.try_recv().as_deref(), Ok(&2)); } #[test] @@ -157,10 +161,10 @@ fn bounded_resubscribe_keeps_the_original_receivers_backlog() { let mut rx2 = rx.resubscribe(); tx.try_send(3).unwrap(); - assert_eq!(rx2.try_recv(), Ok(3)); - assert_eq!(rx.try_recv(), Ok(1)); - assert_eq!(rx.try_recv(), Ok(2)); - assert_eq!(rx.try_recv(), Ok(3)); + assert_eq!(rx2.try_recv().as_deref(), Ok(&3)); + assert_eq!(rx.try_recv().as_deref(), Ok(&1)); + assert_eq!(rx.try_recv().as_deref(), Ok(&2)); + assert_eq!(rx.try_recv().as_deref(), Ok(&3)); } #[test] @@ -175,11 +179,81 @@ fn bounded_unread_message_count_tracks_each_receiver() { assert_eq!(rx1.unread_message_count(), 2); assert_eq!(rx2.unread_message_count(), 2); - assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.try_recv().as_deref(), Ok(&1)); assert_eq!(rx1.unread_message_count(), 1); assert_eq!(rx2.unread_message_count(), 2); } +#[tokio::test] +async fn bounded_shares_one_arc_across_receivers() { + // Not `Clone`: fanning out never needs to copy the payload. + struct Payload(u64); + + let (tx, mut rx1) = bounded(4); + let mut rx2 = tx.subscribe(); + let mut rx3 = tx.subscribe(); + tx.send(Payload(7)).await; + + let first = rx1.recv().await.unwrap(); + let second = rx2.try_recv().unwrap(); + let third = rx3.recv().await.unwrap(); + assert!(Arc::ptr_eq(&first, &second)); + assert!(Arc::ptr_eq(&first, &third)); + assert_eq!(first.0, 7); + assert_eq!(tx.retained_message_count(), 0); +} + +#[tokio::test] +async fn bounded_shared_and_cloned_receives_advance_the_same_cursor() { + let (tx, mut rx) = bounded(4); + for value in 1..=4 { + tx.send(value).await; + } + + assert_eq!(rx.recv().await.as_deref(), Ok(&1)); + assert_eq!(rx.recv_cloned().await, Ok(2)); + assert_eq!(rx.try_recv().as_deref(), Ok(&3)); + assert_eq!(rx.try_recv_cloned(), Ok(4)); + assert_eq!(rx.unread_message_count(), 0); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); +} + +#[test] +fn bounded_sole_receiver_takes_messages_without_cloning() { + static CLONES: AtomicUsize = AtomicUsize::new(0); + + struct CountClone(u32); + + impl Clone for CountClone { + fn clone(&self) -> Self { + CLONES.fetch_add(1, Ordering::Relaxed); + Self(self.0) + } + } + + let (tx, mut rx) = bounded(4); + for i in 0..8 { + tx.try_send(CountClone(i)).unwrap(); + assert_eq!(FutureExt::block_on(rx.recv_cloned()).unwrap().0, i); + } + assert_eq!(CLONES.load(Ordering::Relaxed), 0); + + // A second receiver means the payload is shared, so one of the two reads has to clone it. + let mut second = tx.subscribe(); + tx.try_send(CountClone(8)).unwrap(); + assert_eq!(FutureExt::block_on(rx.recv_cloned()).unwrap().0, 8); + assert_eq!(FutureExt::block_on(second.recv_cloned()).unwrap().0, 8); + assert_eq!(CLONES.load(Ordering::Relaxed), 1); + + // An `Arc` that outlives the last read keeps the payload shared, so even the read that + // reclaims the slot has to clone. + tx.try_send(CountClone(9)).unwrap(); + let shared = second.try_recv().unwrap(); + assert_eq!(rx.try_recv_cloned().unwrap().0, 9); + assert_eq!(CLONES.load(Ordering::Relaxed), 2); + assert_eq!(shared.0, 9); +} + // --------------------------------------------------------------------------------------------- // Strict capacity // --------------------------------------------------------------------------------------------- @@ -194,8 +268,8 @@ fn try_send_rejects_at_capacity_and_returns_the_value() { // The rejected value is handed back untouched, and nothing was published. assert_eq!(tx.retained_message_count(), 2); - assert_eq!(rx.try_recv(), Ok(1)); - assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv().as_deref(), Ok(&1)); + assert_eq!(rx.try_recv().as_deref(), Ok(&2)); assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } @@ -222,13 +296,13 @@ fn retained_message_count_tracks_the_slowest_receiver() { assert_eq!(tx.retained_message_count(), 2); // Draining one receiver does not release what the other has not read. - assert_eq!(rx1.try_recv(), Ok(1)); - assert_eq!(rx1.try_recv(), Ok(2)); + assert_eq!(rx1.try_recv().as_deref(), Ok(&1)); + assert_eq!(rx1.try_recv().as_deref(), Ok(&2)); assert_eq!(tx.retained_message_count(), 2); - assert_eq!(rx2.try_recv(), Ok(1)); + assert_eq!(rx2.try_recv().as_deref(), Ok(&1)); assert_eq!(tx.retained_message_count(), 1); - assert_eq!(rx2.try_recv(), Ok(2)); + assert_eq!(rx2.try_recv().as_deref(), Ok(&2)); assert_eq!(tx.retained_message_count(), 0); } @@ -246,12 +320,12 @@ fn send_waits_while_the_slowest_subscription_holds_capacity() { assert!(poll_once(send.as_mut()).is_pending()); // The fast receiver draining is not enough while the slow one still retains the message. - assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.try_recv().as_deref(), Ok(&1)); assert!(poll_once(send.as_mut()).is_pending()); - assert_eq!(rx2.try_recv(), Ok(1)); + assert_eq!(rx2.try_recv().as_deref(), Ok(&1)); assert!(poll_once(send.as_mut()).is_ready()); - assert_eq!(rx1.try_recv(), Ok(2)); + assert_eq!(rx1.try_recv().as_deref(), Ok(&2)); } #[test] @@ -269,7 +343,7 @@ fn receive_that_vacates_the_head_wakes_a_blocked_sender() { ); assert_eq!(tracker.count(), 0); - assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(rx.try_recv().as_deref(), Ok(&0)); assert_eq!(tracker.count(), 1); assert!(poll_once(send.as_mut()).is_ready()); } @@ -291,7 +365,7 @@ fn parked_recv_that_reclaims_wakes_a_blocked_sender() { // Reclaim through the `recv` future rather than `try_recv`: it is a separate call site, and a // release wired into only one of them would strand this producer. let mut recv = Box::pin(rx.recv()); - assert_eq!(poll_once(recv.as_mut()), std::task::Poll::Ready(Ok(0))); + assert_eq!(expect_ready(poll_once(recv.as_mut())).as_deref(), Ok(&0)); drop(recv); assert_eq!(tracker.count(), 1); @@ -309,13 +383,13 @@ fn wakes_blocked_senders_as_capacity_frees() { assert!(poll_once(second.as_mut()).is_pending()); // One freed slot admits exactly one producer. - assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(rx.try_recv().as_deref(), Ok(&0)); assert!(poll_once(first.as_mut()).is_ready()); assert!(poll_once(second.as_mut()).is_pending()); - assert_eq!(rx.try_recv(), Ok(1)); + assert_eq!(rx.try_recv().as_deref(), Ok(&1)); assert!(poll_once(second.as_mut()).is_ready()); - assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv().as_deref(), Ok(&2)); } #[test] @@ -325,7 +399,7 @@ fn dropping_a_lagging_receiver_wakes_blocked_senders() { tx.try_send(0).unwrap(); let mut send = Box::pin(tx.send(1)); - assert_eq!(rx1.try_recv(), Ok(0)); + assert_eq!(rx1.try_recv().as_deref(), Ok(&0)); assert!(poll_once(send.as_mut()).is_pending()); // `rx2` is the one holding the backlog; dropping it releases the slot. @@ -402,6 +476,22 @@ fn subscribing_while_producers_are_blocked_does_not_release_capacity() { assert!(poll_once(send.as_mut()).is_pending()); } +#[test] +fn retained_arc_holds_no_capacity() { + let (tx, mut rx) = bounded(1); + tx.try_send(1).unwrap(); + + // Reading releases the slot even though the `Arc` keeps the payload alive. + let retained = rx.try_recv().unwrap(); + assert_eq!(tx.retained_message_count(), 0); + assert_eq!(Arc::strong_count(&retained), 1); + + let mut send = Box::pin(tx.send(2)); + assert!(poll_once(send.as_mut()).is_ready()); + assert_eq!(*retained, 1); + assert_eq!(rx.try_recv().as_deref(), Ok(&2)); +} + // --------------------------------------------------------------------------------------------- // Cancellation // --------------------------------------------------------------------------------------------- @@ -417,9 +507,9 @@ fn cancelled_send_publishes_nothing() { // The cancelled value never entered the committed order. The receiver sees the previously // published value followed by a newly sent value. - assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(rx.try_recv().as_deref(), Ok(&0)); tx.try_send(2).unwrap(); - assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv().as_deref(), Ok(&2)); assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } @@ -433,11 +523,11 @@ fn cancelled_notified_sender_passes_capacity_to_the_next_sender() { assert!(poll_once(first.as_mut()).is_pending()); assert!(poll_once(second.as_mut()).is_pending()); - assert_eq!(rx.try_recv(), Ok(0)); + assert_eq!(rx.try_recv().as_deref(), Ok(&0)); drop(first); assert!(poll_once(second.as_mut()).is_ready()); - assert_eq!(rx.try_recv(), Ok(2)); + assert_eq!(rx.try_recv().as_deref(), Ok(&2)); } #[test] @@ -463,6 +553,22 @@ fn cancelled_recv_releases_its_waker() { assert_eq!(tracker.count(), 0); } +#[test] +fn bounded_cancelled_recv_cloned_leaves_the_cursor_unchanged() { + let (tx, mut rx) = bounded(4); + let (waker, counter) = WakeCounter::new(); + + let mut recv = Box::pin(rx.recv_cloned()); + assert!(poll_with(recv.as_mut(), &waker).is_pending()); + tx.try_send(1).unwrap(); + assert_eq!(counter.count(), 1); + + // The woken future is dropped before it is polled again, so it never advances the cursor. + drop(recv); + assert_eq!(rx.unread_message_count(), 1); + assert_eq!(rx.try_recv_cloned(), Ok(1)); +} + // --------------------------------------------------------------------------------------------- // Disconnection // --------------------------------------------------------------------------------------------- @@ -480,7 +586,7 @@ fn dropping_a_woken_recv_keeps_another_receivers_waiter() { tx.try_send(1).unwrap(); assert_eq!(first.count(), 1); - assert_eq!(rx2.try_recv(), Ok(1)); + assert_eq!(rx2.try_recv().as_deref(), Ok(&1)); let second = Arc::new(WakeCounter::default()); let waker = Waker::from(second.clone()); @@ -502,8 +608,8 @@ async fn bounded_recv_drains_buffered_messages_before_reporting_disconnection() tx.try_send(2).unwrap(); drop(tx); - assert_eq!(rx.recv().await, Ok(1)); - assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await.as_deref(), Ok(&1)); + assert_eq!(rx.recv().await.as_deref(), Ok(&2)); assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); } @@ -577,7 +683,7 @@ fn large_reclaim_notifies_every_blocked_sender() { assert!(poll_once(send.as_mut()).is_ready()); } for value in 40..80 { - assert_eq!(fast.try_recv(), Ok(value)); + assert_eq!(fast.try_recv().as_deref(), Ok(&value)); } } @@ -601,14 +707,14 @@ fn bounded_panicking_clone_leaves_the_channel_consistent() { }) .unwrap(); - // Two receivers share the payload, so `try_recv` has to clone it. + // Two receivers share the payload, so `try_recv_cloned` has to clone it. let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - rx1.try_recv().map(|msg| msg.value) + rx1.try_recv_cloned().map(|msg| msg.value) })); assert!(result.is_err()); - // The panicking `try_recv` call still consumed the message for `rx1`, and left the channel - // usable for both receivers. + // The panicking `try_recv_cloned` call still consumed the message for `rx1`, and left the + // channel usable for both receivers. assert_eq!(rx1.try_recv().unwrap().value, 2); assert_eq!(rx2.try_recv().unwrap().value, 1); assert_eq!(rx2.try_recv().unwrap().value, 2); @@ -616,6 +722,39 @@ fn bounded_panicking_clone_leaves_the_channel_consistent() { assert_eq!(rx1.try_recv().unwrap_err(), TryRecvError::Empty); } +#[test] +fn panicking_clone_in_recv_cloned_does_not_strand_a_blocked_sender() { + let (tx, mut rx1) = bounded(1); + let mut rx2 = tx.subscribe(); + tx.try_send(PanicOnClone { + value: 1, + panic: true, + }) + .unwrap(); + + let (waker, counter) = WakeCounter::new(); + let mut send = Box::pin(tx.send(PanicOnClone { + value: 2, + panic: false, + })); + assert!(poll_with(send.as_mut(), &waker).is_pending()); + + // `rx2` keeps its `Arc`, so the read that reclaims the slot cannot move the payload + // out and has to clone it. + let shared = rx2.try_recv().unwrap(); + let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { + FutureExt::block_on(rx1.recv_cloned()).map(|msg| msg.value) + })); + assert!(result.is_err()); + + // The capacity was released, and the producer woken, before the clone ran. + assert_eq!(counter.count(), 1); + assert!(poll_once(send.as_mut()).is_ready()); + assert_eq!(shared.value, 1); + assert_eq!(rx1.try_recv().unwrap().value, 2); + assert_eq!(rx2.try_recv().unwrap().value, 2); +} + #[test] fn panicking_payload_destructor_still_releases_capacity() { let (tx, mut rx1) = bounded(3); @@ -628,7 +767,7 @@ fn panicking_payload_destructor_still_releases_capacity() { } // `rx1` reads clones, which are disarmed; the armed originals stay retained for `rx2`. for _ in 0..3 { - rx1.try_recv().unwrap(); + rx1.try_recv_cloned().unwrap(); } let mut send = Box::pin(tx.send(PanicOnDrop { armed: false })); @@ -765,7 +904,7 @@ fn bounded_concurrent_producers_commit_one_order_seen_by_every_receiver() { let mut seen = Vec::with_capacity(TOTAL as usize); for _ in 0..TOTAL { seen.push( - FutureExt::block_on(receiver.recv()).expect("sender dropped early"), + *FutureExt::block_on(receiver.recv()).expect("sender dropped early"), ); } seen diff --git a/tests-integration/tests/broadcast_test/unbounded.rs b/tests-integration/tests/broadcast_test/unbounded.rs index 3324e4d9..63bff5f1 100644 --- a/tests-integration/tests/broadcast_test/unbounded.rs +++ b/tests-integration/tests/broadcast_test/unbounded.rs @@ -19,7 +19,6 @@ use std::sync::Arc; use std::sync::atomic::AtomicUsize; use std::sync::atomic::Ordering; use std::task::Context; -use std::task::Poll; use std::task::Waker; use std::thread; @@ -27,6 +26,8 @@ use asyncband::blocking::FutureExt; use asyncband::broadcast::mpmc::*; use tests_integration::WakeCounter; use tests_integration::assert_completes_without_deadlock; +use tests_integration::expect_ready; +use tests_integration::poll_with; /// A payload whose destructor re-enters the channel it was sent through. struct Reentrant { @@ -93,7 +94,7 @@ fn test_try_recv() { // Success tx.send(10); - assert_eq!(rx.try_recv(), Ok(10)); + assert_eq!(rx.try_recv().as_deref(), Ok(&10)); assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); // Disconnected @@ -112,12 +113,12 @@ async fn test_slow_receiver_keeps_every_message() { // The fast receiver draining fully must not reclaim anything the slow one still needs. for i in 0..1024 { - assert_eq!(rx1.recv().await, Ok(i)); + assert_eq!(rx1.recv().await.as_deref(), Ok(&i)); } assert_eq!(tx.retained_message_count(), 1024); for i in 0..1024 { - assert_eq!(rx2.recv().await, Ok(i)); + assert_eq!(rx2.recv().await.as_deref(), Ok(&i)); } assert_eq!(tx.retained_message_count(), 0); } @@ -132,14 +133,14 @@ async fn retained_message_count_tracks_the_slowest_receiver() { assert_eq!(tx.retained_message_count(), 2); // Reclaiming waits for the slowest receiver, message by message. - assert_eq!(rx1.recv().await, Ok(1)); + assert_eq!(rx1.recv().await.as_deref(), Ok(&1)); assert_eq!(tx.retained_message_count(), 2); - assert_eq!(rx2.recv().await, Ok(1)); + assert_eq!(rx2.recv().await.as_deref(), Ok(&1)); assert_eq!(tx.retained_message_count(), 1); - assert_eq!(rx1.recv().await, Ok(2)); + assert_eq!(rx1.recv().await.as_deref(), Ok(&2)); assert_eq!(tx.retained_message_count(), 1); - assert_eq!(rx2.recv().await, Ok(2)); + assert_eq!(rx2.recv().await.as_deref(), Ok(&2)); assert_eq!(tx.retained_message_count(), 0); } @@ -152,7 +153,7 @@ async fn test_dropping_a_lagging_receiver_releases_its_backlog() { tx.send(i); } for i in 0..128 { - assert_eq!(rx1.recv().await, Ok(i)); + assert_eq!(rx1.recv().await.as_deref(), Ok(&i)); } assert_eq!(tx.retained_message_count(), 128); @@ -172,12 +173,12 @@ async fn resubscribe_keeps_the_original_receivers_backlog() { tx.send(3); - assert_eq!(rx2.recv().await, Ok(3)); + assert_eq!(rx2.recv().await.as_deref(), Ok(&3)); assert_eq!(tx.retained_message_count(), 3); - assert_eq!(rx.recv().await, Ok(1)); - assert_eq!(rx.recv().await, Ok(2)); - assert_eq!(rx.recv().await, Ok(3)); + assert_eq!(rx.recv().await.as_deref(), Ok(&1)); + assert_eq!(rx.recv().await.as_deref(), Ok(&2)); + assert_eq!(rx.recv().await.as_deref(), Ok(&3)); assert_eq!(tx.retained_message_count(), 0); } @@ -194,7 +195,7 @@ async fn send_without_receivers_does_not_buffer() { assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); tx.send(3); - assert_eq!(rx.recv().await, Ok(3)); + assert_eq!(rx.recv().await.as_deref(), Ok(&3)); } #[test] @@ -213,11 +214,11 @@ fn unread_message_count_tracks_each_receiver() { assert_eq!(rx1.unread_message_count(), 3); assert_eq!(rx2.unread_message_count(), 1); - assert_eq!(rx2.try_recv(), Ok(3)); + assert_eq!(rx2.try_recv().as_deref(), Ok(&3)); assert_eq!(rx2.unread_message_count(), 0); drop(rx2); - assert_eq!(rx1.try_recv(), Ok(1)); + assert_eq!(rx1.try_recv().as_deref(), Ok(&1)); assert_eq!(rx1.unread_message_count(), 2); } @@ -237,16 +238,58 @@ fn sole_receiver_takes_messages_without_cloning() { let (tx, mut rx) = unbounded(); for i in 0..8 { tx.send(CountClone(i)); - assert_eq!(rx.try_recv().unwrap().0, i); + assert_eq!(rx.try_recv_cloned().unwrap().0, i); } assert_eq!(CLONES.load(Ordering::Relaxed), 0); // A second receiver means the payload is shared, so it has to be cloned again. let mut second = tx.subscribe(); tx.send(CountClone(8)); - assert_eq!(rx.try_recv().unwrap().0, 8); - assert_eq!(second.try_recv().unwrap().0, 8); + assert_eq!(rx.try_recv_cloned().unwrap().0, 8); + assert_eq!(second.try_recv_cloned().unwrap().0, 8); assert_eq!(CLONES.load(Ordering::Relaxed), 1); + + // An `Arc` that outlives the last read keeps the payload shared, so even the read that + // reclaims the slot has to clone. + tx.send(CountClone(9)); + let shared = second.try_recv().unwrap(); + assert_eq!(rx.try_recv_cloned().unwrap().0, 9); + assert_eq!(CLONES.load(Ordering::Relaxed), 2); + assert_eq!(shared.0, 9); +} + +#[tokio::test] +async fn shares_one_arc_across_receivers() { + // Not `Clone`: fanning out never needs to copy the payload. + struct Payload(u64); + + let (tx, mut rx1) = unbounded(); + let mut rx2 = tx.subscribe(); + let mut rx3 = tx.subscribe(); + tx.send(Payload(7)); + + let first = rx1.recv().await.unwrap(); + let second = rx2.try_recv().unwrap(); + let third = rx3.recv().await.unwrap(); + assert!(Arc::ptr_eq(&first, &second)); + assert!(Arc::ptr_eq(&first, &third)); + assert_eq!(first.0, 7); + assert_eq!(tx.retained_message_count(), 0); +} + +#[tokio::test] +async fn shared_and_cloned_receives_advance_the_same_cursor() { + let (tx, mut rx) = unbounded(); + for value in 1..=4 { + tx.send(value); + } + + assert_eq!(rx.recv().await.as_deref(), Ok(&1)); + assert_eq!(rx.recv_cloned().await, Ok(2)); + assert_eq!(rx.try_recv().as_deref(), Ok(&3)); + assert_eq!(rx.try_recv_cloned(), Ok(4)); + assert_eq!(rx.unread_message_count(), 0); + assert_eq!(rx.try_recv(), Err(TryRecvError::Empty)); } #[test] @@ -263,14 +306,14 @@ fn panicking_clone_leaves_the_channel_consistent() { panic: false, }); - // Two receivers share the payload, so `try_recv` has to clone it. + // Two receivers share the payload, so `try_recv_cloned` has to clone it. let result = std::panic::catch_unwind(std::panic::AssertUnwindSafe(|| { - rx1.try_recv().map(|msg| msg.value) + rx1.try_recv_cloned().map(|msg| msg.value) })); assert!(result.is_err()); - // The panicking `try_recv` call still consumed the message for `rx1`, and left the channel - // usable for both receivers. + // The panicking `try_recv_cloned` call still consumed the message for `rx1`, and left the + // channel usable for both receivers. assert_eq!(rx1.try_recv().unwrap().value, 2); assert_eq!(rx2.try_recv().unwrap().value, 1); assert_eq!(rx2.try_recv().unwrap().value, 2); @@ -319,7 +362,10 @@ fn send_wakes_a_parked_receiver_exactly_once() { tx.send(42); assert_eq!(tracker.count(), 1); - assert_eq!(recv.as_mut().poll(&mut context), Poll::Ready(Ok(42))); + assert_eq!( + expect_ready(recv.as_mut().poll(&mut context)).as_deref(), + Ok(&42) + ); } #[test] @@ -339,7 +385,23 @@ fn cancelled_recv_releases_its_waker() { tx.send(()); assert_eq!(tracker.count(), 0); - assert_eq!(rx.try_recv(), Ok(())); + assert_eq!(rx.try_recv().as_deref(), Ok(&())); +} + +#[test] +fn cancelled_recv_cloned_leaves_the_cursor_unchanged() { + let (tx, mut rx) = unbounded(); + let (waker, counter) = WakeCounter::new(); + + let mut recv = Box::pin(rx.recv_cloned()); + assert!(poll_with(recv.as_mut(), &waker).is_pending()); + tx.send(1); + assert_eq!(counter.count(), 1); + + // The woken future is dropped before it is polled again, so it never advances the cursor. + drop(recv); + assert_eq!(rx.unread_message_count(), 1); + assert_eq!(rx.try_recv_cloned(), Ok(1)); } #[test] @@ -355,7 +417,7 @@ fn dropping_a_woken_recv_keeps_another_receivers_waiter() { tx.send(1); assert_eq!(first.count(), 1); - assert_eq!(rx2.try_recv(), Ok(1)); + assert_eq!(rx2.try_recv().as_deref(), Ok(&1)); let second = Arc::new(WakeCounter::default()); let waker = Waker::from(second.clone()); @@ -404,7 +466,10 @@ fn parked_recv_prefers_buffered_messages_over_disconnection() { tx.send(7); drop(tx); - assert_eq!(recv.as_mut().poll(&mut context), Poll::Ready(Ok(7))); + assert_eq!( + expect_ready(recv.as_mut().poll(&mut context)).as_deref(), + Ok(&7) + ); drop(recv); assert_eq!(rx.try_recv(), Err(TryRecvError::Disconnected)); } @@ -417,8 +482,8 @@ async fn recv_drains_buffered_messages_before_reporting_disconnection() { tx.send(2); drop(tx); - assert_eq!(rx.recv().await, Ok(1)); - assert_eq!(rx.recv().await, Ok(2)); + assert_eq!(rx.recv().await.as_deref(), Ok(&1)); + assert_eq!(rx.recv().await.as_deref(), Ok(&2)); assert_eq!(rx.recv().await, Err(RecvError::Disconnected)); } @@ -456,7 +521,7 @@ fn concurrent_senders_deliver_every_message_to_every_receiver() { thread::spawn(move || { let mut seen = vec![]; while let Ok(value) = FutureExt::block_on(receiver.recv()) { - seen.push(value); + seen.push(*value); } seen }) @@ -494,7 +559,7 @@ fn randomized_operations_track_the_reference_model() { let index = rng.below(model.len() as u64) as usize; let (receiver, cursor) = &mut model[index]; if *cursor < tail { - assert_eq!(receiver.try_recv(), Ok(*cursor), "seed {seed}"); + assert_eq!(receiver.try_recv().as_deref(), Ok(&*cursor), "seed {seed}"); *cursor += 1; } else { assert_eq!(receiver.try_recv(), Err(TryRecvError::Empty), "seed {seed}");