1use bytes::{Buf, BufMut, Bytes};
19use krabka_compression::RecordDecompressionPolicy;
20
21use crate::records::{
22 RecordsError, borrowed::RecordBatch as RecordBatchBorrowed, owned::RecordBatch,
23};
24
25#[derive(Debug, Clone, PartialEq, Eq)]
27pub enum RecordsPayload {
28 V2(Vec<RecordBatch>),
30 Raw(Bytes),
33 Legacy(Bytes),
36 #[cfg(any(
45 target_os = "linux",
46 target_os = "macos",
47 target_os = "ios",
48 target_os = "tvos",
49 target_os = "watchos",
50 target_os = "freebsd",
51 target_os = "dragonfly",
52 ))]
53 FileRegions(Vec<crate::records::FileRegion>),
54}
55
56impl RecordsPayload {
57 pub fn from_bytes(bytes: Bytes) -> Result<Self, RecordsError> {
64 Self::from_bytes_with_policy(bytes, RecordDecompressionPolicy::default())
65 }
66
67 pub fn from_bytes_with_policy(
75 bytes: Bytes,
76 policy: RecordDecompressionPolicy,
77 ) -> Result<Self, RecordsError> {
78 if looks_like_v2(&bytes) {
79 let mut cur: &[u8] = &bytes;
80 let mut batches = Vec::new();
81 while !cur.is_empty() {
82 batches.push(RecordBatch::decode_with_policy(&mut cur, policy)?);
83 }
84 Ok(Self::V2(batches))
85 } else {
86 Ok(Self::Legacy(bytes))
87 }
88 }
89
90 #[must_use]
92 pub fn payload_len(&self) -> usize {
93 match self {
94 Self::V2(batches) => batches.iter().map(RecordBatch::encoded_len).sum(),
95 Self::Raw(b) | Self::Legacy(b) => b.len(),
96 #[cfg(any(
97 target_os = "linux",
98 target_os = "macos",
99 target_os = "ios",
100 target_os = "tvos",
101 target_os = "watchos",
102 target_os = "freebsd",
103 target_os = "dragonfly",
104 ))]
105 Self::FileRegions(regions) => regions.iter().map(|r| r.len).sum(),
106 }
107 }
108
109 pub fn encode_to<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
121 match self {
122 Self::V2(batches) => {
123 for b in batches {
124 b.encode(buf)?;
125 }
126 Ok(())
127 }
128 Self::Raw(b) | Self::Legacy(b) => {
129 buf.put_slice(b);
130 Ok(())
131 }
132 #[cfg(any(
133 target_os = "linux",
134 target_os = "macos",
135 target_os = "ios",
136 target_os = "tvos",
137 target_os = "watchos",
138 target_os = "freebsd",
139 target_os = "dragonfly",
140 ))]
141 Self::FileRegions(regions) => {
142 use std::os::unix::fs::FileExt;
143 let mut scratch = vec![0u8; 0];
144 for region in regions {
145 scratch.resize(region.len, 0);
146 let mut filled = 0usize;
147 let mut offset = region.offset;
148 while filled < region.len {
149 match region.file.read_at(&mut scratch[filled..], offset) {
150 Ok(0) => {
151 return Err(RecordsError::RecordParse(
152 "FileRegion read hit EOF before len bytes".into(),
153 ));
154 }
155 Ok(n) => {
156 filled += n;
157 offset += n as u64;
158 }
159 Err(e) if e.kind() == std::io::ErrorKind::Interrupted => {}
160 Err(e) => {
161 return Err(RecordsError::RecordParse(format!(
162 "FileRegion read error: {e}"
163 )));
164 }
165 }
166 }
167 buf.put_slice(&scratch);
168 }
169 Ok(())
170 }
171 }
172 }
173
174 #[must_use]
180 pub fn as_v2(&self) -> Option<&[RecordBatch]> {
181 match self {
182 Self::V2(batches) => Some(batches),
183 #[cfg(any(
184 target_os = "linux",
185 target_os = "macos",
186 target_os = "ios",
187 target_os = "tvos",
188 target_os = "watchos",
189 target_os = "freebsd",
190 target_os = "dragonfly",
191 ))]
192 Self::FileRegions(_) => None,
193 Self::Raw(_) | Self::Legacy(_) => None,
194 }
195 }
196
197 #[must_use]
199 pub fn as_legacy(&self) -> Option<&Bytes> {
200 match self {
201 Self::Legacy(b) => Some(b),
202 #[cfg(any(
203 target_os = "linux",
204 target_os = "macos",
205 target_os = "ios",
206 target_os = "tvos",
207 target_os = "watchos",
208 target_os = "freebsd",
209 target_os = "dragonfly",
210 ))]
211 Self::FileRegions(_) => None,
212 Self::V2(_) | Self::Raw(_) => None,
213 }
214 }
215
216 pub fn from_fetch_bytes(bytes: Bytes) -> Result<Self, RecordsError> {
237 if !looks_like_v2(&bytes) {
238 return Ok(Self::Legacy(bytes));
239 }
240 let mut cur: &[u8] = &bytes;
241 let mut batches = Vec::new();
242 while !cur.is_empty() {
243 match RecordBatch::decode(&mut cur) {
244 Ok(rb) => batches.push(rb),
245 Err(RecordsError::HeaderTooShort { .. } | RecordsError::BodyTooShort { .. }) => {
246 break;
247 }
248 Err(e) => return Err(e),
249 }
250 }
251 Ok(Self::V2(batches))
252 }
253
254 pub fn decode_lenient<B: Buf>(
263 buf: &mut B,
264 _version: i16,
265 ) -> Result<Self, crate::ProtocolError> {
266 let bytes = buf.copy_to_bytes(buf.remaining());
267 Self::from_fetch_bytes(bytes).map_err(Into::into)
268 }
269}
270
271impl From<RecordBatch> for RecordsPayload {
272 fn from(rb: RecordBatch) -> Self {
273 Self::V2(vec![rb])
274 }
275}
276
277impl From<Vec<RecordBatch>> for RecordsPayload {
278 fn from(v: Vec<RecordBatch>) -> Self {
279 Self::V2(v)
280 }
281}
282
283impl Default for RecordsPayload {
284 fn default() -> Self {
285 Self::V2(Vec::new())
286 }
287}
288
289impl crate::Encode for RecordsPayload {
290 fn encode<B: BufMut>(&self, buf: &mut B, _version: i16) -> Result<(), crate::ProtocolError> {
291 self.encode_to(buf).map_err(Into::into)
292 }
293
294 fn encoded_len(&self, _version: i16) -> usize {
295 self.payload_len()
296 }
297}
298
299impl crate::Decode<'_> for RecordsPayload {
300 fn decode<B: Buf>(buf: &mut B, _version: i16) -> Result<Self, crate::ProtocolError> {
301 let bytes = buf.copy_to_bytes(buf.remaining());
305 Self::from_bytes(bytes).map_err(Into::into)
306 }
307}
308
309#[derive(Debug, Clone, PartialEq, Eq)]
311pub enum RecordsPayloadBorrowed<'a> {
312 V2(Vec<RecordBatchBorrowed<'a>>),
313 Legacy(&'a [u8]),
314}
315
316impl<'a> RecordsPayloadBorrowed<'a> {
317 pub fn from_slice(bytes: &'a [u8]) -> Result<Self, RecordsError> {
320 if looks_like_v2(bytes) {
321 let mut cur: &'a [u8] = bytes;
322 let mut batches = Vec::new();
323 while !cur.is_empty() {
324 let rb = <RecordBatchBorrowed<'a> as crate::DecodeBorrow<'a>>::decode_borrow(
325 &mut cur, 0,
326 )
327 .map_err(|e| RecordsError::RecordParse(format!("borrowed v2 decode: {e}")))?;
328 batches.push(rb);
329 }
330 Ok(Self::V2(batches))
331 } else {
332 Ok(Self::Legacy(bytes))
333 }
334 }
335
336 #[must_use]
337 pub fn payload_len(&self) -> usize {
338 match self {
339 Self::V2(batches) => batches
340 .iter()
341 .map(|rb| crate::Encode::encoded_len(rb, 0))
342 .sum(),
343 Self::Legacy(b) => b.len(),
344 }
345 }
346
347 pub fn encode_to<B: BufMut>(&self, buf: &mut B) -> Result<(), RecordsError> {
350 match self {
351 Self::V2(batches) => {
352 for rb in batches {
353 crate::Encode::encode(rb, buf, 0).map_err(|e| {
354 RecordsError::RecordParse(format!("borrowed v2 encode: {e}"))
355 })?;
356 }
357 Ok(())
358 }
359 Self::Legacy(b) => {
360 buf.put_slice(b);
361 Ok(())
362 }
363 }
364 }
365
366 pub fn to_owned(&self) -> Result<RecordsPayload, RecordsError> {
370 match self {
371 Self::V2(batches) => {
372 let mut owned = Vec::with_capacity(batches.len());
373 for rb in batches {
374 owned.push(rb.to_owned()?);
375 }
376 Ok(RecordsPayload::V2(owned))
377 }
378 Self::Legacy(b) => Ok(RecordsPayload::Legacy(Bytes::copy_from_slice(b))),
379 }
380 }
381}
382
383impl Default for RecordsPayloadBorrowed<'_> {
384 fn default() -> Self {
385 Self::V2(Vec::new())
386 }
387}
388
389impl crate::Encode for RecordsPayloadBorrowed<'_> {
390 fn encode<B: BufMut>(&self, buf: &mut B, _version: i16) -> Result<(), crate::ProtocolError> {
391 self.encode_to(buf).map_err(Into::into)
392 }
393
394 fn encoded_len(&self, _version: i16) -> usize {
395 self.payload_len()
396 }
397}
398
399impl<'de> crate::DecodeBorrow<'de> for RecordsPayloadBorrowed<'de> {
400 fn decode_borrow(buf: &mut &'de [u8], _version: i16) -> Result<Self, crate::ProtocolError> {
401 let bytes = std::mem::take(buf);
404 Self::from_slice(bytes).map_err(Into::into)
405 }
406}
407
408#[inline]
419fn looks_like_v2(bytes: &[u8]) -> bool {
420 const MAGIC_OFFSET: usize = 16;
425 bytes.len() > MAGIC_OFFSET && bytes[MAGIC_OFFSET] == 2
426}
427
428#[cfg(test)]
429mod tests {
430
431 use bytes::BytesMut;
432 use krabka_compression::{CompressionError, CompressionType, RecordDecompressionPolicy};
433 use krabka_units::{bytes, fraction};
434
435 use super::*;
436 use crate::records::{Record, RecordBatch};
437
438 fn sample_v2() -> RecordBatch {
439 RecordBatch {
440 base_offset: 42,
441 records: vec![Record {
442 key: Some(Bytes::from_static(b"k")),
443 value: Some(Bytes::from_static(b"v")),
444 ..Default::default()
445 }],
446 ..RecordBatch::default()
447 }
448 }
449
450 #[test]
451 fn from_bytes_dispatches_v2() {
452 let rb = sample_v2();
453 let mut buf = BytesMut::new();
454 rb.encode(&mut buf).unwrap();
455 let p = RecordsPayload::from_bytes(buf.freeze()).unwrap();
456 match p {
457 RecordsPayload::V2(batches) => assert2::assert!(batches == vec![rb]),
458 _ => panic!("expected V2"),
459 }
460 }
461
462 #[test]
463 fn decompression_policy_limits_payload_decode() {
464 let mut rb = sample_v2();
465 rb.records[0].value = Some(Bytes::from(vec![b'x'; 4096]));
466 rb.attributes = rb.attributes.with_compression(CompressionType::Lz4);
467 let mut wire = BytesMut::new();
468 rb.encode(&mut wire).unwrap();
469 let wire = wire.freeze();
470
471 RecordsPayload::from_bytes(wire.clone()).unwrap();
472
473 let policy = RecordDecompressionPolicy::new(fraction(1.0), bytes(1), bytes(32)).unwrap();
474 assert2::assert!(matches!(
475 RecordsPayload::from_bytes_with_policy(wire, policy),
476 Err(RecordsError::Compression(CompressionError::TooLarge {
477 limit: 32
478 }))
479 ));
480 }
481
482 #[test]
483 fn from_bytes_parses_all_batches() {
484 let mut b0 = sample_v2();
486 b0.base_offset = 0;
487 let mut b1 = sample_v2();
488 b1.base_offset = 1;
489 let mut buf = BytesMut::new();
490 b0.encode(&mut buf).unwrap();
491 b1.encode(&mut buf).unwrap();
492 let p = RecordsPayload::from_bytes(buf.freeze()).unwrap();
493 let batches = p.as_v2().expect("v2");
494 assert2::assert!(batches == &[b0, b1][..]);
495 }
496
497 #[test]
498 fn raw_passthrough_roundtrips() {
499 let mut b = sample_v2();
500 b.base_offset = 7;
501 let mut wire = BytesMut::new();
502 b.encode(&mut wire).unwrap();
503 let wire = wire.freeze();
504 let p = RecordsPayload::Raw(wire.clone());
505 assert2::assert!(p.payload_len() == wire.len());
506 let mut out = BytesMut::new();
507 p.encode_to(&mut out).unwrap();
508 assert2::assert!(&out[..] == &wire[..]); assert2::assert!(p.as_v2().is_none()); }
511
512 #[test]
513 fn from_bytes_dispatches_legacy() {
514 let mut buf = vec![0u8; 17];
517 buf[16] = 1;
518 let p = RecordsPayload::from_bytes(Bytes::from(buf.clone())).unwrap();
519 match p {
520 RecordsPayload::Legacy(b) => assert2::assert!(&b[..] == &buf[..]),
521 _ => panic!("expected Legacy"),
522 }
523 }
524
525 #[test]
526 fn roundtrip_v2() {
527 let p: RecordsPayload = sample_v2().into();
528 let mut buf = BytesMut::new();
529 p.encode_to(&mut buf).unwrap();
530 let back = RecordsPayload::from_bytes(buf.freeze()).unwrap();
531 assert2::assert!(p == back);
532 assert2::assert!(p.payload_len() == back.payload_len());
533 }
534
535 #[test]
536 fn encode_decode_via_traits() {
537 let p: RecordsPayload = sample_v2().into();
538 let mut buf = BytesMut::new();
539 <RecordsPayload as crate::Encode>::encode(&p, &mut buf, 0).unwrap();
540 let mut cur: &[u8] = &buf;
541 let back = <RecordsPayload as crate::Decode>::decode(&mut cur, 0).unwrap();
542 assert2::assert!(p == back);
543 }
544
545 #[test]
546 fn borrowed_dispatches() {
547 let rb = sample_v2();
548 let mut buf = BytesMut::new();
549 rb.encode(&mut buf).unwrap();
550 let frozen = buf.freeze();
551 let p = RecordsPayloadBorrowed::from_slice(&frozen).unwrap();
552 assert2::assert!(matches!(p, RecordsPayloadBorrowed::V2(_)));
553 let owned = p.to_owned().unwrap();
554 match owned {
555 RecordsPayload::V2(batches) => assert2::assert!(batches[0].base_offset == 42),
556 _ => panic!("expected V2"),
557 }
558 }
559
560 #[test]
561 fn from_record_batch() {
562 let rb = sample_v2();
563 let p: RecordsPayload = rb.clone().into();
564 assert2::assert!(p.as_v2() == Some(&[rb][..]));
565 assert2::assert!(p.as_legacy().is_none());
566 }
567
568 fn legacy_bytes() -> Bytes {
569 let mut buf = vec![0u8; 24];
570 buf[16] = 1;
571 for (i, b) in (b'a'..=b'h').enumerate() {
572 buf[17 + i % 7] = b;
573 }
574 Bytes::from(buf)
575 }
576
577 #[test]
578 fn legacy_payload_len_and_encode_owned() {
579 let bytes = legacy_bytes();
580 let p = RecordsPayload::from_bytes(bytes.clone()).unwrap();
581 assert2::assert!(p == RecordsPayload::Legacy(bytes.clone()));
582 assert2::assert!(p.payload_len() == bytes.len());
583
584 let mut out = BytesMut::new();
585 p.encode_to(&mut out).unwrap();
586 assert2::assert!(&out[..] == &bytes[..]);
587 }
588
589 #[test]
590 fn legacy_roundtrip_via_traits() {
591 let bytes = legacy_bytes();
592 let p = RecordsPayload::from_bytes(bytes.clone()).unwrap();
593 let mut buf = BytesMut::new();
594 <RecordsPayload as crate::Encode>::encode(&p, &mut buf, 0).unwrap();
595 assert2::assert!(<RecordsPayload as crate::Encode>::encoded_len(&p, 0) == bytes.len());
596 let mut cur: &[u8] = &buf;
597 let back = <RecordsPayload as crate::Decode>::decode(&mut cur, 0).unwrap();
598 assert2::assert!(matches!(back, RecordsPayload::Legacy(_)));
599 assert2::assert!(back.as_legacy().unwrap() == &bytes);
600 }
601
602 #[test]
603 fn owned_default_is_empty_v2() {
604 let p = RecordsPayload::default();
605 assert2::assert!(matches!(p, RecordsPayload::V2(ref v) if v.is_empty()));
606 }
607
608 #[test]
609 fn looks_like_v2_rejects_short_buffer() {
610 let short = Bytes::from_static(&[0u8; 10]);
612 let p = RecordsPayload::from_bytes(short.clone()).unwrap();
613 assert2::assert!(p.as_legacy() == Some(&short));
614 }
615
616 #[test]
617 fn borrowed_legacy_roundtrip() {
618 let bytes = legacy_bytes();
619 let p = RecordsPayloadBorrowed::from_slice(&bytes).unwrap();
620 assert2::assert!(matches!(p, RecordsPayloadBorrowed::Legacy(_)));
621 assert2::assert!(p.payload_len() == bytes.len());
622
623 let mut out = BytesMut::new();
624 p.encode_to(&mut out).unwrap();
625 assert2::assert!(&out[..] == &bytes[..]);
626
627 let owned = p.to_owned().unwrap();
628 match owned {
629 RecordsPayload::Legacy(b) => assert2::assert!(&b[..] == &bytes[..]),
630 _ => panic!("expected Legacy"),
631 }
632 }
633
634 #[test]
635 fn borrowed_v2_payload_len_and_encode() {
636 let rb = sample_v2();
637 let mut buf = BytesMut::new();
638 rb.encode(&mut buf).unwrap();
639 let frozen = buf.freeze();
640 let p = RecordsPayloadBorrowed::from_slice(&frozen).unwrap();
641 assert2::assert!(p.payload_len() == frozen.len());
642
643 let mut out = BytesMut::new();
644 p.encode_to(&mut out).unwrap();
645 assert2::assert!(&out[..] == &frozen[..]);
646 }
647
648 #[test]
649 fn borrowed_encode_decode_via_traits() {
650 let rb = sample_v2();
651 let mut buf = BytesMut::new();
652 rb.encode(&mut buf).unwrap();
653 let frozen = buf.freeze();
654
655 let p = RecordsPayloadBorrowed::from_slice(&frozen).unwrap();
656 let mut out = BytesMut::new();
657 <RecordsPayloadBorrowed as crate::Encode>::encode(&p, &mut out, 0).unwrap();
658 assert2::assert!(
659 <RecordsPayloadBorrowed as crate::Encode>::encoded_len(&p, 0) == frozen.len()
660 );
661
662 let mut cur: &[u8] = &out;
663 let back =
664 <RecordsPayloadBorrowed as crate::DecodeBorrow>::decode_borrow(&mut cur, 0).unwrap();
665 assert2::assert!(matches!(back, RecordsPayloadBorrowed::V2(_)));
666 }
667
668 #[test]
669 fn borrowed_default_is_empty_v2() {
670 let p = RecordsPayloadBorrowed::default();
671 assert2::assert!(matches!(p, RecordsPayloadBorrowed::V2(ref v) if v.is_empty()));
672 }
673
674 #[test]
675 fn from_fetch_bytes_drops_incomplete_trailing_batch() {
676 let mut b0 = sample_v2();
681 b0.base_offset = 0;
682 let mut b1 = sample_v2();
683 b1.base_offset = 1;
684 let mut buf = BytesMut::new();
685 b0.encode(&mut buf).unwrap();
686 b1.encode(&mut buf).unwrap();
687 buf.extend_from_slice(&[0u8; 7]); let p = RecordsPayload::from_fetch_bytes(buf.freeze()).unwrap();
689 let batches = p.as_v2().expect("v2");
690 assert2::assert!(batches == &[b0, b1][..]);
691 }
692
693 #[test]
694 fn from_fetch_bytes_still_errors_on_corrupt_batch() {
695 let rb = sample_v2();
698 let mut buf = BytesMut::new();
699 rb.encode(&mut buf).unwrap();
700 let mut bytes = buf.to_vec();
701 bytes[61] ^= 0xFF;
703 let err = RecordsPayload::from_fetch_bytes(Bytes::from(bytes)).unwrap_err();
704 assert2::assert!(matches!(err, RecordsError::CrcMismatch { .. }));
705 }
706
707 #[test]
708 fn from_fetch_bytes_legacy_passes_through() {
709 let bytes = legacy_bytes();
710 let p = RecordsPayload::from_fetch_bytes(bytes.clone()).unwrap();
711 assert2::assert!(p.as_legacy() == Some(&bytes));
712 }
713
714 #[test]
715 fn from_fetch_bytes_empty_is_empty_v2() {
716 let p = RecordsPayload::from_fetch_bytes(Bytes::new()).unwrap();
719 assert2::assert!(matches!(p, RecordsPayload::Legacy(ref b) if b.is_empty()));
720 }
721}