krabka_client_streams/barrier/
cut.rs

1//! The `__barrier_state` record layouts and the decoded cut.
2//!
3//! A krabka broker injects an epoch-stamped marker into every partition of a
4//! barrier group and then publishes the resulting cut to the internal topic
5//! [`BARRIER_STATE_TOPIC`]. The topic carries three record kinds, and the key
6//! discriminates them. Only the cut kind matters to a client, so
7//! [`decode_barrier_cut`] returns `None` for a group definition, for an
8//! injection start, and for a tombstone.
9//!
10//! The layouts below are frozen. The broker, `krabka-streams-rs`,
11//! `krabka-streams-java`, and `krabka-streams-go` all write and read them:
12//!
13//! ```text
14//! key:
15//!   version  i16 = 0
16//!   kind     i16               0 group, 1 injection start, 2 cut
17//!   group    string
18//!   epoch    i64               -1 for kind 0
19//!
20//! cut value:
21//!   version      i16 = 0
22//!   triggered_at i64
23//!   completed_at i64
24//!   status       i8            0 complete, 1 partial
25//!   topics       i32 [ topic string | partitions i32 [ partition i32 | offset i64 ] ]
26//!   missing      i32 [ topic string | partition i32 ]
27//! ```
28//!
29//! A `string` is an `i16` byte length and then UTF-8 bytes. Every integer is
30//! big-endian, and an `i32` count precedes each array.
31
32use std::collections::{BTreeMap, BTreeSet};
33
34use crate::{
35    error::StreamsClientError, membership::TopicPartition, runtime::iqv2::request::Position,
36};
37
38/// The internal topic a krabka broker publishes barrier state to.
39pub const BARRIER_STATE_TOPIC: &str = "__barrier_state";
40
41/// The record version that both the key and the cut value carry.
42const RECORD_VERSION: i16 = 0;
43
44/// The key kind of a cut record. Kinds `0` and `1` are the group definition and
45/// the injection start, and a client skips both.
46const CUT_KIND: i16 = 2;
47
48/// The part of a record a decode failed on. It names the
49/// [`StreamsClientError::BarrierFormat`] `part` field.
50const KEY_PART: &str = "key";
51const VALUE_PART: &str = "cut value";
52
53/// Whether the coordinator reached every partition of a barrier group.
54///
55/// The coordinator publishes both outcomes. A complete cut names a marker
56/// offset for every partition of the group. A partial cut leaves partitions
57/// out, and those partitions never receive that epoch's marker, so a task that
58/// waits for one waits forever. Only a complete cut is alignable.
59#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
60pub enum CutStatus {
61    /// Every partition of the barrier group received the epoch's marker.
62    Complete,
63    /// One or more partitions never received the epoch's marker.
64    Partial,
65}
66
67impl CutStatus {
68    /// The wire code the coordinator writes for this status.
69    #[must_use]
70    pub fn code(self) -> i8 {
71        match self {
72            Self::Complete => 0,
73            Self::Partial => 1,
74        }
75    }
76
77    /// Maps a wire code onto a status.
78    ///
79    /// # Errors
80    ///
81    /// Returns [`StreamsClientError::BarrierFormat`] when no status carries the
82    /// code.
83    pub fn from_code(code: i8) -> Result<Self, StreamsClientError> {
84        match code {
85            0 => Ok(Self::Complete),
86            1 => Ok(Self::Partial),
87            other => Err(malformed(
88                VALUE_PART,
89                format!("unknown barrier cut status {other}"),
90            )),
91        }
92    }
93}
94
95/// One epoch's cut across every partition of a barrier group.
96///
97/// A cut names an exact position in every input at once: the offset of the
98/// epoch's marker in each partition. The marker is a Kafka control record, so no
99/// consumer receives it and the cut offset never holds a data record. The
100/// records before the cut of a partition are exactly the records with a lower
101/// offset.
102///
103/// [`CutReader`](super::CutReader) decodes cuts from [`BARRIER_STATE_TOPIC`].
104#[derive(Debug, Clone, PartialEq, Eq)]
105pub struct BarrierCut {
106    /// The barrier group the coordinator injected markers for.
107    pub group: String,
108    /// The epoch that names this cut inside the group. A coordinator never
109    /// reuses one.
110    pub epoch: i64,
111    /// The time the injection started, in milliseconds since the Unix epoch.
112    pub triggered_at: i64,
113    /// The time the coordinator published the cut, in milliseconds since the
114    /// Unix epoch.
115    pub completed_at: i64,
116    /// Whether every partition received the epoch's marker.
117    pub status: CutStatus,
118    /// The marker offset of every partition that received one, as topic to
119    /// partition to offset.
120    pub offsets: Position,
121    /// The partitions that never received the epoch's marker. It is empty for a
122    /// complete cut.
123    pub missing: BTreeSet<TopicPartition>,
124}
125
126impl BarrierCut {
127    /// Whether every partition of the group received the epoch's marker.
128    ///
129    /// A runner aligns on a complete cut only.
130    #[must_use]
131    pub fn complete(&self) -> bool {
132        self.status == CutStatus::Complete
133    }
134
135    /// The marker offset of one partition, or `None` when the cut does not name
136    /// it.
137    #[must_use]
138    pub fn offset(&self, topic: &str, partition: i32) -> Option<i64> {
139        self.offsets.offset(topic, partition)
140    }
141
142    /// Whether a partition delivered everything below the cut.
143    ///
144    /// `position` is the next offset the task reads. A partition the cut does
145    /// not name holds nothing back, so it always reports `true`.
146    #[must_use]
147    pub fn reached(&self, topic: &str, partition: i32, position: i64) -> bool {
148        self.offset(topic, partition)
149            .is_none_or(|offset| position >= offset)
150    }
151
152    /// Every partition the cut names, in topic and then partition order.
153    #[must_use]
154    pub fn partitions(&self) -> Vec<TopicPartition> {
155        self.offsets
156            .0
157            .iter()
158            .flat_map(|(topic, partitions)| {
159                partitions.keys().map(|partition| TopicPartition {
160                    topic: topic.clone(),
161                    partition: *partition,
162                })
163            })
164            .collect()
165    }
166}
167
168/// A decoded key of the barrier state topic.
169#[derive(Debug, Clone, PartialEq, Eq)]
170pub(crate) struct BarrierKey {
171    pub kind: i16,
172    pub group: String,
173    pub epoch: i64,
174}
175
176impl BarrierKey {
177    /// Whether this key names a cut record.
178    pub(crate) fn is_cut(&self) -> bool {
179        self.kind == CUT_KIND
180    }
181}
182
183fn malformed(part: &'static str, message: String) -> StreamsClientError {
184    StreamsClientError::BarrierFormat { part, message }
185}
186
187/// Decodes one record of [`BARRIER_STATE_TOPIC`] into a cut.
188///
189/// The result is `None` for a record that carries no cut: a group definition, an
190/// injection start, and a tombstone of any kind.
191///
192/// # Errors
193///
194/// Returns [`StreamsClientError::BarrierFormat`] when the key or the value does
195/// not match the frozen layout. The `part` field names which of the two failed.
196pub fn decode_barrier_cut(
197    key: &[u8],
198    value: Option<&[u8]>,
199) -> Result<Option<BarrierCut>, StreamsClientError> {
200    let key = decode_key(key)?;
201    if !key.is_cut() {
202        return Ok(None);
203    }
204    match value {
205        None => Ok(None),
206        Some(value) => decode_cut_value(&key, value).map(Some),
207    }
208}
209
210/// Decodes the key of any barrier state record.
211pub(crate) fn decode_key(data: &[u8]) -> Result<BarrierKey, StreamsClientError> {
212    let mut reader = Reader::new(data, KEY_PART);
213    let version = reader.i16()?;
214    if version != RECORD_VERSION {
215        return Err(malformed(
216            KEY_PART,
217            format!("unsupported barrier state key version {version}"),
218        ));
219    }
220    let kind = reader.i16()?;
221    let group = reader.string()?;
222    let epoch = reader.i64()?;
223    reader.finish()?;
224    Ok(BarrierKey { kind, group, epoch })
225}
226
227/// Decodes the value of a cut record, with the group and the epoch taken from
228/// the key.
229pub(crate) fn decode_cut_value(
230    key: &BarrierKey,
231    data: &[u8],
232) -> Result<BarrierCut, StreamsClientError> {
233    let mut reader = Reader::new(data, VALUE_PART);
234    let version = reader.i16()?;
235    if version != RECORD_VERSION {
236        return Err(malformed(
237            VALUE_PART,
238            format!("unsupported barrier state cut value version {version}"),
239        ));
240    }
241    let triggered_at = reader.i64()?;
242    let completed_at = reader.i64()?;
243    let status = CutStatus::from_code(reader.i8()?)?;
244    let offsets = reader.cut_offsets()?;
245    let missing = reader.missing_partitions()?;
246    reader.finish()?;
247    Ok(BarrierCut {
248        group: key.group.clone(),
249        epoch: key.epoch,
250        triggered_at,
251        completed_at,
252        status,
253        offsets,
254        missing,
255    })
256}
257
258/// Reads the big-endian layout of the barrier state topic.
259///
260/// Every integer is signed, and a string is an `i16` byte length and then UTF-8
261/// bytes.
262struct Reader<'a> {
263    data: &'a [u8],
264    part: &'static str,
265}
266
267impl<'a> Reader<'a> {
268    fn new(data: &'a [u8], part: &'static str) -> Self {
269        Self { data, part }
270    }
271
272    fn truncated(&self) -> StreamsClientError {
273        malformed(self.part, format!("truncated barrier state {}", self.part))
274    }
275
276    /// Fails when bytes remain after the last field.
277    fn finish(&self) -> Result<(), StreamsClientError> {
278        if self.data.is_empty() {
279            return Ok(());
280        }
281        Err(malformed(
282            self.part,
283            format!("trailing bytes in barrier state {}", self.part),
284        ))
285    }
286
287    fn take(&mut self, count: usize) -> Result<&'a [u8], StreamsClientError> {
288        let Some((head, rest)) = self.data.split_at_checked(count) else {
289            return Err(self.truncated());
290        };
291        self.data = rest;
292        Ok(head)
293    }
294
295    fn i8(&mut self) -> Result<i8, StreamsClientError> {
296        Ok(self.take(1)?[0].cast_signed())
297    }
298
299    fn i16(&mut self) -> Result<i16, StreamsClientError> {
300        let bytes = self.take(2)?;
301        Ok(i16::from_be_bytes(bytes.try_into().expect("two bytes")))
302    }
303
304    fn i32(&mut self) -> Result<i32, StreamsClientError> {
305        let bytes = self.take(4)?;
306        Ok(i32::from_be_bytes(bytes.try_into().expect("four bytes")))
307    }
308
309    fn i64(&mut self) -> Result<i64, StreamsClientError> {
310        let bytes = self.take(8)?;
311        Ok(i64::from_be_bytes(bytes.try_into().expect("eight bytes")))
312    }
313
314    fn string(&mut self) -> Result<String, StreamsClientError> {
315        let length = self.i16()?;
316        let length = usize::try_from(length).map_err(|_error| {
317            malformed(
318                self.part,
319                format!(
320                    "negative string length {length} in barrier state {}",
321                    self.part
322                ),
323            )
324        })?;
325        let bytes = self.take(length)?;
326        String::from_utf8(bytes.to_vec()).map_err(|_error| {
327            malformed(
328                self.part,
329                format!("non-UTF-8 string in barrier state {}", self.part),
330            )
331        })
332    }
333
334    /// Reads an `i32` array length.
335    ///
336    /// A negative length is malformed. A length past the remaining bytes fails
337    /// on the first entry, so no bogus count allocates.
338    fn count(&mut self) -> Result<i32, StreamsClientError> {
339        let count = self.i32()?;
340        if count < 0 {
341            return Err(malformed(
342                self.part,
343                format!(
344                    "negative array length {count} in barrier state {}",
345                    self.part
346                ),
347            ));
348        }
349        Ok(count)
350    }
351
352    fn cut_offsets(&mut self) -> Result<Position, StreamsClientError> {
353        let mut offsets: BTreeMap<String, BTreeMap<i32, i64>> = BTreeMap::new();
354        for _ in 0..self.count()? {
355            let topic = self.string()?;
356            let partitions = self.count()?;
357            let entry = offsets.entry(topic).or_default();
358            for _ in 0..partitions {
359                let partition = self.i32()?;
360                entry.insert(partition, self.i64()?);
361            }
362        }
363        Ok(Position(offsets))
364    }
365
366    fn missing_partitions(&mut self) -> Result<BTreeSet<TopicPartition>, StreamsClientError> {
367        let mut missing = BTreeSet::new();
368        for _ in 0..self.count()? {
369            let topic = self.string()?;
370            missing.insert(TopicPartition {
371                topic,
372                partition: self.i32()?,
373            });
374        }
375        Ok(missing)
376    }
377}
378
379#[cfg(test)]
380mod tests {
381    use assert2::{assert, check};
382
383    use super::*;
384    use crate::barrier::testing::{CutRecord, cut_key, cut_value};
385
386    #[test]
387    fn decodes_a_complete_cut_from_its_frozen_bytes() {
388        let record = CutRecord {
389            group: "transactions".to_owned(),
390            epoch: 7,
391            triggered_at: 100,
392            completed_at: 140,
393            status: CutStatus::Complete,
394            offsets: vec![("orders", 0, 12), ("orders", 1, 5), ("payments", 0, 31)],
395            missing: vec![],
396        };
397        let decoded = decode_barrier_cut(&cut_key(&record), Some(&cut_value(&record)))
398            .unwrap()
399            .unwrap();
400        check!(decoded == record.expected());
401    }
402
403    #[test]
404    fn decodes_a_partial_cut_with_its_missing_partitions() {
405        let record = CutRecord {
406            group: "transactions".to_owned(),
407            epoch: 8,
408            triggered_at: 200,
409            completed_at: 260,
410            status: CutStatus::Partial,
411            offsets: vec![("orders", 0, 40)],
412            missing: vec![("orders", 1)],
413        };
414        let decoded = decode_barrier_cut(&cut_key(&record), Some(&cut_value(&record)))
415            .unwrap()
416            .unwrap();
417        check!(decoded == record.expected());
418        check!(!decoded.complete());
419    }
420
421    #[test]
422    fn skips_other_kinds_and_tombstones() {
423        let record = CutRecord::simple("g", 3, &[("in", 0, 1)]);
424        // Kind 0 is a group definition and kind 1 is an injection start.
425        for kind in [0i16, 1] {
426            let key = crate::barrier::testing::key_of_kind(kind, "g", 3);
427            check!(decode_barrier_cut(&key, Some(&cut_value(&record))).unwrap() == None);
428        }
429        check!(decode_barrier_cut(&cut_key(&record), None).unwrap() == None);
430    }
431
432    #[test]
433    fn cut_answers_offset_reached_and_partitions() {
434        let record = CutRecord::simple("g", 1, &[("in", 0, 10), ("in", 1, 4)]);
435        let cut = record.expected();
436        check!(cut.offset("in", 0) == Some(10));
437        check!(cut.offset("in", 2) == None);
438        check!(!cut.reached("in", 0, 9));
439        check!(cut.reached("in", 0, 10));
440        check!(cut.reached("other", 0, 0));
441        check!(
442            cut.partitions()
443                == vec![
444                    TopicPartition {
445                        topic: "in".to_owned(),
446                        partition: 0
447                    },
448                    TopicPartition {
449                        topic: "in".to_owned(),
450                        partition: 1
451                    },
452                ]
453        );
454    }
455
456    #[test]
457    fn rejects_malformed_records() {
458        let record = CutRecord::simple("g", 1, &[("in", 0, 10)]);
459        let key = cut_key(&record);
460        let value = cut_value(&record);
461
462        let mut bad_key_version = key.to_vec();
463        bad_key_version[1] = 9;
464        assert!(let Err(error) = decode_barrier_cut(&bad_key_version, Some(&value)));
465        check!(
466            error.to_string()
467                == "malformed barrier state key: unsupported barrier state key version 9"
468        );
469
470        assert!(let Err(error) = decode_barrier_cut(&key[..key.len() - 1], Some(&value)));
471        check!(error.to_string().contains("truncated barrier state key"));
472
473        let mut trailing_key = key.to_vec();
474        trailing_key.push(0);
475        assert!(let Err(error) = decode_barrier_cut(&trailing_key, Some(&value)));
476        check!(
477            error
478                .to_string()
479                .contains("trailing bytes in barrier state key")
480        );
481
482        let mut bad_status = value.to_vec();
483        bad_status[18] = 7;
484        assert!(let Err(error) = decode_barrier_cut(&key, Some(&bad_status)));
485        check!(error.to_string().contains("unknown barrier cut status 7"));
486
487        assert!(let Err(error) = decode_barrier_cut(&key, Some(&value[..value.len() - 2])));
488        check!(
489            error
490                .to_string()
491                .contains("truncated barrier state cut value")
492        );
493
494        let mut negative_count = value.to_vec();
495        // The topics array length sits right after version, the two timestamps,
496        // and the status byte.
497        negative_count[19] = 0xFF;
498        assert!(let Err(error) = decode_barrier_cut(&key, Some(&negative_count)));
499        check!(error.to_string().contains("negative array length"));
500    }
501
502    #[test]
503    fn status_codes_round_trip() {
504        for status in [CutStatus::Complete, CutStatus::Partial] {
505            check!(CutStatus::from_code(status.code()).unwrap() == status);
506        }
507        assert!(let Err(error) = CutStatus::from_code(2));
508        check!(error.to_string().contains("unknown barrier cut status 2"));
509    }
510}
krabka_client_streams/barrier/cut.rs

krabka_client_streams/barrier/
cut.rs
