sync: add len and is_empty methods to mpsc senders

Author: tqwewe

tokio/src/sync/mpsc/bounded.rs

@@ -515,7 +515,6 @@ impl<T> Receiver<T> {
     ///     tx.send(0).await.unwrap();
     ///     assert!(!rx.is_empty());
     /// }
-    ///
     /// ```
     pub fn is_empty(&self) -> bool {
         self.chan.is_empty()
@@ -1060,6 +1059,46 @@ impl<T> Sender<T> {
         self.chan.is_closed()
     }
 
+    /// Checks if a channel is empty.
+    ///
+    /// This method returns `true` if the channel has no messages.
+    ///
+    /// # Examples
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, rx) = mpsc::unbounded_channel();
+    ///     assert!(rx.is_empty());
+    ///
+    ///     tx.send(0).await.unwrap();
+    ///     assert!(!rx.is_empty());
+    /// }
+    /// ```
+    pub fn is_empty(&self) -> bool {
+        self.chan.semaphore().bound - self.chan.semaphore().semaphore.available_permits() == 0
+    }
+
+    /// Returns the number of messages in the channel.
+    ///
+    /// # Examples
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, rx) = mpsc::unbounded_channel();
+    ///     assert_eq!(0, rx.len());
+    ///
+    ///     tx.send(0).await.unwrap();
+    ///     assert_eq!(1, rx.len());
+    /// }
+    /// ```
+    pub fn len(&self) -> usize {
+        self.chan.semaphore().bound - self.chan.semaphore().semaphore.available_permits()
+    }
+
     /// Waits for channel capacity. Once capacity to send one message is
     /// available, it is reserved for the caller.
     ///

tokio/src/sync/mpsc/chan.rs

@@ -44,6 +44,8 @@ pub(crate) trait Semaphore {
 
     fn add_permits(&self, n: usize);
 
+    fn available_permits(&self) -> usize;
+
     fn close(&self);
 
     fn is_closed(&self) -> bool;
@@ -193,6 +195,14 @@ impl<T, S> Tx<T, S> {
 }
 
 impl<T, S: Semaphore> Tx<T, S> {
+    pub(super) fn is_empty(&self) -> bool {
+        self.inner.semaphore.available_permits() == 0
+    }
+
+    pub(super) fn len(&self) -> usize {
+        self.inner.semaphore.available_permits()
+    }
+
     pub(crate) fn is_closed(&self) -> bool {
         self.inner.semaphore.is_closed()
     }
@@ -576,6 +586,10 @@ impl Semaphore for bounded::Semaphore {
         self.semaphore.release(n)
     }
 
+    fn available_permits(&self) -> usize {
+        self.semaphore.available_permits()
+    }
+
     fn is_idle(&self) -> bool {
         self.semaphore.available_permits() == self.bound
     }
@@ -610,6 +624,10 @@ impl Semaphore for unbounded::Semaphore {
         }
     }
 
+    fn available_permits(&self) -> usize {
+        self.0.load(Acquire) >> 1
+    }
+
     fn is_idle(&self) -> bool {
         self.0.load(Acquire) >> 1 == 0
     }

tokio/src/sync/mpsc/unbounded.rs

@@ -382,7 +382,6 @@ impl<T> UnboundedReceiver<T> {
     ///     tx.send(0).unwrap();
     ///     assert!(!rx.is_empty());
     /// }
-    ///
     /// ```
     pub fn is_empty(&self) -> bool {
         self.chan.is_empty()
@@ -577,6 +576,46 @@ impl<T> UnboundedSender<T> {
         }
     }
 
+    /// Checks if a channel is empty.
+    ///
+    /// This method returns `true` if the channel has no messages.
+    ///
+    /// # Examples
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, rx) = mpsc::unbounded_channel();
+    ///     assert!(tx.is_empty());
+    ///
+    ///     tx.send(0).unwrap();
+    ///     assert!(!tx.is_empty());
+    /// }
+    /// ```
+    pub fn is_empty(&self) -> bool {
+        self.chan.is_empty()
+    }
+
+    /// Returns the number of messages in the channel.
+    ///
+    /// # Examples
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, rx) = mpsc::unbounded_channel();
+    ///     assert_eq!(0, tx.len());
+    ///
+    ///     tx.send(0).unwrap();
+    ///     assert_eq!(1, tx.len());
+    /// }
+    /// ```
+    pub fn len(&self) -> usize {
+        self.chan.len()
+    }
+
     /// Completes when the receiver has dropped.
     ///
     /// This allows the producers to get notified when interest in the produced

tokio/tests/sync_mpsc.rs

@@ -1039,6 +1039,64 @@ async fn test_tx_capacity() {
     assert_eq!(tx.max_capacity(), 10);
 }
 
+#[tokio::test]
+async fn test_bounded_tx_len() {
+    let (tx, mut rx) = mpsc::channel::<()>(10);
+
+    // initially len should be 0
+    assert_eq!(tx.len(), 0);
+    assert!(tx.is_empty());
+
+    // queue one message, and len should be 1
+    tx.send(()).await.unwrap();
+    assert_eq!(tx.len(), 1);
+    assert!(!tx.is_empty());
+
+    // queue a second message, and len should be 2
+    tx.send(()).await.unwrap();
+    assert_eq!(tx.len(), 2);
+    assert!(!tx.is_empty());
+
+    // consume a message, and len should be 1
+    let _ = rx.recv().await;
+    assert_eq!(tx.len(), 1);
+    assert!(!tx.is_empty());
+
+    // consume a final message, and len should be 0
+    let _ = rx.recv().await;
+    assert_eq!(tx.len(), 0);
+    assert!(tx.is_empty());
+}
+
+#[tokio::test]
+async fn test_unbounded_tx_len() {
+    let (tx, mut rx) = mpsc::unbounded_channel();
+
+    // initially len should be 0
+    assert_eq!(tx.len(), 0);
+    assert!(tx.is_empty());
+
+    // queue one message, and len should be 1
+    tx.send(()).unwrap();
+    assert_eq!(tx.len(), 1);
+    assert!(!tx.is_empty());
+
+    // queue a second message, and len should be 2
+    tx.send(()).unwrap();
+    assert_eq!(tx.len(), 2);
+    assert!(!tx.is_empty());
+
+    // consume a message, and len should be 1
+    let _ = rx.recv().await;
+    assert_eq!(tx.len(), 1);
+    assert!(!tx.is_empty());
+
+    // consume a final message, and len should be 0
+    let _ = rx.recv().await;
+    assert_eq!(tx.len(), 0);
+    assert!(tx.is_empty());
+}
+
 #[tokio::test]
 async fn test_rx_is_closed_when_calling_close_with_sender() {
     // is_closed should return true after calling close but still has a sender