diff --git a/tokio/src/sync/mpsc/bounded.rs b/tokio/src/sync/mpsc/bounded.rs index c2a2f061872..47d7938158a 100644 --- a/tokio/src/sync/mpsc/bounded.rs +++ b/tokio/src/sync/mpsc/bounded.rs @@ -1,3 +1,4 @@ +use crate::loom::sync::Arc; use crate::sync::batch_semaphore::{self as semaphore, TryAcquireError}; use crate::sync::mpsc::chan; use crate::sync::mpsc::error::{SendError, TryRecvError, TrySendError}; @@ -22,6 +23,40 @@ pub struct Sender { chan: chan::Tx, } +/// A sender that does not prevent the channel from being closed. +/// +/// If all [`Sender`] instances of a channel were dropped and only `WeakSender` +/// instances remain, the channel is closed. +/// +/// In order to send messages, the `WeakSender` needs to be upgraded using +/// [`WeakSender::upgrade`], which returns `Option`. It returns `None` +/// if all `Sender`s have been dropped, and otherwise it returns a `Sender`. +/// +/// [`Sender`]: Sender +/// [`WeakSender::upgrade`]: WeakSender::upgrade +/// +/// #Examples +/// +/// ``` +/// use tokio::sync::mpsc::channel; +/// +/// #[tokio::main] +/// async fn main() { +/// let (tx, _rx) = channel::(15); +/// let tx_weak = tx.downgrade(); +/// +/// // Upgrading will succeed because `tx` still exists. +/// assert!(tx_weak.upgrade().is_some()); +/// +/// // If we drop `tx`, then it will fail. +/// drop(tx); +/// assert!(tx_weak.clone().upgrade().is_none()); +/// } +/// ``` +pub struct WeakSender { + chan: Arc>, +} + /// Permits to send one value into the channel. /// /// `Permit` values are returned by [`Sender::reserve()`] and [`Sender::try_reserve()`] @@ -991,6 +1026,16 @@ impl Sender { pub fn capacity(&self) -> usize { self.chan.semaphore().0.available_permits() } + + /// Converts the `Sender` to a [`WeakSender`] that does not count + /// towards RAII semantics, i.e. if all `Sender` instances of the + /// channel were dropped and only `WeakSender` instances remain, + /// the channel is closed. + pub fn downgrade(&self) -> WeakSender { + WeakSender { + chan: self.chan.downgrade(), + } + } } impl Clone for Sender { @@ -1009,6 +1054,29 @@ impl fmt::Debug for Sender { } } +impl Clone for WeakSender { + fn clone(&self) -> Self { + WeakSender { + chan: self.chan.clone(), + } + } +} + +impl WeakSender { + /// Tries to convert a WeakSender into a [`Sender`]. This will return `Some` + /// if there are other `Sender` instances alive and the channel wasn't + /// previously dropped, otherwise `None` is returned. + pub fn upgrade(&self) -> Option> { + chan::Tx::upgrade(self.chan.clone()).map(Sender::new) + } +} + +impl fmt::Debug for WeakSender { + fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result { + fmt.debug_struct("WeakSender").finish() + } +} + // ===== impl Permit ===== impl Permit<'_, T> { diff --git a/tokio/src/sync/mpsc/chan.rs b/tokio/src/sync/mpsc/chan.rs index 1228cfb6c6f..a10ffb7d797 100644 --- a/tokio/src/sync/mpsc/chan.rs +++ b/tokio/src/sync/mpsc/chan.rs @@ -10,9 +10,10 @@ use crate::sync::notify::Notify; use std::fmt; use std::process; -use std::sync::atomic::Ordering::{AcqRel, Relaxed}; +use std::sync::atomic::Ordering::{AcqRel, Acquire, Relaxed, Release}; use std::task::Poll::{Pending, Ready}; use std::task::{Context, Poll}; +use std::usize; /// Channel sender. pub(crate) struct Tx { @@ -46,7 +47,7 @@ pub(crate) trait Semaphore { fn is_closed(&self) -> bool; } -struct Chan { +pub(super) struct Chan { /// Notifies all tasks listening for the receiver being dropped. notify_rx_closed: Notify, @@ -129,6 +130,30 @@ impl Tx { Tx { inner: chan } } + pub(super) fn downgrade(&self) -> Arc> { + self.inner.clone() + } + + // Returns the upgraded channel or None if the upgrade failed. + pub(super) fn upgrade(chan: Arc>) -> Option { + let mut tx_count = chan.tx_count.load(Acquire); + + loop { + if tx_count == 0 { + // channel is closed + return None; + } + + match chan + .tx_count + .compare_exchange_weak(tx_count, tx_count + 1, AcqRel, Acquire) + { + Ok(_) => return Some(Tx { inner: chan }), + Err(prev_count) => tx_count = prev_count, + } + } + } + pub(super) fn semaphore(&self) -> &S { &self.inner.semaphore } @@ -378,9 +403,6 @@ impl Semaphore for (crate::sync::batch_semaphore::Semaphore, usize) { // ===== impl Semaphore for AtomicUsize ===== -use std::sync::atomic::Ordering::{Acquire, Release}; -use std::usize; - impl Semaphore for AtomicUsize { fn add_permit(&self) { let prev = self.fetch_sub(2, Release); diff --git a/tokio/src/sync/mpsc/mod.rs b/tokio/src/sync/mpsc/mod.rs index b1513a9da51..d37779ca761 100644 --- a/tokio/src/sync/mpsc/mod.rs +++ b/tokio/src/sync/mpsc/mod.rs @@ -90,7 +90,7 @@ pub(super) mod block; mod bounded; -pub use self::bounded::{channel, OwnedPermit, Permit, Receiver, Sender}; +pub use self::bounded::{channel, OwnedPermit, Permit, Receiver, Sender, WeakSender}; mod chan; diff --git a/tokio/tests/sync_mpsc.rs b/tokio/tests/sync_mpsc.rs index 99c22931e77..aebdf8548d8 100644 --- a/tokio/tests/sync_mpsc.rs +++ b/tokio/tests/sync_mpsc.rs @@ -10,10 +10,13 @@ use wasm_bindgen_test::wasm_bindgen_test as maybe_tokio_test; #[cfg(not(tokio_wasm_not_wasi))] use tokio::test as maybe_tokio_test; -use tokio::sync::mpsc; use tokio::sync::mpsc::error::{TryRecvError, TrySendError}; +use tokio::sync::mpsc::{self, channel}; +use tokio::sync::oneshot; use tokio_test::*; +use std::sync::atomic::AtomicUsize; +use std::sync::atomic::Ordering::{Acquire, Release}; use std::sync::Arc; #[cfg(not(tokio_wasm))] @@ -657,3 +660,264 @@ fn recv_timeout_panic() { let (tx, _rx) = mpsc::channel(5); tx.send_timeout(10, Duration::from_secs(1)).now_or_never(); } + +#[tokio::test] +async fn weak_sender() { + let (tx, mut rx) = channel(11); + + let tx_weak = tokio::spawn(async move { + let tx_weak = tx.clone().downgrade(); + + for i in 0..10 { + if tx.send(i).await.is_err() { + return None; + } + } + + let tx2 = tx_weak + .upgrade() + .expect("expected to be able to upgrade tx_weak"); + let _ = tx2.send(20).await; + let tx_weak = tx2.downgrade(); + + Some(tx_weak) + }) + .await + .unwrap(); + + for i in 0..12 { + let recvd = rx.recv().await; + + match recvd { + Some(msg) => { + if i == 10 { + assert_eq!(msg, 20); + } + } + None => { + assert_eq!(i, 11); + break; + } + } + } + + let tx_weak = tx_weak.unwrap(); + let upgraded = tx_weak.upgrade(); + assert!(upgraded.is_none()); +} + +#[tokio::test] +async fn actor_weak_sender() { + pub struct MyActor { + receiver: mpsc::Receiver, + sender: mpsc::WeakSender, + next_id: u32, + pub received_self_msg: bool, + } + + enum ActorMessage { + GetUniqueId { respond_to: oneshot::Sender }, + SelfMessage {}, + } + + impl MyActor { + fn new( + receiver: mpsc::Receiver, + sender: mpsc::WeakSender, + ) -> Self { + MyActor { + receiver, + sender, + next_id: 0, + received_self_msg: false, + } + } + + fn handle_message(&mut self, msg: ActorMessage) { + match msg { + ActorMessage::GetUniqueId { respond_to } => { + self.next_id += 1; + + // The `let _ =` ignores any errors when sending. + // + // This can happen if the `select!` macro is used + // to cancel waiting for the response. + let _ = respond_to.send(self.next_id); + } + ActorMessage::SelfMessage { .. } => { + self.received_self_msg = true; + } + } + } + + async fn send_message_to_self(&mut self) { + let msg = ActorMessage::SelfMessage {}; + + let sender = self.sender.clone(); + + // cannot move self.sender here + if let Some(sender) = sender.upgrade() { + let _ = sender.send(msg).await; + self.sender = sender.downgrade(); + } + } + + async fn run(&mut self) { + let mut i = 0; + while let Some(msg) = self.receiver.recv().await { + self.handle_message(msg); + + if i == 0 { + self.send_message_to_self().await; + } + + i += 1 + } + + assert!(self.received_self_msg); + } + } + + #[derive(Clone)] + pub struct MyActorHandle { + sender: mpsc::Sender, + } + + impl MyActorHandle { + pub fn new() -> (Self, MyActor) { + let (sender, receiver) = mpsc::channel(8); + let actor = MyActor::new(receiver, sender.clone().downgrade()); + + (Self { sender }, actor) + } + + pub async fn get_unique_id(&self) -> u32 { + let (send, recv) = oneshot::channel(); + let msg = ActorMessage::GetUniqueId { respond_to: send }; + + // Ignore send errors. If this send fails, so does the + // recv.await below. There's no reason to check the + // failure twice. + let _ = self.sender.send(msg).await; + recv.await.expect("Actor task has been killed") + } + } + + let (handle, mut actor) = MyActorHandle::new(); + + let actor_handle = tokio::spawn(async move { actor.run().await }); + + let _ = tokio::spawn(async move { + let _ = handle.get_unique_id().await; + drop(handle); + }) + .await; + + let _ = actor_handle.await; +} + +static NUM_DROPPED: AtomicUsize = AtomicUsize::new(0); + +#[derive(Debug)] +struct Msg; + +impl Drop for Msg { + fn drop(&mut self) { + NUM_DROPPED.fetch_add(1, Release); + } +} + +// Tests that no pending messages are put onto the channel after `Rx` was +// dropped. +// +// Note: After the introduction of `WeakSender`, which internally +// used `Arc` and doesn't call a drop of the channel after the last strong +// `Sender` was dropped while more than one `WeakSender` remains, we want to +// ensure that no messages are kept in the channel, which were sent after +// the receiver was dropped. +#[tokio::test] +async fn test_msgs_dropped_on_rx_drop() { + let (tx, mut rx) = mpsc::channel(3); + + let _ = tx.send(Msg {}).await.unwrap(); + let _ = tx.send(Msg {}).await.unwrap(); + + // This msg will be pending and should be dropped when `rx` is dropped + let sent_fut = tx.send(Msg {}); + + let _ = rx.recv().await.unwrap(); + let _ = rx.recv().await.unwrap(); + + let _ = sent_fut.await.unwrap(); + + drop(rx); + + assert_eq!(NUM_DROPPED.load(Acquire), 3); + + // This msg will not be put onto `Tx` list anymore, since `Rx` is closed. + assert!(tx.send(Msg {}).await.is_err()); + + assert_eq!(NUM_DROPPED.load(Acquire), 4); +} + +// Tests that a `WeakSender` is upgradeable when other `Sender`s exist. +#[tokio::test] +async fn downgrade_upgrade_sender_success() { + let (tx, _rx) = mpsc::channel::(1); + let weak_tx = tx.downgrade(); + assert!(weak_tx.upgrade().is_some()); +} + +// Tests that a `WeakSender` fails to upgrade when no other `Sender` exists. +#[tokio::test] +async fn downgrade_upgrade_sender_failure() { + let (tx, _rx) = mpsc::channel::(1); + let weak_tx = tx.downgrade(); + drop(tx); + assert!(weak_tx.upgrade().is_none()); +} + +// Tests that a `WeakSender` cannot be upgraded after a `Sender` was dropped, +// which existed at the time of the `downgrade` call. +#[tokio::test] +async fn downgrade_drop_upgrade() { + let (tx, _rx) = mpsc::channel::(1); + + // the cloned `Tx` is dropped right away + let weak_tx = tx.clone().downgrade(); + drop(tx); + assert!(weak_tx.upgrade().is_none()); +} + +// Tests that we can upgrade a weak sender with an outstanding permit +// but no other strong senders. +#[tokio::test] +async fn downgrade_get_permit_upgrade_no_senders() { + let (tx, _rx) = mpsc::channel::(1); + let weak_tx = tx.downgrade(); + let _permit = tx.reserve_owned().await.unwrap(); + assert!(weak_tx.upgrade().is_some()); +} + +// Tests that you can downgrade and upgrade a sender with an outstanding permit +// but no other senders left. +#[tokio::test] +async fn downgrade_upgrade_get_permit_no_senders() { + let (tx, _rx) = mpsc::channel::(1); + let tx2 = tx.clone(); + let _permit = tx.reserve_owned().await.unwrap(); + let weak_tx = tx2.downgrade(); + drop(tx2); + assert!(weak_tx.upgrade().is_some()); +} + +// Tests that `downgrade` does not change the `tx_count` of the channel. +#[tokio::test] +async fn test_tx_count_weak_sender() { + let (tx, _rx) = mpsc::channel::(1); + let tx_weak = tx.downgrade(); + let tx_weak2 = tx.downgrade(); + drop(tx); + + assert!(tx_weak.upgrade().is_none() && tx_weak2.upgrade().is_none()); +}