sync: expose strong and weak counts of mpsc sender handles

tokio/src/sync/mpsc/bounded.rs

@@ -1409,6 +1409,16 @@ impl<T> Sender<T> {
     pub fn max_capacity(&self) -> usize {
         self.chan.semaphore().bound
     }
+
+    /// Returns the number of [`Sender`] handles.
+    pub fn strong_count(&self) -> usize {
+        self.chan.strong_count()
+    }
+
+    /// Returns the number of [`WeakSender`] handles.
+    pub fn weak_count(&self) -> usize {
+        self.chan.weak_count()
+    }
 }
 
 impl<T> Clone for Sender<T> {
@@ -1435,13 +1445,29 @@ impl<T> Clone for WeakSender<T> {
     }
 }
 
+impl<T> Drop for WeakSender<T> {
+    fn drop(&mut self) {
+        self.chan.decrement_weak_count();
+    }
+}
+
 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)
     }
+
+    /// Returns the number of [`Sender`] handles.
+    pub fn strong_count(&self) -> usize {
+        self.chan.strong_count()
+    }
+
+    /// Returns the number of [`WeakSender`] handles.
+    pub fn weak_count(&self) -> usize {
+        self.chan.weak_count()
+    }
 }
 
 impl<T> fmt::Debug for WeakSender<T> {

tokio/src/sync/mpsc/chan.rs

@@ -66,6 +66,9 @@ pub(super) struct Chan<T, S> {
     /// When this drops to zero, the send half of the channel is closed.
     tx_count: AtomicUsize,
 
+    /// Tracks the number of outstanding weak sender handles.
+    tx_weak_count: AtomicUsize,
+
     /// Only accessed by `Rx` handle.
     rx_fields: UnsafeCell<RxFields<T>>,
 }
@@ -115,6 +118,7 @@ pub(crate) fn channel<T, S: Semaphore>(semaphore: S) -> (Tx<T, S>, Rx<T, S>) {
         semaphore,
         rx_waker: CachePadded::new(AtomicWaker::new()),
         tx_count: AtomicUsize::new(1),
+        tx_weak_count: AtomicUsize::new(0),
         rx_fields: UnsafeCell::new(RxFields {
             list: rx,
             rx_closed: false,
@@ -131,7 +135,17 @@ impl<T, S> Tx<T, S> {
         Tx { inner: chan }
     }
 
+    pub(super) fn strong_count(&self) -> usize {
+        self.inner.tx_count.load(Acquire)
+    }
+
+    pub(super) fn weak_count(&self) -> usize {
+        self.inner.tx_weak_count.load(Relaxed)
+    }
+
     pub(super) fn downgrade(&self) -> Arc<Chan<T, S>> {
+        self.inner.tx_weak_count.fetch_add(1, Relaxed);
+
         self.inner.clone()
     }
 
@@ -452,6 +466,18 @@ impl<T, S> Chan<T, S> {
         // Notify the rx task
         self.rx_waker.wake();
     }
+
+    pub(super) fn decrement_weak_count(&self) {
+        self.tx_weak_count.fetch_sub(1, Relaxed);
+    }
+
+    pub(super) fn strong_count(&self) -> usize {
+        self.tx_count.load(Acquire)
+    }
+
+    pub(super) fn weak_count(&self) -> usize {
+        self.tx_weak_count.load(Relaxed)
+    }
 }
 
 impl<T, S> Drop for Chan<T, S> {

tokio/src/sync/mpsc/unbounded.rs

@@ -578,6 +578,16 @@ impl<T> UnboundedSender<T> {
             chan: self.chan.downgrade(),
         }
     }
+
+    /// Returns the number of [`UnboundedSender`] handles.
+    pub fn strong_count(&self) -> usize {
+        self.chan.strong_count()
+    }
+
+    /// Returns the number of [`WeakUnboundedSender`] handles.
+    pub fn weak_count(&self) -> usize {
+        self.chan.weak_count()
+    }
 }
 
 impl<T> Clone for WeakUnboundedSender<T> {
@@ -588,13 +598,29 @@ impl<T> Clone for WeakUnboundedSender<T> {
     }
 }
 
+impl<T> Drop for WeakUnboundedSender<T> {
+    fn drop(&mut self) {
+        self.chan.decrement_weak_count();
+    }
+}
+
 impl<T> WeakUnboundedSender<T> {
     /// Tries to convert a `WeakUnboundedSender` into an [`UnboundedSender`].
     /// 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<UnboundedSender<T>> {
         chan::Tx::upgrade(self.chan.clone()).map(UnboundedSender::new)
     }
+
+    /// Returns the number of [`UnboundedSender`] handles.
+    pub fn strong_count(&self) -> usize {
+        self.chan.strong_count()
+    }
+
+    /// Returns the number of [`WeakUnboundedSender`] handles.
+    pub fn weak_count(&self) -> usize {
+        self.chan.weak_count()
+    }
 }
 
 impl<T> fmt::Debug for WeakUnboundedSender<T> {

tokio/tests/sync_mpsc_weak.rs

@@ -7,8 +7,9 @@ use wasm_bindgen_test::wasm_bindgen_test as test;
 
 use std::sync::atomic::AtomicUsize;
 use std::sync::atomic::Ordering::{Acquire, Release};
+use std::sync::Arc;
 use tokio::sync::mpsc::{self, channel, unbounded_channel};
-use tokio::sync::oneshot;
+use tokio::sync::{oneshot, Notify};
 
 #[tokio::test]
 async fn weak_sender() {
@@ -511,3 +512,129 @@ fn test_tx_count_weak_unbounded_sender() {
 
     assert!(tx_weak.upgrade().is_none() && tx_weak2.upgrade().is_none());
 }
+
+#[tokio::test]
+async fn strong_and_weak_count() {
+    let (tx, _rx) = mpsc::channel::<()>(1);
+
+    let first_downgrade = Arc::new(Notify::new());
+    let second_downgrade = Arc::new(Notify::new());
+
+    let task_handle = {
+        let tx = tx.clone();
+        let first_downgrade = first_downgrade.clone();
+        let second_downgrade = second_downgrade.clone();
+
+        tokio::spawn(async move {
+            let weak = tx.downgrade();
+
+            assert_eq!(tx.weak_count(), 1);
+            assert_eq!(tx.strong_count(), 2);
+            assert_eq!(weak.weak_count(), 1);
+            assert_eq!(weak.strong_count(), 2);
+
+            first_downgrade.notify_one();
+
+            second_downgrade.notified().await;
+
+            assert_eq!(tx.weak_count(), 2);
+            assert_eq!(tx.strong_count(), 2);
+            assert_eq!(weak.weak_count(), 2);
+            assert_eq!(weak.strong_count(), 2);
+
+            drop(weak);
+
+            assert_eq!(tx.weak_count(), 1);
+            assert_eq!(tx.strong_count(), 2);
+
+            drop(tx);
+        })
+    };
+
+    first_downgrade.notified().await;
+
+    assert_eq!(tx.weak_count(), 1);
+    assert_eq!(tx.strong_count(), 2);
+
+    let weak = tx.downgrade();
+
+    assert_eq!(tx.weak_count(), 2);
+    assert_eq!(tx.strong_count(), 2);
+    assert_eq!(weak.weak_count(), 2);
+    assert_eq!(weak.strong_count(), 2);
+
+    second_downgrade.notify_one();
+
+    task_handle.await.unwrap();
+
+    assert_eq!(tx.weak_count(), 1);
+    assert_eq!(tx.strong_count(), 1);
+
+    drop(weak);
+
+    assert_eq!(tx.weak_count(), 0);
+    assert_eq!(tx.strong_count(), 1);
+}
+
+#[tokio::test]
+async fn unbounded_strong_and_weak_count() {
+    let (tx, _rx) = mpsc::unbounded_channel::<()>();
+
+    let first_downgrade = Arc::new(Notify::new());
+    let second_downgrade = Arc::new(Notify::new());
+
+    let task_handle = {
+        let tx = tx.clone();
+        let first_downgrade = first_downgrade.clone();
+        let second_downgrade = second_downgrade.clone();
+
+        tokio::spawn(async move {
+            let weak = tx.downgrade();
+
+            assert_eq!(tx.weak_count(), 1);
+            assert_eq!(tx.strong_count(), 2);
+            assert_eq!(weak.weak_count(), 1);
+            assert_eq!(weak.strong_count(), 2);
+
+            first_downgrade.notify_one();
+
+            second_downgrade.notified().await;
+
+            assert_eq!(tx.weak_count(), 2);
+            assert_eq!(tx.strong_count(), 2);
+            assert_eq!(weak.weak_count(), 2);
+            assert_eq!(weak.strong_count(), 2);
+
+            drop(weak);
+
+            assert_eq!(tx.weak_count(), 1);
+            assert_eq!(tx.strong_count(), 2);
+
+            drop(tx);
+        })
+    };
+
+    first_downgrade.notified().await;
+
+    assert_eq!(tx.weak_count(), 1);
+    assert_eq!(tx.strong_count(), 2);
+
+    let weak = tx.downgrade();
+
+    assert_eq!(tx.weak_count(), 2);
+    assert_eq!(tx.strong_count(), 2);
+    assert_eq!(weak.weak_count(), 2);
+    assert_eq!(weak.strong_count(), 2);
+
+    second_downgrade.notify_one();
+
+    task_handle.await.unwrap();
+
+    assert_eq!(tx.weak_count(), 1);
+    assert_eq!(tx.strong_count(), 1);
+
+    drop(weak);
+
+    assert_eq!(tx.weak_count(), 0);
+    assert_eq!(tx.strong_count(), 1);
+}