kafrust_protocol/api/
fetch.rs

1use crate::api::produce::RecordBatchHeader;
2use crate::codec::{DecodeLimits, Decoder, Encoder};
3use crate::error::{Error, Result};
4use crate::header::RequestHeader;
5use crate::record_batch::{decompress_record_batch_records_with_limit, RecordBatchCompression};
6
7pub const API_KEY: i16 = 1;
8
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct FetchRequestV2 {
11    pub correlation_id: i32,
12    pub client_id: Option<String>,
13    pub replica_id: i32,
14    pub max_wait_ms: i32,
15    pub min_bytes: i32,
16    pub topics: Vec<FetchTopicV2>,
17}
18
19#[derive(Debug, Clone, PartialEq, Eq)]
20pub struct FetchRequestV4 {
21    pub correlation_id: i32,
22    pub client_id: Option<String>,
23    pub replica_id: i32,
24    pub max_wait_ms: i32,
25    pub min_bytes: i32,
26    pub max_bytes: i32,
27    pub isolation_level: i8,
28    pub topics: Vec<FetchTopicV2>,
29}
30
31/// Fetch request version 11 with rack-aware read selection.
32///
33/// Version 11 keeps the non-flexible wire format used by the existing direct
34/// consumer while adding fetch-session fields and the consumer rack ID.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct FetchRequestV11 {
37    pub correlation_id: i32,
38    pub client_id: Option<String>,
39    pub replica_id: i32,
40    pub max_wait_ms: i32,
41    pub min_bytes: i32,
42    pub max_bytes: i32,
43    pub isolation_level: i8,
44    pub session_id: i32,
45    pub session_epoch: i32,
46    pub topics: Vec<FetchTopicV11>,
47    pub forgotten_topics: Vec<FetchForgottenTopicV11>,
48    pub rack_id: String,
49}
50
51/// Fetch request version 12 with flexible encoding and rack-aware read selection.
52#[derive(Debug, Clone, PartialEq, Eq)]
53pub struct FetchRequestV12 {
54    pub correlation_id: i32,
55    pub client_id: Option<String>,
56    pub replica_id: i32,
57    pub max_wait_ms: i32,
58    pub min_bytes: i32,
59    pub max_bytes: i32,
60    pub isolation_level: i8,
61    pub session_id: i32,
62    pub session_epoch: i32,
63    pub topics: Vec<FetchTopicV12>,
64    pub forgotten_topics: Vec<FetchForgottenTopicV12>,
65    pub rack_id: String,
66}
67
68impl FetchRequestV12 {
69    pub fn encode(&self) -> Result<Vec<u8>> {
70        let mut encoder = Encoder::new();
71        RequestHeader {
72            api_key: API_KEY,
73            api_version: 12,
74            correlation_id: self.correlation_id,
75            client_id: self.client_id.clone(),
76        }
77        .encode_v2(&mut encoder)?;
78        encoder.write_i32(self.replica_id);
79        encoder.write_i32(self.max_wait_ms);
80        encoder.write_i32(self.min_bytes);
81        encoder.write_i32(self.max_bytes);
82        encoder.write_i8(self.isolation_level);
83        encoder.write_i32(self.session_id);
84        encoder.write_i32(self.session_epoch);
85        encoder.write_compact_array(Some(self.topics.as_slice()), |encoder, topic| {
86            topic.encode(encoder)
87        })?;
88        encoder.write_compact_array(Some(self.forgotten_topics.as_slice()), |encoder, topic| {
89            topic.encode(encoder)
90        })?;
91        encoder.write_compact_string(&self.rack_id)?;
92        encoder.write_empty_tagged_fields();
93        Ok(encoder.into_bytes())
94    }
95}
96
97#[derive(Debug, Clone, PartialEq, Eq)]
98pub struct FetchTopicV12 {
99    pub name: String,
100    pub partitions: Vec<FetchPartitionV12>,
101}
102
103impl FetchTopicV12 {
104    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
105        encoder.write_compact_string(&self.name)?;
106        encoder.write_compact_array(Some(self.partitions.as_slice()), |encoder, partition| {
107            partition.encode(encoder)
108        })?;
109        encoder.write_empty_tagged_fields();
110        Ok(())
111    }
112}
113
114#[derive(Debug, Clone, PartialEq, Eq)]
115pub struct FetchPartitionV12 {
116    pub partition_index: i32,
117    pub current_leader_epoch: i32,
118    pub fetch_offset: i64,
119    pub last_fetched_epoch: i32,
120    pub log_start_offset: i64,
121    pub max_bytes: i32,
122}
123
124impl FetchPartitionV12 {
125    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
126        encoder.write_i32(self.partition_index);
127        encoder.write_i32(self.current_leader_epoch);
128        encoder.write_i64(self.fetch_offset);
129        encoder.write_i32(self.last_fetched_epoch);
130        encoder.write_i64(self.log_start_offset);
131        encoder.write_i32(self.max_bytes);
132        encoder.write_empty_tagged_fields();
133        Ok(())
134    }
135}
136
137#[derive(Debug, Clone, PartialEq, Eq)]
138pub struct FetchForgottenTopicV12 {
139    pub name: String,
140    pub partitions: Vec<i32>,
141}
142
143impl FetchForgottenTopicV12 {
144    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
145        encoder.write_compact_string(&self.name)?;
146        encoder.write_compact_array(Some(self.partitions.as_slice()), |encoder, partition| {
147            encoder.write_i32(*partition);
148            Ok(())
149        })?;
150        encoder.write_empty_tagged_fields();
151        Ok(())
152    }
153}
154
155impl FetchRequestV11 {
156    pub fn encode(&self) -> Result<Vec<u8>> {
157        let mut encoder = Encoder::new();
158        RequestHeader {
159            api_key: API_KEY,
160            api_version: 11,
161            correlation_id: self.correlation_id,
162            client_id: self.client_id.clone(),
163        }
164        .encode_v1(&mut encoder)?;
165        encoder.write_i32(self.replica_id);
166        encoder.write_i32(self.max_wait_ms);
167        encoder.write_i32(self.min_bytes);
168        encoder.write_i32(self.max_bytes);
169        encoder.write_i8(self.isolation_level);
170        encoder.write_i32(self.session_id);
171        encoder.write_i32(self.session_epoch);
172        encoder.write_array(Some(self.topics.as_slice()), |encoder, topic| {
173            topic.encode(encoder)
174        })?;
175        encoder.write_array(Some(self.forgotten_topics.as_slice()), |encoder, topic| {
176            topic.encode(encoder)
177        })?;
178        encoder.write_string(&self.rack_id)?;
179        Ok(encoder.into_bytes())
180    }
181}
182
183#[derive(Debug, Clone, PartialEq, Eq)]
184pub struct FetchTopicV11 {
185    pub name: String,
186    pub partitions: Vec<FetchPartitionV11>,
187}
188
189impl FetchTopicV11 {
190    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
191        encoder.write_string(&self.name)?;
192        encoder.write_array(Some(self.partitions.as_slice()), |encoder, partition| {
193            partition.encode(encoder)
194        })
195    }
196}
197
198#[derive(Debug, Clone, PartialEq, Eq)]
199pub struct FetchPartitionV11 {
200    pub partition_index: i32,
201    pub current_leader_epoch: i32,
202    pub fetch_offset: i64,
203    pub log_start_offset: i64,
204    pub max_bytes: i32,
205}
206
207impl FetchPartitionV11 {
208    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
209        encoder.write_i32(self.partition_index);
210        encoder.write_i32(self.current_leader_epoch);
211        encoder.write_i64(self.fetch_offset);
212        encoder.write_i64(self.log_start_offset);
213        encoder.write_i32(self.max_bytes);
214        Ok(())
215    }
216}
217
218#[derive(Debug, Clone, PartialEq, Eq)]
219pub struct FetchForgottenTopicV11 {
220    pub name: String,
221    pub partitions: Vec<i32>,
222}
223
224impl FetchForgottenTopicV11 {
225    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
226        encoder.write_string(&self.name)?;
227        encoder.write_array(Some(self.partitions.as_slice()), |encoder, partition| {
228            encoder.write_i32(*partition);
229            Ok(())
230        })
231    }
232}
233
234impl FetchRequestV4 {
235    pub fn encode(&self) -> Result<Vec<u8>> {
236        let mut encoder = Encoder::new();
237        RequestHeader {
238            api_key: API_KEY,
239            api_version: 4,
240            correlation_id: self.correlation_id,
241            client_id: self.client_id.clone(),
242        }
243        .encode_v1(&mut encoder)?;
244        encoder.write_i32(self.replica_id);
245        encoder.write_i32(self.max_wait_ms);
246        encoder.write_i32(self.min_bytes);
247        encoder.write_i32(self.max_bytes);
248        encoder.write_i8(self.isolation_level);
249        encoder.write_array(Some(self.topics.as_slice()), |encoder, topic| {
250            topic.encode(encoder)
251        })?;
252        Ok(encoder.into_bytes())
253    }
254}
255
256impl FetchRequestV2 {
257    pub fn encode(&self) -> Result<Vec<u8>> {
258        let mut encoder = Encoder::new();
259        RequestHeader {
260            api_key: API_KEY,
261            api_version: 2,
262            correlation_id: self.correlation_id,
263            client_id: self.client_id.clone(),
264        }
265        .encode_v1(&mut encoder)?;
266        encoder.write_i32(self.replica_id);
267        encoder.write_i32(self.max_wait_ms);
268        encoder.write_i32(self.min_bytes);
269        encoder.write_array(Some(self.topics.as_slice()), |encoder, topic| {
270            topic.encode(encoder)
271        })?;
272        Ok(encoder.into_bytes())
273    }
274}
275
276#[derive(Debug, Clone, PartialEq, Eq)]
277pub struct FetchTopicV2 {
278    pub name: String,
279    pub partitions: Vec<FetchPartitionV2>,
280}
281
282impl FetchTopicV2 {
283    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
284        encoder.write_string(&self.name)?;
285        encoder.write_array(Some(self.partitions.as_slice()), |encoder, partition| {
286            partition.encode(encoder)
287        })
288    }
289}
290
291#[derive(Debug, Clone, PartialEq, Eq)]
292pub struct FetchPartitionV2 {
293    pub partition_index: i32,
294    pub fetch_offset: i64,
295    pub max_bytes: i32,
296}
297
298impl FetchPartitionV2 {
299    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
300        encoder.write_i32(self.partition_index);
301        encoder.write_i64(self.fetch_offset);
302        encoder.write_i32(self.max_bytes);
303        Ok(())
304    }
305}
306
307#[derive(Debug, Clone, PartialEq, Eq)]
308pub struct FetchResponseV2 {
309    pub throttle_time_ms: i32,
310    pub responses: Vec<FetchTopicResponseV2>,
311}
312
313#[derive(Debug, Clone, PartialEq, Eq)]
314pub struct FetchResponseV4 {
315    pub throttle_time_ms: i32,
316    pub responses: Vec<FetchTopicResponseV4>,
317}
318
319/// Fetch response version 11 with broker-selected read replicas.
320#[derive(Debug, Clone, PartialEq, Eq)]
321pub struct FetchResponseV11 {
322    pub throttle_time_ms: i32,
323    pub error_code: i16,
324    pub session_id: i32,
325    pub responses: Vec<FetchTopicResponseV11>,
326}
327
328/// Fetch response version 12 with flexible encoding and broker-selected reads.
329#[derive(Debug, Clone, PartialEq, Eq)]
330pub struct FetchResponseV12 {
331    pub throttle_time_ms: i32,
332    pub error_code: i16,
333    pub session_id: i32,
334    pub responses: Vec<FetchTopicResponseV12>,
335}
336
337impl FetchResponseV12 {
338    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
339        Ok(Self {
340            throttle_time_ms: decoder.read_i32()?,
341            error_code: decoder.read_i16()?,
342            session_id: decoder.read_i32()?,
343            responses: decoder
344                .read_compact_array("fetch responses", FetchTopicResponseV12::decode)?
345                .unwrap_or_default(),
346        })
347        .and_then(|response| {
348            decoder.read_tagged_fields()?;
349            Ok(response)
350        })
351    }
352}
353
354#[derive(Debug, Clone, PartialEq, Eq)]
355pub struct FetchTopicResponseV12 {
356    pub name: String,
357    pub partitions: Vec<FetchPartitionResponseV12>,
358}
359
360impl FetchTopicResponseV12 {
361    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
362        let name = decoder.read_compact_string()?;
363        let partitions = decoder
364            .read_compact_array(
365                "fetch partition responses",
366                FetchPartitionResponseV12::decode,
367            )?
368            .unwrap_or_default();
369        decoder.read_tagged_fields()?;
370        Ok(Self { name, partitions })
371    }
372}
373
374#[derive(Debug, Clone, PartialEq, Eq)]
375pub struct FetchPartitionResponseV12 {
376    pub partition_index: i32,
377    pub error_code: i16,
378    pub high_watermark: i64,
379    pub last_stable_offset: i64,
380    pub log_start_offset: i64,
381    pub aborted_transactions: Vec<AbortedTransactionV12>,
382    pub preferred_read_replica: i32,
383    pub records: Vec<MessageSetRecord>,
384}
385
386impl FetchPartitionResponseV12 {
387    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
388        let limits = decoder.limits();
389        let partition_index = decoder.read_i32()?;
390        let error_code = decoder.read_i16()?;
391        let high_watermark = decoder.read_i64()?;
392        let last_stable_offset = decoder.read_i64()?;
393        let log_start_offset = decoder.read_i64()?;
394        let aborted_transactions = decoder
395            .read_compact_array("aborted transactions", AbortedTransactionV12::decode)?
396            .unwrap_or_default();
397        let preferred_read_replica = decoder.read_i32()?;
398        let records = decoder
399            .read_compact_nullable_bytes()?
400            .map(|bytes| decode_message_set(&bytes, limits))
401            .transpose()?
402            .unwrap_or_default();
403        decoder.read_tagged_fields()?;
404        Ok(Self {
405            partition_index,
406            error_code,
407            high_watermark,
408            last_stable_offset,
409            log_start_offset,
410            aborted_transactions,
411            preferred_read_replica,
412            records,
413        })
414    }
415}
416
417#[derive(Debug, Clone, PartialEq, Eq)]
418pub struct AbortedTransactionV12 {
419    pub producer_id: i64,
420    pub first_offset: i64,
421}
422
423impl AbortedTransactionV12 {
424    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
425        let producer_id = decoder.read_i64()?;
426        let first_offset = decoder.read_i64()?;
427        decoder.read_tagged_fields()?;
428        Ok(Self {
429            producer_id,
430            first_offset,
431        })
432    }
433}
434
435impl FetchResponseV11 {
436    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
437        Ok(Self {
438            throttle_time_ms: decoder.read_i32()?,
439            error_code: decoder.read_i16()?,
440            session_id: decoder.read_i32()?,
441            responses: decoder
442                .read_array("fetch responses", FetchTopicResponseV11::decode)?
443                .unwrap_or_default(),
444        })
445    }
446}
447
448#[derive(Debug, Clone, PartialEq, Eq)]
449pub struct FetchTopicResponseV11 {
450    pub name: String,
451    pub partitions: Vec<FetchPartitionResponseV11>,
452}
453
454impl FetchTopicResponseV11 {
455    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
456        Ok(Self {
457            name: decoder.read_string()?,
458            partitions: decoder
459                .read_array(
460                    "fetch partition responses",
461                    FetchPartitionResponseV11::decode,
462                )?
463                .unwrap_or_default(),
464        })
465    }
466}
467
468#[derive(Debug, Clone, PartialEq, Eq)]
469pub struct FetchPartitionResponseV11 {
470    pub partition_index: i32,
471    pub error_code: i16,
472    pub high_watermark: i64,
473    pub last_stable_offset: i64,
474    pub log_start_offset: i64,
475    pub aborted_transactions: Vec<AbortedTransactionV4>,
476    pub preferred_read_replica: i32,
477    pub records: Vec<MessageSetRecord>,
478}
479
480impl FetchPartitionResponseV11 {
481    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
482        let limits = decoder.limits();
483        Ok(Self {
484            partition_index: decoder.read_i32()?,
485            error_code: decoder.read_i16()?,
486            high_watermark: decoder.read_i64()?,
487            last_stable_offset: decoder.read_i64()?,
488            log_start_offset: decoder.read_i64()?,
489            aborted_transactions: decoder
490                .read_array("aborted transactions", AbortedTransactionV4::decode)?
491                .unwrap_or_default(),
492            preferred_read_replica: decoder.read_i32()?,
493            records: decode_message_set(&decoder.read_bytes()?, limits)?,
494        })
495    }
496}
497
498impl FetchResponseV4 {
499    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
500        Ok(Self {
501            throttle_time_ms: decoder.read_i32()?,
502            responses: decoder
503                .read_array("fetch responses", FetchTopicResponseV4::decode)?
504                .unwrap_or_default(),
505        })
506    }
507}
508
509#[derive(Debug, Clone, PartialEq, Eq)]
510pub struct FetchTopicResponseV4 {
511    pub name: String,
512    pub partitions: Vec<FetchPartitionResponseV4>,
513}
514
515impl FetchTopicResponseV4 {
516    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
517        Ok(Self {
518            name: decoder.read_string()?,
519            partitions: decoder
520                .read_array(
521                    "fetch partition responses",
522                    FetchPartitionResponseV4::decode,
523                )?
524                .unwrap_or_default(),
525        })
526    }
527}
528
529#[derive(Debug, Clone, PartialEq, Eq)]
530pub struct FetchPartitionResponseV4 {
531    pub partition_index: i32,
532    pub error_code: i16,
533    pub high_watermark: i64,
534    pub last_stable_offset: i64,
535    pub aborted_transactions: Vec<AbortedTransactionV4>,
536    pub records: Vec<MessageSetRecord>,
537}
538
539impl FetchPartitionResponseV4 {
540    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
541        let limits = decoder.limits();
542        Ok(Self {
543            partition_index: decoder.read_i32()?,
544            error_code: decoder.read_i16()?,
545            high_watermark: decoder.read_i64()?,
546            last_stable_offset: decoder.read_i64()?,
547            aborted_transactions: decoder
548                .read_array("aborted transactions", AbortedTransactionV4::decode)?
549                .unwrap_or_default(),
550            records: decode_message_set(&decoder.read_bytes()?, limits)?,
551        })
552    }
553}
554
555#[derive(Debug, Clone, PartialEq, Eq)]
556pub struct AbortedTransactionV4 {
557    pub producer_id: i64,
558    pub first_offset: i64,
559}
560
561impl AbortedTransactionV4 {
562    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
563        Ok(Self {
564            producer_id: decoder.read_i64()?,
565            first_offset: decoder.read_i64()?,
566        })
567    }
568}
569
570impl FetchResponseV2 {
571    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
572        Ok(Self {
573            throttle_time_ms: decoder.read_i32()?,
574            responses: decoder
575                .read_array("fetch responses", FetchTopicResponseV2::decode)?
576                .unwrap_or_default(),
577        })
578    }
579}
580
581#[derive(Debug, Clone, PartialEq, Eq)]
582pub struct FetchTopicResponseV2 {
583    pub name: String,
584    pub partitions: Vec<FetchPartitionResponseV2>,
585}
586
587impl FetchTopicResponseV2 {
588    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
589        Ok(Self {
590            name: decoder.read_string()?,
591            partitions: decoder
592                .read_array(
593                    "fetch partition responses",
594                    FetchPartitionResponseV2::decode,
595                )?
596                .unwrap_or_default(),
597        })
598    }
599}
600
601#[derive(Debug, Clone, PartialEq, Eq)]
602pub struct FetchPartitionResponseV2 {
603    pub partition_index: i32,
604    pub error_code: i16,
605    pub high_watermark: i64,
606    pub records: Vec<MessageSetRecord>,
607}
608
609impl FetchPartitionResponseV2 {
610    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
611        let limits = decoder.limits();
612        Ok(Self {
613            partition_index: decoder.read_i32()?,
614            error_code: decoder.read_i16()?,
615            high_watermark: decoder.read_i64()?,
616            records: decode_message_set(&decoder.read_bytes()?, limits)?,
617        })
618    }
619}
620
621#[derive(Debug, Clone, PartialEq, Eq)]
622pub struct MessageSetRecord {
623    pub offset: i64,
624    pub timestamp_ms: i64,
625    pub key: Option<Vec<u8>>,
626    pub value: Option<Vec<u8>>,
627    pub headers: Vec<RecordBatchHeader>,
628    pub producer_id: Option<i64>,
629    pub transactional: bool,
630    pub control: bool,
631}
632
633fn decode_message_set(bytes: &[u8], limits: DecodeLimits) -> Result<Vec<MessageSetRecord>> {
634    let mut decoder = Decoder::with_limits(bytes, limits);
635    let mut records = Vec::new();
636
637    while decoder.remaining() >= 12 {
638        let offset = decoder.read_i64()?;
639        let message_size = decoder.read_i32()?;
640        if message_size < 0 {
641            return Err(Error::NegativeLength {
642                kind: "message",
643                length: message_size,
644            });
645        }
646        let message_size =
647            usize::try_from(message_size).map_err(|_| Error::LengthOverflow("message"))?;
648        // Fetch responses may end with a partial trailing message set entry.
649        if decoder.remaining() < message_size {
650            break;
651        }
652        let message = decoder.read_exact(message_size)?;
653        let decoded = decode_message_or_batch(offset, message, limits)?;
654        let total = records
655            .len()
656            .checked_add(decoded.len())
657            .ok_or(Error::LengthOverflow("fetch records"))?;
658        decoder.ensure_collection_length("fetch records", total)?;
659        records.extend(decoded);
660    }
661
662    Ok(records)
663}
664
665fn decode_message_or_batch(
666    offset: i64,
667    bytes: &[u8],
668    limits: DecodeLimits,
669) -> Result<Vec<MessageSetRecord>> {
670    match bytes.get(4).copied() {
671        Some(2) => decode_record_batch(offset, bytes, limits),
672        _ => Ok(vec![decode_message(offset, bytes, limits)?]),
673    }
674}
675
676fn decode_message(offset: i64, bytes: &[u8], limits: DecodeLimits) -> Result<MessageSetRecord> {
677    let mut decoder = Decoder::with_limits(bytes, limits);
678    let _crc = decoder.read_i32()?;
679    let magic = decoder.read_i8()?;
680    let _attributes = decoder.read_i8()?;
681    let timestamp_ms = match magic {
682        0 => -1,
683        1 => decoder.read_i64()?,
684        _ => {
685            return Err(Error::UnsupportedVersion {
686                kind: "message magic",
687                version: i16::from(magic),
688            })
689        }
690    };
691    let key = decoder.read_nullable_bytes()?;
692    let value = decoder.read_nullable_bytes()?;
693
694    Ok(MessageSetRecord {
695        offset,
696        timestamp_ms,
697        key,
698        value,
699        headers: Vec::new(),
700        producer_id: None,
701        transactional: false,
702        control: false,
703    })
704}
705
706fn decode_record_batch(
707    base_offset: i64,
708    bytes: &[u8],
709    limits: DecodeLimits,
710) -> Result<Vec<MessageSetRecord>> {
711    let mut decoder = Decoder::with_limits(bytes, limits);
712    let _partition_leader_epoch = decoder.read_i32()?;
713    let magic = decoder.read_i8()?;
714    if magic != 2 {
715        return Err(Error::UnsupportedVersion {
716            kind: "record batch magic",
717            version: i16::from(magic),
718        });
719    }
720    let _crc = decoder.read_i32()?;
721    let attributes = decoder.read_i16()?;
722    let compression = RecordBatchCompression::from_attributes(attributes)?;
723    let _last_offset_delta = decoder.read_i32()?;
724    let base_timestamp = decoder.read_i64()?;
725    let _max_timestamp = decoder.read_i64()?;
726    let producer_id = decoder.read_i64()?;
727    let _producer_epoch = decoder.read_i16()?;
728    let _base_sequence = decoder.read_i32()?;
729    let record_count = decoder.read_i32()?;
730    if record_count < 0 {
731        return Err(Error::NegativeLength {
732            kind: "record batch records",
733            length: record_count,
734        });
735    }
736
737    let record_count =
738        usize::try_from(record_count).map_err(|_| Error::LengthOverflow("record batch records"))?;
739    decoder.ensure_collection_length("record batch records", record_count)?;
740    let record_bytes = if compression.is_compressed() {
741        let compressed = decoder.read_exact(decoder.remaining())?;
742        decompress_record_batch_records_with_limit(
743            compression,
744            compressed,
745            limits.max_decompressed_record_bytes(),
746        )?
747    } else {
748        if decoder.remaining() > limits.max_decompressed_record_bytes() {
749            return Err(Error::LimitExceeded {
750                kind: "decompressed record batch bytes",
751                actual: decoder.remaining(),
752                max: limits.max_decompressed_record_bytes(),
753            });
754        }
755        decoder.read_exact(decoder.remaining())?.to_vec()
756    };
757    let mut record_decoder = Decoder::with_limits(&record_bytes, limits);
758    let mut records = Vec::with_capacity(record_count);
759    for _ in 0..record_count {
760        let record_length = record_decoder.read_varint()?;
761        if record_length < 0 {
762            return Err(Error::NegativeLength {
763                kind: "record",
764                length: record_length,
765            });
766        }
767        let record_length =
768            usize::try_from(record_length).map_err(|_| Error::LengthOverflow("record"))?;
769        let record_bytes = record_decoder.read_exact(record_length)?;
770        records.push(decode_record(
771            base_offset,
772            base_timestamp,
773            producer_id,
774            attributes,
775            record_bytes,
776            limits,
777        )?);
778    }
779
780    Ok(records)
781}
782
783fn decode_record(
784    base_offset: i64,
785    base_timestamp: i64,
786    producer_id: i64,
787    batch_attributes: i16,
788    bytes: &[u8],
789    limits: DecodeLimits,
790) -> Result<MessageSetRecord> {
791    let mut decoder = Decoder::with_limits(bytes, limits);
792    let _attributes = decoder.read_i8()?;
793    let timestamp_delta = decoder.read_varlong()?;
794    let offset_delta = decoder.read_varint()?;
795    let key = decoder.read_varint_nullable_bytes()?;
796    let value = decoder.read_varint_nullable_bytes()?;
797    let header_count = decoder.read_varint()?;
798    if header_count < 0 {
799        return Err(Error::NegativeLength {
800            kind: "record headers",
801            length: header_count,
802        });
803    }
804    let header_count =
805        usize::try_from(header_count).map_err(|_| Error::LengthOverflow("record headers"))?;
806    decoder.ensure_collection_length("record headers", header_count)?;
807    let mut headers = Vec::with_capacity(header_count);
808    for _ in 0..header_count {
809        let header_key =
810            String::from_utf8(decoder.read_varint_bytes()?).map_err(|_| Error::InvalidUtf8)?;
811        let header_value = decoder.read_varint_nullable_bytes()?;
812        headers.push(RecordBatchHeader::new(header_key, header_value));
813    }
814
815    Ok(MessageSetRecord {
816        offset: base_offset.saturating_add(i64::from(offset_delta)),
817        timestamp_ms: base_timestamp.saturating_add(timestamp_delta),
818        key,
819        value,
820        headers,
821        producer_id: (producer_id >= 0).then_some(producer_id),
822        transactional: batch_attributes & 0x10 != 0,
823        control: batch_attributes & 0x20 != 0,
824    })
825}
826
827#[cfg(test)]
828#[allow(clippy::unwrap_used)]
829mod tests {
830    use super::{
831        FetchPartitionV11, FetchPartitionV12, FetchPartitionV2, FetchRequestV11, FetchRequestV12,
832        FetchRequestV2, FetchRequestV4, FetchResponseV11, FetchResponseV12, FetchResponseV2,
833        FetchResponseV4, FetchTopicV11, FetchTopicV12, FetchTopicV2, MessageSetRecord,
834        RecordBatchHeader,
835    };
836    use crate::codec::{Decoder, Encoder};
837
838    #[test]
839    fn encodes_fetch_request_v2() {
840        let request = FetchRequestV2 {
841            correlation_id: 7,
842            client_id: Some("kafrust".to_owned()),
843            replica_id: -1,
844            max_wait_ms: 500,
845            min_bytes: 1,
846            topics: vec![FetchTopicV2 {
847                name: "orders".to_owned(),
848                partitions: vec![FetchPartitionV2 {
849                    partition_index: 0,
850                    fetch_offset: 42,
851                    max_bytes: 1_048_576,
852                }],
853            }],
854        };
855
856        let bytes = request.encode().unwrap();
857        assert_eq!(&bytes[0..4], &[0, 1, 0, 2]);
858        assert!(bytes.len() > 40);
859    }
860
861    #[test]
862    fn encodes_fetch_request_v4() {
863        let request = FetchRequestV4 {
864            correlation_id: 8,
865            client_id: Some("kafrust".to_owned()),
866            replica_id: -1,
867            max_wait_ms: 500,
868            min_bytes: 1,
869            max_bytes: 1_048_576,
870            isolation_level: 0,
871            topics: vec![FetchTopicV2 {
872                name: "orders".to_owned(),
873                partitions: vec![FetchPartitionV2 {
874                    partition_index: 0,
875                    fetch_offset: 42,
876                    max_bytes: 1_048_576,
877                }],
878            }],
879        };
880
881        let bytes = request.encode().unwrap();
882        assert_eq!(&bytes[0..4], &[0, 1, 0, 4]);
883        assert_eq!(&bytes[4..8], &[0, 0, 0, 8]);
884        assert!(bytes.len() > 45);
885    }
886
887    #[test]
888    fn encodes_fetch_request_v11_with_rack_and_fetch_session_fields() {
889        let request = FetchRequestV11 {
890            correlation_id: 9,
891            client_id: Some("kafrust".to_owned()),
892            replica_id: -1,
893            max_wait_ms: 500,
894            min_bytes: 1,
895            max_bytes: 1_048_576,
896            isolation_level: 1,
897            session_id: 0,
898            session_epoch: 0,
899            topics: vec![FetchTopicV11 {
900                name: "orders".to_owned(),
901                partitions: vec![FetchPartitionV11 {
902                    partition_index: 0,
903                    current_leader_epoch: -1,
904                    fetch_offset: 42,
905                    log_start_offset: -1,
906                    max_bytes: 1_048_576,
907                }],
908            }],
909            forgotten_topics: Vec::new(),
910            rack_id: "rack-a".to_owned(),
911        };
912
913        let bytes = request.encode().unwrap();
914        let mut decoder = Decoder::new(&bytes);
915        assert_eq!(decoder.read_i16().unwrap(), 1);
916        assert_eq!(decoder.read_i16().unwrap(), 11);
917        assert_eq!(decoder.read_i32().unwrap(), 9);
918        assert_eq!(
919            decoder.read_nullable_string().unwrap().as_deref(),
920            Some("kafrust")
921        );
922        assert_eq!(decoder.read_i32().unwrap(), -1);
923        assert_eq!(decoder.read_i32().unwrap(), 500);
924        assert_eq!(decoder.read_i32().unwrap(), 1);
925        assert_eq!(decoder.read_i32().unwrap(), 1_048_576);
926        assert_eq!(decoder.read_i8().unwrap(), 1);
927        assert_eq!(decoder.read_i32().unwrap(), 0);
928        assert_eq!(decoder.read_i32().unwrap(), 0);
929        assert_eq!(decoder.read_i32().unwrap(), 1);
930        assert_eq!(decoder.read_string().unwrap(), "orders");
931        assert_eq!(decoder.read_i32().unwrap(), 1);
932        assert_eq!(decoder.read_i32().unwrap(), 0);
933        assert_eq!(decoder.read_i32().unwrap(), -1);
934        assert_eq!(decoder.read_i64().unwrap(), 42);
935        assert_eq!(decoder.read_i64().unwrap(), -1);
936        assert_eq!(decoder.read_i32().unwrap(), 1_048_576);
937        assert_eq!(decoder.read_i32().unwrap(), 0);
938        assert_eq!(decoder.read_string().unwrap(), "rack-a");
939        assert!(decoder.is_empty());
940    }
941
942    #[test]
943    fn decodes_fetch_response_v11_with_preferred_read_replica() {
944        let mut bytes = Encoder::new();
945        bytes.write_i32(3);
946        bytes.write_i16(0);
947        bytes.write_i32(17);
948        bytes.write_i32(1);
949        bytes.write_string("orders").unwrap();
950        bytes.write_i32(1);
951        bytes.write_i32(0);
952        bytes.write_i16(0);
953        bytes.write_i64(43);
954        bytes.write_i64(42);
955        bytes.write_i64(40);
956        bytes.write_i32(0);
957        bytes.write_i32(2);
958        bytes.write_bytes(&[]).unwrap();
959
960        let bytes = bytes.into_bytes();
961        let mut decoder = Decoder::new(&bytes);
962        let response = FetchResponseV11::decode_body(&mut decoder).unwrap();
963        let partition = &response.responses[0].partitions[0];
964
965        assert_eq!(response.throttle_time_ms, 3);
966        assert_eq!(response.session_id, 17);
967        assert_eq!(partition.log_start_offset, 40);
968        assert_eq!(partition.preferred_read_replica, 2);
969        assert!(partition.records.is_empty());
970        assert!(decoder.is_empty());
971    }
972
973    #[test]
974    fn encodes_fetch_request_v12_with_flexible_rack_fields() {
975        let request = FetchRequestV12 {
976            correlation_id: 10,
977            client_id: Some("kafrust".to_owned()),
978            replica_id: -1,
979            max_wait_ms: 500,
980            min_bytes: 1,
981            max_bytes: 1_048_576,
982            isolation_level: 1,
983            session_id: 0,
984            session_epoch: 0,
985            topics: vec![FetchTopicV12 {
986                name: "orders".to_owned(),
987                partitions: vec![FetchPartitionV12 {
988                    partition_index: 0,
989                    current_leader_epoch: -1,
990                    fetch_offset: 42,
991                    last_fetched_epoch: -1,
992                    log_start_offset: -1,
993                    max_bytes: 1_048_576,
994                }],
995            }],
996            forgotten_topics: Vec::new(),
997            rack_id: "rack-a".to_owned(),
998        };
999
1000        let bytes = request.encode().unwrap();
1001        let mut decoder = Decoder::new(&bytes);
1002        assert_eq!(decoder.read_i16().unwrap(), 1);
1003        assert_eq!(decoder.read_i16().unwrap(), 12);
1004        assert_eq!(decoder.read_i32().unwrap(), 10);
1005        assert_eq!(
1006            decoder.read_nullable_string().unwrap().as_deref(),
1007            Some("kafrust")
1008        );
1009        assert_eq!(decoder.read_unsigned_varint().unwrap(), 0);
1010        assert_eq!(decoder.read_i32().unwrap(), -1);
1011        assert_eq!(decoder.read_i32().unwrap(), 500);
1012        assert_eq!(decoder.read_i32().unwrap(), 1);
1013        assert_eq!(decoder.read_i32().unwrap(), 1_048_576);
1014        assert_eq!(decoder.read_i8().unwrap(), 1);
1015        assert_eq!(decoder.read_i32().unwrap(), 0);
1016        assert_eq!(decoder.read_i32().unwrap(), 0);
1017        assert_eq!(decoder.read_unsigned_varint().unwrap(), 2);
1018        assert_eq!(decoder.read_compact_string().unwrap(), "orders");
1019        assert_eq!(decoder.read_unsigned_varint().unwrap(), 2);
1020        assert_eq!(decoder.read_i32().unwrap(), 0);
1021        assert_eq!(decoder.read_i32().unwrap(), -1);
1022        assert_eq!(decoder.read_i64().unwrap(), 42);
1023        assert_eq!(decoder.read_i32().unwrap(), -1);
1024        assert_eq!(decoder.read_i64().unwrap(), -1);
1025        assert_eq!(decoder.read_i32().unwrap(), 1_048_576);
1026        assert_eq!(decoder.read_unsigned_varint().unwrap(), 0);
1027        assert_eq!(decoder.read_unsigned_varint().unwrap(), 0);
1028        assert_eq!(decoder.read_unsigned_varint().unwrap(), 1);
1029        assert_eq!(decoder.read_compact_string().unwrap(), "rack-a");
1030        assert_eq!(decoder.read_unsigned_varint().unwrap(), 0);
1031        assert!(decoder.is_empty());
1032    }
1033
1034    #[test]
1035    fn decodes_fetch_response_v12_with_preferred_read_replica() {
1036        let mut bytes = Encoder::new();
1037        bytes.write_i32(3);
1038        bytes.write_i16(0);
1039        bytes.write_i32(17);
1040        bytes.write_unsigned_varint(2);
1041        bytes.write_compact_string("orders").unwrap();
1042        bytes.write_unsigned_varint(2);
1043        bytes.write_i32(0);
1044        bytes.write_i16(0);
1045        bytes.write_i64(43);
1046        bytes.write_i64(42);
1047        bytes.write_i64(40);
1048        bytes.write_unsigned_varint(2);
1049        bytes.write_i64(7);
1050        bytes.write_i64(40);
1051        bytes.write_unsigned_varint(0);
1052        bytes.write_i32(2);
1053        bytes.write_compact_nullable_bytes(Some(&[])).unwrap();
1054        bytes.write_unsigned_varint(0);
1055        bytes.write_unsigned_varint(0);
1056        bytes.write_unsigned_varint(0);
1057
1058        let bytes = bytes.into_bytes();
1059        let mut decoder = Decoder::new(&bytes);
1060        let response = FetchResponseV12::decode_body(&mut decoder).unwrap();
1061        let partition = &response.responses[0].partitions[0];
1062
1063        assert_eq!(response.throttle_time_ms, 3);
1064        assert_eq!(response.session_id, 17);
1065        assert_eq!(partition.log_start_offset, 40);
1066        assert_eq!(partition.preferred_read_replica, 2);
1067        assert_eq!(partition.aborted_transactions[0].producer_id, 7);
1068        assert!(partition.records.is_empty());
1069        assert!(decoder.is_empty());
1070    }
1071
1072    #[test]
1073    fn decodes_fetch_response_v4_with_aborted_transaction() {
1074        let mut bytes = Encoder::new();
1075        bytes.write_i32(0);
1076        bytes.write_i32(1);
1077        bytes.write_string("orders").unwrap();
1078        bytes.write_i32(1);
1079        bytes.write_i32(0);
1080        bytes.write_i16(0);
1081        bytes.write_i64(43);
1082        bytes.write_i64(42);
1083        bytes.write_i32(1);
1084        bytes.write_i64(7);
1085        bytes.write_i64(40);
1086        bytes.write_bytes(&[]).unwrap();
1087        let bytes = bytes.into_bytes();
1088
1089        let mut decoder = Decoder::new(&bytes);
1090        let response = FetchResponseV4::decode_body(&mut decoder).unwrap();
1091        let partition = &response.responses[0].partitions[0];
1092
1093        assert_eq!(partition.high_watermark, 43);
1094        assert_eq!(partition.last_stable_offset, 42);
1095        assert_eq!(partition.aborted_transactions[0].producer_id, 7);
1096        assert_eq!(partition.aborted_transactions[0].first_offset, 40);
1097        assert!(partition.records.is_empty());
1098        assert!(decoder.is_empty());
1099    }
1100
1101    #[test]
1102    fn decodes_fetch_response_v2_with_message_set() {
1103        let mut message = Encoder::new();
1104        message.write_i32(0);
1105        message.write_i8(1);
1106        message.write_i8(0);
1107        message.write_i64(123);
1108        message.write_nullable_bytes(Some(b"order-1")).unwrap();
1109        message.write_nullable_bytes(Some(b"created")).unwrap();
1110        let message = message.into_bytes();
1111
1112        let mut set = Encoder::new();
1113        set.write_i64(42);
1114        set.write_i32(i32::try_from(message.len()).unwrap());
1115        set.write_raw(&message);
1116        let set = set.into_bytes();
1117
1118        let mut bytes = Encoder::new();
1119        bytes.write_i32(0);
1120        bytes.write_i32(1);
1121        bytes.write_string("orders").unwrap();
1122        bytes.write_i32(1);
1123        bytes.write_i32(0);
1124        bytes.write_i16(0);
1125        bytes.write_i64(43);
1126        bytes.write_bytes(&set).unwrap();
1127        let bytes = bytes.into_bytes();
1128
1129        let mut decoder = Decoder::new(&bytes);
1130        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
1131        let record = MessageSetRecord {
1132            offset: 42,
1133            timestamp_ms: 123,
1134            key: Some(b"order-1".to_vec()),
1135            value: Some(b"created".to_vec()),
1136            headers: Vec::new(),
1137            producer_id: None,
1138            transactional: false,
1139            control: false,
1140        };
1141
1142        assert_eq!(response.throttle_time_ms, 0);
1143        assert_eq!(response.responses[0].partitions[0].high_watermark, 43);
1144        assert_eq!(response.responses[0].partitions[0].records, vec![record]);
1145        assert!(decoder.is_empty());
1146    }
1147
1148    #[test]
1149    fn decodes_fetch_response_v2_ignores_partial_trailing_message_set_entry() {
1150        let mut message = Encoder::new();
1151        message.write_i32(0);
1152        message.write_i8(1);
1153        message.write_i8(0);
1154        message.write_i64(123);
1155        message.write_nullable_bytes(Some(b"order-1")).unwrap();
1156        message.write_nullable_bytes(Some(b"created")).unwrap();
1157        let message = message.into_bytes();
1158
1159        let mut set = Encoder::new();
1160        set.write_i64(42);
1161        set.write_i32(i32::try_from(message.len()).unwrap());
1162        set.write_raw(&message);
1163        set.write_i64(-1);
1164        set.write_i32(61);
1165        set.write_raw(&[0; 22]);
1166        let set = set.into_bytes();
1167
1168        let mut bytes = Encoder::new();
1169        bytes.write_i32(0);
1170        bytes.write_i32(1);
1171        bytes.write_string("orders").unwrap();
1172        bytes.write_i32(1);
1173        bytes.write_i32(0);
1174        bytes.write_i16(0);
1175        bytes.write_i64(43);
1176        bytes.write_bytes(&set).unwrap();
1177        let bytes = bytes.into_bytes();
1178
1179        let mut decoder = Decoder::new(&bytes);
1180        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
1181
1182        assert_eq!(response.responses[0].partitions[0].records.len(), 1);
1183        assert_eq!(response.responses[0].partitions[0].records[0].offset, 42);
1184        assert!(decoder.is_empty());
1185    }
1186
1187    #[test]
1188    fn decodes_fetch_response_v2_with_record_batch() {
1189        let mut record = Vec::new();
1190        record.push(0);
1191        write_varlong(&mut record, 5);
1192        write_varint(&mut record, 0);
1193        write_varint(&mut record, 7);
1194        record.extend_from_slice(b"order-1");
1195        write_varint(&mut record, 7);
1196        record.extend_from_slice(b"created");
1197        write_varint(&mut record, 2);
1198        write_varint(&mut record, 6);
1199        record.extend_from_slice(b"source");
1200        write_varint(&mut record, 8);
1201        record.extend_from_slice(b"checkout");
1202        write_varint(&mut record, 9);
1203        record.extend_from_slice(b"tombstone");
1204        write_varint(&mut record, -1);
1205
1206        let mut batch = Encoder::new();
1207        batch.write_i32(0);
1208        batch.write_i8(2);
1209        batch.write_i32(0);
1210        batch.write_i16(0x10);
1211        batch.write_i32(0);
1212        batch.write_i64(1_000);
1213        batch.write_i64(1_005);
1214        batch.write_i64(7);
1215        batch.write_i16(-1);
1216        batch.write_i32(-1);
1217        batch.write_i32(1);
1218        let mut encoded_record = Vec::new();
1219        write_varint(&mut encoded_record, i32::try_from(record.len()).unwrap());
1220        encoded_record.extend_from_slice(&record);
1221        batch.write_raw(&encoded_record);
1222        let batch = batch.into_bytes();
1223
1224        let mut set = Encoder::new();
1225        set.write_i64(42);
1226        set.write_i32(i32::try_from(batch.len()).unwrap());
1227        set.write_raw(&batch);
1228        let set = set.into_bytes();
1229
1230        let mut bytes = Encoder::new();
1231        bytes.write_i32(0);
1232        bytes.write_i32(1);
1233        bytes.write_string("orders").unwrap();
1234        bytes.write_i32(1);
1235        bytes.write_i32(0);
1236        bytes.write_i16(0);
1237        bytes.write_i64(43);
1238        bytes.write_bytes(&set).unwrap();
1239        let bytes = bytes.into_bytes();
1240
1241        let mut decoder = Decoder::new(&bytes);
1242        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
1243        let record = MessageSetRecord {
1244            offset: 42,
1245            timestamp_ms: 1_005,
1246            key: Some(b"order-1".to_vec()),
1247            value: Some(b"created".to_vec()),
1248            headers: vec![
1249                RecordBatchHeader::new("source", Some(b"checkout".to_vec())),
1250                RecordBatchHeader::new("tombstone", None),
1251            ],
1252            producer_id: Some(7),
1253            transactional: true,
1254            control: false,
1255        };
1256
1257        assert_eq!(response.responses[0].partitions[0].records, vec![record]);
1258        assert!(decoder.is_empty());
1259    }
1260
1261    fn write_varint(output: &mut Vec<u8>, value: i32) {
1262        write_unsigned_varint(output, u64::from(((value << 1) ^ (value >> 31)) as u32));
1263    }
1264
1265    fn write_varlong(output: &mut Vec<u8>, value: i64) {
1266        write_unsigned_varint(output, ((value << 1) ^ (value >> 63)) as u64);
1267    }
1268
1269    fn write_unsigned_varint(output: &mut Vec<u8>, mut value: u64) {
1270        loop {
1271            let mut byte = (value & 0x7f) as u8;
1272            value >>= 7;
1273            if value != 0 {
1274                byte |= 0x80;
1275            }
1276            output.push(byte);
1277            if value == 0 {
1278                break;
1279            }
1280        }
1281    }
1282}
kafrust_protocol/api/fetch.rs

kafrust_protocol/api/
fetch.rs
