diff --git a/tokio/src/sync/mpsc/block.rs b/tokio/src/sync/mpsc/block.rs index 927c45664..7d7e59afd 100644 --- a/tokio/src/sync/mpsc/block.rs +++ b/tokio/src/sync/mpsc/block.rs @@ -211,11 +211,6 @@ impl Block { 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. /// diff --git a/tokio/src/sync/mpsc/list.rs b/tokio/src/sync/mpsc/list.rs index 118bac856..1f3d23839 100644 --- a/tokio/src/sync/mpsc/list.rs +++ b/tokio/src/sync/mpsc/list.rs @@ -224,15 +224,6 @@ impl Tx { 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 fmt::Debug for Tx { @@ -256,11 +247,58 @@ impl Rx { 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, 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) -> 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. diff --git a/tokio/src/sync/mpsc/mod.rs b/tokio/src/sync/mpsc/mod.rs index 3df612ca4..e724f62a7 100644 --- a/tokio/src/sync/mpsc/mod.rs +++ b/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; diff --git a/tokio/src/sync/tests/loom_mpsc.rs b/tokio/src/sync/tests/loom_mpsc.rs index 039b87a77..3e6c84798 100644 --- a/tokio/src/sync/tests/loom_mpsc.rs +++ b/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); +}