Skip to main content

proxima_primitives/pipe/
batch_source.rs

1//! `BatchSource` — the multi-consumer (`&self`) owned-batch source.
2//!
3//! The algebra's third source flavor, filling the gap between the two
4//! single-consumer sources:
5//! - [`DrainSource`](crate::pipe::drain_source::DrainSource) — `&mut self`, zero-copy
6//!   borrow-visitor (the `*DK` frame relay).
7//! - [`FanIn`](crate::pipe::fan_in::FanIn)'s `UnpinPipe` sources — `&self`, owned async pull.
8//! - **`BatchSource`** — `&self`, owned batch pull, **safe for many concurrent
9//!   drainers**. A lock-free MPMC ring ([`proxima_core::ring::BoundedQueue`]) is
10//!   the canonical impl: several threads can each pull a disjoint batch at once,
11//!   which is what lets a drainer partition work across cores and lift the
12//!   single-drainer ceiling. The single-consumer sources cannot express that —
13//!   `&mut self` is one owner by construction.
14//!
15//! A consumer drives it by handing a scratch slice to fill: `drain_batch(&mut
16//! [Item]) -> count`. No alloc, no waker, no `&mut self` — legal at T0.
17
18use proxima_core::ring::{BoundedQueue, RingStorage};
19
20/// A source drained in owned batches through a shared `&self`, so multiple
21/// consumers may pull disjoint batches concurrently.
22pub trait BatchSource {
23    /// The owned item pulled from the source.
24    type Item;
25
26    /// Pull up to `out.len()` items FIFO into `out`, returning the count filled.
27    /// Safe to call from several threads at once — each caller gets a disjoint
28    /// batch (the impl linearises the dequeue).
29    fn drain_batch(&self, out: &mut [Self::Item]) -> usize;
30
31    /// A snapshot lower bound on the items available (a racing producer just
32    /// defers to the next pull). Consumers size their scratch buffer from it.
33    fn len(&self) -> usize;
34
35    /// True when the snapshot shows nothing to pull.
36    fn is_empty(&self) -> bool {
37        self.len() == 0
38    }
39}
40
41impl<R: RingStorage> BatchSource for BoundedQueue<R> {
42    type Item = R::Item;
43
44    #[inline]
45    fn drain_batch(&self, out: &mut [R::Item]) -> usize {
46        self.drain_into(out)
47    }
48
49    fn len(&self) -> usize {
50        BoundedQueue::len(self)
51    }
52}
53
54#[cfg(test)]
55#[allow(clippy::unwrap_used, clippy::expect_used)]
56mod tests {
57    use super::*;
58    use proxima_core::ring::{FailMode, StaticBoundedQueue};
59
60    #[test]
61    fn bounded_queue_is_a_batch_source() {
62        let queue = StaticBoundedQueue::<u32, 4>::new(FailMode::DropNewest);
63        for value in 1..=3 {
64            queue.enqueue(value);
65        }
66        // drive it through the trait, not the inherent method — proving the
67        // recorder's drain can be generic over any BatchSource.
68        fn drain_via_trait<Source: BatchSource<Item = u32>>(
69            source: &Source,
70            out: &mut [u32],
71        ) -> usize {
72            source.drain_batch(out)
73        }
74        assert_eq!(queue.len(), 3);
75        assert!(!queue.is_empty());
76        let mut out = [0u32; 8];
77        assert_eq!(drain_via_trait(&queue, &mut out), 3);
78        assert_eq!(&out[..3], &[1, 2, 3]);
79        assert!(queue.is_empty());
80    }
81}