sync: Implement Clone for watch::Sender

tokio/src/sync/tests/loom_watch.rs

@@ -88,3 +88,23 @@ fn wait_for_returns_correct_value() {
         jh.join().unwrap();
     });
 }
+
+#[test]
+fn multiple_sender_drop_concurrently() {
+    loom::model(move || {
+        let (tx1, rx) = watch::channel(0);
+        let tx2 = tx1.clone();
+
+        let jh = thread::spawn(move || {
+            drop(tx2);
+        });
+        assert!(rx.has_changed().is_ok());
+
+        drop(tx1);
+
+        jh.join().unwrap();
+
+        // Check if all sender are dropped and closed flag is set.
+        assert!(rx.has_changed().is_err());
+    });
+}

tokio/src/sync/watch.rs

@@ -114,7 +114,7 @@
 use crate::sync::notify::Notify;
 
 use crate::loom::sync::atomic::AtomicUsize;
-use crate::loom::sync::atomic::Ordering::Relaxed;
+use crate::loom::sync::atomic::Ordering::{AcqRel, Relaxed};
 use crate::loom::sync::{Arc, RwLock, RwLockReadGuard};
 use std::fmt;
 use std::mem;
@@ -146,6 +146,16 @@ pub struct Sender<T> {
     shared: Arc<Shared<T>>,
 }
 
+impl<T> Clone for Sender<T> {
+    fn clone(&self) -> Self {
+        self.shared.ref_count_tx.fetch_add(1, AcqRel);
+
+        Self {
+            shared: self.shared.clone(),
+        }
+    }
+}
+
 /// Returns a reference to the inner value.
 ///
 /// Outstanding borrows hold a read lock on the inner value. This means that
@@ -238,6 +248,9 @@ struct Shared<T> {
     /// Tracks the number of `Receiver` instances.
     ref_count_rx: AtomicUsize,
 
+    /// Tracks the number of `Sender` instances.
+    ref_count_tx: AtomicUsize,
+
     /// Notifies waiting receivers that the value changed.
     notify_rx: big_notify::BigNotify,
 
@@ -354,12 +367,12 @@ mod big_notify {
     }
 }
 
-use self::state::{AtomicState, Version};
+use self::state::{AtomicState, Version, CLOSED_BIT};
 mod state {
     use crate::loom::sync::atomic::AtomicUsize;
     use crate::loom::sync::atomic::Ordering;
 
-    const CLOSED_BIT: usize = 1;
+    pub(super) const CLOSED_BIT: usize = 1;
 
     // Using 2 as the step size preserves the `CLOSED_BIT`.
     const STEP_SIZE: usize = 2;
@@ -374,7 +387,7 @@ mod state {
     /// The CLOSED bit tracks whether the Sender has been dropped. Dropping all
     /// receivers does not set it.
     #[derive(Copy, Clone, Debug)]
-    pub(super) struct StateSnapshot(usize);
+    pub(super) struct StateSnapshot(pub(super) usize);
 
     /// The state stored in an atomic integer.
     ///
@@ -383,7 +396,7 @@ mod state {
     /// current state. This ensures that written values are seen by
     /// the `Receiver`s for a proper handover.
     #[derive(Debug)]
-    pub(super) struct AtomicState(AtomicUsize);
+    pub(super) struct AtomicState(pub(super) AtomicUsize);
 
     impl Version {
         /// Decrements the version.
@@ -436,11 +449,6 @@ mod state {
             // method.
             self.0.fetch_add(STEP_SIZE, Ordering::Release);
         }
-
-        /// Set the closed bit in the state.
-        pub(super) fn set_closed(&self) {
-            self.0.fetch_or(CLOSED_BIT, Ordering::Release);
-        }
     }
 }
 
@@ -485,6 +493,7 @@ pub fn channel<T>(init: T) -> (Sender<T>, Receiver<T>) {
         value: RwLock::new(init),
         state: AtomicState::new(),
         ref_count_rx: AtomicUsize::new(1),
+        ref_count_tx: AtomicUsize::new(1),
         notify_rx: big_notify::BigNotify::new(),
         notify_tx: Notify::new(),
     });
@@ -1302,8 +1311,29 @@ impl<T> Sender<T> {
 
 impl<T> Drop for Sender<T> {
     fn drop(&mut self) {
-        self.shared.state.set_closed();
-        self.shared.notify_rx.notify_waiters();
+        if self.shared.ref_count_tx.fetch_sub(1, AcqRel) == 1 {
+            // Now that the sender's ref count is zero, so we attempt to close the sender.
+
+            // Fetch the latest state.
+            let current_state = self.shared.state.load();
+            let currrent_state_inner = current_state.0;
+
+            if self
+                .shared
+                .state
+                .0
+                .compare_exchange(
+                    currrent_state_inner,
+                    currrent_state_inner | CLOSED_BIT,
+                    AcqRel,
+                    Relaxed,
+                )
+                .is_ok()
+            {
+                // Only when last sender is dropped, we notify the receivers.
+                self.shared.notify_rx.notify_waiters();
+            }
+        }
     }
 }
 

tokio/tests/sync_watch.rs

@@ -7,7 +7,9 @@ use wasm_bindgen_test::wasm_bindgen_test as test;
 
 use tokio::sync::watch;
 use tokio_test::task::spawn;
-use tokio_test::{assert_pending, assert_ready, assert_ready_err, assert_ready_ok};
+use tokio_test::{
+    assert_pending, assert_ready, assert_ready_eq, assert_ready_err, assert_ready_ok,
+};
 
 #[test]
 fn single_rx_recv() {
@@ -332,3 +334,37 @@ fn send_modify_panic() {
     assert_ready_ok!(task.poll());
     assert_eq!(*rx.borrow_and_update(), "three");
 }
+
+#[tokio::test]
+async fn multiple_sender() {
+    let (tx1, mut rx) = watch::channel(0);
+    let tx2 = tx1.clone();
+
+    let mut t = spawn(async {
+        rx.changed().await.unwrap();
+        let v1 = *rx.borrow_and_update();
+        rx.changed().await.unwrap();
+        let v2 = *rx.borrow_and_update();
+        (v1, v2)
+    });
+
+    tx1.send(1).unwrap();
+    assert_pending!(t.poll());
+    tx2.send(2).unwrap();
+    assert_ready_eq!(t.poll(), (1, 2));
+}
+
+#[tokio::test]
+async fn reciever_is_notified_when_last_sender_is_dropped() {
+    let (tx1, mut rx) = watch::channel(0);
+    let tx2 = tx1.clone();
+
+    let mut t = spawn(rx.changed());
+    assert_pending!(t.poll());
+
+    drop(tx1);
+    assert!(!t.is_woken());
+    drop(tx2);
+
+    assert!(t.is_woken());
+}