Skip to main content

krabka_replicator/mm2/
mod.rs

1//! MirrorMaker-2 byte-compatible record codecs.
2mod checkpoint;
3mod heartbeat;
4mod offset_sync;
5
6pub use checkpoint::Checkpoint;
7pub use heartbeat::Heartbeat;
8pub use offset_sync::OffsetSync;
9
10use crate::error::ReplicatorError;
11
12pub(crate) const VERSION: i16 = 0;
13
14pub(crate) struct Writer(Vec<u8>);
15
16impl Writer {
17    /// Versioned buffer: emits the `[int16 version]` header first.
18    ///
19    /// The real JVM MM2 codecs version only the record *value* (via
20    /// `HEADER_SCHEMA`), never the key (`serializeKey` writes `KEY_SCHEMA`
21    /// alone). Use [`Writer::new`] for values, [`Writer::keyless`] for keys.
22    pub fn new() -> Self {
23        let mut w = Vec::new();
24        w.extend_from_slice(&VERSION.to_be_bytes());
25        Self(w)
26    }
27
28    /// Versionless buffer for record keys (no `[int16 version]` header), as
29    /// produced by the JVM MM2 `serializeKey()` path.
30    pub fn keyless() -> Self {
31        Self(Vec::new())
32    }
33
34    pub fn string(&mut self, s: &str) -> &mut Self {
35        let len = i16::try_from(s.len()).expect("string < 32k");
36        self.0.extend_from_slice(&len.to_be_bytes());
37        self.0.extend_from_slice(s.as_bytes());
38        self
39    }
40
41    pub fn i32(&mut self, v: i32) -> &mut Self {
42        self.0.extend_from_slice(&v.to_be_bytes());
43        self
44    }
45
46    pub fn i64(&mut self, v: i64) -> &mut Self {
47        self.0.extend_from_slice(&v.to_be_bytes());
48        self
49    }
50
51    #[must_use]
52    pub fn finish(&self) -> Vec<u8> {
53        self.0.clone()
54    }
55}
56
57pub(crate) struct Reader<'a> {
58    buf: &'a [u8],
59    pos: usize,
60}
61
62impl<'a> Reader<'a> {
63    /// Reader over a versioned buffer (record *value*): consumes and validates
64    /// the leading `[int16 version]` header.
65    pub fn new(buf: &'a [u8]) -> Result<Self, ReplicatorError> {
66        if buf.len() < 2 || i16::from_be_bytes([buf[0], buf[1]]) != VERSION {
67            return Err(ReplicatorError::Codec(
68                "bad/unsupported MM2 record version".into(),
69            ));
70        }
71        Ok(Self { buf, pos: 2 })
72    }
73
74    /// Reader over a versionless buffer (record *key*): the JVM MM2
75    /// `serializeKey()` path emits no version header.
76    pub fn keyless(buf: &'a [u8]) -> Self {
77        Self { buf, pos: 0 }
78    }
79
80    fn take(&mut self, n: usize) -> Result<&'a [u8], ReplicatorError> {
81        let end = self.pos + n;
82        if end > self.buf.len() {
83            return Err(ReplicatorError::Codec("short MM2 record".into()));
84        }
85        let s = &self.buf[self.pos..end];
86        self.pos = end;
87        Ok(s)
88    }
89
90    pub fn string(&mut self) -> Result<String, ReplicatorError> {
91        let raw = self.take(2)?;
92        // SAFETY: take(2) guarantees exactly 2 bytes.
93        let len_i16 = i16::from_be_bytes([raw[0], raw[1]]);
94        let len = usize::try_from(len_i16)
95            .map_err(|_| ReplicatorError::Codec("negative MM2 string length".into()))?;
96        Ok(String::from_utf8_lossy(self.take(len)?).into_owned())
97    }
98
99    pub fn i32(&mut self) -> Result<i32, ReplicatorError> {
100        let raw = self.take(4)?;
101        Ok(i32::from_be_bytes([raw[0], raw[1], raw[2], raw[3]]))
102    }
103
104    pub fn i64(&mut self) -> Result<i64, ReplicatorError> {
105        let raw = self.take(8)?;
106        Ok(i64::from_be_bytes([
107            raw[0], raw[1], raw[2], raw[3], raw[4], raw[5], raw[6], raw[7],
108        ]))
109    }
110}
111
112#[cfg(test)]
113mod tests {
114    use super::*;
115    use crate::ids::{DownstreamOffset, PartitionIndex, UpstreamOffset};
116
117    #[test]
118    fn heartbeat_key_value_bytes_exact() {
119        let hb = Heartbeat {
120            source: "us-east".into(),
121            target: "eu-west".into(),
122            timestamp_ms: 100,
123        };
124        // Key is versionless (JVM MM2 serializeKey writes no version header).
125        assert2::assert!(hb.key_bytes() == b"\x00\x07us-east\x00\x07eu-west".to_vec());
126        // value = version(i16=0x0000) + timestamp(i64=100=0x0000_0000_0000_0064) = 10 bytes
127        assert2::assert!(hb.value_bytes() == b"\x00\x00\x00\x00\x00\x00\x00\x00\x00\x64".to_vec());
128        let (k, v) = (hb.key_bytes(), hb.value_bytes());
129        assert2::assert!(Heartbeat::from_bytes(&k, &v).unwrap() == hb);
130    }
131
132    #[test]
133    fn offset_sync_bytes_roundtrip() {
134        let os = OffsetSync {
135            topic: "orders".into(),
136            partition: PartitionIndex(7),
137            upstream: UpstreamOffset(1000),
138            downstream: DownstreamOffset(742),
139        };
140        assert2::assert!(OffsetSync::from_bytes(&os.key_bytes(), &os.value_bytes()).unwrap() == os);
141        // Key is versionless (JVM MM2 serializeKey writes no version header).
142        assert2::assert!(os.key_bytes() == b"\x00\x06orders\x00\x00\x00\x07".to_vec());
143    }
144
145    #[test]
146    fn internal_topic_names_match_mm2() {
147        // MM2-compatible internal topic naming, pinned exactly.
148        assert2::assert!(Checkpoint::topic_name("us-east") == "us-east.checkpoints.internal");
149        assert2::assert!(OffsetSync::topic_name("us-east") == "mm2-offset-syncs.us-east.internal");
150    }
151
152    #[test]
153    fn checkpoint_bytes_roundtrip() {
154        let c = Checkpoint {
155            group: "analytics".into(),
156            topic: "orders".into(),
157            partition: PartitionIndex(7),
158            upstream: UpstreamOffset(1000),
159            downstream: DownstreamOffset(742),
160            metadata: String::new(),
161        };
162        assert2::assert!(Checkpoint::from_bytes(&c.key_bytes(), &c.value_bytes()).unwrap() == c);
163    }
164}