kafrust_protocol/api/
consumer_group_heartbeat.rs

1use crate::codec::{Decoder, Encoder};
2use crate::error::Result;
3use crate::header::RequestHeader;
4
5/// Kafka ConsumerGroupHeartbeat API key.
6pub const API_KEY: i16 = 68;
7
8/// Topic partitions owned by a member in the KIP-848 consumer group protocol.
9#[derive(Debug, Clone, PartialEq, Eq)]
10pub struct ConsumerGroupHeartbeatTopicPartitions {
11    pub topic_id: [u8; 16],
12    pub partitions: Vec<i32>,
13}
14
15impl ConsumerGroupHeartbeatTopicPartitions {
16    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
17        encoder.write_uuid(&self.topic_id);
18        encoder.write_compact_array(Some(&self.partitions), |encoder, partition| {
19            encoder.write_i32(*partition);
20            Ok(())
21        })?;
22        encoder.write_empty_tagged_fields();
23        Ok(())
24    }
25
26    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
27        let topic_id = decoder.read_uuid()?;
28        let partitions = decoder
29            .read_compact_array("consumer group heartbeat partitions", |decoder| {
30                decoder.read_i32()
31            })?
32            .unwrap_or_default();
33        decoder.read_tagged_fields()?;
34        Ok(Self {
35            topic_id,
36            partitions,
37        })
38    }
39}
40
41/// KIP-848 ConsumerGroupHeartbeat v0 request.
42#[derive(Debug, Clone, PartialEq, Eq)]
43pub struct ConsumerGroupHeartbeatRequestV0 {
44    pub correlation_id: i32,
45    pub client_id: Option<String>,
46    pub group_id: String,
47    pub member_id: String,
48    pub member_epoch: i32,
49    pub instance_id: Option<String>,
50    pub rack_id: Option<String>,
51    pub rebalance_timeout_ms: i32,
52    pub subscribed_topic_names: Option<Vec<String>>,
53    pub server_assignor: Option<String>,
54    pub topic_partitions: Option<Vec<ConsumerGroupHeartbeatTopicPartitions>>,
55}
56
57impl ConsumerGroupHeartbeatRequestV0 {
58    pub fn encode(&self) -> Result<Vec<u8>> {
59        let mut encoder = Encoder::new();
60        RequestHeader {
61            api_key: API_KEY,
62            api_version: 0,
63            correlation_id: self.correlation_id,
64            client_id: self.client_id.clone(),
65        }
66        .encode_v2(&mut encoder)?;
67        encoder.write_compact_string(&self.group_id)?;
68        encoder.write_compact_string(&self.member_id)?;
69        encoder.write_i32(self.member_epoch);
70        encoder.write_compact_nullable_string(self.instance_id.as_deref())?;
71        encoder.write_compact_nullable_string(self.rack_id.as_deref())?;
72        encoder.write_i32(self.rebalance_timeout_ms);
73        encoder.write_compact_array(self.subscribed_topic_names.as_deref(), |encoder, topic| {
74            encoder.write_compact_string(topic)
75        })?;
76        encoder.write_compact_nullable_string(self.server_assignor.as_deref())?;
77        encoder.write_compact_array(self.topic_partitions.as_deref(), |encoder, topic| {
78            topic.encode(encoder)
79        })?;
80        encoder.write_empty_tagged_fields();
81        Ok(encoder.into_bytes())
82    }
83}
84
85/// KIP-848 ConsumerGroupHeartbeat v1 request.
86///
87/// Version 1 keeps the v0 flexible wire shape and adds the nullable topic
88/// subscription regular expression introduced by KIP-848. It also lets the
89/// consumer provide a stable member ID, which is represented by the existing
90/// `member_id` field.
91#[derive(Debug, Clone, PartialEq, Eq)]
92pub struct ConsumerGroupHeartbeatRequestV1 {
93    pub correlation_id: i32,
94    pub client_id: Option<String>,
95    pub group_id: String,
96    pub member_id: String,
97    pub member_epoch: i32,
98    pub instance_id: Option<String>,
99    pub rack_id: Option<String>,
100    pub rebalance_timeout_ms: i32,
101    pub subscribed_topic_names: Option<Vec<String>>,
102    pub subscribed_topic_regex: Option<String>,
103    pub server_assignor: Option<String>,
104    pub topic_partitions: Option<Vec<ConsumerGroupHeartbeatTopicPartitions>>,
105}
106
107impl ConsumerGroupHeartbeatRequestV1 {
108    pub fn encode(&self) -> Result<Vec<u8>> {
109        let mut encoder = Encoder::new();
110        RequestHeader {
111            api_key: API_KEY,
112            api_version: 1,
113            correlation_id: self.correlation_id,
114            client_id: self.client_id.clone(),
115        }
116        .encode_v2(&mut encoder)?;
117        encoder.write_compact_string(&self.group_id)?;
118        encoder.write_compact_string(&self.member_id)?;
119        encoder.write_i32(self.member_epoch);
120        encoder.write_compact_nullable_string(self.instance_id.as_deref())?;
121        encoder.write_compact_nullable_string(self.rack_id.as_deref())?;
122        encoder.write_i32(self.rebalance_timeout_ms);
123        encoder.write_compact_array(self.subscribed_topic_names.as_deref(), |encoder, topic| {
124            encoder.write_compact_string(topic)
125        })?;
126        encoder.write_compact_nullable_string(self.subscribed_topic_regex.as_deref())?;
127        encoder.write_compact_nullable_string(self.server_assignor.as_deref())?;
128        encoder.write_compact_array(self.topic_partitions.as_deref(), |encoder, topic| {
129            topic.encode(encoder)
130        })?;
131        encoder.write_empty_tagged_fields();
132        Ok(encoder.into_bytes())
133    }
134}
135
136/// KIP-848 ConsumerGroupHeartbeat v0 response.
137#[derive(Debug, Clone, PartialEq, Eq)]
138pub struct ConsumerGroupHeartbeatResponseV0 {
139    pub throttle_time_ms: i32,
140    pub error_code: i16,
141    pub error_message: Option<String>,
142    pub member_id: Option<String>,
143    pub member_epoch: i32,
144    pub heartbeat_interval_ms: i32,
145    pub assignment: Option<Vec<ConsumerGroupHeartbeatTopicPartitions>>,
146}
147
148impl ConsumerGroupHeartbeatResponseV0 {
149    pub fn decode_body(decoder: &mut Decoder<'_>) -> Result<Self> {
150        let throttle_time_ms = decoder.read_i32()?;
151        let error_code = decoder.read_i16()?;
152        let error_message = decoder.read_compact_nullable_string()?;
153        let member_id = decoder.read_compact_nullable_string()?;
154        let member_epoch = decoder.read_i32()?;
155        let heartbeat_interval_ms = decoder.read_i32()?;
156        // Assignment is a nullable struct, not a compact nullable array. The
157        // struct contains the compact topic-partitions array and its tags.
158        let assignment = match decoder.read_i8()? {
159            -1 => None,
160            1 => {
161                let topic_partitions = decoder
162                    .read_compact_array("consumer group heartbeat assignment", |decoder| {
163                        ConsumerGroupHeartbeatTopicPartitions::decode(decoder)
164                    })?
165                    .unwrap_or_default();
166                decoder.read_tagged_fields()?;
167                Some(topic_partitions)
168            }
169            marker => return Err(crate::error::Error::InvalidNullableStruct(marker)),
170        };
171        decoder.read_tagged_fields()?;
172        Ok(Self {
173            throttle_time_ms,
174            error_code,
175            error_message,
176            member_id,
177            member_epoch,
178            heartbeat_interval_ms,
179            assignment,
180        })
181    }
182}
183
184/// ConsumerGroupHeartbeat v1 has the same response wire shape as v0.
185pub type ConsumerGroupHeartbeatResponseV1 = ConsumerGroupHeartbeatResponseV0;
186
187#[cfg(test)]
188#[allow(clippy::unwrap_used)]
189mod tests {
190    use super::{
191        ConsumerGroupHeartbeatRequestV0, ConsumerGroupHeartbeatRequestV1,
192        ConsumerGroupHeartbeatResponseV0, ConsumerGroupHeartbeatResponseV1,
193        ConsumerGroupHeartbeatTopicPartitions,
194    };
195    use crate::codec::{Decoder, Encoder};
196
197    #[test]
198    fn encodes_kip_848_heartbeat_request_with_subscription_and_assignment() {
199        let request = ConsumerGroupHeartbeatRequestV0 {
200            correlation_id: 23,
201            client_id: Some("kafrust".to_owned()),
202            group_id: "orders-group".to_owned(),
203            member_id: "member-a".to_owned(),
204            member_epoch: 4,
205            instance_id: Some("instance-a".to_owned()),
206            rack_id: None,
207            rebalance_timeout_ms: 30_000,
208            subscribed_topic_names: Some(vec!["orders".to_owned()]),
209            server_assignor: Some("uniform".to_owned()),
210            topic_partitions: Some(vec![ConsumerGroupHeartbeatTopicPartitions {
211                topic_id: [1; 16],
212                partitions: vec![0, 2],
213            }]),
214        };
215
216        let encoded = request.encode().unwrap();
217        assert_eq!(&encoded[0..4], &[0, 68, 0, 0]);
218        assert_eq!(&encoded[4..12], &[0, 0, 0, 23, 0, 7, b'k', b'a']);
219        assert!(encoded.windows(12).any(|bytes| bytes == b"orders-group"));
220        assert!(encoded.windows(16).any(|bytes| bytes == [1; 16]));
221        assert_eq!(*encoded.last().unwrap(), 0);
222    }
223
224    #[test]
225    fn decodes_kip_848_heartbeat_response_and_assignment() {
226        let mut bytes = Encoder::new();
227        bytes.write_i32(12);
228        bytes.write_i16(0);
229        bytes.write_compact_nullable_string(None).unwrap();
230        bytes
231            .write_compact_nullable_string(Some("member-a"))
232            .unwrap();
233        bytes.write_i32(5);
234        bytes.write_i32(2500);
235        bytes.write_i8(1);
236        bytes
237            .write_compact_array(
238                Some(&[ConsumerGroupHeartbeatTopicPartitions {
239                    topic_id: [2; 16],
240                    partitions: vec![1, 3],
241                }]),
242                |encoder, topic| topic.encode(encoder),
243            )
244            .unwrap();
245        bytes.write_empty_tagged_fields();
246        bytes.write_empty_tagged_fields();
247
248        let bytes = bytes.into_bytes();
249        let mut decoder = Decoder::new(&bytes);
250        let response = ConsumerGroupHeartbeatResponseV0::decode_body(&mut decoder).unwrap();
251
252        assert_eq!(response.throttle_time_ms, 12);
253        assert_eq!(response.error_code, 0);
254        assert_eq!(response.member_id.as_deref(), Some("member-a"));
255        assert_eq!(response.member_epoch, 5);
256        assert_eq!(response.heartbeat_interval_ms, 2500);
257        assert_eq!(response.assignment.unwrap()[0].partitions, vec![1, 3]);
258        assert!(decoder.is_empty());
259    }
260
261    #[test]
262    fn decodes_kip_848_heartbeat_response_with_null_assignment_struct() {
263        let mut bytes = Encoder::new();
264        bytes.write_i32(0);
265        bytes.write_i16(0);
266        bytes.write_compact_nullable_string(None).unwrap();
267        bytes
268            .write_compact_nullable_string(Some("member-a"))
269            .unwrap();
270        bytes.write_i32(1);
271        bytes.write_i32(2500);
272        bytes.write_i8(-1);
273        bytes.write_empty_tagged_fields();
274
275        let bytes = bytes.into_bytes();
276        let mut decoder = Decoder::new(&bytes);
277        let response = ConsumerGroupHeartbeatResponseV0::decode_body(&mut decoder).unwrap();
278
279        assert!(response.assignment.is_none());
280        assert!(decoder.is_empty());
281    }
282
283    #[test]
284    fn encodes_nullable_kip_848_fields_as_null_arrays() {
285        let request = ConsumerGroupHeartbeatRequestV0 {
286            correlation_id: 1,
287            client_id: None,
288            group_id: "g".to_owned(),
289            member_id: "m".to_owned(),
290            member_epoch: -1,
291            instance_id: None,
292            rack_id: None,
293            rebalance_timeout_ms: -1,
294            subscribed_topic_names: None,
295            server_assignor: None,
296            topic_partitions: None,
297        };
298
299        let encoded = request.encode().unwrap();
300        assert_eq!(&encoded[0..10], &[0, 68, 0, 0, 0, 0, 0, 1, 0xff, 0xff]);
301        assert_eq!(encoded[10], 0); // request-header tagged fields
302        assert!(encoded.ends_with(&[0, 0]));
303    }
304
305    #[test]
306    fn encodes_kip_848_heartbeat_v1_with_topic_regex() {
307        let request = ConsumerGroupHeartbeatRequestV1 {
308            correlation_id: 9,
309            client_id: Some("kafrust".to_owned()),
310            group_id: "orders-group".to_owned(),
311            member_id: "member-a".to_owned(),
312            member_epoch: 2,
313            instance_id: None,
314            rack_id: Some("rack-a".to_owned()),
315            rebalance_timeout_ms: 30_000,
316            subscribed_topic_names: None,
317            subscribed_topic_regex: Some("orders-.*".to_owned()),
318            server_assignor: Some("uniform".to_owned()),
319            topic_partitions: None,
320        };
321
322        let encoded = request.encode().unwrap();
323        assert_eq!(&encoded[0..4], &[0, 68, 0, 1]);
324        assert!(encoded.windows(9).any(|bytes| bytes == b"orders-.*"));
325        assert!(encoded.ends_with(&[0, 0]));
326    }
327
328    #[test]
329    fn v1_response_alias_decodes_the_v0_wire_shape() {
330        let mut bytes = Encoder::new();
331        bytes.write_i32(0);
332        bytes.write_i16(0);
333        bytes.write_compact_nullable_string(None).unwrap();
334        bytes
335            .write_compact_nullable_string(Some("member-a"))
336            .unwrap();
337        bytes.write_i32(3);
338        bytes.write_i32(2500);
339        bytes.write_i8(-1);
340        bytes.write_empty_tagged_fields();
341
342        let bytes = bytes.into_bytes();
343        let mut decoder = Decoder::new(&bytes);
344        let response = ConsumerGroupHeartbeatResponseV1::decode_body(&mut decoder).unwrap();
345        assert_eq!(response.member_epoch, 3);
346        assert!(response.assignment.is_none());
347        assert!(decoder.is_empty());
348    }
349}
kafrust_protocol/api/consumer_group_heartbeat.rs

kafrust_protocol/api/
consumer_group_heartbeat.rs
