sync: fix underflow in mpsc channel len() (#8062)

tokio/src/sync/mpsc/block.rs

@@ -211,11 +211,6 @@ impl<T> Block<T> {
         self.header.ready_slots.fetch_or(TX_CLOSED, Release);
     }
 
-    pub(crate) unsafe fn is_closed(&self) -> bool {
-        let ready_bits = self.header.ready_slots.load(Acquire);
-        is_tx_closed(ready_bits)
-    }
-
     /// Resets the block to a blank state. This enables reusing blocks in the
     /// channel.
     ///

tokio/src/sync/mpsc/list.rs

@@ -224,15 +224,6 @@ impl<T> Tx<T> {
             let _ = Box::from_raw(block.as_ptr());
         }
     }
-
-    pub(crate) fn is_closed(&self) -> bool {
-        let tail = self.block_tail.load(Acquire);
-
-        unsafe {
-            let tail_block = &*tail;
-            tail_block.is_closed()
-        }
-    }
 }
 
 impl<T> fmt::Debug for Tx<T> {
@@ -256,11 +247,58 @@ impl<T> Rx<T> {
         self.len(tx) == 0
     }
 
+    // Guaranteed to return true if `slot_index` is the fake message sent on channel close.
+    // Guaranteed to return false if `slot_index` is a fully sent message.
+    //
+    // For messages that are partially sent, may return either true or false.
+    fn is_maybe_closed(&self, tx: &Tx<T>, slot_index: usize) -> bool {
+        let start_index = block::start_index(slot_index);
+
+        let tail = tx.block_tail.load(Acquire);
+        // SAFETY: Only the receiver frees blocks, so since we are the receiver, this will not be
+        // freed right now.
+        let tail_ref = unsafe { &*tail };
+        if tail_ref.is_at_index(start_index) {
+            return !tail_ref.has_value(slot_index);
+        }
+
+        // This method is optimized for checking whether the last value is present, so most of the
+        // time it is in `block_tail`. However, this isn't always the case since it's possible
+        // that the list was grown with an empty block, in which case `block_tail` points one block
+        // too far. To handle this case, we walk the list from the head.
+        let mut block_ptr = Some(self.head);
+
+        while let Some(block) = block_ptr {
+            // SAFETY: Only the receiver frees blocks, so since we are the receiver, this will not
+            // be freed right now.
+            let block_ref = unsafe { block.as_ref() };
+            if block_ref.is_at_index(start_index) {
+                return !block_ref.has_value(slot_index);
+            }
+            block_ptr = block_ref.load_next(Acquire);
+        }
+        true
+    }
+
     pub(crate) fn len(&self, tx: &Tx<T>) -> usize {
-        // When all the senders are dropped, there will be a last block in the tail position,
-        // but it will be closed
         let tail_position = tx.tail_position.load(Acquire);
-        tail_position - self.index - (tx.is_closed() as usize)
+        let mut len = tail_position.wrapping_sub(self.index);
+        debug_assert!(0 <= len as isize);
+        if len == 0 {
+            return 0;
+        }
+        // There are messages present in the queue. However, it's possible that the last message is
+        // a fake "closed" message that we do not wish to count. To avoid counting it, we do not
+        // count the last message if the ready bit is unset.
+        //
+        // Note that it is also possible for the ready bit to be unset on a normal message, but
+        // this happens only if that message is currently being sent *right now* in parallel on
+        // another thread. That is okay because it is optional to count messages that are currently
+        // being sent.
+        if self.is_maybe_closed(tx, tail_position.wrapping_sub(1)) {
+            len -= 1;
+        }
+        len
     }
 
     /// Pops the next value off the queue.

tokio/src/sync/mpsc/mod.rs

@@ -135,10 +135,10 @@ pub mod error;
 /// This value must be a power of 2. It also must be smaller than the number of
 /// bits in `usize`.
 #[cfg(all(target_pointer_width = "64", not(loom)))]
-const BLOCK_CAP: usize = 32;
+pub(crate) const BLOCK_CAP: usize = 32;
 
 #[cfg(all(not(target_pointer_width = "64"), not(loom)))]
-const BLOCK_CAP: usize = 16;
+pub(crate) const BLOCK_CAP: usize = 16;
 
 #[cfg(loom)]
-const BLOCK_CAP: usize = 2;
+pub(crate) const BLOCK_CAP: usize = 2;

tokio/src/sync/tests/loom_mpsc.rs

@@ -1,4 +1,4 @@
-use crate::sync::mpsc;
+use crate::sync::mpsc::{self, BLOCK_CAP};
 
 use loom::future::block_on;
 use loom::sync::Arc;
@@ -222,3 +222,56 @@ fn nonempty_after_send() {
         join.join().unwrap();
     });
 }
+
+#[test]
+fn is_empty_during_close() {
+    loom::model(|| {
+        let (tx, rx) = mpsc::channel::<()>(1);
+
+        let th1 = thread::spawn(move || {
+            assert!(rx.is_empty());
+        });
+
+        drop(tx);
+
+        th1.join().unwrap();
+    });
+}
+
+fn len_during_close_helper(n: usize) {
+    loom::model(move || {
+        let (tx, rx) = mpsc::channel::<()>(n + 1);
+
+        for _ in 0..n {
+            tx.try_send(()).unwrap();
+        }
+
+        let th1 = thread::spawn(move || {
+            assert_eq!(rx.len(), n);
+        });
+
+        drop(tx);
+
+        th1.join().unwrap();
+    });
+}
+
+#[test]
+fn len_during_close_0() {
+    len_during_close_helper(0);
+}
+
+#[test]
+fn len_during_close_1() {
+    len_during_close_helper(1);
+}
+
+#[test]
+fn len_during_close_block_cap() {
+    len_during_close_helper(BLOCK_CAP);
+}
+
+#[test]
+fn len_during_close_block_cap_plus_1() {
+    len_during_close_helper(BLOCK_CAP + 1);
+}