etl_core/sink/
queue.rs

1//! Bounded per-shard chunk queues: the pipeline→sink handoff.
2//!
3//! Senders live on pipeline threads and only ever `try_send` (the
4//! backpressure invariant — never block a poll loop); receivers live in
5//! shard worker tasks on the I/O runtime.
6
7use super::EncodedChunk;
8use tokio::sync::mpsc;
9
10/// A rejected chunk, handed back so the terminal stage can park it and
11/// report `Blocked` upstream.
12#[derive(Debug)]
13pub struct ChunkSendError(pub EncodedChunk);
14
15/// Sending side of every shard queue, shared by pipeline threads.
16#[derive(Clone, Debug)]
17pub struct ShardQueues {
18    senders: Vec<mpsc::Sender<EncodedChunk>>,
19    capacity: usize,
20}
21
22impl ShardQueues {
23    /// Number of shards.
24    #[must_use]
25    pub fn num_shards(&self) -> usize {
26        self.senders.len()
27    }
28
29    /// Configured per-shard capacity (in chunks).
30    #[must_use]
31    pub fn capacity(&self) -> usize {
32        self.capacity
33    }
34
35    /// Non-blocking send to `shard`. On `Err` the chunk comes back and the
36    /// caller applies backpressure. A closed queue (sink shut down) also
37    /// returns the chunk; the driver observes shutdown separately.
38    pub fn try_send(&self, shard: usize, chunk: EncodedChunk) -> Result<(), ChunkSendError> {
39        match self.senders[shard].try_send(chunk) {
40            Ok(()) => Ok(()),
41            Err(mpsc::error::TrySendError::Full(c) | mpsc::error::TrySendError::Closed(c)) => {
42                Err(ChunkSendError(c))
43            }
44        }
45    }
46
47    /// Whether every shard queue is below `ratio` of its capacity —
48    /// the resume condition the backpressure controller asks about.
49    #[must_use]
50    pub fn all_below(&self, ratio: f64) -> bool {
51        let threshold = (self.capacity as f64 * ratio) as usize;
52        self.senders
53            .iter()
54            .all(|s| self.capacity - s.capacity() <= threshold)
55    }
56}
57
58/// Build the queues: one bounded channel per shard. Returns the shared
59/// sender handle and the per-shard receivers for the workers.
60#[must_use]
61pub fn shard_queues(
62    num_shards: usize,
63    capacity: usize,
64) -> (ShardQueues, Vec<mpsc::Receiver<EncodedChunk>>) {
65    assert!(num_shards > 0, "a sink needs at least one shard");
66    assert!(capacity > 0, "shard queues need non-zero capacity");
67    let (senders, receivers): (Vec<_>, Vec<_>) =
68        (0..num_shards).map(|_| mpsc::channel(capacity)).unzip();
69    (ShardQueues { senders, capacity }, receivers)
70}
71
72#[cfg(all(test, not(loom)))]
73mod tests {
74    use super::*;
75    use bytes::Bytes;
76
77    fn chunk() -> EncodedChunk {
78        EncodedChunk {
79            oldest_ingest: std::time::Instant::now(),
80            oldest_event_ms: 0,
81            frame: Bytes::from_static(b"x"),
82            rows: 1,
83            acks: crate::checkpoint::AckSet::new(),
84        }
85    }
86
87    #[test]
88    fn try_send_never_blocks_and_returns_the_chunk_when_full() {
89        let (q, mut rx) = shard_queues(1, 2);
90        assert!(q.try_send(0, chunk()).is_ok());
91        assert!(q.try_send(0, chunk()).is_ok());
92        let ChunkSendError(returned) = q.try_send(0, chunk()).unwrap_err();
93        assert_eq!(returned.rows, 1);
94        assert!(rx[0].try_recv().is_ok());
95        assert!(q.try_send(0, chunk()).is_ok(), "capacity freed");
96        let _ = rx;
97    }
98
99    #[test]
100    fn dropping_a_receiver_with_queued_chunks_fails_their_acks() {
101        use crate::checkpoint::{AckRef, AckStatus};
102        let (q, rx) = shard_queues(1, 4);
103        let (ack, ack_rx) = AckRef::test_pair();
104        let mut c = chunk();
105        c.acks = vec![ack.clone()].into();
106        drop(ack);
107        q.try_send(0, c).expect("queued");
108        drop(rx); // sink torn down with the chunk still queued
109        assert_eq!(
110            ack_rx.try_recv().expect("resolved").status,
111            AckStatus::Failed,
112            "chunks lost in a dropped queue must fail their batches"
113        );
114    }
115
116    #[test]
117    fn closed_queue_hands_the_chunk_back() {
118        let (q, rx) = shard_queues(1, 1);
119        drop(rx);
120        assert!(q.try_send(0, chunk()).is_err());
121    }
122
123    #[test]
124    fn all_below_reflects_fill_ratio() {
125        let (q, _rx) = shard_queues(2, 4);
126        assert!(q.all_below(0.5));
127        q.try_send(0, chunk()).unwrap();
128        q.try_send(0, chunk()).unwrap();
129        q.try_send(0, chunk()).unwrap();
130        assert!(!q.all_below(0.5), "shard 0 is 75% full");
131    }
132}
etl_core/sink/queue.rs

etl_core/sink/
queue.rs
