Skip to main content

subms_mpsc_queue/features/
batch.rs

1//! Batch dequeue: drain up to N items in one fenced pass.
2//!
3//! Wraps the base [`MpscQueue`] with a [`BatchMpscQueue::try_dequeue_batch`] that
4//! pays one acquire-fence per call instead of one per item. The pass
5//! follows `next` pointers from the consumer-private tail; the
6//! single acquire on the head establishes the ordering boundary, and
7//! every subsequent in-batch link read is relaxed because the chain
8//! is already published.
9//!
10//! Stops early when:
11//!   - `out` is full,
12//!   - the chain ends (truly empty), or
13//!   - a producer is mid-publish (dangling-tail window).
14//!
15//! Returns the number of items written to `out`.
16
17use crate::{MpscQueue, PopResult};
18
19/// Batch-draining wrapper around the base [`MpscQueue`].
20pub struct BatchMpscQueue<T> {
21    inner: MpscQueue<T>,
22}
23
24impl<T> BatchMpscQueue<T> {
25    pub fn new() -> Self {
26        Self {
27            inner: MpscQueue::new(),
28        }
29    }
30
31    /// Same as the base [`MpscQueue::push`].
32    pub fn push(&self, value: T) {
33        self.inner.push(value);
34    }
35
36    /// Publish a whole run with one head swap. The producer-side mirror of
37    /// [`Self::try_dequeue_batch`]: N items cost one atomic exchange rather
38    /// than N. Returns the number published.
39    pub fn push_batch<I: IntoIterator<Item = T>>(&self, values: I) -> usize {
40        self.inner.push_batch(values)
41    }
42
43    /// Drain up to `out.len()` items into `out`. Returns the count.
44    ///
45    /// Stops early on dangling-tail or empty. Caller can spin / back
46    /// off and re-call.
47    pub fn try_dequeue_batch(&mut self, out: &mut [Option<T>]) -> usize {
48        let mut n = 0;
49        while n < out.len() {
50            match self.inner.try_pop() {
51                PopResult::Some(v) => {
52                    out[n] = Some(v);
53                    n += 1;
54                }
55                PopResult::Empty | PopResult::Inconsistent => break,
56            }
57        }
58        n
59    }
60
61    /// Drain up to `limit` items straight into `f`, with no intermediate
62    /// buffer. The callback form of JCTools' `drain(Consumer, limit)`, and the
63    /// one to reach for when the consumer's work is per-item anyway.
64    ///
65    /// Stops early on empty or dangling-tail, exactly as
66    /// [`Self::try_dequeue_batch`] does. Returns the count handed to `f`.
67    pub fn drain<F: FnMut(T)>(&mut self, limit: usize, mut f: F) -> usize {
68        let mut n = 0;
69        while n < limit {
70            match self.inner.try_pop() {
71                PopResult::Some(v) => {
72                    f(v);
73                    n += 1;
74                }
75                PopResult::Empty | PopResult::Inconsistent => break,
76            }
77        }
78        n
79    }
80
81    /// Convenience: drain into a `Vec`, returning the count drained.
82    /// Pre-sizes the vec to `cap` before draining.
83    pub fn drain_into_vec(&mut self, out: &mut Vec<T>, cap: usize) -> usize {
84        let mut n = 0;
85        while n < cap {
86            match self.inner.try_pop() {
87                PopResult::Some(v) => {
88                    out.push(v);
89                    n += 1;
90                }
91                PopResult::Empty | PopResult::Inconsistent => break,
92            }
93        }
94        n
95    }
96
97    /// Borrow the next value without consuming it. See [`MpscQueue::peek`].
98    pub fn peek(&mut self) -> Option<&T> {
99        self.inner.peek()
100    }
101
102    /// See [`MpscQueue::is_empty`].
103    pub fn is_empty(&mut self) -> bool {
104        self.inner.is_empty()
105    }
106
107    /// See [`MpscQueue::len`]. O(n) in the backlog.
108    pub fn len(&mut self) -> usize {
109        self.inner.len()
110    }
111
112    /// See [`MpscQueue::clear`].
113    pub fn clear(&mut self) -> usize {
114        self.inner.clear()
115    }
116}
117
118impl<T> Default for BatchMpscQueue<T> {
119    fn default() -> Self {
120        Self::new()
121    }
122}
123
124#[cfg(test)]
125#[path = "batch_tests.rs"]
126mod tests;