sync: implement Weak version of mpsc::Sender (#4595)

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<T> {
     chan: chan::Tx<T, Semaphore>,
 }
 
+/// 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<Sender>`. 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::<i32>(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<T> {
+    chan: Arc<chan::Chan<T, Semaphore>>,
+}
+
 /// Permits to send one value into the channel.
 ///
 /// `Permit` values are returned by [`Sender::reserve()`] and [`Sender::try_reserve()`]
@@ -991,6 +1026,16 @@ impl<T> Sender<T> {
     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<T> {
+        WeakSender {
+            chan: self.chan.downgrade(),
+        }
+    }
 }
 
 impl<T> Clone for Sender<T> {
@@ -1009,6 +1054,29 @@ impl<T> fmt::Debug for Sender<T> {
     }
 }
 
+impl<T> Clone for WeakSender<T> {
+    fn clone(&self) -> Self {
+        WeakSender {
+            chan: self.chan.clone(),
+        }
+    }
+}
+
+impl<T> WeakSender<T> {
+    /// 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<Sender<T>> {
+        chan::Tx::upgrade(self.chan.clone()).map(Sender::new)
+    }
+}
+
+impl<T> fmt::Debug for WeakSender<T> {
+    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
+        fmt.debug_struct("WeakSender").finish()
+    }
+}
+
 // ===== impl Permit =====
 
 impl<T> Permit<'_, T> {

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<T, S> {
@@ -46,7 +47,7 @@ pub(crate) trait Semaphore {
     fn is_closed(&self) -> bool;
 }
 
-struct Chan<T, S> {
+pub(super) struct Chan<T, S> {
     /// Notifies all tasks listening for the receiver being dropped.
     notify_rx_closed: Notify,
 
@@ -129,6 +130,30 @@ impl<T, S> Tx<T, S> {
         Tx { inner: chan }
     }
 
+    pub(super) fn downgrade(&self) -> Arc<Chan<T, S>> {
+        self.inner.clone()
+    }
+
+    // Returns the upgraded channel or None if the upgrade failed.
+    pub(super) fn upgrade(chan: Arc<Chan<T, S>>) -> Option<Self> {
+        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);

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;
 

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<ActorMessage>,
+        sender: mpsc::WeakSender<ActorMessage>,
+        next_id: u32,
+        pub received_self_msg: bool,
+    }
+
+    enum ActorMessage {
+        GetUniqueId { respond_to: oneshot::Sender<u32> },
+        SelfMessage {},
+    }
+
+    impl MyActor {
+        fn new(
+            receiver: mpsc::Receiver<ActorMessage>,
+            sender: mpsc::WeakSender<ActorMessage>,
+        ) -> 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<ActorMessage>,
+    }
+
+    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::<i32>(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::<i32>(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::<i32>(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::<i32>(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::<i32>(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::<i32>(1);
+    let tx_weak = tx.downgrade();
+    let tx_weak2 = tx.downgrade();
+    drop(tx);
+
+    assert!(tx_weak.upgrade().is_none() && tx_weak2.upgrade().is_none());
+}