krabka_records_legacy/set/
encode.rs

1//! Encoding a v0/v1 `MessageSet`, flat or wrapped in a single compressed
2//! outer message.
3//!
4//! The two entry points are separate because the wire layouts are: a flat set
5//! writes one outer message per record, while a compressed set writes exactly
6//! one outer message whose value is a complete inner `MessageSet`. The
7//! wrapper's own offset and timestamp, and whether the inner offsets are
8//! absolute or relative, are the KIP-32 conventions that differ between v0
9//! and v1.
10
11use {
12    super::ParsedRecord,
13    crate::{
14        error::LegacyRecordsError,
15        message::{Magic, Message, attrs_with_compression},
16    },
17    bytes::{BufMut, BytesMut},
18    krabka_compression::CompressionType,
19};
20
21fn message_for_record(record: ParsedRecord, magic: Magic) -> Message {
22    Message {
23        magic,
24        attributes: 0,
25        timestamp: match magic {
26            Magic::V0 => None,
27            Magic::V1 => Some(record.timestamp.unwrap_or(-1)),
28        },
29        key: record.key,
30        value: record.value,
31    }
32}
33
34/// Encode a flat `MessageSet` of magic `magic` into `buf`.
35///
36/// This function writes one outer message per record. Use it to emit an
37/// uncompressed batch.
38pub fn encode_flat_message_set<B: BufMut, I: IntoIterator<Item = ParsedRecord>>(
39    records: I,
40    magic: Magic,
41    buf: &mut B,
42) {
43    for r in records {
44        let offset = r.offset;
45        let msg = message_for_record(r, magic);
46        let msg_len = msg.encoded_len();
47        buf.put_i64(offset.0);
48        // Safe: legacy messages are well-bounded; capping at i32::MAX is
49        // sufficient for any realistic batch.
50        buf.put_i32(i32::try_from(msg_len).unwrap_or(i32::MAX));
51        msg.encode_into(buf);
52    }
53}
54
55/// Encode a `MessageSet` wrapped in a single compressed outer message.
56///
57/// The inner set is uncompressed and contains one message per record. It
58/// follows the KIP-32 conventions: v1 inner offsets are relative to the
59/// first record, preserving gaps left by compaction.
60/// # Errors
61/// Returns `Malformed` if `codec` is `CompressionType::Zstd`, which v0/v1 cannot represent.
62/// Returns `Malformed` if a v1 relative offset exceeds `i64`.
63/// Returns a compression error if the codec cannot compress the inner set.
64pub fn encode_compressed_message_set<B: BufMut>(
65    records: &[ParsedRecord],
66    magic: Magic,
67    codec: CompressionType,
68    buf: &mut B,
69) -> Result<(), LegacyRecordsError> {
70    assert2::assert!(
71        codec != CompressionType::None,
72        "use encode_flat_message_set for uncompressed"
73    );
74    if matches!(codec, CompressionType::Zstd) {
75        return Err(LegacyRecordsError::Malformed(
76            "zstd compression not representable in v0/v1".into(),
77        ));
78    }
79    let (Some(first), Some(last)) = (records.first(), records.last()) else {
80        // Nothing to wrap. Encode as a zero-message wrapper would be
81        // ambiguous; emit nothing instead.
82        return Ok(());
83    };
84
85    // Build inner uncompressed MessageSet.
86    let mut inner = BytesMut::new();
87    let base_offset = first.offset.0;
88    for r in records {
89        let inner_offset = match magic {
90            Magic::V0 => r.offset.0,
91            // KIP-32 preserves the original distance between retained records.
92            Magic::V1 => r.offset.0.checked_sub(base_offset).ok_or_else(|| {
93                LegacyRecordsError::Malformed("relative inner offset exceeds i64".into())
94            })?,
95        };
96        let msg = message_for_record(r.clone(), magic);
97        let msg_len = msg.encoded_len();
98        inner.put_i64(inner_offset);
99        inner.put_i32(i32::try_from(msg_len).unwrap_or(i32::MAX));
100        msg.encode_into(&mut inner);
101    }
102
103    // Compress.
104    let compressed = krabka_compression::compress(codec, &inner)?;
105
106    // Wrapper message.
107    let wrapper_attributes = attrs_with_compression(0, codec);
108    let wrapper_timestamp = match magic {
109        Magic::V0 => None,
110        Magic::V1 => Some(
111            records
112                .iter()
113                .filter_map(|r| r.timestamp)
114                .max()
115                .unwrap_or(-1),
116        ),
117    };
118    let wrapper = Message {
119        magic,
120        attributes: wrapper_attributes,
121        timestamp: wrapper_timestamp,
122        key: None,
123        value: Some(compressed),
124    };
125    let wrapper_len = wrapper.encoded_len();
126
127    // Wrapper offset: v0 = 0 (per Kafka convention pre-KIP-32),
128    // v1 = absolute offset of last inner record.
129    let wrapper_offset = match magic {
130        Magic::V0 => 0,
131        Magic::V1 => last.offset.0,
132    };
133    buf.put_i64(wrapper_offset);
134    buf.put_i32(i32::try_from(wrapper_len).unwrap_or(i32::MAX));
135    wrapper.encode_into(buf);
136    Ok(())
137}
138
139#[cfg(test)]
140mod tests {
141    use {
142        super::*,
143        crate::set::{
144            decode_message_set,
145            test_support::{compressed_bytes, decode_bytes, sample_records_v0, sample_records_v1},
146        },
147        bytes::{Buf, Bytes},
148        krabka_ids::Offset,
149    };
150
151    #[test]
152    fn compressed_v1_writes_relative_sparse_offsets() {
153        let mut records = sample_records_v1();
154        records[0].offset = Offset(100);
155        records[1].offset = Offset(102);
156        records[2].offset = Offset(105);
157        let wire = compressed_bytes(&records, Magic::V1, CompressionType::Gzip);
158        let mut cur = &wire[..];
159        assert2::assert!(cur.get_i64() == 105);
160        let size = usize::try_from(cur.get_i32()).unwrap();
161        let wrapper = Message::decode_from(&mut cur, size).unwrap();
162        let inner = krabka_compression::decompress(
163            CompressionType::Gzip,
164            &wrapper.value.unwrap(),
165            krabka_units::mebibytes(1),
166        )
167        .unwrap();
168        let decoded = decode_bytes(&inner);
169        let mut expected = records.clone();
170        for (record, offset) in expected.iter_mut().zip([0, 2, 5]) {
171            record.offset = Offset(offset);
172        }
173        assert2::assert!(decoded == expected);
174        assert2::assert!(decode_bytes(&wire) == records);
175    }
176
177    #[test]
178    fn compressed_v1_rejects_unrepresentable_relative_span() {
179        let mut records = sample_records_v1();
180        records[0].offset = Offset(i64::MIN);
181        records[1].offset = Offset(0);
182        records[2].offset = Offset(i64::MAX);
183        let mut wire = BytesMut::new();
184        assert2::assert!(matches!(
185            encode_compressed_message_set(&records, Magic::V1, CompressionType::Gzip, &mut wire),
186            Err(LegacyRecordsError::Malformed(_))
187        ));
188        assert2::assert!(wire.is_empty());
189    }
190
191    #[yare::parameterized(
192        flat_v0 = { Magic::V0, None },
193        flat_v1 = { Magic::V1, None },
194        gzip_v0 = { Magic::V0, Some(CompressionType::Gzip) },
195        gzip_v1 = { Magic::V1, Some(CompressionType::Gzip) },
196        snappy_v1 = { Magic::V1, Some(CompressionType::Snappy) },
197    )]
198    fn message_set_roundtrips(magic: Magic, codec: Option<CompressionType>) {
199        let records = match magic {
200            Magic::V0 => sample_records_v0(),
201            Magic::V1 => sample_records_v1(),
202        };
203        let mut buffer = BytesMut::new();
204        if let Some(codec) = codec {
205            encode_compressed_message_set(&records, magic, codec, &mut buffer).unwrap();
206        } else {
207            encode_flat_message_set(records.clone(), magic, &mut buffer);
208        }
209
210        let decoded = decode_message_set(&mut &buffer[..], buffer.len()).unwrap();
211        assert2::assert!(decoded == records);
212    }
213
214    #[test]
215    fn flat_v1_missing_timestamp_encodes_minus_one() {
216        let recs = vec![ParsedRecord {
217            offset: Offset(7),
218            timestamp: None,
219            key: None,
220            value: Some(Bytes::from_static(b"v")),
221        }];
222        let mut buf = BytesMut::new();
223        encode_flat_message_set(recs, Magic::V1, &mut buf);
224        let decoded = decode_bytes(&buf);
225        assert2::assert!(
226            decoded
227                == vec![ParsedRecord {
228                    offset: Offset(7),
229                    timestamp: Some(-1),
230                    key: None,
231                    value: Some(Bytes::from_static(b"v")),
232                }]
233        );
234    }
235
236    #[test]
237    fn compressed_v1_missing_timestamps_default_to_minus_one() {
238        // Records with no timestamps: inner messages encode ts = -1, and the
239        // wrapper's own timestamp (max over records, none present) is -1.
240        let recs = vec![ParsedRecord {
241            offset: Offset(9),
242            timestamp: None,
243            key: None,
244            value: Some(Bytes::from_static(b"v")),
245        }];
246        let buf = compressed_bytes(&recs, Magic::V1, CompressionType::Gzip);
247
248        // Inspect the raw wrapper message's own timestamp before unwrapping.
249        let mut cur: &[u8] = &buf[..];
250        let _wrapper_offset = cur.get_i64();
251        let wrapper_size = usize::try_from(cur.get_i32()).unwrap();
252        let wrapper = Message::decode_from(&mut cur, wrapper_size).unwrap();
253        // The inner record's timestamp survives the unwrap as -1.
254        let decoded = decode_bytes(&buf);
255        assert2::assert!(wrapper.timestamp == Some(-1));
256        assert2::assert!(decoded[0].timestamp == Some(-1));
257    }
258}
krabka_records_legacy/set/encode.rs

krabka_records_legacy/set/
encode.rs
