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    /// Kafka's partition leader epoch for the enclosing RecordBatch.
625    /// Legacy MessageSet records do not carry an epoch and use `-1`.
626    pub leader_epoch: i32,
627    pub timestamp_ms: i64,
628    pub key: Option<Vec<u8>>,
629    pub value: Option<Vec<u8>>,
630    pub headers: Vec<RecordBatchHeader>,
631    pub producer_id: Option<i64>,
632    pub transactional: bool,
633    pub control: bool,
634}
635
636fn decode_message_set(bytes: &[u8], limits: DecodeLimits) -> Result<Vec<MessageSetRecord>> {
637    let mut decoder = Decoder::with_limits(bytes, limits);
638    let mut records = Vec::new();
639
640    while decoder.remaining() >= 12 {
641        let offset = decoder.read_i64()?;
642        let message_size = decoder.read_i32()?;
643        if message_size < 0 {
644            return Err(Error::NegativeLength {
645                kind: "message",
646                length: message_size,
647            });
648        }
649        let message_size =
650            usize::try_from(message_size).map_err(|_| Error::LengthOverflow("message"))?;
651        // Fetch responses may end with a partial trailing message set entry.
652        if decoder.remaining() < message_size {
653            break;
654        }
655        let message = decoder.read_exact(message_size)?;
656        let decoded = decode_message_or_batch(offset, message, limits)?;
657        let total = records
658            .len()
659            .checked_add(decoded.len())
660            .ok_or(Error::LengthOverflow("fetch records"))?;
661        decoder.ensure_collection_length("fetch records", total)?;
662        records.extend(decoded);
663    }
664
665    Ok(records)
666}
667
668fn decode_message_or_batch(
669    offset: i64,
670    bytes: &[u8],
671    limits: DecodeLimits,
672) -> Result<Vec<MessageSetRecord>> {
673    match bytes.get(4).copied() {
674        Some(2) => decode_record_batch(offset, bytes, limits),
675        _ => Ok(vec![decode_message(offset, bytes, limits)?]),
676    }
677}
678
679fn decode_message(offset: i64, bytes: &[u8], limits: DecodeLimits) -> Result<MessageSetRecord> {
680    let mut decoder = Decoder::with_limits(bytes, limits);
681    let _crc = decoder.read_i32()?;
682    let magic = decoder.read_i8()?;
683    let _attributes = decoder.read_i8()?;
684    let timestamp_ms = match magic {
685        0 => -1,
686        1 => decoder.read_i64()?,
687        _ => {
688            return Err(Error::UnsupportedVersion {
689                kind: "message magic",
690                version: i16::from(magic),
691            })
692        }
693    };
694    let key = decoder.read_nullable_bytes()?;
695    let value = decoder.read_nullable_bytes()?;
696
697    Ok(MessageSetRecord {
698        offset,
699        leader_epoch: -1,
700        timestamp_ms,
701        key,
702        value,
703        headers: Vec::new(),
704        producer_id: None,
705        transactional: false,
706        control: false,
707    })
708}
709
710fn decode_record_batch(
711    base_offset: i64,
712    bytes: &[u8],
713    limits: DecodeLimits,
714) -> Result<Vec<MessageSetRecord>> {
715    let mut decoder = Decoder::with_limits(bytes, limits);
716    let partition_leader_epoch = decoder.read_i32()?;
717    let magic = decoder.read_i8()?;
718    if magic != 2 {
719        return Err(Error::UnsupportedVersion {
720            kind: "record batch magic",
721            version: i16::from(magic),
722        });
723    }
724    let _crc = decoder.read_i32()?;
725    let attributes = decoder.read_i16()?;
726    let compression = RecordBatchCompression::from_attributes(attributes)?;
727    let _last_offset_delta = decoder.read_i32()?;
728    let base_timestamp = decoder.read_i64()?;
729    let _max_timestamp = decoder.read_i64()?;
730    let producer_id = decoder.read_i64()?;
731    let _producer_epoch = decoder.read_i16()?;
732    let _base_sequence = decoder.read_i32()?;
733    let record_count = decoder.read_i32()?;
734    if record_count < 0 {
735        return Err(Error::NegativeLength {
736            kind: "record batch records",
737            length: record_count,
738        });
739    }
740
741    let record_count =
742        usize::try_from(record_count).map_err(|_| Error::LengthOverflow("record batch records"))?;
743    decoder.ensure_collection_length("record batch records", record_count)?;
744    let record_bytes = if compression.is_compressed() {
745        let compressed = decoder.read_exact(decoder.remaining())?;
746        decompress_record_batch_records_with_limit(
747            compression,
748            compressed,
749            limits.max_decompressed_record_bytes(),
750        )?
751    } else {
752        if decoder.remaining() > limits.max_decompressed_record_bytes() {
753            return Err(Error::LimitExceeded {
754                kind: "decompressed record batch bytes",
755                actual: decoder.remaining(),
756                max: limits.max_decompressed_record_bytes(),
757            });
758        }
759        decoder.read_exact(decoder.remaining())?.to_vec()
760    };
761    let mut record_decoder = Decoder::with_limits(&record_bytes, limits);
762    let mut records = Vec::with_capacity(record_count);
763    for _ in 0..record_count {
764        let record_length = record_decoder.read_varint()?;
765        if record_length < 0 {
766            return Err(Error::NegativeLength {
767                kind: "record",
768                length: record_length,
769            });
770        }
771        let record_length =
772            usize::try_from(record_length).map_err(|_| Error::LengthOverflow("record"))?;
773        let record_bytes = record_decoder.read_exact(record_length)?;
774        records.push(decode_record(
775            base_offset,
776            partition_leader_epoch,
777            base_timestamp,
778            producer_id,
779            attributes,
780            record_bytes,
781            limits,
782        )?);
783    }
784
785    Ok(records)
786}
787
788fn decode_record(
789    base_offset: i64,
790    partition_leader_epoch: i32,
791    base_timestamp: i64,
792    producer_id: i64,
793    batch_attributes: i16,
794    bytes: &[u8],
795    limits: DecodeLimits,
796) -> Result<MessageSetRecord> {
797    let mut decoder = Decoder::with_limits(bytes, limits);
798    let _attributes = decoder.read_i8()?;
799    let timestamp_delta = decoder.read_varlong()?;
800    let offset_delta = decoder.read_varint()?;
801    let key = decoder.read_varint_nullable_bytes()?;
802    let value = decoder.read_varint_nullable_bytes()?;
803    let header_count = decoder.read_varint()?;
804    if header_count < 0 {
805        return Err(Error::NegativeLength {
806            kind: "record headers",
807            length: header_count,
808        });
809    }
810    let header_count =
811        usize::try_from(header_count).map_err(|_| Error::LengthOverflow("record headers"))?;
812    decoder.ensure_collection_length("record headers", header_count)?;
813    let mut headers = Vec::with_capacity(header_count);
814    for _ in 0..header_count {
815        let header_key =
816            String::from_utf8(decoder.read_varint_bytes()?).map_err(|_| Error::InvalidUtf8)?;
817        let header_value = decoder.read_varint_nullable_bytes()?;
818        headers.push(RecordBatchHeader::new(header_key, header_value));
819    }
820
821    Ok(MessageSetRecord {
822        offset: base_offset.saturating_add(i64::from(offset_delta)),
823        leader_epoch: partition_leader_epoch,
824        timestamp_ms: base_timestamp.saturating_add(timestamp_delta),
825        key,
826        value,
827        headers,
828        producer_id: (producer_id >= 0).then_some(producer_id),
829        transactional: batch_attributes & 0x10 != 0,
830        control: batch_attributes & 0x20 != 0,
831    })
832}
833
834#[cfg(test)]
835#[allow(clippy::unwrap_used)]
836mod tests {
837    use super::{
838        FetchPartitionV11, FetchPartitionV12, FetchPartitionV2, FetchRequestV11, FetchRequestV12,
839        FetchRequestV2, FetchRequestV4, FetchResponseV11, FetchResponseV12, FetchResponseV2,
840        FetchResponseV4, FetchTopicV11, FetchTopicV12, FetchTopicV2, MessageSetRecord,
841        RecordBatchHeader,
842    };
843    use crate::codec::{Decoder, Encoder};
844
845    #[test]
846    fn encodes_fetch_request_v2() {
847        let request = FetchRequestV2 {
848            correlation_id: 7,
849            client_id: Some("kafrust".to_owned()),
850            replica_id: -1,
851            max_wait_ms: 500,
852            min_bytes: 1,
853            topics: vec![FetchTopicV2 {
854                name: "orders".to_owned(),
855                partitions: vec![FetchPartitionV2 {
856                    partition_index: 0,
857                    fetch_offset: 42,
858                    max_bytes: 1_048_576,
859                }],
860            }],
861        };
862
863        let bytes = request.encode().unwrap();
864        assert_eq!(&bytes[0..4], &[0, 1, 0, 2]);
865        assert!(bytes.len() > 40);
866    }
867
868    #[test]
869    fn encodes_fetch_request_v4() {
870        let request = FetchRequestV4 {
871            correlation_id: 8,
872            client_id: Some("kafrust".to_owned()),
873            replica_id: -1,
874            max_wait_ms: 500,
875            min_bytes: 1,
876            max_bytes: 1_048_576,
877            isolation_level: 0,
878            topics: vec![FetchTopicV2 {
879                name: "orders".to_owned(),
880                partitions: vec![FetchPartitionV2 {
881                    partition_index: 0,
882                    fetch_offset: 42,
883                    max_bytes: 1_048_576,
884                }],
885            }],
886        };
887
888        let bytes = request.encode().unwrap();
889        assert_eq!(&bytes[0..4], &[0, 1, 0, 4]);
890        assert_eq!(&bytes[4..8], &[0, 0, 0, 8]);
891        assert!(bytes.len() > 45);
892    }
893
894    #[test]
895    fn encodes_fetch_request_v11_with_rack_and_fetch_session_fields() {
896        let request = FetchRequestV11 {
897            correlation_id: 9,
898            client_id: Some("kafrust".to_owned()),
899            replica_id: -1,
900            max_wait_ms: 500,
901            min_bytes: 1,
902            max_bytes: 1_048_576,
903            isolation_level: 1,
904            session_id: 0,
905            session_epoch: 0,
906            topics: vec![FetchTopicV11 {
907                name: "orders".to_owned(),
908                partitions: vec![FetchPartitionV11 {
909                    partition_index: 0,
910                    current_leader_epoch: -1,
911                    fetch_offset: 42,
912                    log_start_offset: -1,
913                    max_bytes: 1_048_576,
914                }],
915            }],
916            forgotten_topics: Vec::new(),
917            rack_id: "rack-a".to_owned(),
918        };
919
920        let bytes = request.encode().unwrap();
921        let mut decoder = Decoder::new(&bytes);
922        assert_eq!(decoder.read_i16().unwrap(), 1);
923        assert_eq!(decoder.read_i16().unwrap(), 11);
924        assert_eq!(decoder.read_i32().unwrap(), 9);
925        assert_eq!(
926            decoder.read_nullable_string().unwrap().as_deref(),
927            Some("kafrust")
928        );
929        assert_eq!(decoder.read_i32().unwrap(), -1);
930        assert_eq!(decoder.read_i32().unwrap(), 500);
931        assert_eq!(decoder.read_i32().unwrap(), 1);
932        assert_eq!(decoder.read_i32().unwrap(), 1_048_576);
933        assert_eq!(decoder.read_i8().unwrap(), 1);
934        assert_eq!(decoder.read_i32().unwrap(), 0);
935        assert_eq!(decoder.read_i32().unwrap(), 0);
936        assert_eq!(decoder.read_i32().unwrap(), 1);
937        assert_eq!(decoder.read_string().unwrap(), "orders");
938        assert_eq!(decoder.read_i32().unwrap(), 1);
939        assert_eq!(decoder.read_i32().unwrap(), 0);
940        assert_eq!(decoder.read_i32().unwrap(), -1);
941        assert_eq!(decoder.read_i64().unwrap(), 42);
942        assert_eq!(decoder.read_i64().unwrap(), -1);
943        assert_eq!(decoder.read_i32().unwrap(), 1_048_576);
944        assert_eq!(decoder.read_i32().unwrap(), 0);
945        assert_eq!(decoder.read_string().unwrap(), "rack-a");
946        assert!(decoder.is_empty());
947    }
948
949    #[test]
950    fn decodes_fetch_response_v11_with_preferred_read_replica() {
951        let mut bytes = Encoder::new();
952        bytes.write_i32(3);
953        bytes.write_i16(0);
954        bytes.write_i32(17);
955        bytes.write_i32(1);
956        bytes.write_string("orders").unwrap();
957        bytes.write_i32(1);
958        bytes.write_i32(0);
959        bytes.write_i16(0);
960        bytes.write_i64(43);
961        bytes.write_i64(42);
962        bytes.write_i64(40);
963        bytes.write_i32(0);
964        bytes.write_i32(2);
965        bytes.write_bytes(&[]).unwrap();
966
967        let bytes = bytes.into_bytes();
968        let mut decoder = Decoder::new(&bytes);
969        let response = FetchResponseV11::decode_body(&mut decoder).unwrap();
970        let partition = &response.responses[0].partitions[0];
971
972        assert_eq!(response.throttle_time_ms, 3);
973        assert_eq!(response.session_id, 17);
974        assert_eq!(partition.log_start_offset, 40);
975        assert_eq!(partition.preferred_read_replica, 2);
976        assert!(partition.records.is_empty());
977        assert!(decoder.is_empty());
978    }
979
980    #[test]
981    fn encodes_fetch_request_v12_with_flexible_rack_fields() {
982        let request = FetchRequestV12 {
983            correlation_id: 10,
984            client_id: Some("kafrust".to_owned()),
985            replica_id: -1,
986            max_wait_ms: 500,
987            min_bytes: 1,
988            max_bytes: 1_048_576,
989            isolation_level: 1,
990            session_id: 0,
991            session_epoch: 0,
992            topics: vec![FetchTopicV12 {
993                name: "orders".to_owned(),
994                partitions: vec![FetchPartitionV12 {
995                    partition_index: 0,
996                    current_leader_epoch: -1,
997                    fetch_offset: 42,
998                    last_fetched_epoch: -1,
999                    log_start_offset: -1,
1000                    max_bytes: 1_048_576,
1001                }],
1002            }],
1003            forgotten_topics: Vec::new(),
1004            rack_id: "rack-a".to_owned(),
1005        };
1006
1007        let bytes = request.encode().unwrap();
1008        let mut decoder = Decoder::new(&bytes);
1009        assert_eq!(decoder.read_i16().unwrap(), 1);
1010        assert_eq!(decoder.read_i16().unwrap(), 12);
1011        assert_eq!(decoder.read_i32().unwrap(), 10);
1012        assert_eq!(
1013            decoder.read_nullable_string().unwrap().as_deref(),
1014            Some("kafrust")
1015        );
1016        assert_eq!(decoder.read_unsigned_varint().unwrap(), 0);
1017        assert_eq!(decoder.read_i32().unwrap(), -1);
1018        assert_eq!(decoder.read_i32().unwrap(), 500);
1019        assert_eq!(decoder.read_i32().unwrap(), 1);
1020        assert_eq!(decoder.read_i32().unwrap(), 1_048_576);
1021        assert_eq!(decoder.read_i8().unwrap(), 1);
1022        assert_eq!(decoder.read_i32().unwrap(), 0);
1023        assert_eq!(decoder.read_i32().unwrap(), 0);
1024        assert_eq!(decoder.read_unsigned_varint().unwrap(), 2);
1025        assert_eq!(decoder.read_compact_string().unwrap(), "orders");
1026        assert_eq!(decoder.read_unsigned_varint().unwrap(), 2);
1027        assert_eq!(decoder.read_i32().unwrap(), 0);
1028        assert_eq!(decoder.read_i32().unwrap(), -1);
1029        assert_eq!(decoder.read_i64().unwrap(), 42);
1030        assert_eq!(decoder.read_i32().unwrap(), -1);
1031        assert_eq!(decoder.read_i64().unwrap(), -1);
1032        assert_eq!(decoder.read_i32().unwrap(), 1_048_576);
1033        assert_eq!(decoder.read_unsigned_varint().unwrap(), 0);
1034        assert_eq!(decoder.read_unsigned_varint().unwrap(), 0);
1035        assert_eq!(decoder.read_unsigned_varint().unwrap(), 1);
1036        assert_eq!(decoder.read_compact_string().unwrap(), "rack-a");
1037        assert_eq!(decoder.read_unsigned_varint().unwrap(), 0);
1038        assert!(decoder.is_empty());
1039    }
1040
1041    #[test]
1042    fn decodes_fetch_response_v12_with_preferred_read_replica() {
1043        let mut bytes = Encoder::new();
1044        bytes.write_i32(3);
1045        bytes.write_i16(0);
1046        bytes.write_i32(17);
1047        bytes.write_unsigned_varint(2);
1048        bytes.write_compact_string("orders").unwrap();
1049        bytes.write_unsigned_varint(2);
1050        bytes.write_i32(0);
1051        bytes.write_i16(0);
1052        bytes.write_i64(43);
1053        bytes.write_i64(42);
1054        bytes.write_i64(40);
1055        bytes.write_unsigned_varint(2);
1056        bytes.write_i64(7);
1057        bytes.write_i64(40);
1058        bytes.write_unsigned_varint(0);
1059        bytes.write_i32(2);
1060        bytes.write_compact_nullable_bytes(Some(&[])).unwrap();
1061        bytes.write_unsigned_varint(0);
1062        bytes.write_unsigned_varint(0);
1063        bytes.write_unsigned_varint(0);
1064
1065        let bytes = bytes.into_bytes();
1066        let mut decoder = Decoder::new(&bytes);
1067        let response = FetchResponseV12::decode_body(&mut decoder).unwrap();
1068        let partition = &response.responses[0].partitions[0];
1069
1070        assert_eq!(response.throttle_time_ms, 3);
1071        assert_eq!(response.session_id, 17);
1072        assert_eq!(partition.log_start_offset, 40);
1073        assert_eq!(partition.preferred_read_replica, 2);
1074        assert_eq!(partition.aborted_transactions[0].producer_id, 7);
1075        assert!(partition.records.is_empty());
1076        assert!(decoder.is_empty());
1077    }
1078
1079    #[test]
1080    fn decodes_fetch_response_v4_with_aborted_transaction() {
1081        let mut bytes = Encoder::new();
1082        bytes.write_i32(0);
1083        bytes.write_i32(1);
1084        bytes.write_string("orders").unwrap();
1085        bytes.write_i32(1);
1086        bytes.write_i32(0);
1087        bytes.write_i16(0);
1088        bytes.write_i64(43);
1089        bytes.write_i64(42);
1090        bytes.write_i32(1);
1091        bytes.write_i64(7);
1092        bytes.write_i64(40);
1093        bytes.write_bytes(&[]).unwrap();
1094        let bytes = bytes.into_bytes();
1095
1096        let mut decoder = Decoder::new(&bytes);
1097        let response = FetchResponseV4::decode_body(&mut decoder).unwrap();
1098        let partition = &response.responses[0].partitions[0];
1099
1100        assert_eq!(partition.high_watermark, 43);
1101        assert_eq!(partition.last_stable_offset, 42);
1102        assert_eq!(partition.aborted_transactions[0].producer_id, 7);
1103        assert_eq!(partition.aborted_transactions[0].first_offset, 40);
1104        assert!(partition.records.is_empty());
1105        assert!(decoder.is_empty());
1106    }
1107
1108    #[test]
1109    fn decodes_fetch_response_v2_with_message_set() {
1110        let mut message = Encoder::new();
1111        message.write_i32(0);
1112        message.write_i8(1);
1113        message.write_i8(0);
1114        message.write_i64(123);
1115        message.write_nullable_bytes(Some(b"order-1")).unwrap();
1116        message.write_nullable_bytes(Some(b"created")).unwrap();
1117        let message = message.into_bytes();
1118
1119        let mut set = Encoder::new();
1120        set.write_i64(42);
1121        set.write_i32(i32::try_from(message.len()).unwrap());
1122        set.write_raw(&message);
1123        let set = set.into_bytes();
1124
1125        let mut bytes = Encoder::new();
1126        bytes.write_i32(0);
1127        bytes.write_i32(1);
1128        bytes.write_string("orders").unwrap();
1129        bytes.write_i32(1);
1130        bytes.write_i32(0);
1131        bytes.write_i16(0);
1132        bytes.write_i64(43);
1133        bytes.write_bytes(&set).unwrap();
1134        let bytes = bytes.into_bytes();
1135
1136        let mut decoder = Decoder::new(&bytes);
1137        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
1138        let record = MessageSetRecord {
1139            offset: 42,
1140            leader_epoch: -1,
1141            timestamp_ms: 123,
1142            key: Some(b"order-1".to_vec()),
1143            value: Some(b"created".to_vec()),
1144            headers: Vec::new(),
1145            producer_id: None,
1146            transactional: false,
1147            control: false,
1148        };
1149
1150        assert_eq!(response.throttle_time_ms, 0);
1151        assert_eq!(response.responses[0].partitions[0].high_watermark, 43);
1152        assert_eq!(response.responses[0].partitions[0].records, vec![record]);
1153        assert!(decoder.is_empty());
1154    }
1155
1156    #[test]
1157    fn decodes_fetch_response_v2_ignores_partial_trailing_message_set_entry() {
1158        let mut message = Encoder::new();
1159        message.write_i32(0);
1160        message.write_i8(1);
1161        message.write_i8(0);
1162        message.write_i64(123);
1163        message.write_nullable_bytes(Some(b"order-1")).unwrap();
1164        message.write_nullable_bytes(Some(b"created")).unwrap();
1165        let message = message.into_bytes();
1166
1167        let mut set = Encoder::new();
1168        set.write_i64(42);
1169        set.write_i32(i32::try_from(message.len()).unwrap());
1170        set.write_raw(&message);
1171        set.write_i64(-1);
1172        set.write_i32(61);
1173        set.write_raw(&[0; 22]);
1174        let set = set.into_bytes();
1175
1176        let mut bytes = Encoder::new();
1177        bytes.write_i32(0);
1178        bytes.write_i32(1);
1179        bytes.write_string("orders").unwrap();
1180        bytes.write_i32(1);
1181        bytes.write_i32(0);
1182        bytes.write_i16(0);
1183        bytes.write_i64(43);
1184        bytes.write_bytes(&set).unwrap();
1185        let bytes = bytes.into_bytes();
1186
1187        let mut decoder = Decoder::new(&bytes);
1188        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
1189
1190        assert_eq!(response.responses[0].partitions[0].records.len(), 1);
1191        assert_eq!(response.responses[0].partitions[0].records[0].offset, 42);
1192        assert!(decoder.is_empty());
1193    }
1194
1195    #[test]
1196    fn decodes_fetch_response_v2_with_record_batch() {
1197        let mut record = Vec::new();
1198        record.push(0);
1199        write_varlong(&mut record, 5);
1200        write_varint(&mut record, 0);
1201        write_varint(&mut record, 7);
1202        record.extend_from_slice(b"order-1");
1203        write_varint(&mut record, 7);
1204        record.extend_from_slice(b"created");
1205        write_varint(&mut record, 2);
1206        write_varint(&mut record, 6);
1207        record.extend_from_slice(b"source");
1208        write_varint(&mut record, 8);
1209        record.extend_from_slice(b"checkout");
1210        write_varint(&mut record, 9);
1211        record.extend_from_slice(b"tombstone");
1212        write_varint(&mut record, -1);
1213
1214        let mut batch = Encoder::new();
1215        batch.write_i32(0);
1216        batch.write_i8(2);
1217        batch.write_i32(0);
1218        batch.write_i16(0x10);
1219        batch.write_i32(0);
1220        batch.write_i64(1_000);
1221        batch.write_i64(1_005);
1222        batch.write_i64(7);
1223        batch.write_i16(-1);
1224        batch.write_i32(-1);
1225        batch.write_i32(1);
1226        let mut encoded_record = Vec::new();
1227        write_varint(&mut encoded_record, i32::try_from(record.len()).unwrap());
1228        encoded_record.extend_from_slice(&record);
1229        batch.write_raw(&encoded_record);
1230        let batch = batch.into_bytes();
1231
1232        let mut set = Encoder::new();
1233        set.write_i64(42);
1234        set.write_i32(i32::try_from(batch.len()).unwrap());
1235        set.write_raw(&batch);
1236        let set = set.into_bytes();
1237
1238        let mut bytes = Encoder::new();
1239        bytes.write_i32(0);
1240        bytes.write_i32(1);
1241        bytes.write_string("orders").unwrap();
1242        bytes.write_i32(1);
1243        bytes.write_i32(0);
1244        bytes.write_i16(0);
1245        bytes.write_i64(43);
1246        bytes.write_bytes(&set).unwrap();
1247        let bytes = bytes.into_bytes();
1248
1249        let mut decoder = Decoder::new(&bytes);
1250        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
1251        let record = MessageSetRecord {
1252            offset: 42,
1253            leader_epoch: 0,
1254            timestamp_ms: 1_005,
1255            key: Some(b"order-1".to_vec()),
1256            value: Some(b"created".to_vec()),
1257            headers: vec![
1258                RecordBatchHeader::new("source", Some(b"checkout".to_vec())),
1259                RecordBatchHeader::new("tombstone", None),
1260            ],
1261            producer_id: Some(7),
1262            transactional: true,
1263            control: false,
1264        };
1265
1266        assert_eq!(response.responses[0].partitions[0].records, vec![record]);
1267        assert!(decoder.is_empty());
1268    }
1269
1270    fn write_varint(output: &mut Vec<u8>, value: i32) {
1271        write_unsigned_varint(output, u64::from(((value << 1) ^ (value >> 31)) as u32));
1272    }
1273
1274    fn write_varlong(output: &mut Vec<u8>, value: i64) {
1275        write_unsigned_varint(output, ((value << 1) ^ (value >> 63)) as u64);
1276    }
1277
1278    fn write_unsigned_varint(output: &mut Vec<u8>, mut value: u64) {
1279        loop {
1280            let mut byte = (value & 0x7f) as u8;
1281            value >>= 7;
1282            if value != 0 {
1283                byte |= 0x80;
1284            }
1285            output.push(byte);
1286            if value == 0 {
1287                break;
1288            }
1289        }
1290    }
1291}
kafrust_protocol/api/fetch.rs

kafrust_protocol/api/
fetch.rs
