krabka_client_streams/barrier/
reader.rs

1//! Reads published cuts from the internal `__barrier_state` topic.
2
3use std::{
4    collections::{BTreeMap, HashMap},
5    sync::Arc,
6};
7
8use crate::{
9    barrier::cut::{BARRIER_STATE_TOPIC, BarrierCut, decode_cut_value, decode_key},
10    error::StreamsClientError,
11    runtime::io::{FetchedRec, IsolationLevel, RecordFetcher},
12};
13
14/// Reads the cuts a barrier coordinator published to `__barrier_state`.
15///
16/// The reader drives the runtime's own [`RecordFetcher`] with a manual read of
17/// every partition of the topic, so it joins no consumer group and needs no new
18/// broker RPC. The first read starts at offset `0` of each partition. A later
19/// read continues where the previous one stopped, so it costs one fetch per
20/// partition of the records the coordinator published since.
21///
22/// The reader keeps complete cuts only. A partial cut names partitions that
23/// never received the epoch's marker, so a task that waited for one would wait
24/// forever. That is why the coordinator publishes partial cuts at all: a reader
25/// skips the epoch instead of stalling on markers that never arrive.
26///
27/// # Examples
28///
29/// ```no_run
30/// use std::sync::Arc;
31///
32/// use krabka_client_streams::{barrier::CutReader, runtime::RecordFetcher};
33///
34/// # async fn read(fetcher: Arc<dyn RecordFetcher>) -> Result<(), Box<dyn std::error::Error>> {
35/// let mut reader = CutReader::new(fetcher);
36/// if let Some(cut) = reader.latest_complete_cut("transactions").await? {
37///     println!("epoch {} at {:?}", cut.epoch, cut.offsets);
38/// }
39/// # Ok(()) }
40/// ```
41pub struct CutReader {
42    fetcher: Arc<dyn RecordFetcher>,
43    topic: String,
44    /// The partitions of the topic, discovered on the first read.
45    partitions: Option<Vec<i32>>,
46    /// The next offset to read in each partition.
47    offsets: HashMap<i32, i64>,
48    /// Complete cuts, as barrier group to epoch to cut.
49    cuts: HashMap<String, BTreeMap<i64, BarrierCut>>,
50}
51
52impl CutReader {
53    /// Creates a reader over the `__barrier_state` topic.
54    #[must_use]
55    pub fn new(fetcher: Arc<dyn RecordFetcher>) -> Self {
56        Self::with_topic(fetcher, BARRIER_STATE_TOPIC)
57    }
58
59    /// Creates a reader over another topic than `__barrier_state`.
60    #[must_use]
61    pub fn with_topic(fetcher: Arc<dyn RecordFetcher>, topic: impl Into<String>) -> Self {
62        Self {
63            fetcher,
64            topic: topic.into(),
65            partitions: None,
66            offsets: HashMap::new(),
67            cuts: HashMap::new(),
68        }
69    }
70
71    /// The complete cut of a barrier group with the highest epoch.
72    ///
73    /// The result is `None` when the topic holds no complete cut of that group.
74    ///
75    /// # Errors
76    ///
77    /// Returns the fetcher's error, or
78    /// [`StreamsClientError::BarrierFormat`](crate::StreamsClientError::BarrierFormat)
79    /// when a record does not match the frozen layout.
80    pub async fn latest_complete_cut(
81        &mut self,
82        group: &str,
83    ) -> Result<Option<BarrierCut>, StreamsClientError> {
84        self.refresh().await?;
85        Ok(self
86            .cuts
87            .get(group)
88            .and_then(|by_epoch| by_epoch.values().next_back())
89            .cloned())
90    }
91
92    /// The complete cuts of a barrier group above `epoch`, oldest first.
93    ///
94    /// Pass `-1` to get every complete cut. Partial cuts are left out.
95    ///
96    /// # Errors
97    ///
98    /// Returns the fetcher's error, or
99    /// [`StreamsClientError::BarrierFormat`](crate::StreamsClientError::BarrierFormat)
100    /// when a record does not match the frozen layout.
101    pub async fn complete_cuts_after(
102        &mut self,
103        group: &str,
104        epoch: i64,
105    ) -> Result<Vec<BarrierCut>, StreamsClientError> {
106        self.refresh().await?;
107        Ok(self.cuts.get(group).map_or_else(Vec::new, |by_epoch| {
108            by_epoch
109                .range((epoch + 1)..)
110                .map(|(_epoch, cut)| cut.clone())
111                .collect()
112        }))
113    }
114
115    /// The complete cut of one epoch, or `None` when the topic holds no such
116    /// cut.
117    ///
118    /// # Errors
119    ///
120    /// Returns the fetcher's error, or
121    /// [`StreamsClientError::BarrierFormat`](crate::StreamsClientError::BarrierFormat)
122    /// when a record does not match the frozen layout.
123    pub async fn complete_cut_at(
124        &mut self,
125        group: &str,
126        epoch: i64,
127    ) -> Result<Option<BarrierCut>, StreamsClientError> {
128        self.refresh().await?;
129        Ok(self
130            .cuts
131            .get(group)
132            .and_then(|by_epoch| by_epoch.get(&epoch))
133            .cloned())
134    }
135
136    /// Reads every partition of the topic up to its end.
137    ///
138    /// The first call discovers the partitions and starts at offset `0`. A later
139    /// call resumes at the offset the previous one left.
140    async fn refresh(&mut self) -> Result<(), StreamsClientError> {
141        if self.partitions.is_none() {
142            let mut partitions = self.fetcher.partitions(&self.topic).await?;
143            partitions.sort_unstable();
144            self.partitions = Some(partitions);
145        }
146        let partitions = self.partitions.clone().unwrap_or_default();
147        for partition in partitions {
148            let mut offset = *self.offsets.entry(partition).or_insert(0);
149            loop {
150                let batch = self
151                    .fetcher
152                    .fetch(
153                        &self.topic,
154                        partition,
155                        offset,
156                        IsolationLevel::ReadUncommitted,
157                    )
158                    .await?;
159                if batch.records.is_empty() {
160                    break;
161                }
162                let mut advanced = false;
163                for record in &batch.records {
164                    self.accept(record)?;
165                    let next = record.offset + 1;
166                    if next > offset {
167                        offset = next;
168                        advanced = true;
169                    }
170                }
171                // Stop when no record moved the offset, so a fetcher that
172                // replays one batch cannot spin here.
173                if !advanced {
174                    break;
175                }
176            }
177            self.offsets.insert(partition, offset);
178        }
179        Ok(())
180    }
181
182    /// Applies one record of the topic.
183    ///
184    /// A record of another kind is skipped. A tombstone and a partial cut both
185    /// drop the epoch, so neither can be adopted later.
186    fn accept(&mut self, record: &FetchedRec) -> Result<(), StreamsClientError> {
187        let Some(key) = record.key.as_deref() else {
188            return Ok(()); // A record with no key carries no barrier state.
189        };
190        let key = decode_key(key)?;
191        if !key.is_cut() {
192            return Ok(());
193        }
194        let value = record.value.as_deref().filter(|value| !value.is_empty());
195        let Some(value) = value else {
196            self.forget(&key.group, key.epoch);
197            return Ok(());
198        };
199        let cut = decode_cut_value(&key, value)?;
200        if cut.complete() {
201            self.cuts
202                .entry(key.group)
203                .or_default()
204                .insert(key.epoch, cut);
205        } else {
206            self.forget(&key.group, key.epoch);
207        }
208        Ok(())
209    }
210
211    fn forget(&mut self, group: &str, epoch: i64) {
212        if let Some(by_epoch) = self.cuts.get_mut(group) {
213            by_epoch.remove(&epoch);
214        }
215    }
216}
217
218#[cfg(test)]
219mod tests {
220    use std::sync::Mutex as StdMutex;
221
222    use assert2::{assert, check};
223    use bytes::Bytes;
224
225    use super::*;
226    use crate::{
227        barrier::testing::{CutRecord, cut_key, cut_value, key_of_kind},
228        runtime::io::FetchBatch,
229    };
230
231    /// A fetcher that serves one scripted batch per `(partition, offset)` of the
232    /// barrier state topic and an empty batch for anything else.
233    #[derive(Default)]
234    struct CutFetcher {
235        partitions: Vec<i32>,
236        batches: StdMutex<HashMap<(i32, i64), FetchBatch>>,
237    }
238
239    impl CutFetcher {
240        /// Serves `records` from partition `0` at offset `0`.
241        fn one_batch(records: Vec<FetchedRec>) -> Self {
242            let mut fetcher = Self {
243                partitions: vec![0],
244                batches: StdMutex::new(HashMap::new()),
245            };
246            fetcher.partitions = vec![0];
247            fetcher.script(0, 0, records);
248            fetcher
249        }
250
251        /// Adds the batch that one partition serves at one offset.
252        fn script(&self, partition: i32, offset: i64, records: Vec<FetchedRec>) {
253            self.batches
254                .lock()
255                .unwrap()
256                .insert((partition, offset), FetchBatch { records });
257        }
258    }
259
260    #[async_trait::async_trait]
261    impl RecordFetcher for CutFetcher {
262        async fn fetch(
263            &self,
264            _topic: &str,
265            partition: i32,
266            offset: i64,
267            _isolation: IsolationLevel,
268        ) -> Result<FetchBatch, StreamsClientError> {
269            Ok(self
270                .batches
271                .lock()
272                .unwrap()
273                .remove(&(partition, offset))
274                .unwrap_or_default())
275        }
276
277        async fn partitions(&self, _topic: &str) -> Result<Vec<i32>, StreamsClientError> {
278            Ok(self.partitions.clone())
279        }
280    }
281
282    fn record(offset: i64, key: Bytes, value: Option<Bytes>) -> FetchedRec {
283        FetchedRec {
284            offset,
285            key: Some(key),
286            value,
287            timestamp: -1,
288        }
289    }
290
291    fn cut_record(offset: i64, cut: &CutRecord) -> FetchedRec {
292        record(offset, cut_key(cut), Some(cut_value(cut)))
293    }
294
295    #[tokio::test]
296    async fn reads_complete_cuts_in_epoch_order() {
297        let first = CutRecord::simple("g", 1, &[("in", 0, 10)]);
298        let second = CutRecord::simple("g", 2, &[("in", 0, 20)]);
299        let other_group = CutRecord::simple("other", 3, &[("in", 0, 30)]);
300        let fetcher = Arc::new(CutFetcher::one_batch(vec![
301            cut_record(0, &first),
302            cut_record(1, &second),
303            cut_record(2, &other_group),
304        ]));
305        let mut reader = CutReader::new(fetcher);
306
307        check!(
308            reader.complete_cuts_after("g", -1).await.unwrap()
309                == vec![first.expected(), second.expected()]
310        );
311        check!(reader.complete_cuts_after("g", 1).await.unwrap() == vec![second.expected()]);
312        check!(reader.latest_complete_cut("g").await.unwrap() == Some(second.expected()));
313        check!(reader.complete_cut_at("g", 1).await.unwrap() == Some(first.expected()));
314        check!(reader.complete_cut_at("g", 9).await.unwrap() == None);
315        check!(reader.latest_complete_cut("absent").await.unwrap() == None);
316    }
317
318    #[tokio::test]
319    async fn never_adopts_a_partial_cut() {
320        let partial = CutRecord::simple("g", 1, &[("in", 0, 10)]).partial(&[("in", 1)]);
321        let complete = CutRecord::simple("g", 2, &[("in", 0, 20), ("in", 1, 7)]);
322        let fetcher = Arc::new(CutFetcher::one_batch(vec![
323            cut_record(0, &partial),
324            cut_record(1, &complete),
325        ]));
326        let mut reader = CutReader::new(fetcher);
327
328        check!(reader.complete_cuts_after("g", -1).await.unwrap() == vec![complete.expected()]);
329        check!(reader.complete_cut_at("g", 1).await.unwrap() == None);
330    }
331
332    #[tokio::test]
333    async fn a_republished_partial_cut_drops_the_epoch() {
334        let complete = CutRecord::simple("g", 1, &[("in", 0, 10)]);
335        let partial = CutRecord::simple("g", 1, &[("in", 0, 10)]).partial(&[("in", 1)]);
336        let fetcher = Arc::new(CutFetcher::one_batch(vec![
337            cut_record(0, &complete),
338            cut_record(1, &partial),
339        ]));
340        let mut reader = CutReader::new(fetcher);
341        check!(reader.complete_cuts_after("g", -1).await.unwrap() == vec![]);
342    }
343
344    #[tokio::test]
345    async fn skips_other_kinds_tombstones_and_keyless_records() {
346        let cut = CutRecord::simple("g", 1, &[("in", 0, 10)]);
347        let fetcher = Arc::new(CutFetcher::one_batch(vec![
348            record(
349                0,
350                key_of_kind(0, "g", -1),
351                Some(Bytes::from_static(b"junk")),
352            ),
353            record(1, key_of_kind(1, "g", 1), Some(Bytes::from_static(b"junk"))),
354            FetchedRec {
355                offset: 2,
356                key: None,
357                value: Some(Bytes::from_static(b"junk")),
358                timestamp: -1,
359            },
360            cut_record(3, &cut),
361            record(4, cut_key(&cut), None),
362        ]));
363        let mut reader = CutReader::new(fetcher);
364        // The group record, the injection start, and the keyless record are
365        // skipped; the tombstone drops the cut the batch published before it.
366        check!(reader.complete_cuts_after("g", -1).await.unwrap() == vec![]);
367    }
368
369    #[tokio::test]
370    async fn resumes_at_the_offset_the_previous_read_left() {
371        let first = CutRecord::simple("g", 1, &[("in", 0, 10)]);
372        let second = CutRecord::simple("g", 2, &[("in", 0, 20)]);
373        let fetcher = Arc::new(CutFetcher::one_batch(vec![cut_record(0, &first)]));
374        let mut reader = CutReader::new(Arc::clone(&fetcher) as Arc<dyn RecordFetcher>);
375        check!(reader.complete_cuts_after("g", -1).await.unwrap() == vec![first.expected()]);
376
377        // The coordinator publishes another cut at offset 1. The next read picks
378        // it up from there and keeps the cut the first read already saw.
379        fetcher.script(0, 1, vec![cut_record(1, &second)]);
380        check!(
381            reader.complete_cuts_after("g", -1).await.unwrap()
382                == vec![first.expected(), second.expected()]
383        );
384    }
385
386    #[tokio::test]
387    async fn reads_every_partition_of_the_topic() {
388        let first = CutRecord::simple("g", 1, &[("in", 0, 10)]);
389        let second = CutRecord::simple("g", 2, &[("in", 0, 20)]);
390        let fetcher = Arc::new(CutFetcher {
391            partitions: vec![0, 1],
392            batches: StdMutex::new(HashMap::new()),
393        });
394        fetcher.script(0, 0, vec![cut_record(0, &first)]);
395        fetcher.script(1, 0, vec![cut_record(0, &second)]);
396        let mut reader = CutReader::new(fetcher);
397        check!(
398            reader.complete_cuts_after("g", -1).await.unwrap()
399                == vec![first.expected(), second.expected()]
400        );
401    }
402
403    #[tokio::test]
404    async fn a_malformed_record_fails_the_read() {
405        let fetcher = Arc::new(CutFetcher::one_batch(vec![record(
406            0,
407            Bytes::from_static(b"\x00\x00"),
408            Some(Bytes::from_static(b"x")),
409        )]));
410        let mut reader = CutReader::new(fetcher);
411        assert!(let Err(error) = reader.latest_complete_cut("g").await);
412        check!(error.to_string().contains("truncated barrier state key"));
413    }
414}
krabka_client_streams/barrier/reader.rs

krabka_client_streams/barrier/
reader.rs
