Add is_closed, is_empty and len to mpsc::Receiver and mpsc::UnboundedReceiver

tokio/src/sync/mpsc/block.rs

@@ -168,6 +168,20 @@ impl<T> Block<T> {
         Some(Read::Value(value.assume_init()))
     }
 
+    /// Returns true if the sender half of the list is closed and there is no
+    /// value in the slot to be consumed
+    ///
+    /// # Safety
+    ///
+    /// To maintain safety, the caller must ensure:
+    ///
+    /// * No concurrent access to the slot.
+    pub(crate) fn is_closed_and_empty(&self, slot_index: usize) -> bool {
+        let offset = offset(slot_index);
+        let ready_bits = self.header.ready_slots.load(Acquire);
+        !is_ready(ready_bits, offset) && is_tx_closed(ready_bits)
+    }
+
     /// Writes a value to the block at the given offset.
     ///
     /// # Safety

tokio/src/sync/mpsc/bounded.rs

@@ -463,6 +463,42 @@ impl<T> Receiver<T> {
         self.chan.close();
     }
 
+    /// Checks if a channel is closed.
+    ///
+    /// This method returns `true` if the channel has been closed and there are
+    /// no remaining messages in the channel's buffer. The channel is closed
+    /// when all [`Sender`] have been dropped, or when [`Receiver::close`] is called.
+    ///
+    /// [`Sender`]: crate::sync::mpsc::Sender
+    /// [`Receiver::close`]: crate::sync::mpsc::Receiver::close
+    ///
+    /// # Examples
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, mut rx) = mpsc::channel(10);
+    ///     assert!(!rx.is_closed());
+    ///
+    ///     tx.send(100).await.unwrap();
+    ///     tx.send(200).await.unwrap();
+    ///
+    ///     rx.close();
+    ///     assert!(!rx.is_closed());
+    ///     assert_eq!(rx.recv().await, Some(100));
+    ///
+    ///     assert!(!rx.is_closed());
+    ///     assert_eq!(rx.recv().await, Some(200));
+    ///     assert!(rx.recv().await.is_none());
+    ///     
+    ///     assert!(rx.is_closed());
+    /// }
+    /// ```
+    pub fn is_closed(&mut self) -> bool {
+        self.chan.is_closed_and_empty()
+    }
+
     /// Polls to receive the next message on this channel.
     ///
     /// This method returns:

tokio/src/sync/mpsc/chan.rs

@@ -241,6 +241,24 @@ impl<T, S: Semaphore> Rx<T, S> {
         self.inner.notify_rx_closed.notify_waiters();
     }
 
+    pub(crate) fn is_closed_and_empty(&mut self) -> bool {
+        // There two internal states that can represent a closed channel that has no messages
+        //
+        //  1. When `close` is called.
+        //  In this case, the inner semaphore will be closed, and the channel has no
+        //  outstanding messages if the list is at the tail position.
+        //
+        //  2. When all senders are dropped.
+        //  In this case, the semaphore remains unclosed, and the `index` in the list won't
+        //  reach the tail position. It is necessary to check the list if the last block is
+        //  `closed` and has no value available.
+        self.inner.rx_fields.with_mut(|rx_fields_ptr| {
+            let rx_fields = unsafe { &mut *rx_fields_ptr };
+            (self.inner.semaphore.is_closed() && rx_fields.list.is_at_tail_position(&self.inner.tx))
+                || rx_fields.list.is_closed_and_empty(&self.inner.tx)
+        })
+    }
+
     /// Receive the next value
     pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
         use super::block::Read;

tokio/src/sync/mpsc/list.rs

@@ -230,6 +230,23 @@ impl<T> fmt::Debug for Tx<T> {
 }
 
 impl<T> Rx<T> {
+    pub(crate) fn is_closed_and_empty(&mut self, tx: &Tx<T>) -> bool {
+        // Advance `head`, if needed
+        if self.try_advancing_head() {
+            self.reclaim_blocks(tx);
+        }
+
+        unsafe {
+            let block = self.head.as_ref();
+            block.is_closed_and_empty(self.index)
+        }
+    }
+
+    pub(crate) fn is_at_tail_position(&mut self, tx: &Tx<T>) -> bool {
+        let tail_position = tx.tail_position.load(Acquire);
+        tail_position == self.index
+    }
+
     /// Pops the next value off the queue.
     pub(crate) fn pop(&mut self, tx: &Tx<T>) -> Option<block::Read<T>> {
         // Advance `head`, if needed

tokio/src/sync/mpsc/unbounded.rs

@@ -330,6 +330,42 @@ impl<T> UnboundedReceiver<T> {
         self.chan.close();
     }
 
+    /// Checks if a channel is closed.
+    ///
+    /// This method returns `true` if the channel has been closed and there are
+    /// no remaining messages in the channel's buffer. The channel is closed
+    /// when all [`UnboundedSender`] have been dropped, or when [`UnboundedReceiver::close`] is called.
+    ///
+    /// [`UnboundedSender`]: crate::sync::mpsc::UnboundedSender
+    /// [`UnboundedReceiver::close`]: crate::sync::mpsc::UnboundedReceiver::close
+    ///
+    /// # Examples
+    /// ```
+    /// use tokio::sync::mpsc;
+    ///
+    /// #[tokio::main]
+    /// async fn main() {
+    ///     let (tx, mut rx) = mpsc::unbounded_channel();
+    ///     assert!(!rx.is_closed());
+    ///
+    ///     let _ = tx.send(100);
+    ///     let _ = tx.send(200);
+    ///
+    ///     rx.close();
+    ///     assert!(!rx.is_closed());
+    ///     assert_eq!(rx.recv().await, Some(100));
+    ///
+    ///     assert!(!rx.is_closed());
+    ///     assert_eq!(rx.recv().await, Some(200));
+    ///     assert!(rx.recv().await.is_none());
+    ///     
+    ///     assert!(rx.is_closed());
+    /// }
+    /// ```
+    pub fn is_closed(&mut self) -> bool {
+        self.chan.is_closed_and_empty()
+    }
+
     /// Polls to receive the next message on this channel.
     ///
     /// This method returns:

tokio/tests/sync_mpsc.rs

@@ -1017,4 +1017,81 @@ async fn test_tx_capacity() {
     assert_eq!(tx.max_capacity(), 10);
 }
 
+#[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
+    let (_tx, mut rx) = mpsc::channel::<()>(10);
+    rx.close();
+
+    assert!(rx.is_closed());
+}
+
+#[tokio::test]
+async fn test_rx_is_closed_when_dropping_all_senders() {
+    // is_closed should return true after dropping all senders
+    let (tx, mut rx) = mpsc::channel::<()>(10);
+    let another_tx = tx.clone();
+    let task = tokio::spawn(async move {
+        drop(another_tx);
+    });
+
+    drop(tx);
+    let _ = task.await;
+
+    assert!(rx.is_closed());
+}
+
+#[tokio::test]
+async fn test_rx_is_not_closed_when_there_are_senders() {
+    // is_closed should return false when there is a sender
+    let (_tx, mut rx) = mpsc::channel::<()>(10);
+    assert!(!rx.is_closed());
+}
+
+#[tokio::test]
+async fn test_rx_is_not_closed_when_there_are_senders_and_buffer_filled() {
+    // is_closed should return false when there is a sender, even if enough messages have been sent to fill the channel
+    let (tx, mut rx) = mpsc::channel(10);
+    for i in 0..10 {
+        assert!(tx.send(i).await.is_ok());
+    }
+    assert!(!rx.is_closed());
+}
+
+#[tokio::test]
+async fn test_rx_is_not_closed_when_there_are_messages_but_not_senders() {
+    // is_closed should return false when there is a permit (but no senders)
+    let (tx, mut rx) = mpsc::channel(10);
+    for i in 0..10 {
+        assert!(tx.send(i).await.is_ok());
+    }
+    drop(tx);
+    assert!(!rx.is_closed());
+}
+
+#[tokio::test]
+async fn test_rx_is_not_closed_when_there_are_messages_and_close_is_called() {
+    // is_closed should return false when there is a permit (but no senders)
+    let (tx, mut rx) = mpsc::channel(10);
+    for i in 0..10 {
+        assert!(tx.send(i).await.is_ok());
+    }
+    rx.close();
+    assert!(!rx.is_closed());
+}
+
+#[tokio::test]
+async fn test_rx_is_closed_after_consuming_messages() {
+    // is_closed should return false when there is a permit (but no senders)
+    let (tx, mut rx) = mpsc::channel(10);
+    for i in 0..10 {
+        assert!(tx.send(i).await.is_ok());
+    }
+    drop(tx);
+
+    assert!(!rx.is_closed());
+    while (rx.recv().await).is_some() {}
+    assert!(rx.is_closed());
+}
+
 fn is_debug<T: fmt::Debug>(_: &T) {}

tokio/tests/sync_mpsc_weak.rs

@@ -511,3 +511,12 @@ fn test_tx_count_weak_unbounded_sender() {
 
     assert!(tx_weak.upgrade().is_none() && tx_weak2.upgrade().is_none());
 }
+
+#[tokio::test]
+async fn test_rx_is_closed_when_dropping_all_senders_except_weak_senders() {
+    // is_closed should return true after dropping all senders except for a weak sender
+    let (tx, mut rx) = mpsc::channel::<()>(10);
+    let _weak_sender = tx.clone().downgrade();
+    drop(tx);
+    assert!(rx.is_closed());
+}