diff --git a/tokio/src/sync/tests/loom_watch.rs b/tokio/src/sync/tests/loom_watch.rs index 51589cd80..0f69ae7e9 100644 --- a/tokio/src/sync/tests/loom_watch.rs +++ b/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()); + }); +} diff --git a/tokio/src/sync/watch.rs b/tokio/src/sync/watch.rs index 3979b0720..a813a829b 100644 --- a/tokio/src/sync/watch.rs +++ b/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 { shared: Arc>, } +impl Clone for Sender { + fn clone(&self) -> Self { + self.shared.ref_count_tx.fetch_add(1, Relaxed); + + 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 { /// 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, @@ -485,6 +498,7 @@ pub fn channel(init: T) -> (Sender, Receiver) { 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 +1316,10 @@ impl Sender { impl Drop for Sender { 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 { + self.shared.state.set_closed(); + self.shared.notify_rx.notify_waiters(); + } } } diff --git a/tokio/tests/sync_watch.rs b/tokio/tests/sync_watch.rs index a5b229f7d..7957c7da2 100644 --- a/tokio/tests/sync_watch.rs +++ b/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()); +}