1use bytes::{Buf, BufMut, Bytes, BytesMut};
4use krabka_compression::RecordDecompressionPolicy;
5use krabka_units::{ByteSize, convert::ByteSizeExt as _};
6use zerocopy::FromBytes as _;
7
8use crate::{
9 primitives::varint::{
10 get_varint, get_varlong, put_varint, put_varlong, varint_len, varlong_len,
11 },
12 records::{
13 RecordsError,
14 borrowed::parse_one_record,
15 crc::{crc32c, crc32c_append},
16 header::{Attributes, CRC_COVERAGE_START, HEADER_LEN, RecordBatchHeader},
17 },
18};
19
20#[derive(Debug, Clone, PartialEq, Eq, Default)]
21pub struct RecordHeader {
22 pub key: String,
23 pub value: Option<Bytes>,
24}
25
26#[derive(Debug, Clone, PartialEq, Eq, Default)]
27pub struct Record {
28 pub attributes: i8,
29 pub timestamp_delta: i64,
30 pub offset_delta: i32,
31 pub key: Option<Bytes>,
32 pub value: Option<Bytes>,
33 pub headers: Vec<RecordHeader>,
34}
35
36#[derive(Debug, Clone, PartialEq, Eq)]
37pub struct RecordBatch {
38 pub base_offset: i64,
39 pub partition_leader_epoch: i32,
40 pub attributes: Attributes,
41 pub last_offset_delta: i32,
42 pub base_timestamp: i64,
43 pub max_timestamp: i64,
44 pub producer_id: i64,
45 pub producer_epoch: i16,
46 pub base_sequence: i32,
47 pub records: Vec<Record>,
48}
49
50impl Default for RecordBatch {
51 fn default() -> Self {
52 Self {
53 base_offset: 0,
54 partition_leader_epoch: 0,
55 attributes: Attributes::default(),
56 last_offset_delta: 0,
57 base_timestamp: 0,
58 max_timestamp: 0,
59 producer_id: -1, producer_epoch: -1,
61 base_sequence: -1,
62 records: Vec::new(),
63 }
64 }
65}
66
67impl Record {
68 pub fn encode<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
74 let body_len = self.body_len();
75 put_varlong(
76 buf,
77 i64::try_from(body_len)
78 .map_err(|_| RecordsError::RecordParse("record body length overflow".into()))?,
79 );
80 self.encode_body(buf)
81 }
82
83 pub fn encoded_len(&self) -> usize {
85 let body = self.body_len();
86 let body_i64 = i64::try_from(body).unwrap_or(i64::MAX);
89 varlong_len(body_i64) + body
90 }
91
92 #[must_use]
100 pub fn body_len(&self) -> usize {
101 body_len_of(
102 self.timestamp_delta,
103 self.offset_delta,
104 self.key.as_ref().map(Bytes::len),
105 self.value.as_ref().map(Bytes::len),
106 self.headers
107 .iter()
108 .map(|h| (h.key.len(), h.value.as_ref().map(Bytes::len))),
109 )
110 }
111
112 fn encode_body<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
113 const SMALL_KEY_LEN: usize = 32;
115 const PREFIX_LEN: usize = 21 + SMALL_KEY_LEN;
116 let mut prefix = [0_u8; PREFIX_LEN];
117 let mut fields = prefix.as_mut_slice();
118 fields.put_i8(self.attributes);
119 put_varlong(&mut fields, self.timestamp_delta);
120 put_varint(&mut fields, self.offset_delta);
121 match &self.key {
122 None => put_varint(&mut fields, -1),
123 Some(k) => {
124 let Ok(len) = i32::try_from(k.len()) else {
125 let unused = fields.len();
127 buf.put_slice(written_prefix(&prefix, unused)?);
128 return Err(RecordsError::RecordParse(
129 "record key length overflow".into(),
130 ));
131 };
132 put_varint(&mut fields, len);
133 }
134 }
135 if let Some(key) = &self.key
136 && key.len() <= SMALL_KEY_LEN
137 {
138 fields.put_slice(key);
139 }
140 let unused = fields.len();
141 buf.put_slice(written_prefix(&prefix, unused)?);
142 if let Some(key) = &self.key
143 && key.len() > SMALL_KEY_LEN
144 {
145 buf.put_slice(key);
146 }
147 match &self.value {
148 None => put_varint(buf, -1),
149 Some(v) => {
150 put_varint(
151 buf,
152 i32::try_from(v.len()).map_err(|_| {
153 RecordsError::RecordParse("record value length overflow".into())
154 })?,
155 );
156 buf.put_slice(v);
157 }
158 }
159 put_varint(
160 buf,
161 i32::try_from(self.headers.len())
162 .map_err(|_| RecordsError::RecordParse("record header count overflow".into()))?,
163 );
164 for h in &self.headers {
165 let key_bytes = h.key.as_bytes();
166 put_varint(
167 buf,
168 i32::try_from(key_bytes.len())
169 .map_err(|_| RecordsError::RecordParse("header key length overflow".into()))?,
170 );
171 buf.put_slice(key_bytes);
172 match &h.value {
173 None => put_varint(buf, -1),
174 Some(v) => {
175 put_varint(
176 buf,
177 i32::try_from(v.len()).map_err(|_| {
178 RecordsError::RecordParse("header value length overflow".into())
179 })?,
180 );
181 buf.put_slice(v);
182 }
183 }
184 }
185 Ok(())
186 }
187
188 pub fn decode<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
196 let body_len = get_varlong(buf)
197 .map_err(|e| RecordsError::RecordParse(format!("record length: {e}")))?;
198 let body_len = usize::try_from(body_len).map_err(|_| {
199 RecordsError::RecordParse(format!("record length negative or too large: {body_len}"))
200 })?;
201 if buf.remaining() < body_len {
202 return Err(RecordsError::BodyTooShort {
203 needed: body_len - buf.remaining(),
204 });
205 }
206 let mut body = buf.take(body_len);
209 let r = Self::decode_body(&mut body)?;
210 if body.has_remaining() {
212 return Err(RecordsError::RecordParse(format!(
213 "trailing bytes inside record (left={})",
214 body.remaining()
215 )));
216 }
217 Ok(r)
218 }
219
220 fn decode_body<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
221 if buf.remaining() == 0 {
222 return Err(RecordsError::RecordParse("record body empty".into()));
223 }
224 let attributes = buf.get_i8();
225 let timestamp_delta = get_varlong(buf)
226 .map_err(|e| RecordsError::RecordParse(format!("timestamp_delta: {e}")))?;
227 let offset_delta =
228 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("offset_delta: {e}")))?;
229
230 let key = decode_nullable_bytes(buf, "key")?;
231 let value = decode_nullable_bytes(buf, "value")?;
232
233 let header_count =
234 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("header_count: {e}")))?;
235 let Ok(header_count_usize) = usize::try_from(header_count) else {
236 return Err(RecordsError::RecordParse(format!(
237 "negative header count {header_count}"
238 )));
239 };
240 let mut headers = Vec::with_capacity(header_count_usize.min(buf.remaining()));
245 for i in 0..header_count {
246 headers.push(
247 decode_record_header(buf)
248 .map_err(|e| RecordsError::RecordParse(format!("header[{i}]: {e}")))?,
249 );
250 }
251
252 Ok(Self {
253 attributes,
254 timestamp_delta,
255 offset_delta,
256 key,
257 value,
258 headers,
259 })
260 }
261}
262
263pub fn body_len_of(
267 timestamp_delta: i64,
268 offset_delta: i32,
269 key: Option<usize>,
270 value: Option<usize>,
271 headers: impl ExactSizeIterator<Item = (usize, Option<usize>)>,
272) -> usize {
273 let field = |len: Option<usize>| {
274 len.map_or_else(
275 || varint_len(-1),
276 |len| varint_len(i32::try_from(len).unwrap_or(i32::MAX)) + len,
277 )
278 };
279 let header_count = varint_len(i32::try_from(headers.len()).unwrap_or(i32::MAX));
280 let headers: usize = headers
281 .map(|(key, value)| field(Some(key)) + field(value))
282 .sum();
283 1 + varlong_len(timestamp_delta)
285 + varint_len(offset_delta)
286 + field(key)
287 + field(value)
288 + header_count
289 + headers
290}
291
292fn decode_nullable_bytes<B: Buf>(buf: &mut B, label: &str) -> Result<Option<Bytes>, RecordsError> {
293 let len =
294 get_varint(buf).map_err(|e| RecordsError::RecordParse(format!("{label} length: {e}")))?;
295 let Ok(n) = usize::try_from(len) else {
297 return Ok(None);
298 };
299 if buf.remaining() < n {
300 return Err(RecordsError::BodyTooShort {
301 needed: n - buf.remaining(),
302 });
303 }
304 let mut v = vec![0u8; n];
305 buf.copy_to_slice(&mut v);
306 Ok(Some(Bytes::from(v)))
307}
308
309fn written_prefix(prefix: &[u8], unused: usize) -> Result<&[u8], RecordsError> {
312 prefix
313 .len()
314 .checked_sub(unused)
315 .and_then(|used| prefix.get(..used))
316 .ok_or_else(|| RecordsError::RecordParse("record prefix overflow".into()))
317}
318
319fn decode_record_header<B: Buf>(buf: &mut B) -> Result<RecordHeader, String> {
320 let key_len = get_varint(buf).map_err(|e| format!("key length: {e}"))?;
321 let Ok(n) = usize::try_from(key_len) else {
322 return Err(format!("non-nullable key has negative length {key_len}"));
323 };
324 if buf.remaining() < n {
325 return Err(format!("key truncated (need {} more)", n - buf.remaining()));
326 }
327 let mut kv = vec![0u8; n];
328 buf.copy_to_slice(&mut kv);
329 let key = String::from_utf8(kv).map_err(|e| format!("key utf-8: {e}"))?;
330
331 let value_len = get_varint(buf).map_err(|e| format!("value length: {e}"))?;
332 let value = if let Ok(n) = usize::try_from(value_len) {
334 if buf.remaining() < n {
335 return Err(format!(
336 "value truncated (need {} more)",
337 n - buf.remaining()
338 ));
339 }
340 let mut vv = vec![0u8; n];
341 buf.copy_to_slice(&mut vv);
342 Some(Bytes::from(vv))
343 } else {
344 None
345 };
346
347 Ok(RecordHeader { key, value })
348}
349
350#[cfg(test)]
351mod record_tests {
352
353 use std::iter;
354
355 use bytes::BytesMut;
356
357 use super::*;
358
359 fn fixture_minimal_record() -> Record {
360 Record {
361 attributes: 0,
362 timestamp_delta: 0,
363 offset_delta: 0,
364 key: None,
365 value: None,
366 headers: vec![],
367 }
368 }
369
370 fn fixture_keyed_record() -> Record {
371 Record {
372 attributes: 0,
373 timestamp_delta: 17,
374 offset_delta: 2,
375 key: Some(Bytes::from_static(b"the-key")),
376 value: Some(Bytes::from_static(b"hello kafka")),
377 headers: vec![
378 RecordHeader {
379 key: "trace-id".to_owned(),
380 value: Some(Bytes::from_static(b"abc")),
381 },
382 RecordHeader {
383 key: "null-val".to_owned(),
384 value: None,
385 },
386 ],
387 }
388 }
389
390 fn fixture_large_payload_record() -> Record {
391 Record {
392 attributes: 0,
393 timestamp_delta: 1_000_000,
394 offset_delta: 999,
395 key: Some(Bytes::from(vec![b'k'; 128])),
396 value: Some(Bytes::from(vec![b'v'; 4096])),
397 headers: vec![],
398 }
399 }
400
401 #[test]
402 fn record_prefix_golden_bytes_across_destination_splits() {
403 let mut cases = vec![
404 (Record::default(), vec![0x0c, 0, 0, 0, 1, 1, 0]),
405 (
406 Record {
407 attributes: i8::MIN,
408 timestamp_delta: i64::MIN,
409 offset_delta: i32::MIN,
410 key: Some(Bytes::from_static(b"k")),
411 value: Some(Bytes::from_static(b"v")),
412 headers: vec![],
413 },
414 vec![
415 0x2a, 0x80, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0x01, 0xff,
416 0xff, 0xff, 0xff, 0x0f, 0x02, b'k', 0x02, b'v', 0,
417 ],
418 ),
419 (
420 Record {
421 offset_delta: 129,
422 key: Some(Bytes::from_static(b"k")),
423 value: Some(Bytes::from_static(b"v")),
424 ..Default::default()
425 },
426 vec![0x12, 0, 0, 0x82, 0x02, 0x02, b'k', 0x02, b'v', 0],
427 ),
428 ];
429 for (len, length_prefix, key_prefix) in [
432 (0, vec![12], vec![0]),
433 (32, vec![76], vec![64]),
434 (33, vec![78], vec![66]),
435 (63, vec![138, 1], vec![126]),
436 (64, vec![142, 1], vec![128, 1]),
437 ] {
438 let mut expected = length_prefix;
439 expected.extend_from_slice(&[0, 0, 0]);
440 expected.extend(key_prefix);
441 expected.extend(iter::repeat_n(b'k', len));
442 expected.extend_from_slice(&[1, 0]);
443 cases.push((
444 Record {
445 key: Some(Bytes::from(vec![b'k'; len])),
446 ..Default::default()
447 },
448 expected,
449 ));
450 }
451 for (record, expected) in cases {
452 for split in 0..=expected.len() {
453 let mut encoded = vec![0xcd; expected.len()];
454 let (first, second) = encoded.split_at_mut(split);
455 let mut destination = first.chain_mut(second);
456 record.encode(&mut destination).unwrap();
457 assert2::assert!(destination.remaining_mut() == 0);
458 assert2::assert!(encoded == expected);
459 let mut cursor = encoded.as_slice();
460 assert2::assert!(Record::decode(&mut cursor).unwrap() == record);
461 assert2::assert!(cursor.is_empty());
462 }
463 }
464 }
465
466 #[test]
467 fn record_roundtrip_cases() {
468 type TestCase1<'a> = (&'a str, fn() -> Record);
469 let cases: [TestCase1<'_>; 3] = [
470 ("minimal", fixture_minimal_record),
471 ("keyed with headers", fixture_keyed_record),
472 ("large payload", fixture_large_payload_record),
473 ];
474 for (_case, fixture) in cases {
475 let record = fixture();
476 let mut buf = BytesMut::new();
477 record.encode(&mut buf).unwrap();
478 assert2::assert!(buf.len() == record.encoded_len());
479
480 let mut cur: &[u8] = &buf[..];
481 let decoded = Record::decode(&mut cur).unwrap();
482 assert2::assert!((decoded, cur.is_empty()) == (record, true));
483 }
484 }
485
486 #[test]
495 fn body_len_is_the_size_the_length_prefix_declares() {
496 for (value_len, body, prefix) in [
497 (0, 6, 1),
498 (57, 63, 1),
499 (58, 64, 2),
500 (63, 69, 2),
501 (64, 71, 2),
502 (8_184, 8_191, 2),
503 (8_185, 8_192, 3),
504 (8_192, 8_200, 3),
505 (70_000, 70_008, 3),
506 ] {
507 let record = Record {
508 value: Some(Bytes::from(vec![7_u8; value_len])),
509 ..Default::default()
510 };
511 let mut buf = BytesMut::new();
512 record.encode(&mut buf).unwrap();
513
514 let mut cur: &[u8] = &buf[..];
515 let declared = get_varlong(&mut cur).unwrap();
516 let case = format!("value_len={value_len}");
517 assert2::check!(record.body_len() == body, "{case}");
518 assert2::check!(record.encoded_len() == prefix + body, "{case}");
519 assert2::check!(buf.len() == prefix + body, "{case}");
520 assert2::check!(declared == i64::try_from(body).unwrap(), "{case}");
521 assert2::check!(buf.len() - cur.len() == prefix, "{case}");
522 }
523 }
524
525 #[test]
529 fn body_len_counts_every_field() {
530 type Fixture<'a> = (&'a str, fn() -> Record, usize);
531 let cases: [Fixture<'_>; 3] = [
532 ("minimal", fixture_minimal_record, 6),
534 ("keyed with headers", fixture_keyed_record, 47),
538 ("large payload", fixture_large_payload_record, 4_235),
541 ];
542 for (case, fixture, body) in cases {
543 let record = fixture();
544 let mut buf = BytesMut::new();
545 record.encode(&mut buf).unwrap();
546 let declared = get_varlong(&mut &buf[..]).unwrap();
547 assert2::check!(record.body_len() == body, "{case}");
548 assert2::check!(declared == i64::try_from(body).unwrap(), "{case}");
549 }
550 }
551
552 #[test]
553 fn decode_rejects_negative_header_count() {
554 let mut buf = BytesMut::new();
555 put_varlong(&mut buf, 6); buf.put_i8(0); put_varlong(&mut buf, 0); put_varint(&mut buf, 0); put_varint(&mut buf, -1); put_varint(&mut buf, -1); put_varint(&mut buf, -1); let mut cur: &[u8] = &buf[..];
566 match Record::decode(&mut cur) {
567 Err(RecordsError::RecordParse(msg)) => {
568 assert2::assert!(msg.contains("negative header count"));
569 }
570 other => panic!("expected RecordParse, got {other:?}"),
571 }
572 }
573
574 #[test]
575 fn decode_huge_header_count_does_not_overallocate() {
576 let mut inner = BytesMut::new();
581 inner.put_i8(0); put_varlong(&mut inner, 0); put_varint(&mut inner, 0); put_varint(&mut inner, -1); put_varint(&mut inner, -1); put_varint(&mut inner, 1_000_000_000); let mut buf = BytesMut::new();
589 put_varlong(&mut buf, i64::try_from(inner.len()).unwrap());
590 buf.extend_from_slice(&inner);
591
592 let mut cur: &[u8] = &buf[..];
593 assert2::assert!(Record::decode(&mut cur).is_err());
595 }
596}
597
598impl RecordBatch {
599 #[must_use]
603 pub fn delete_horizon_ms(&self) -> Option<i64> {
604 self.attributes
605 .has_delete_horizon()
606 .then_some(self.base_timestamp)
607 }
608
609 #[must_use]
615 pub fn with_delete_horizon(mut self, horizon_ms: i64) -> Self {
616 let old_base = self.base_timestamp;
617 for r in &mut self.records {
618 let abs = old_base.saturating_add(r.timestamp_delta);
623 r.timestamp_delta = abs.saturating_sub(horizon_ms);
624 }
625 self.base_timestamp = horizon_ms;
626 self.attributes = self.attributes.with_delete_horizon(true);
627 self
628 }
629
630 pub fn decode_batch<B: Buf>(buf: &mut B) -> Result<Self, RecordsError> {
640 Self::decode_with_policy(buf, RecordDecompressionPolicy::default())
641 }
642
643 pub fn decode_with_policy<B: Buf>(
650 buf: &mut B,
651 policy: RecordDecompressionPolicy,
652 ) -> Result<Self, RecordsError> {
653 if buf.remaining() < HEADER_LEN {
655 return Err(RecordsError::HeaderTooShort {
656 needed: HEADER_LEN - buf.remaining(),
657 });
658 }
659 let contiguous_crc = contiguous_covered_crc(buf.chunk());
662 let mut hdr_bytes = [0u8; HEADER_LEN];
664 buf.copy_to_slice(&mut hdr_bytes);
665
666 let hdr = RecordBatchHeader::ref_from_bytes(&hdr_bytes[..])
667 .map_err(|_| RecordsError::ZerocopyFailure)?;
668
669 if hdr.magic != 2 {
670 return Err(RecordsError::UnsupportedMagic { found: hdr.magic });
671 }
672
673 let body_len = declared_body_len(hdr.batch_length.get()).ok_or_else(|| {
674 RecordsError::RecordParse("negative or oversized batch_length".into())
675 })?;
676
677 if buf.remaining() < body_len {
678 return Err(RecordsError::BodyTooShort {
679 needed: body_len - buf.remaining(),
680 });
681 }
682
683 let expected_crc = hdr.crc.get();
687 let verify_crc = |body: &[u8]| {
688 let computed = contiguous_crc.unwrap_or_else(|| {
689 crc32c_append(crc32c(&hdr_bytes[CRC_COVERAGE_START..HEADER_LEN]), body)
690 });
691 if computed == expected_crc {
692 Ok(())
693 } else {
694 Err(RecordsError::CrcMismatch {
695 expected: expected_crc,
696 computed,
697 })
698 }
699 };
700
701 let attributes = Attributes(hdr.attributes.get());
702 let codec = attributes.compression();
703
704 let decompress = |body: &[u8]| -> Result<Bytes, RecordsError> {
705 verify_crc(body)?;
706 Ok(krabka_compression::decompress(
707 codec,
708 body,
709 policy.output_limit(ByteSize::from_bytes(body.len() as u64)),
710 )?)
711 };
712
713 let body_for_records = if codec != krabka_compression::CompressionType::None
719 && let Some(chunk) = buf.chunk().get(..body_len)
720 {
721 let decoded = decompress(chunk);
722 buf.advance(body_len);
723 decoded?
724 } else {
725 let mut body = vec![0u8; body_len];
726 buf.copy_to_slice(&mut body);
727 if codec == krabka_compression::CompressionType::None {
728 verify_crc(&body)?;
729 Bytes::from(body)
730 } else {
731 decompress(&body)?
732 }
733 };
734
735 let count = hdr.records_count.get();
737 let Ok(count_usize) = usize::try_from(count) else {
738 return Err(RecordsError::RecordParse(format!(
739 "negative records_count {count}"
740 )));
741 };
742 let mut body_cur: &[u8] = &body_for_records[..];
743 let record_capacity = count_usize.min(body_for_records.len());
747 let mut records = Vec::with_capacity(record_capacity);
748 for i in 0..count {
749 let r = parse_one_record(&mut body_cur)
756 .map_err(|e| RecordsError::RecordParse(format!("record[{i}]: {e}")))?;
757 records.push(Record {
758 attributes: r.attributes,
759 timestamp_delta: r.timestamp_delta,
760 offset_delta: r.offset_delta,
761 key: r.key.map(|s| body_for_records.slice_ref(s)),
762 value: r.value.map(|s| body_for_records.slice_ref(s)),
763 headers: r
764 .headers
765 .into_iter()
766 .map(|h| RecordHeader {
767 key: h.key.to_owned(),
768 value: h.value.map(|s| body_for_records.slice_ref(s)),
769 })
770 .collect(),
771 });
772 }
773 if !body_cur.is_empty() {
774 return Err(RecordsError::RecordParse(format!(
775 "trailing bytes after records (left={})",
776 body_cur.len()
777 )));
778 }
779
780 Ok(Self {
781 base_offset: hdr.base_offset.get(),
782 partition_leader_epoch: hdr.partition_leader_epoch.get(),
783 attributes,
784 last_offset_delta: hdr.last_offset_delta.get(),
785 base_timestamp: hdr.base_timestamp.get(),
786 max_timestamp: hdr.max_timestamp.get(),
787 producer_id: hdr.producer_id.get(),
788 producer_epoch: hdr.producer_epoch.get(),
789 base_sequence: hdr.base_sequence.get(),
790 records,
791 })
792 }
793
794 pub fn encode_batch<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
803 self.encode_with_compression_level(buf, None)
804 }
805
806 pub fn encode_to_bytes(&self) -> Result<Bytes, RecordsError> {
815 if self.attributes.compression() == krabka_compression::CompressionType::None {
816 return Ok(self.encode_uncompressed()?.freeze());
817 }
818 let mut wire = BytesMut::with_capacity(self.encoded_batch_len());
819 self.encode_batch(&mut wire)?;
820 Ok(wire.freeze())
821 }
822
823 fn encode_uncompressed(&self) -> Result<BytesMut, RecordsError> {
827 let mut wire = BytesMut::with_capacity(self.encoded_batch_len());
828 wire.resize(HEADER_LEN, 0);
829 for record in &self.records {
830 record.encode(&mut wire)?;
831 }
832 let body_len = wire.len() - HEADER_LEN;
833 let batch_length = batch_length_for(body_len)?;
834 let covered_head = self.covered_head()?;
835 let short = || RecordsError::RecordParse("batch buffer shorter than its header".into());
836 let (prefix, covered) = wire
837 .split_at_mut_checked(CRC_COVERAGE_START)
838 .ok_or_else(short)?;
839 let (covered_head_dst, _) = covered
840 .split_first_chunk_mut::<{ HEADER_LEN - CRC_COVERAGE_START }>()
841 .ok_or_else(short)?;
842 *covered_head_dst = covered_head;
843 let crc = crc32c(covered);
844 self.put_header_prefix(&mut &mut *prefix, batch_length, crc);
845 Ok(wire)
846 }
847
848 pub fn encode_with_compression_level<B: BufMut>(
862 &self,
863 buf: &mut B,
864 level: Option<i32>,
865 ) -> Result<(), RecordsError> {
866 let codec = self.attributes.compression();
867 if level.is_none() && codec == krabka_compression::CompressionType::None {
868 buf.put_slice(&self.encode_uncompressed()?);
869 return Ok(());
870 }
871
872 let mut raw_body =
874 BytesMut::with_capacity(self.records.iter().map(Record::encoded_len).sum());
875 for r in &self.records {
876 r.encode(&mut raw_body)?;
877 }
878 let raw_body = raw_body.freeze();
879
880 let body: Bytes = match level {
882 Some(level) => krabka_compression::compress_with_level(codec, &raw_body, level)?,
883 None => krabka_compression::compress(codec, &raw_body)?,
884 };
885
886 self.encode_header(buf, &body)?;
887 buf.put_slice(&body);
888 Ok(())
889 }
890
891 fn encode_header<B: BufMut>(&self, buf: &mut B, body: &[u8]) -> Result<(), RecordsError> {
895 let batch_length = batch_length_for(body.len())?;
896 let covered_head = self.covered_head()?;
897 let crc = crc32c_append(crc32c(&covered_head), body);
898 self.put_header_prefix(buf, batch_length, crc);
899 buf.put_slice(&covered_head);
900 Ok(())
901 }
902
903 fn covered_head(&self) -> Result<[u8; HEADER_LEN - CRC_COVERAGE_START], RecordsError> {
905 let mut covered_head = [0_u8; HEADER_LEN - CRC_COVERAGE_START];
906 let mut covered = covered_head.as_mut_slice();
907 covered.put_i16(self.attributes.0);
908 covered.put_i32(self.last_offset_delta);
909 covered.put_i64(self.base_timestamp);
910 covered.put_i64(self.max_timestamp);
911 covered.put_i64(self.producer_id);
912 covered.put_i16(self.producer_epoch);
913 covered.put_i32(self.base_sequence);
914 covered.put_i32(
915 i32::try_from(self.records.len())
916 .map_err(|_| RecordsError::RecordParse("records_count exceeds i32".into()))?,
917 );
918 Ok(covered_head)
919 }
920
921 fn put_header_prefix<B: BufMut>(&self, buf: &mut B, batch_length: i32, crc: u32) {
924 buf.put_i64(self.base_offset);
925 buf.put_i32(batch_length);
926 buf.put_i32(self.partition_leader_epoch);
927 buf.put_i8(2); buf.put_u32(crc);
929 }
930
931 pub fn encoded_batch_len(&self) -> usize {
937 let body: usize = self.records.iter().map(Record::encoded_len).sum();
938 HEADER_LEN + body
939 }
940}
941
942const HEADER_TAIL_LEN: i32 = 49;
948
949fn declared_body_len(batch_length: i32) -> Option<usize> {
952 let n = i32::checked_sub(batch_length, HEADER_TAIL_LEN)?;
953 usize::try_from(n).ok()
954}
955
956fn batch_length_for(body_len: usize) -> Result<i32, RecordsError> {
958 i32::try_from(body_len)
959 .ok()
960 .and_then(|body_len| HEADER_TAIL_LEN.checked_add(body_len))
961 .ok_or_else(|| RecordsError::RecordParse("body length exceeds i32".into()))
962}
963
964fn contiguous_covered_crc(chunk: &[u8]) -> Option<u32> {
967 let hdr = RecordBatchHeader::ref_from_bytes(chunk.get(..HEADER_LEN)?).ok()?;
968 if hdr.magic != 2 {
969 return None;
970 }
971 let body_len = declared_body_len(hdr.batch_length.get())?;
972 let covered = chunk.get(CRC_COVERAGE_START..HEADER_LEN.checked_add(body_len)?)?;
973 Some(crc32c(covered))
974}
975
976#[cfg(test)]
977mod batch_tests {
978 use assert2::check;
979 use krabka_compression::{CompressionError, CompressionType, RecordDecompressionPolicy};
980 use krabka_units::{bytes, fraction};
981
982 use super::*;
983
984 fn fixture_empty_batch() -> RecordBatch {
985 RecordBatch::default()
986 }
987
988 fn fixture_single_record_batch() -> RecordBatch {
989 RecordBatch {
990 records: vec![Record {
991 key: Some(Bytes::from_static(b"k1")),
992 value: Some(Bytes::from_static(b"v1")),
993 ..Default::default()
994 }],
995 ..RecordBatch::default()
996 }
997 }
998
999 fn fixture_multi_record_batch() -> RecordBatch {
1000 RecordBatch {
1001 base_offset: 42,
1002 partition_leader_epoch: 5,
1003 last_offset_delta: 2,
1004 base_timestamp: 1_700_000_000,
1005 max_timestamp: 1_700_000_500,
1006 producer_id: 100,
1007 producer_epoch: 3,
1008 base_sequence: 7,
1009 records: vec![
1010 Record {
1011 offset_delta: 0,
1012 timestamp_delta: 0,
1013 key: Some(Bytes::from_static(b"a")),
1014 value: Some(Bytes::from_static(b"1")),
1015 ..Default::default()
1016 },
1017 Record {
1018 offset_delta: 1,
1019 timestamp_delta: 100,
1020 key: Some(Bytes::from_static(b"b")),
1021 value: Some(Bytes::from_static(b"2")),
1022 ..Default::default()
1023 },
1024 Record {
1025 offset_delta: 2,
1026 timestamp_delta: 500,
1027 key: None,
1028 value: Some(Bytes::from_static(b"3")),
1029 headers: vec![RecordHeader {
1030 key: "h".to_owned(),
1031 value: Some(Bytes::from_static(b"hv")),
1032 }],
1033 ..Default::default()
1034 },
1035 ],
1036 ..RecordBatch::default()
1037 }
1038 }
1039
1040 #[test]
1041 fn uncompressed_roundtrip_cases() {
1042 type TestCase2 = (&'static str, fn() -> RecordBatch);
1043 let cases: [TestCase2; 3] = [
1044 ("empty", fixture_empty_batch),
1045 ("single", fixture_single_record_batch),
1046 ("multiple", fixture_multi_record_batch),
1047 ];
1048 for (_case, fixture) in cases {
1049 let mut batch = fixture();
1050 batch.attributes = batch.attributes.with_compression(CompressionType::None);
1051 let mut buf = BytesMut::new();
1052 batch.encode_batch(&mut buf).unwrap();
1053 assert2::assert!(buf.len() == batch.encoded_batch_len());
1054 let mut cur: &[u8] = &buf[..];
1055 let decoded = RecordBatch::decode_batch(&mut cur).unwrap();
1056 assert2::assert!((decoded, cur.is_empty()) == (batch, true));
1057 }
1058 }
1059
1060 #[test]
1063 fn encodings_agree_and_store_the_reference_crc() {
1064 for codec in [
1065 CompressionType::None,
1066 CompressionType::Gzip,
1067 CompressionType::Lz4,
1068 ] {
1069 for fixture in [
1070 fixture_empty_batch,
1071 fixture_single_record_batch,
1072 fixture_multi_record_batch,
1073 ] {
1074 let mut batch = fixture();
1075 batch.attributes = batch.attributes.with_compression(codec);
1076 let mut encoded = BytesMut::new();
1077 batch.encode_batch(&mut encoded).unwrap();
1078 let to_bytes = batch.encode_to_bytes().unwrap();
1079 let stored = u32::from_be_bytes(encoded[17..21].try_into().unwrap());
1080 assert2::assert!(
1081 (
1082 to_bytes[..] == encoded[..],
1083 stored == crc32c::crc32c(&encoded[CRC_COVERAGE_START..]),
1084 ) == (true, true),
1085 "{codec:?}"
1086 );
1087 }
1088 }
1089 }
1090
1091 #[test]
1096 fn fragmented_input_decodes_like_contiguous_input() {
1097 for codec in [CompressionType::None, CompressionType::Lz4] {
1098 let mut batch = fixture_multi_record_batch();
1099 batch.attributes = batch.attributes.with_compression(codec);
1100 let mut encoded = BytesMut::new();
1101 batch.encode_batch(&mut encoded).unwrap();
1102 let mut corrupt = encoded.to_vec();
1103 *corrupt.last_mut().unwrap() ^= 0xFF;
1104 let contiguous_error = RecordBatch::decode_batch(&mut &corrupt[..]).unwrap_err();
1105 let RecordsError::CrcMismatch { expected, computed } = contiguous_error else {
1106 panic!("{codec:?}: {contiguous_error:?}");
1107 };
1108 assert2::assert!(computed == crc32c::crc32c(&corrupt[CRC_COVERAGE_START..]));
1109
1110 for split in [0, 10, 30, HEADER_LEN, HEADER_LEN + 1, encoded.len() - 1] {
1111 let mut fragmented = (&encoded[..split]).chain(&encoded[split..]);
1112 let decoded = RecordBatch::decode_batch(&mut fragmented).unwrap();
1113 assert2::assert!(
1114 (decoded == batch, fragmented.has_remaining()) == (true, false),
1115 "{codec:?} split at {split}"
1116 );
1117
1118 let mut fragmented = (&corrupt[..split]).chain(&corrupt[split..]);
1119 let error = RecordBatch::decode_batch(&mut fragmented).unwrap_err();
1120 assert2::assert!(
1121 matches!(
1122 error,
1123 RecordsError::CrcMismatch { expected: e, computed: c }
1124 if (e, c) == (expected, computed)
1125 ),
1126 "{codec:?} split at {split}: {error:?}"
1127 );
1128 }
1129 }
1130 }
1131
1132 #[test]
1137 fn encode_with_compression_level_compresses_at_the_level() {
1138 let value: Vec<u8> = (0..16 * 1024u32)
1139 .map(|i| u8::try_from((i * 7 + i / 13) % 251).unwrap_or(0))
1140 .collect();
1141 let mut batch = RecordBatch {
1142 records: vec![Record {
1143 value: Some(Bytes::from(value)),
1144 ..Default::default()
1145 }],
1146 ..RecordBatch::default()
1147 };
1148 batch.attributes = batch.attributes.with_compression(CompressionType::Gzip);
1149 let encoded = |level: Option<i32>| {
1150 let mut buf = BytesMut::new();
1151 batch
1152 .encode_with_compression_level(&mut buf, level)
1153 .map(|()| buf.freeze())
1154 };
1155 let default = encoded(None).unwrap();
1156 let fastest = encoded(Some(1)).unwrap();
1157 let best = encoded(Some(9)).unwrap();
1158 let decoded = [&fastest, &best]
1159 .map(|bytes| RecordBatch::decode_batch(&mut &bytes[..]).unwrap() == batch);
1160
1161 assert2::assert!(
1162 (
1163 encoded(Some(-1)).unwrap() == default,
1164 fastest != best,
1165 decoded,
1166 ) == (true, true, [true, true])
1167 );
1168 assert2::assert!(matches!(
1169 encoded(Some(10)),
1170 Err(RecordsError::Compression(CompressionError::InvalidLevel {
1171 level: 10,
1172 ..
1173 }))
1174 ));
1175 }
1176
1177 #[test]
1178 fn rejects_pre_v2_magic() {
1179 let mut buf = BytesMut::new();
1180 buf.put_i64(0); buf.put_i32(49); buf.put_i32(0); buf.put_i8(1); buf.put_u32(0); for _ in 21..HEADER_LEN {
1186 buf.put_u8(0);
1187 }
1188 let mut cur: &[u8] = &buf[..];
1189 assert2::assert!(matches!(
1190 RecordBatch::decode_batch(&mut cur),
1191 Err(RecordsError::UnsupportedMagic { found: 1 })
1192 ));
1193 }
1194
1195 #[test]
1196 fn rejects_bad_crc() {
1197 let b = fixture_single_record_batch();
1198 let mut buf = BytesMut::new();
1199 b.encode_batch(&mut buf).unwrap();
1200 buf[17] ^= 0xFF;
1202 let mut cur: &[u8] = &buf[..];
1203 assert2::assert!(matches!(
1204 RecordBatch::decode_batch(&mut cur),
1205 Err(RecordsError::CrcMismatch { .. })
1206 ));
1207 }
1208
1209 #[test]
1210 fn compressed_roundtrip_cases() {
1211 for (_case, codec) in [
1212 ("gzip", CompressionType::Gzip),
1213 ("snappy", CompressionType::Snappy),
1214 ("lz4", CompressionType::Lz4),
1215 ("zstd", CompressionType::Zstd),
1216 ] {
1217 let mut batch = fixture_multi_record_batch();
1218 batch.attributes = batch.attributes.with_compression(codec);
1219 let mut buf = BytesMut::new();
1220 batch.encode_batch(&mut buf).unwrap();
1221 let mut cur: &[u8] = &buf[..];
1222 let decoded = RecordBatch::decode_batch(&mut cur).unwrap();
1223 assert2::assert!((decoded, cur.is_empty()) == (batch, true));
1224 }
1225 }
1226
1227 #[test]
1228 fn decompression_policy_limits_owned_decode() {
1229 let mut batch = fixture_single_record_batch();
1230 batch.records[0].value = Some(Bytes::from(vec![b'x'; 4096]));
1231 batch.attributes = batch.attributes.with_compression(CompressionType::Lz4);
1232 let mut wire = BytesMut::new();
1233 batch.encode_batch(&mut wire).unwrap();
1234
1235 let mut default_cur = &wire[..];
1236 RecordBatch::decode_batch(&mut default_cur).unwrap();
1237
1238 let policy = RecordDecompressionPolicy::new(fraction(1.0), bytes(1), bytes(32)).unwrap();
1239 let mut limited_cur = &wire[..];
1240 assert2::assert!(matches!(
1241 RecordBatch::decode_with_policy(&mut limited_cur, policy),
1242 Err(RecordsError::Compression(CompressionError::TooLarge {
1243 limit: 32
1244 }))
1245 ));
1246 }
1247
1248 #[test]
1249 fn with_delete_horizon_stamps_and_preserves_record_timestamps() {
1250 let b = RecordBatch {
1252 base_timestamp: 1000,
1253 records: vec![
1254 Record {
1255 timestamp_delta: 0,
1256 ..Default::default()
1257 },
1258 Record {
1259 timestamp_delta: 5,
1260 ..Default::default()
1261 },
1262 ],
1263 ..RecordBatch::default()
1264 };
1265
1266 let stamped = b.with_delete_horizon(9999);
1267
1268 check!(stamped.attributes.has_delete_horizon());
1269 check!(stamped.base_timestamp == 9999);
1270 check!(stamped.delete_horizon_ms() == Some(9999));
1271
1272 let absolutes: Vec<i64> = stamped
1274 .records
1275 .iter()
1276 .map(|r| stamped.base_timestamp + r.timestamp_delta)
1277 .collect();
1278 assert2::assert!(absolutes == vec![1000, 1005]);
1279 }
1280
1281 #[test]
1282 fn delete_horizon_round_trips_through_encode_decode() {
1283 let b = RecordBatch {
1285 base_timestamp: 1000,
1286 last_offset_delta: 1,
1287 records: vec![
1288 Record {
1289 timestamp_delta: 0,
1290 offset_delta: 0,
1291 key: Some(Bytes::from_static(b"k1")),
1292 value: Some(Bytes::from_static(b"v1")),
1293 ..Default::default()
1294 },
1295 Record {
1296 timestamp_delta: 5,
1297 offset_delta: 1,
1298 key: Some(Bytes::from_static(b"k2")),
1299 value: Some(Bytes::from_static(b"v2")),
1300 ..Default::default()
1301 },
1302 ],
1303 ..RecordBatch::default()
1304 }
1305 .with_delete_horizon(9999);
1306
1307 let mut buf = BytesMut::new();
1308 b.encode_batch(&mut buf).unwrap();
1309
1310 let mut cur: &[u8] = &buf[..];
1311 let decoded = RecordBatch::decode_batch(&mut cur).unwrap();
1312 assert2::assert!(cur.is_empty());
1313
1314 assert2::assert!(decoded.delete_horizon_ms() == Some(9999));
1315
1316 let absolutes: Vec<i64> = decoded
1317 .records
1318 .iter()
1319 .map(|r| decoded.base_timestamp + r.timestamp_delta)
1320 .collect();
1321 assert2::assert!(absolutes == vec![1000, 1005]);
1322 }
1323
1324 #[test]
1325 fn decode_huge_records_count_does_not_overallocate() {
1326 let mut b = fixture_empty_batch();
1331 b.attributes = b.attributes.with_compression(CompressionType::None);
1332 let mut buf = BytesMut::new();
1333 b.encode_batch(&mut buf).unwrap();
1334
1335 let rc_off = HEADER_LEN - 4;
1337 buf[rc_off..HEADER_LEN].copy_from_slice(&1_000_000_000i32.to_be_bytes());
1338
1339 let computed = crc32c(&buf[CRC_COVERAGE_START..]);
1341 buf[17..21].copy_from_slice(&computed.to_be_bytes());
1342
1343 let mut cur: &[u8] = &buf[..];
1344 assert2::assert!(RecordBatch::decode_batch(&mut cur).is_err());
1346 }
1347}
1348
1349impl crate::Encode for RecordBatch {
1350 fn encode<B: BufMut>(&self, buf: &mut B, _version: i16) -> Result<(), crate::ProtocolError> {
1351 self.encode_batch(buf).map_err(Into::into)
1352 }
1353
1354 fn encoded_len(&self, _version: i16) -> usize {
1355 self.encoded_batch_len()
1356 }
1357}
1358
1359impl crate::Decode<'_> for RecordBatch {
1360 fn decode<B: Buf>(buf: &mut B, _version: i16) -> Result<Self, crate::ProtocolError> {
1361 Self::decode_batch(buf).map_err(Into::into)
1362 }
1363}