kafrust_protocol/api/
produce.rs

1use crate::codec::{Decoder, Encoder};
2use crate::error::{Error, Result};
3use crate::header::RequestHeader;
4use crate::record_batch::{compress_record_batch_records, RecordBatchCompression};
5
6pub const API_KEY: i16 = 0;
7
8#[derive(Debug, Clone, PartialEq, Eq)]
9pub struct ProduceRequestV2 {
10    pub correlation_id: i32,
11    pub client_id: Option<String>,
12    pub acks: i16,
13    pub timeout_ms: i32,
14    pub topics: Vec<ProduceTopicV2>,
15}
16
17#[derive(Debug, Clone, PartialEq, Eq)]
18pub struct ProduceRequestV3 {
19    pub correlation_id: i32,
20    pub client_id: Option<String>,
21    pub transactional_id: Option<String>,
22    pub acks: i16,
23    pub timeout_ms: i32,
24    pub topics: Vec<ProduceTopicV3>,
25}
26
27#[derive(Debug, Clone, PartialEq, Eq)]
28pub struct ProduceRequestV7 {
29    pub correlation_id: i32,
30    pub client_id: Option<String>,
31    pub transactional_id: Option<String>,
32    pub acks: i16,
33    pub timeout_ms: i32,
34    pub topics: Vec<ProduceTopicV3>,
35}
36
37impl ProduceRequestV3 {
38    pub fn encode(&self) -> Result<Vec<u8>> {
39        encode_record_batch_request(
40            3,
41            self.correlation_id,
42            self.client_id.clone(),
43            self.transactional_id.as_deref(),
44            self.acks,
45            self.timeout_ms,
46            &self.topics,
47        )
48    }
49}
50
51impl ProduceRequestV7 {
52    pub fn encode(&self) -> Result<Vec<u8>> {
53        encode_record_batch_request(
54            7,
55            self.correlation_id,
56            self.client_id.clone(),
57            self.transactional_id.as_deref(),
58            self.acks,
59            self.timeout_ms,
60            &self.topics,
61        )
62    }
63}
64
65fn encode_record_batch_request(
66    api_version: i16,
67    correlation_id: i32,
68    client_id: Option<String>,
69    transactional_id: Option<&str>,
70    acks: i16,
71    timeout_ms: i32,
72    topics: &[ProduceTopicV3],
73) -> Result<Vec<u8>> {
74    let mut encoder = Encoder::new();
75    RequestHeader {
76        api_key: API_KEY,
77        api_version,
78        correlation_id,
79        client_id,
80    }
81    .encode_v1(&mut encoder)?;
82    encoder.write_nullable_string(transactional_id)?;
83    encoder.write_i16(acks);
84    encoder.write_i32(timeout_ms);
85    encoder.write_array(Some(topics), |encoder, topic| {
86        topic.encode(encoder, transactional_id.is_some())
87    })?;
88    Ok(encoder.into_bytes())
89}
90
91impl ProduceRequestV2 {
92    pub fn encode(&self) -> Result<Vec<u8>> {
93        let mut encoder = Encoder::new();
94        RequestHeader {
95            api_key: API_KEY,
96            api_version: 2,
97            correlation_id: self.correlation_id,
98            client_id: self.client_id.clone(),
99        }
100        .encode_v1(&mut encoder)?;
101        encoder.write_i16(self.acks);
102        encoder.write_i32(self.timeout_ms);
103        encoder.write_array(Some(self.topics.as_slice()), |encoder, topic| {
104            topic.encode(encoder)
105        })?;
106        Ok(encoder.into_bytes())
107    }
108}
109
110#[derive(Debug, Clone, PartialEq, Eq)]
111pub struct ProduceTopicV2 {
112    pub name: String,
113    pub partitions: Vec<ProducePartitionV2>,
114}
115
116#[derive(Debug, Clone, PartialEq, Eq)]
117pub struct ProduceTopicV3 {
118    pub name: String,
119    pub partitions: Vec<ProducePartitionV3>,
120}
121
122impl ProduceTopicV3 {
123    fn encode(&self, encoder: &mut Encoder, transactional: bool) -> Result<()> {
124        encoder.write_string(&self.name)?;
125        encoder.write_array(Some(self.partitions.as_slice()), |encoder, partition| {
126            partition.encode(encoder, transactional)
127        })
128    }
129}
130
131impl ProduceTopicV2 {
132    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
133        encoder.write_string(&self.name)?;
134        encoder.write_array(Some(self.partitions.as_slice()), |encoder, partition| {
135            partition.encode(encoder)
136        })
137    }
138}
139
140#[derive(Debug, Clone, PartialEq, Eq)]
141pub struct ProducePartitionV2 {
142    pub partition_index: i32,
143    pub records: Vec<MessageSetMessage>,
144}
145
146#[derive(Debug, Clone, PartialEq, Eq)]
147pub struct ProducePartitionV3 {
148    pub partition_index: i32,
149    pub records: Vec<RecordBatchMessage>,
150    pub compression: RecordBatchCompression,
151    pub identity: RecordBatchIdentity,
152}
153
154#[derive(Debug, Clone, Copy, PartialEq, Eq)]
155pub struct RecordBatchIdentity {
156    pub producer_id: i64,
157    pub producer_epoch: i16,
158    pub base_sequence: i32,
159}
160
161impl RecordBatchIdentity {
162    pub const NON_IDEMPOTENT: Self = Self {
163        producer_id: -1,
164        producer_epoch: -1,
165        base_sequence: -1,
166    };
167}
168
169impl ProducePartitionV3 {
170    fn encode(&self, encoder: &mut Encoder, transactional: bool) -> Result<()> {
171        encoder.write_i32(self.partition_index);
172        let record_set = encode_record_batch_set_with_compression_identity_and_transaction(
173            &self.records,
174            self.compression,
175            self.identity,
176            transactional,
177        )?;
178        encoder.write_bytes(&record_set)
179    }
180}
181
182impl ProducePartitionV2 {
183    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
184        encoder.write_i32(self.partition_index);
185        let record_set = encode_message_set(&self.records)?;
186        encoder.write_bytes(&record_set)
187    }
188}
189
190/// Returns the encoded byte length of a Produce v2 message set.
191pub fn encoded_message_set_len(records: &[MessageSetMessage]) -> Result<usize> {
192    Ok(encode_message_set(records)?.len())
193}
194
195/// Returns the encoded byte length of a Produce v3 record batch set.
196pub fn encoded_record_batch_set_len(records: &[RecordBatchMessage]) -> Result<usize> {
197    Ok(encode_record_batch_set(records)?.len())
198}
199
200/// Returns the encoded byte length of a Produce v3 record batch set.
201pub fn encoded_record_batch_set_len_with_compression(
202    records: &[RecordBatchMessage],
203    compression: RecordBatchCompression,
204) -> Result<usize> {
205    encoded_record_batch_set_len_with_compression_and_identity(
206        records,
207        compression,
208        RecordBatchIdentity::NON_IDEMPOTENT,
209    )
210}
211
212pub fn encoded_record_batch_set_len_with_compression_and_identity(
213    records: &[RecordBatchMessage],
214    compression: RecordBatchCompression,
215    identity: RecordBatchIdentity,
216) -> Result<usize> {
217    Ok(
218        encode_record_batch_set_with_compression_and_identity(records, compression, identity)?
219            .len(),
220    )
221}
222
223#[derive(Debug, Clone, PartialEq, Eq)]
224pub struct MessageSetMessage {
225    pub key: Option<Vec<u8>>,
226    pub value: Option<Vec<u8>>,
227    pub timestamp_ms: i64,
228}
229
230impl MessageSetMessage {
231    pub fn new(key: Option<Vec<u8>>, value: Option<Vec<u8>>, timestamp_ms: i64) -> Self {
232        Self {
233            key,
234            value,
235            timestamp_ms,
236        }
237    }
238}
239
240#[derive(Debug, Clone, PartialEq, Eq)]
241pub struct RecordBatchHeader {
242    pub key: String,
243    pub value: Option<Vec<u8>>,
244}
245
246impl RecordBatchHeader {
247    pub fn new(key: impl Into<String>, value: Option<Vec<u8>>) -> Self {
248        Self {
249            key: key.into(),
250            value,
251        }
252    }
253}
254
255#[derive(Debug, Clone, PartialEq, Eq)]
256pub struct RecordBatchMessage {
257    pub key: Option<Vec<u8>>,
258    pub value: Option<Vec<u8>>,
259    pub timestamp_ms: i64,
260    pub headers: Vec<RecordBatchHeader>,
261}
262
263impl RecordBatchMessage {
264    pub fn new(key: Option<Vec<u8>>, value: Option<Vec<u8>>, timestamp_ms: i64) -> Self {
265        Self {
266            key,
267            value,
268            timestamp_ms,
269            headers: Vec::new(),
270        }
271    }
272
273    pub fn header(mut self, key: impl Into<String>, value: Option<Vec<u8>>) -> Self {
274        self.headers.push(RecordBatchHeader::new(key, value));
275        self
276    }
277}
278
279#[derive(Debug, Clone, PartialEq, Eq)]
280pub struct ProduceResponseV2 {
281    pub responses: Vec<ProduceTopicResponseV2>,
282    pub throttle_time_ms: i32,
283}
284
285impl ProduceResponseV2 {
286    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
287        Ok(Self {
288            responses: decoder
289                .read_array("produce responses", ProduceTopicResponseV2::decode)?
290                .unwrap_or_default(),
291            throttle_time_ms: decoder.read_i32()?,
292        })
293    }
294}
295
296#[derive(Debug, Clone, PartialEq, Eq)]
297pub struct ProduceTopicResponseV2 {
298    pub name: String,
299    pub partitions: Vec<ProducePartitionResponseV2>,
300}
301
302impl ProduceTopicResponseV2 {
303    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
304        Ok(Self {
305            name: decoder.read_string()?,
306            partitions: decoder
307                .read_array(
308                    "produce partition responses",
309                    ProducePartitionResponseV2::decode,
310                )?
311                .unwrap_or_default(),
312        })
313    }
314}
315
316#[derive(Debug, Clone, PartialEq, Eq)]
317pub struct ProducePartitionResponseV2 {
318    pub partition_index: i32,
319    pub error_code: i16,
320    pub base_offset: i64,
321    pub log_append_time_ms: i64,
322}
323
324impl ProducePartitionResponseV2 {
325    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
326        Ok(Self {
327            partition_index: decoder.read_i32()?,
328            error_code: decoder.read_i16()?,
329            base_offset: decoder.read_i64()?,
330            log_append_time_ms: decoder.read_i64()?,
331        })
332    }
333}
334
335#[derive(Debug, Clone, PartialEq, Eq)]
336pub struct ProduceResponseV7 {
337    pub responses: Vec<ProduceTopicResponseV7>,
338    pub throttle_time_ms: i32,
339}
340
341impl ProduceResponseV7 {
342    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
343        Ok(Self {
344            responses: decoder
345                .read_array("produce responses", ProduceTopicResponseV7::decode)?
346                .unwrap_or_default(),
347            throttle_time_ms: decoder.read_i32()?,
348        })
349    }
350}
351
352#[derive(Debug, Clone, PartialEq, Eq)]
353pub struct ProduceTopicResponseV7 {
354    pub name: String,
355    pub partitions: Vec<ProducePartitionResponseV7>,
356}
357
358impl ProduceTopicResponseV7 {
359    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
360        Ok(Self {
361            name: decoder.read_string()?,
362            partitions: decoder
363                .read_array(
364                    "produce partition responses",
365                    ProducePartitionResponseV7::decode,
366                )?
367                .unwrap_or_default(),
368        })
369    }
370}
371
372#[derive(Debug, Clone, PartialEq, Eq)]
373pub struct ProducePartitionResponseV7 {
374    pub partition_index: i32,
375    pub error_code: i16,
376    pub base_offset: i64,
377    pub log_append_time_ms: i64,
378    pub log_start_offset: i64,
379}
380
381impl ProducePartitionResponseV7 {
382    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
383        Ok(Self {
384            partition_index: decoder.read_i32()?,
385            error_code: decoder.read_i16()?,
386            base_offset: decoder.read_i64()?,
387            log_append_time_ms: decoder.read_i64()?,
388            log_start_offset: decoder.read_i64()?,
389        })
390    }
391}
392
393fn encode_message_set(records: &[MessageSetMessage]) -> Result<Vec<u8>> {
394    let mut set = Encoder::new();
395    for record in records {
396        let message = encode_message(record)?;
397        set.write_i64(0);
398        set.write_i32(i32::try_from(message.len()).map_err(|_| Error::LengthOverflow("message"))?);
399        set.write_raw(&message);
400    }
401    Ok(set.into_bytes())
402}
403
404fn encode_message(record: &MessageSetMessage) -> Result<Vec<u8>> {
405    let mut body = Encoder::new();
406    body.write_i8(1);
407    body.write_i8(0);
408    body.write_i64(record.timestamp_ms);
409    body.write_nullable_bytes(record.key.as_deref())?;
410    body.write_nullable_bytes(record.value.as_deref())?;
411    let body = body.into_bytes();
412
413    let mut message = Encoder::new();
414    message.write_i32(crc32_ieee(&body) as i32);
415    message.write_raw(&body);
416    Ok(message.into_bytes())
417}
418
419fn encode_record_batch_set(records: &[RecordBatchMessage]) -> Result<Vec<u8>> {
420    encode_record_batch_set_with_compression(records, RecordBatchCompression::None)
421}
422
423fn encode_record_batch_set_with_compression(
424    records: &[RecordBatchMessage],
425    compression: RecordBatchCompression,
426) -> Result<Vec<u8>> {
427    encode_record_batch_set_with_compression_and_identity(
428        records,
429        compression,
430        RecordBatchIdentity::NON_IDEMPOTENT,
431    )
432}
433
434fn encode_record_batch_set_with_compression_and_identity(
435    records: &[RecordBatchMessage],
436    compression: RecordBatchCompression,
437    identity: RecordBatchIdentity,
438) -> Result<Vec<u8>> {
439    encode_record_batch_set_with_compression_identity_and_transaction(
440        records,
441        compression,
442        identity,
443        false,
444    )
445}
446
447fn encode_record_batch_set_with_compression_identity_and_transaction(
448    records: &[RecordBatchMessage],
449    compression: RecordBatchCompression,
450    identity: RecordBatchIdentity,
451    transactional: bool,
452) -> Result<Vec<u8>> {
453    let base_timestamp = records
454        .first()
455        .map(|record| record.timestamp_ms)
456        .unwrap_or_default();
457    let max_timestamp = records
458        .iter()
459        .map(|record| record.timestamp_ms)
460        .max()
461        .unwrap_or(base_timestamp);
462    let last_offset_delta = records
463        .len()
464        .checked_sub(1)
465        .map(|delta| i32::try_from(delta).map_err(|_| Error::LengthOverflow("record batch")))
466        .transpose()?
467        .unwrap_or_default();
468
469    let record_count =
470        i32::try_from(records.len()).map_err(|_| Error::LengthOverflow("record batch records"))?;
471    let mut record_bytes = Encoder::new();
472    for (offset_delta, record) in records.iter().enumerate() {
473        let encoded = encode_record(record, base_timestamp, offset_delta)?;
474        record_bytes.write_varint(
475            i32::try_from(encoded.len()).map_err(|_| Error::LengthOverflow("record"))?,
476        );
477        record_bytes.write_raw(&encoded);
478    }
479    let record_bytes = compress_record_batch_records(compression, &record_bytes.into_bytes())?;
480
481    let mut crc_payload = Encoder::new();
482    let attributes = compression.attributes() | if transactional { 0x10 } else { 0 };
483    crc_payload.write_i16(attributes);
484    crc_payload.write_i32(last_offset_delta);
485    crc_payload.write_i64(base_timestamp);
486    crc_payload.write_i64(max_timestamp);
487    crc_payload.write_i64(identity.producer_id);
488    crc_payload.write_i16(identity.producer_epoch);
489    crc_payload.write_i32(identity.base_sequence);
490    crc_payload.write_i32(record_count);
491    crc_payload.write_raw(&record_bytes);
492    let crc_payload = crc_payload.into_bytes();
493
494    let mut batch = Encoder::new();
495    batch.write_i32(0);
496    batch.write_i8(2);
497    batch.write_i32(crc32c(&crc_payload) as i32);
498    batch.write_raw(&crc_payload);
499    let batch = batch.into_bytes();
500
501    let mut set = Encoder::new();
502    set.write_i64(0);
503    set.write_i32(i32::try_from(batch.len()).map_err(|_| Error::LengthOverflow("record batch"))?);
504    set.write_raw(&batch);
505    Ok(set.into_bytes())
506}
507
508fn encode_record(
509    record: &RecordBatchMessage,
510    base_timestamp: i64,
511    offset_delta: usize,
512) -> Result<Vec<u8>> {
513    let mut encoder = Encoder::new();
514    encoder.write_i8(0);
515    encoder.write_varlong(record.timestamp_ms.saturating_sub(base_timestamp));
516    encoder.write_varint(
517        i32::try_from(offset_delta).map_err(|_| Error::LengthOverflow("record offset delta"))?,
518    );
519    encoder.write_varint_nullable_bytes(record.key.as_deref())?;
520    encoder.write_varint_nullable_bytes(record.value.as_deref())?;
521    encoder.write_varint(
522        i32::try_from(record.headers.len()).map_err(|_| Error::LengthOverflow("record headers"))?,
523    );
524    for header in &record.headers {
525        encoder.write_varint_bytes(header.key.as_bytes())?;
526        encoder.write_varint_nullable_bytes(header.value.as_deref())?;
527    }
528    Ok(encoder.into_bytes())
529}
530
531fn crc32_ieee(bytes: &[u8]) -> u32 {
532    crc32_with_table(bytes, &CRC32_IEEE_TABLE)
533}
534
535fn crc32c(bytes: &[u8]) -> u32 {
536    crc32_with_table(bytes, &CRC32C_TABLE)
537}
538
539const CRC32_IEEE_TABLE: [u32; 256] = crc32_table(0xedb8_8320);
540const CRC32C_TABLE: [u32; 256] = crc32_table(0x82f6_3b78);
541
542const fn crc32_table(polynomial: u32) -> [u32; 256] {
543    let mut table = [0; 256];
544    let mut index = 0;
545    while index < table.len() {
546        let mut value = index as u32;
547        let mut bit = 0;
548        while bit < 8 {
549            let mask = 0u32.wrapping_sub(value & 1);
550            value = (value >> 1) ^ (polynomial & mask);
551            bit += 1;
552        }
553        table[index] = value;
554        index += 1;
555    }
556    table
557}
558
559fn crc32_with_table(bytes: &[u8], table: &[u32; 256]) -> u32 {
560    let mut crc = 0xffff_ffffu32;
561    for byte in bytes {
562        let index = usize::from((crc as u8) ^ byte);
563        crc = (crc >> 8) ^ table[index];
564    }
565    !crc
566}
567
568#[cfg(test)]
569#[allow(clippy::unwrap_used)]
570mod tests {
571    use super::{
572        crc32_ieee, crc32c, encode_message_set, encode_record_batch_set,
573        encode_record_batch_set_with_compression,
574        encode_record_batch_set_with_compression_and_identity,
575        encode_record_batch_set_with_compression_identity_and_transaction, encoded_message_set_len,
576        encoded_record_batch_set_len, MessageSetMessage, ProducePartitionV2, ProducePartitionV3,
577        ProduceRequestV2, ProduceRequestV3, ProduceRequestV7, ProduceResponseV2, ProduceResponseV7,
578        ProduceTopicV2, ProduceTopicV3, RecordBatchIdentity, RecordBatchMessage,
579    };
580    use crate::codec::{DecodeLimits, Decoder};
581    use crate::record_batch::RecordBatchCompression;
582    use crate::{api::fetch::FetchResponseV2, codec::Encoder};
583
584    #[test]
585    fn crc_implementations_match_standard_check_vectors() {
586        assert_eq!(crc32_ieee(b"123456789"), 0xcbf4_3926);
587        assert_eq!(crc32c(b"123456789"), 0xe306_9283);
588    }
589
590    #[test]
591    fn encodes_produce_request_v2() {
592        let request = ProduceRequestV2 {
593            correlation_id: 5,
594            client_id: Some("kafrust".to_owned()),
595            acks: 1,
596            timeout_ms: 30_000,
597            topics: vec![ProduceTopicV2 {
598                name: "orders".to_owned(),
599                partitions: vec![ProducePartitionV2 {
600                    partition_index: 0,
601                    records: vec![MessageSetMessage::new(
602                        Some(b"order-1".to_vec()),
603                        Some(b"created".to_vec()),
604                        0,
605                    )],
606                }],
607            }],
608        };
609
610        let bytes = request.encode().unwrap();
611        assert_eq!(
612            &bytes[0..17],
613            &[0, 0, 0, 2, 0, 0, 0, 5, 0, 7, b'k', b'a', b'f', b'r', b'u', b's', b't',]
614        );
615        assert!(bytes.len() > 60);
616    }
617
618    #[test]
619    fn encodes_idempotent_record_batch_identity() {
620        let set = encode_record_batch_set_with_compression_and_identity(
621            &[RecordBatchMessage::new(
622                Some(b"order-1".to_vec()),
623                Some(b"created".to_vec()),
624                1_000,
625            )],
626            RecordBatchCompression::None,
627            RecordBatchIdentity {
628                producer_id: 42,
629                producer_epoch: 3,
630                base_sequence: 7,
631            },
632        )
633        .unwrap();
634
635        assert_eq!(&set[43..51], &42_i64.to_be_bytes());
636        assert_eq!(&set[51..53], &3_i16.to_be_bytes());
637        assert_eq!(&set[53..57], &7_i32.to_be_bytes());
638    }
639
640    #[test]
641    fn encodes_transactional_record_batch_attribute() {
642        let set = encode_record_batch_set_with_compression_identity_and_transaction(
643            &[RecordBatchMessage::new(
644                None,
645                Some(b"created".to_vec()),
646                1_000,
647            )],
648            RecordBatchCompression::None,
649            RecordBatchIdentity {
650                producer_id: 42,
651                producer_epoch: 3,
652                base_sequence: 0,
653            },
654            true,
655        )
656        .unwrap();
657
658        assert_eq!(&set[21..23], &0x10_i16.to_be_bytes());
659    }
660
661    #[test]
662    fn encodes_produce_request_v3_with_record_batch() {
663        let request = ProduceRequestV3 {
664            correlation_id: 5,
665            client_id: Some("kafrust".to_owned()),
666            transactional_id: None,
667            acks: 1,
668            timeout_ms: 30_000,
669            topics: vec![ProduceTopicV3 {
670                name: "orders".to_owned(),
671                partitions: vec![ProducePartitionV3 {
672                    partition_index: 0,
673                    compression: RecordBatchCompression::None,
674                    identity: RecordBatchIdentity::NON_IDEMPOTENT,
675                    records: vec![RecordBatchMessage::new(
676                        Some(b"order-1".to_vec()),
677                        Some(b"created".to_vec()),
678                        1_000,
679                    )
680                    .header("source", Some(b"checkout".to_vec()))],
681                }],
682            }],
683        };
684
685        let bytes = request.encode().unwrap();
686
687        assert_eq!(&bytes[0..4], &[0, 0, 0, 3]);
688        assert!(bytes.len() > 80);
689    }
690
691    #[test]
692    fn record_batch_encoding_roundtrips_through_fetch_decoder() {
693        let record_set = encode_record_batch_set(&[RecordBatchMessage::new(
694            Some(b"order-1".to_vec()),
695            Some(b"created".to_vec()),
696            1_000,
697        )
698        .header("source", Some(b"checkout".to_vec()))])
699        .unwrap();
700
701        let mut bytes = Encoder::new();
702        bytes.write_i32(0);
703        bytes.write_i32(1);
704        bytes.write_string("orders").unwrap();
705        bytes.write_i32(1);
706        bytes.write_i32(0);
707        bytes.write_i16(0);
708        bytes.write_i64(43);
709        bytes.write_bytes(&record_set).unwrap();
710        let bytes = bytes.into_bytes();
711
712        let mut decoder = Decoder::new(&bytes);
713        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
714        let record = &response.responses[0].partitions[0].records[0];
715
716        assert_eq!(record.offset, 0);
717        assert_eq!(record.timestamp_ms, 1_000);
718        assert_eq!(record.key.as_deref(), Some(&b"order-1"[..]));
719        assert_eq!(record.value.as_deref(), Some(&b"created"[..]));
720        assert!(decoder.is_empty());
721    }
722
723    #[test]
724    fn gzip_record_batch_encoding_roundtrips_through_fetch_decoder() {
725        let record_set =
726            encode_record_batch_set_with_compression(
727                &[RecordBatchMessage::new(
728                    Some(b"order-1".to_vec()),
729                    Some(b"created".to_vec()),
730                    1_000,
731                )
732                .header("source", Some(b"checkout".to_vec()))],
733                RecordBatchCompression::Gzip,
734            )
735            .unwrap();
736
737        let mut bytes = Encoder::new();
738        bytes.write_i32(0);
739        bytes.write_i32(1);
740        bytes.write_string("orders").unwrap();
741        bytes.write_i32(1);
742        bytes.write_i32(0);
743        bytes.write_i16(0);
744        bytes.write_i64(43);
745        bytes.write_bytes(&record_set).unwrap();
746        let bytes = bytes.into_bytes();
747
748        let mut decoder = Decoder::new(&bytes);
749        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
750        let record = &response.responses[0].partitions[0].records[0];
751
752        assert_eq!(record.offset, 0);
753        assert_eq!(record.timestamp_ms, 1_000);
754        assert_eq!(record.key.as_deref(), Some(&b"order-1"[..]));
755        assert_eq!(record.value.as_deref(), Some(&b"created"[..]));
756        assert!(decoder.is_empty());
757    }
758
759    #[test]
760    fn snappy_record_batch_encoding_roundtrips_through_fetch_decoder() {
761        let record_set =
762            encode_record_batch_set_with_compression(
763                &[RecordBatchMessage::new(
764                    Some(b"order-1".to_vec()),
765                    Some(b"created".to_vec()),
766                    1_000,
767                )
768                .header("source", Some(b"checkout".to_vec()))],
769                RecordBatchCompression::Snappy,
770            )
771            .unwrap();
772
773        let mut bytes = Encoder::new();
774        bytes.write_i32(0);
775        bytes.write_i32(1);
776        bytes.write_string("orders").unwrap();
777        bytes.write_i32(1);
778        bytes.write_i32(0);
779        bytes.write_i16(0);
780        bytes.write_i64(43);
781        bytes.write_bytes(&record_set).unwrap();
782        let bytes = bytes.into_bytes();
783
784        let mut decoder = Decoder::new(&bytes);
785        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
786        let record = &response.responses[0].partitions[0].records[0];
787
788        assert_eq!(record.offset, 0);
789        assert_eq!(record.timestamp_ms, 1_000);
790        assert_eq!(record.key.as_deref(), Some(&b"order-1"[..]));
791        assert_eq!(record.value.as_deref(), Some(&b"created"[..]));
792        assert!(decoder.is_empty());
793    }
794
795    #[test]
796    fn lz4_record_batch_encoding_roundtrips_through_fetch_decoder() {
797        let record_set =
798            encode_record_batch_set_with_compression(
799                &[RecordBatchMessage::new(
800                    Some(b"order-1".to_vec()),
801                    Some(b"created".to_vec()),
802                    1_000,
803                )
804                .header("source", Some(b"checkout".to_vec()))],
805                RecordBatchCompression::Lz4,
806            )
807            .unwrap();
808
809        let mut bytes = Encoder::new();
810        bytes.write_i32(0);
811        bytes.write_i32(1);
812        bytes.write_string("orders").unwrap();
813        bytes.write_i32(1);
814        bytes.write_i32(0);
815        bytes.write_i16(0);
816        bytes.write_i64(43);
817        bytes.write_bytes(&record_set).unwrap();
818        let bytes = bytes.into_bytes();
819
820        let mut decoder = Decoder::new(&bytes);
821        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
822        let record = &response.responses[0].partitions[0].records[0];
823
824        assert_eq!(record.offset, 0);
825        assert_eq!(record.timestamp_ms, 1_000);
826        assert_eq!(record.key.as_deref(), Some(&b"order-1"[..]));
827        assert_eq!(record.value.as_deref(), Some(&b"created"[..]));
828        assert!(decoder.is_empty());
829    }
830
831    #[test]
832    fn encodes_produce_request_v7_with_record_batch() {
833        let request = ProduceRequestV7 {
834            correlation_id: 5,
835            client_id: Some("kafrust".to_owned()),
836            transactional_id: None,
837            acks: 1,
838            timeout_ms: 30_000,
839            topics: vec![ProduceTopicV3 {
840                name: "orders".to_owned(),
841                partitions: vec![ProducePartitionV3 {
842                    partition_index: 0,
843                    compression: RecordBatchCompression::Zstd,
844                    identity: RecordBatchIdentity::NON_IDEMPOTENT,
845                    records: vec![RecordBatchMessage::new(
846                        Some(b"order-1".to_vec()),
847                        Some(b"created".to_vec()),
848                        1_000,
849                    )],
850                }],
851            }],
852        };
853
854        let bytes = request.encode().unwrap();
855
856        assert_eq!(&bytes[0..4], &[0, 0, 0, 7]);
857        assert!(bytes.len() > 70);
858    }
859
860    #[test]
861    fn zstd_record_batch_encoding_roundtrips_through_fetch_decoder() {
862        let record_set =
863            encode_record_batch_set_with_compression(
864                &[RecordBatchMessage::new(
865                    Some(b"order-1".to_vec()),
866                    Some(b"created".to_vec()),
867                    1_000,
868                )
869                .header("source", Some(b"checkout".to_vec()))],
870                RecordBatchCompression::Zstd,
871            )
872            .unwrap();
873
874        let mut bytes = Encoder::new();
875        bytes.write_i32(0);
876        bytes.write_i32(1);
877        bytes.write_string("orders").unwrap();
878        bytes.write_i32(1);
879        bytes.write_i32(0);
880        bytes.write_i16(0);
881        bytes.write_i64(43);
882        bytes.write_bytes(&record_set).unwrap();
883        let bytes = bytes.into_bytes();
884
885        let mut decoder = Decoder::new(&bytes);
886        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
887        let record = &response.responses[0].partitions[0].records[0];
888
889        assert_eq!(record.offset, 0);
890        assert_eq!(record.timestamp_ms, 1_000);
891        assert_eq!(record.key.as_deref(), Some(&b"order-1"[..]));
892        assert_eq!(record.value.as_deref(), Some(&b"created"[..]));
893        assert!(decoder.is_empty());
894    }
895
896    #[test]
897    fn fetch_decoder_applies_custom_decompression_limit_to_record_batch() {
898        let record_set = encode_record_batch_set_with_compression(
899            &[RecordBatchMessage::new(
900                Some(b"order-1".to_vec()),
901                Some(vec![b'x'; 1024]),
902                1_000,
903            )],
904            RecordBatchCompression::Zstd,
905        )
906        .unwrap();
907
908        let mut bytes = Encoder::new();
909        bytes.write_i32(0);
910        bytes.write_i32(1);
911        bytes.write_string("orders").unwrap();
912        bytes.write_i32(1);
913        bytes.write_i32(0);
914        bytes.write_i16(0);
915        bytes.write_i64(43);
916        bytes.write_bytes(&record_set).unwrap();
917        let bytes = bytes.into_bytes();
918        let limits = DecodeLimits::new().with_max_decompressed_record_bytes(64);
919        let mut decoder = Decoder::with_limits(&bytes, limits);
920
921        assert!(matches!(
922            FetchResponseV2::decode_body(&mut decoder),
923            Err(crate::Error::LimitExceeded {
924                kind: "decompressed record batch bytes",
925                max: 64,
926                ..
927            })
928        ));
929    }
930
931    #[test]
932    fn reports_message_set_encoded_len() {
933        let records = [MessageSetMessage::new(
934            Some(b"order-1".to_vec()),
935            Some(b"created".to_vec()),
936            1_000,
937        )];
938
939        assert_eq!(
940            encoded_message_set_len(&records).unwrap(),
941            encode_message_set(&records).unwrap().len()
942        );
943    }
944
945    #[test]
946    fn reports_record_batch_set_encoded_len() {
947        let records =
948            [
949                RecordBatchMessage::new(
950                    Some(b"order-1".to_vec()),
951                    Some(b"created".to_vec()),
952                    1_000,
953                )
954                .header("source", Some(b"checkout".to_vec())),
955            ];
956
957        assert_eq!(
958            encoded_record_batch_set_len(&records).unwrap(),
959            encode_record_batch_set(&records).unwrap().len()
960        );
961    }
962
963    #[test]
964    fn decodes_produce_response_v2() {
965        let bytes = [
966            0, 0, 0, 1, // topic response count
967            0, 6, b'o', b'r', b'd', b'e', b'r', b's', // topic
968            0, 0, 0, 1, // partition response count
969            0, 0, 0, 0, // partition
970            0, 0, // error code
971            0, 0, 0, 0, 0, 0, 0, 42, // base offset
972            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // log append time -1
973            0, 0, 0, 0, // throttle time
974        ];
975        let mut decoder = Decoder::new(&bytes);
976        let response = ProduceResponseV2::decode_body(&mut decoder).unwrap();
977
978        assert_eq!(response.throttle_time_ms, 0);
979        assert_eq!(response.responses[0].name, "orders");
980        assert_eq!(response.responses[0].partitions[0].partition_index, 0);
981        assert_eq!(response.responses[0].partitions[0].error_code, 0);
982        assert_eq!(response.responses[0].partitions[0].base_offset, 42);
983        assert!(decoder.is_empty());
984    }
985
986    #[test]
987    fn decodes_produce_response_v7() {
988        let bytes = [
989            0, 0, 0, 1, // topic response count
990            0, 6, b'o', b'r', b'd', b'e', b'r', b's', // topic
991            0, 0, 0, 1, // partition response count
992            0, 0, 0, 0, // partition
993            0, 0, // error code
994            0, 0, 0, 0, 0, 0, 0, 42, // base offset
995            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // log append time -1
996            0, 0, 0, 0, 0, 0, 0, 7, // log start offset
997            0, 0, 0, 0, // throttle time
998        ];
999        let mut decoder = Decoder::new(&bytes);
1000        let response = ProduceResponseV7::decode_body(&mut decoder).unwrap();
1001
1002        assert_eq!(response.throttle_time_ms, 0);
1003        assert_eq!(response.responses[0].name, "orders");
1004        assert_eq!(response.responses[0].partitions[0].base_offset, 42);
1005        assert_eq!(response.responses[0].partitions[0].log_start_offset, 7);
1006        assert!(decoder.is_empty());
1007    }
1008}
kafrust_protocol/api/produce.rs

kafrust_protocol/api/
produce.rs
