kafrust_protocol/api/
produce.rs

1use crate::codec::{Decoder, Encoder};
2use crate::error::{Error, Result};
3use crate::header::RequestHeader;
4
5pub const API_KEY: i16 = 0;
6
7#[derive(Debug, Clone, PartialEq, Eq)]
8pub struct ProduceRequestV2 {
9    pub correlation_id: i32,
10    pub client_id: Option<String>,
11    pub acks: i16,
12    pub timeout_ms: i32,
13    pub topics: Vec<ProduceTopicV2>,
14}
15
16#[derive(Debug, Clone, PartialEq, Eq)]
17pub struct ProduceRequestV3 {
18    pub correlation_id: i32,
19    pub client_id: Option<String>,
20    pub transactional_id: Option<String>,
21    pub acks: i16,
22    pub timeout_ms: i32,
23    pub topics: Vec<ProduceTopicV3>,
24}
25
26impl ProduceRequestV3 {
27    pub fn encode(&self) -> Result<Vec<u8>> {
28        let mut encoder = Encoder::new();
29        RequestHeader {
30            api_key: API_KEY,
31            api_version: 3,
32            correlation_id: self.correlation_id,
33            client_id: self.client_id.clone(),
34        }
35        .encode_v1(&mut encoder)?;
36        encoder.write_nullable_string(self.transactional_id.as_deref())?;
37        encoder.write_i16(self.acks);
38        encoder.write_i32(self.timeout_ms);
39        encoder.write_array(Some(self.topics.as_slice()), |encoder, topic| {
40            topic.encode(encoder)
41        })?;
42        Ok(encoder.into_bytes())
43    }
44}
45
46impl ProduceRequestV2 {
47    pub fn encode(&self) -> Result<Vec<u8>> {
48        let mut encoder = Encoder::new();
49        RequestHeader {
50            api_key: API_KEY,
51            api_version: 2,
52            correlation_id: self.correlation_id,
53            client_id: self.client_id.clone(),
54        }
55        .encode_v1(&mut encoder)?;
56        encoder.write_i16(self.acks);
57        encoder.write_i32(self.timeout_ms);
58        encoder.write_array(Some(self.topics.as_slice()), |encoder, topic| {
59            topic.encode(encoder)
60        })?;
61        Ok(encoder.into_bytes())
62    }
63}
64
65#[derive(Debug, Clone, PartialEq, Eq)]
66pub struct ProduceTopicV2 {
67    pub name: String,
68    pub partitions: Vec<ProducePartitionV2>,
69}
70
71#[derive(Debug, Clone, PartialEq, Eq)]
72pub struct ProduceTopicV3 {
73    pub name: String,
74    pub partitions: Vec<ProducePartitionV3>,
75}
76
77impl ProduceTopicV3 {
78    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
79        encoder.write_string(&self.name)?;
80        encoder.write_array(Some(self.partitions.as_slice()), |encoder, partition| {
81            partition.encode(encoder)
82        })
83    }
84}
85
86impl ProduceTopicV2 {
87    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
88        encoder.write_string(&self.name)?;
89        encoder.write_array(Some(self.partitions.as_slice()), |encoder, partition| {
90            partition.encode(encoder)
91        })
92    }
93}
94
95#[derive(Debug, Clone, PartialEq, Eq)]
96pub struct ProducePartitionV2 {
97    pub partition_index: i32,
98    pub records: Vec<MessageSetMessage>,
99}
100
101#[derive(Debug, Clone, PartialEq, Eq)]
102pub struct ProducePartitionV3 {
103    pub partition_index: i32,
104    pub records: Vec<RecordBatchMessage>,
105}
106
107impl ProducePartitionV3 {
108    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
109        encoder.write_i32(self.partition_index);
110        let record_set = encode_record_batch_set(&self.records)?;
111        encoder.write_bytes(&record_set)
112    }
113}
114
115impl ProducePartitionV2 {
116    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
117        encoder.write_i32(self.partition_index);
118        let record_set = encode_message_set(&self.records)?;
119        encoder.write_bytes(&record_set)
120    }
121}
122
123#[derive(Debug, Clone, PartialEq, Eq)]
124pub struct MessageSetMessage {
125    pub key: Option<Vec<u8>>,
126    pub value: Option<Vec<u8>>,
127    pub timestamp_ms: i64,
128}
129
130impl MessageSetMessage {
131    pub fn new(key: Option<Vec<u8>>, value: Option<Vec<u8>>, timestamp_ms: i64) -> Self {
132        Self {
133            key,
134            value,
135            timestamp_ms,
136        }
137    }
138}
139
140#[derive(Debug, Clone, PartialEq, Eq)]
141pub struct RecordBatchHeader {
142    pub key: String,
143    pub value: Option<Vec<u8>>,
144}
145
146impl RecordBatchHeader {
147    pub fn new(key: impl Into<String>, value: Option<Vec<u8>>) -> Self {
148        Self {
149            key: key.into(),
150            value,
151        }
152    }
153}
154
155#[derive(Debug, Clone, PartialEq, Eq)]
156pub struct RecordBatchMessage {
157    pub key: Option<Vec<u8>>,
158    pub value: Option<Vec<u8>>,
159    pub timestamp_ms: i64,
160    pub headers: Vec<RecordBatchHeader>,
161}
162
163impl RecordBatchMessage {
164    pub fn new(key: Option<Vec<u8>>, value: Option<Vec<u8>>, timestamp_ms: i64) -> Self {
165        Self {
166            key,
167            value,
168            timestamp_ms,
169            headers: Vec::new(),
170        }
171    }
172
173    pub fn header(mut self, key: impl Into<String>, value: Option<Vec<u8>>) -> Self {
174        self.headers.push(RecordBatchHeader::new(key, value));
175        self
176    }
177}
178
179#[derive(Debug, Clone, PartialEq, Eq)]
180pub struct ProduceResponseV2 {
181    pub responses: Vec<ProduceTopicResponseV2>,
182    pub throttle_time_ms: i32,
183}
184
185impl ProduceResponseV2 {
186    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
187        Ok(Self {
188            responses: decoder
189                .read_array("produce responses", ProduceTopicResponseV2::decode)?
190                .unwrap_or_default(),
191            throttle_time_ms: decoder.read_i32()?,
192        })
193    }
194}
195
196#[derive(Debug, Clone, PartialEq, Eq)]
197pub struct ProduceTopicResponseV2 {
198    pub name: String,
199    pub partitions: Vec<ProducePartitionResponseV2>,
200}
201
202impl ProduceTopicResponseV2 {
203    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
204        Ok(Self {
205            name: decoder.read_string()?,
206            partitions: decoder
207                .read_array(
208                    "produce partition responses",
209                    ProducePartitionResponseV2::decode,
210                )?
211                .unwrap_or_default(),
212        })
213    }
214}
215
216#[derive(Debug, Clone, PartialEq, Eq)]
217pub struct ProducePartitionResponseV2 {
218    pub partition_index: i32,
219    pub error_code: i16,
220    pub base_offset: i64,
221    pub log_append_time_ms: i64,
222}
223
224impl ProducePartitionResponseV2 {
225    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
226        Ok(Self {
227            partition_index: decoder.read_i32()?,
228            error_code: decoder.read_i16()?,
229            base_offset: decoder.read_i64()?,
230            log_append_time_ms: decoder.read_i64()?,
231        })
232    }
233}
234
235fn encode_message_set(records: &[MessageSetMessage]) -> Result<Vec<u8>> {
236    let mut set = Encoder::new();
237    for record in records {
238        let message = encode_message(record)?;
239        set.write_i64(0);
240        set.write_i32(i32::try_from(message.len()).map_err(|_| Error::LengthOverflow("message"))?);
241        set.write_raw(&message);
242    }
243    Ok(set.into_bytes())
244}
245
246fn encode_message(record: &MessageSetMessage) -> Result<Vec<u8>> {
247    let mut body = Encoder::new();
248    body.write_i8(1);
249    body.write_i8(0);
250    body.write_i64(record.timestamp_ms);
251    body.write_nullable_bytes(record.key.as_deref())?;
252    body.write_nullable_bytes(record.value.as_deref())?;
253    let body = body.into_bytes();
254
255    let mut message = Encoder::new();
256    message.write_i32(crc32_ieee(&body) as i32);
257    message.write_raw(&body);
258    Ok(message.into_bytes())
259}
260
261fn encode_record_batch_set(records: &[RecordBatchMessage]) -> Result<Vec<u8>> {
262    let base_timestamp = records
263        .first()
264        .map(|record| record.timestamp_ms)
265        .unwrap_or_default();
266    let max_timestamp = records
267        .iter()
268        .map(|record| record.timestamp_ms)
269        .max()
270        .unwrap_or(base_timestamp);
271    let last_offset_delta = records
272        .len()
273        .checked_sub(1)
274        .map(|delta| i32::try_from(delta).map_err(|_| Error::LengthOverflow("record batch")))
275        .transpose()?
276        .unwrap_or_default();
277
278    let mut record_bytes = Encoder::new();
279    record_bytes.write_i32(
280        i32::try_from(records.len()).map_err(|_| Error::LengthOverflow("record batch records"))?,
281    );
282    for (offset_delta, record) in records.iter().enumerate() {
283        let encoded = encode_record(record, base_timestamp, offset_delta)?;
284        record_bytes.write_varint(
285            i32::try_from(encoded.len()).map_err(|_| Error::LengthOverflow("record"))?,
286        );
287        record_bytes.write_raw(&encoded);
288    }
289
290    let mut crc_payload = Encoder::new();
291    crc_payload.write_i16(0);
292    crc_payload.write_i32(last_offset_delta);
293    crc_payload.write_i64(base_timestamp);
294    crc_payload.write_i64(max_timestamp);
295    crc_payload.write_i64(-1);
296    crc_payload.write_i16(-1);
297    crc_payload.write_i32(-1);
298    crc_payload.write_raw(&record_bytes.into_bytes());
299    let crc_payload = crc_payload.into_bytes();
300
301    let mut batch = Encoder::new();
302    batch.write_i32(0);
303    batch.write_i8(2);
304    batch.write_i32(crc32c(&crc_payload) as i32);
305    batch.write_raw(&crc_payload);
306    let batch = batch.into_bytes();
307
308    let mut set = Encoder::new();
309    set.write_i64(0);
310    set.write_i32(i32::try_from(batch.len()).map_err(|_| Error::LengthOverflow("record batch"))?);
311    set.write_raw(&batch);
312    Ok(set.into_bytes())
313}
314
315fn encode_record(
316    record: &RecordBatchMessage,
317    base_timestamp: i64,
318    offset_delta: usize,
319) -> Result<Vec<u8>> {
320    let mut encoder = Encoder::new();
321    encoder.write_i8(0);
322    encoder.write_varlong(record.timestamp_ms.saturating_sub(base_timestamp));
323    encoder.write_varint(
324        i32::try_from(offset_delta).map_err(|_| Error::LengthOverflow("record offset delta"))?,
325    );
326    encoder.write_varint_nullable_bytes(record.key.as_deref())?;
327    encoder.write_varint_nullable_bytes(record.value.as_deref())?;
328    encoder.write_varint(
329        i32::try_from(record.headers.len()).map_err(|_| Error::LengthOverflow("record headers"))?,
330    );
331    for header in &record.headers {
332        encoder.write_varint_bytes(header.key.as_bytes())?;
333        encoder.write_varint_nullable_bytes(header.value.as_deref())?;
334    }
335    Ok(encoder.into_bytes())
336}
337
338fn crc32_ieee(bytes: &[u8]) -> u32 {
339    let mut crc = 0xffff_ffffu32;
340    for byte in bytes {
341        crc ^= u32::from(*byte);
342        for _ in 0..8 {
343            let mask = 0u32.wrapping_sub(crc & 1);
344            crc = (crc >> 1) ^ (0xedb8_8320 & mask);
345        }
346    }
347    !crc
348}
349
350fn crc32c(bytes: &[u8]) -> u32 {
351    let mut crc = 0xffff_ffffu32;
352    for byte in bytes {
353        crc ^= u32::from(*byte);
354        for _ in 0..8 {
355            let mask = 0u32.wrapping_sub(crc & 1);
356            crc = (crc >> 1) ^ (0x82f6_3b78 & mask);
357        }
358    }
359    !crc
360}
361
362#[cfg(test)]
363#[allow(clippy::unwrap_used)]
364mod tests {
365    use super::{
366        encode_record_batch_set, MessageSetMessage, ProducePartitionV2, ProducePartitionV3,
367        ProduceRequestV2, ProduceRequestV3, ProduceResponseV2, ProduceTopicV2, ProduceTopicV3,
368        RecordBatchMessage,
369    };
370    use crate::codec::Decoder;
371    use crate::{api::fetch::FetchResponseV2, codec::Encoder};
372
373    #[test]
374    fn encodes_produce_request_v2() {
375        let request = ProduceRequestV2 {
376            correlation_id: 5,
377            client_id: Some("kafrust".to_owned()),
378            acks: 1,
379            timeout_ms: 30_000,
380            topics: vec![ProduceTopicV2 {
381                name: "orders".to_owned(),
382                partitions: vec![ProducePartitionV2 {
383                    partition_index: 0,
384                    records: vec![MessageSetMessage::new(
385                        Some(b"order-1".to_vec()),
386                        Some(b"created".to_vec()),
387                        0,
388                    )],
389                }],
390            }],
391        };
392
393        let bytes = request.encode().unwrap();
394        assert_eq!(
395            &bytes[0..17],
396            &[0, 0, 0, 2, 0, 0, 0, 5, 0, 7, b'k', b'a', b'f', b'r', b'u', b's', b't',]
397        );
398        assert!(bytes.len() > 60);
399    }
400
401    #[test]
402    fn encodes_produce_request_v3_with_record_batch() {
403        let request = ProduceRequestV3 {
404            correlation_id: 5,
405            client_id: Some("kafrust".to_owned()),
406            transactional_id: None,
407            acks: 1,
408            timeout_ms: 30_000,
409            topics: vec![ProduceTopicV3 {
410                name: "orders".to_owned(),
411                partitions: vec![ProducePartitionV3 {
412                    partition_index: 0,
413                    records: vec![RecordBatchMessage::new(
414                        Some(b"order-1".to_vec()),
415                        Some(b"created".to_vec()),
416                        1_000,
417                    )
418                    .header("source", Some(b"checkout".to_vec()))],
419                }],
420            }],
421        };
422
423        let bytes = request.encode().unwrap();
424
425        assert_eq!(&bytes[0..4], &[0, 0, 0, 3]);
426        assert!(bytes.len() > 80);
427    }
428
429    #[test]
430    fn record_batch_encoding_roundtrips_through_fetch_decoder() {
431        let record_set = encode_record_batch_set(&[RecordBatchMessage::new(
432            Some(b"order-1".to_vec()),
433            Some(b"created".to_vec()),
434            1_000,
435        )
436        .header("source", Some(b"checkout".to_vec()))])
437        .unwrap();
438
439        let mut bytes = Encoder::new();
440        bytes.write_i32(0);
441        bytes.write_i32(1);
442        bytes.write_string("orders").unwrap();
443        bytes.write_i32(1);
444        bytes.write_i32(0);
445        bytes.write_i16(0);
446        bytes.write_i64(43);
447        bytes.write_bytes(&record_set).unwrap();
448        let bytes = bytes.into_bytes();
449
450        let mut decoder = Decoder::new(&bytes);
451        let response = FetchResponseV2::decode_body(&mut decoder).unwrap();
452        let record = &response.responses[0].partitions[0].records[0];
453
454        assert_eq!(record.offset, 0);
455        assert_eq!(record.timestamp_ms, 1_000);
456        assert_eq!(record.key.as_deref(), Some(&b"order-1"[..]));
457        assert_eq!(record.value.as_deref(), Some(&b"created"[..]));
458        assert!(decoder.is_empty());
459    }
460
461    #[test]
462    fn decodes_produce_response_v2() {
463        let bytes = [
464            0, 0, 0, 1, // topic response count
465            0, 6, b'o', b'r', b'd', b'e', b'r', b's', // topic
466            0, 0, 0, 1, // partition response count
467            0, 0, 0, 0, // partition
468            0, 0, // error code
469            0, 0, 0, 0, 0, 0, 0, 42, // base offset
470            0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, 0xff, // log append time -1
471            0, 0, 0, 0, // throttle time
472        ];
473        let mut decoder = Decoder::new(&bytes);
474        let response = ProduceResponseV2::decode_body(&mut decoder).unwrap();
475
476        assert_eq!(response.throttle_time_ms, 0);
477        assert_eq!(response.responses[0].name, "orders");
478        assert_eq!(response.responses[0].partitions[0].partition_index, 0);
479        assert_eq!(response.responses[0].partitions[0].error_code, 0);
480        assert_eq!(response.responses[0].partitions[0].base_offset, 42);
481        assert!(decoder.is_empty());
482    }
483}
kafrust_protocol/api/produce.rs

kafrust_protocol/api/
produce.rs
