Skip to content

Commit 3ce4720

Browse files
balliegojrballiegojr
authored
sync: add is_closed, is_empty, and len to mpsc receivers (#6348)
Fixes: #4638
1 parent 8342e4b commit 3ce4720

File tree

8 files changed

+661
-0
lines changed

8 files changed

+661
-0
lines changed

‎tokio/src/sync/mpsc/block.rs

+18
Original file line numberDiff line numberDiff line change
@@ -168,6 +168,19 @@ impl<T> Block<T> {
168168
Some(Read::Value(value.assume_init()))
169169
}
170170

171+
/// Returns true if there is a value in the slot to be consumed
172+
///
173+
/// # Safety
174+
///
175+
/// To maintain safety, the caller must ensure:
176+
///
177+
/// * No concurrent access to the slot.
178+
pub(crate) fn has_value(&self, slot_index: usize) -> bool {
179+
let offset = offset(slot_index);
180+
let ready_bits = self.header.ready_slots.load(Acquire);
181+
is_ready(ready_bits, offset)
182+
}
183+
171184
/// Writes a value to the block at the given offset.
172185
///
173186
/// # Safety
@@ -195,6 +208,11 @@ impl<T> Block<T> {
195208
self.header.ready_slots.fetch_or(TX_CLOSED, Release);
196209
}
197210

211+
pub(crate) unsafe fn is_closed(&self) -> bool {
212+
let ready_bits = self.header.ready_slots.load(Acquire);
213+
is_tx_closed(ready_bits)
214+
}
215+
198216
/// Resets the block to a blank state. This enables reusing blocks in the
199217
/// channel.
200218
///

‎tokio/src/sync/mpsc/bounded.rs

+67
Original file line numberDiff line numberDiff line change
@@ -463,6 +463,73 @@ impl<T> Receiver<T> {
463463
self.chan.close();
464464
}
465465

466+
/// Checks if a channel is closed.
467+
///
468+
/// This method returns `true` if the channel has been closed. The channel is closed
469+
/// when all [`Sender`] have been dropped, or when [`Receiver::close`] is called.
470+
///
471+
/// [`Sender`]: crate::sync::mpsc::Sender
472+
/// [`Receiver::close`]: crate::sync::mpsc::Receiver::close
473+
///
474+
/// # Examples
475+
/// ```
476+
/// use tokio::sync::mpsc;
477+
///
478+
/// #[tokio::main]
479+
/// async fn main() {
480+
/// let (_tx, mut rx) = mpsc::channel::<()>(10);
481+
/// assert!(!rx.is_closed());
482+
///
483+
/// rx.close();
484+
///
485+
/// assert!(rx.is_closed());
486+
/// }
487+
/// ```
488+
pub fn is_closed(&self) -> bool {
489+
self.chan.is_closed()
490+
}
491+
492+
/// Checks if a channel is empty.
493+
///
494+
/// This method returns `true` if the channel has no messages.
495+
///
496+
/// # Examples
497+
/// ```
498+
/// use tokio::sync::mpsc;
499+
///
500+
/// #[tokio::main]
501+
/// async fn main() {
502+
/// let (tx, rx) = mpsc::channel(10);
503+
/// assert!(rx.is_empty());
504+
///
505+
/// tx.send(0).await.unwrap();
506+
/// assert!(!rx.is_empty());
507+
/// }
508+
///
509+
/// ```
510+
pub fn is_empty(&self) -> bool {
511+
self.chan.is_empty()
512+
}
513+
514+
/// Returns the number of messages in the channel.
515+
///
516+
/// # Examples
517+
/// ```
518+
/// use tokio::sync::mpsc;
519+
///
520+
/// #[tokio::main]
521+
/// async fn main() {
522+
/// let (tx, rx) = mpsc::channel(10);
523+
/// assert_eq!(0, rx.len());
524+
///
525+
/// tx.send(0).await.unwrap();
526+
/// assert_eq!(1, rx.len());
527+
/// }
528+
/// ```
529+
pub fn len(&self) -> usize {
530+
self.chan.len()
531+
}
532+
466533
/// Polls to receive the next message on this channel.
467534
///
468535
/// This method returns:

‎tokio/src/sync/mpsc/chan.rs

+27
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,33 @@ impl<T, S: Semaphore> Rx<T, S> {
255255
self.inner.notify_rx_closed.notify_waiters();
256256
}
257257

258+
pub(crate) fn is_closed(&self) -> bool {
259+
// There two internal states that can represent a closed channel
260+
//
261+
// 1. When `close` is called.
262+
// In this case, the inner semaphore will be closed.
263+
//
264+
// 2. When all senders are dropped.
265+
// In this case, the semaphore remains unclosed, and the `index` in the list won't
266+
// reach the tail position. It is necessary to check the list if the last block is
267+
// `closed`.
268+
self.inner.semaphore.is_closed() || self.inner.tx_count.load(Acquire) == 0
269+
}
270+
271+
pub(crate) fn is_empty(&self) -> bool {
272+
self.inner.rx_fields.with(|rx_fields_ptr| {
273+
let rx_fields = unsafe { &*rx_fields_ptr };
274+
rx_fields.list.is_empty(&self.inner.tx)
275+
})
276+
}
277+
278+
pub(crate) fn len(&self) -> usize {
279+
self.inner.rx_fields.with(|rx_fields_ptr| {
280+
let rx_fields = unsafe { &*rx_fields_ptr };
281+
rx_fields.list.len(&self.inner.tx)
282+
})
283+
}
284+
258285
/// Receive the next value
259286
pub(crate) fn recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<T>> {
260287
use super::block::Read;

‎tokio/src/sync/mpsc/list.rs

+27
Original file line numberDiff line numberDiff line change
@@ -218,6 +218,15 @@ impl<T> Tx<T> {
218218
let _ = Box::from_raw(block.as_ptr());
219219
}
220220
}
221+
222+
pub(crate) fn is_closed(&self) -> bool {
223+
let tail = self.block_tail.load(Acquire);
224+
225+
unsafe {
226+
let tail_block = &*tail;
227+
tail_block.is_closed()
228+
}
229+
}
221230
}
222231

223232
impl<T> fmt::Debug for Tx<T> {
@@ -230,6 +239,24 @@ impl<T> fmt::Debug for Tx<T> {
230239
}
231240

232241
impl<T> Rx<T> {
242+
pub(crate) fn is_empty(&self, tx: &Tx<T>) -> bool {
243+
let block = unsafe { self.head.as_ref() };
244+
if block.has_value(self.index) {
245+
return false;
246+
}
247+
248+
// It is possible that a block has no value "now" but the list is still not empty.
249+
// To be sure, it is necessary to check the length of the list.
250+
self.len(tx) == 0
251+
}
252+
253+
pub(crate) fn len(&self, tx: &Tx<T>) -> usize {
254+
// When all the senders are dropped, there will be a last block in the tail position,
255+
// but it will be closed
256+
let tail_position = tx.tail_position.load(Acquire);
257+
tail_position - self.index - (tx.is_closed() as usize)
258+
}
259+
233260
/// Pops the next value off the queue.
234261
pub(crate) fn pop(&mut self, tx: &Tx<T>) -> Option<block::Read<T>> {
235262
// Advance `head`, if needed

‎tokio/src/sync/mpsc/unbounded.rs

+67
Original file line numberDiff line numberDiff line change
@@ -330,6 +330,73 @@ impl<T> UnboundedReceiver<T> {
330330
self.chan.close();
331331
}
332332

333+
/// Checks if a channel is closed.
334+
///
335+
/// This method returns `true` if the channel has been closed. The channel is closed
336+
/// when all [`UnboundedSender`] have been dropped, or when [`UnboundedReceiver::close`] is called.
337+
///
338+
/// [`UnboundedSender`]: crate::sync::mpsc::UnboundedSender
339+
/// [`UnboundedReceiver::close`]: crate::sync::mpsc::UnboundedReceiver::close
340+
///
341+
/// # Examples
342+
/// ```
343+
/// use tokio::sync::mpsc;
344+
///
345+
/// #[tokio::main]
346+
/// async fn main() {
347+
/// let (_tx, mut rx) = mpsc::unbounded_channel::<()>();
348+
/// assert!(!rx.is_closed());
349+
///
350+
/// rx.close();
351+
///
352+
/// assert!(rx.is_closed());
353+
/// }
354+
/// ```
355+
pub fn is_closed(&self) -> bool {
356+
self.chan.is_closed()
357+
}
358+
359+
/// Checks if a channel is empty.
360+
///
361+
/// This method returns `true` if the channel has no messages.
362+
///
363+
/// # Examples
364+
/// ```
365+
/// use tokio::sync::mpsc;
366+
///
367+
/// #[tokio::main]
368+
/// async fn main() {
369+
/// let (tx, rx) = mpsc::unbounded_channel();
370+
/// assert!(rx.is_empty());
371+
///
372+
/// tx.send(0).unwrap();
373+
/// assert!(!rx.is_empty());
374+
/// }
375+
///
376+
/// ```
377+
pub fn is_empty(&self) -> bool {
378+
self.chan.is_empty()
379+
}
380+
381+
/// Returns the number of messages in the channel.
382+
///
383+
/// # Examples
384+
/// ```
385+
/// use tokio::sync::mpsc;
386+
///
387+
/// #[tokio::main]
388+
/// async fn main() {
389+
/// let (tx, rx) = mpsc::unbounded_channel();
390+
/// assert_eq!(0, rx.len());
391+
///
392+
/// tx.send(0).unwrap();
393+
/// assert_eq!(1, rx.len());
394+
/// }
395+
/// ```
396+
pub fn len(&self) -> usize {
397+
self.chan.len()
398+
}
399+
333400
/// Polls to receive the next message on this channel.
334401
///
335402
/// This method returns:

‎tokio/src/sync/tests/loom_mpsc.rs

+34
Original file line numberDiff line numberDiff line change
@@ -188,3 +188,37 @@ fn try_recv() {
188188
}
189189
});
190190
}
191+
192+
#[test]
193+
fn len_nonzero_after_send() {
194+
loom::model(|| {
195+
let (send, recv) = mpsc::channel(10);
196+
let send2 = send.clone();
197+
198+
let join = thread::spawn(move || {
199+
block_on(send2.send("message2")).unwrap();
200+
});
201+
202+
block_on(send.send("message1")).unwrap();
203+
assert!(recv.len() != 0);
204+
205+
join.join().unwrap();
206+
});
207+
}
208+
209+
#[test]
210+
fn nonempty_after_send() {
211+
loom::model(|| {
212+
let (send, recv) = mpsc::channel(10);
213+
let send2 = send.clone();
214+
215+
let join = thread::spawn(move || {
216+
block_on(send2.send("message2")).unwrap();
217+
});
218+
219+
block_on(send.send("message1")).unwrap();
220+
assert!(!recv.is_empty());
221+
222+
join.join().unwrap();
223+
});
224+
}

0 commit comments

Comments
 (0)