Skip to main content

kafrust_protocol/
consumer_group.rs

1use crate::codec::{Decoder, Encoder};
2use crate::error::{Error, Result};
3
4#[derive(Debug, Clone, PartialEq, Eq)]
5pub struct ConsumerProtocolSubscriptionV0 {
6    pub topics: Vec<String>,
7    pub user_data: Option<Vec<u8>>,
8}
9
10impl ConsumerProtocolSubscriptionV0 {
11    pub fn encode(&self) -> Result<Vec<u8>> {
12        let mut encoder = Encoder::new();
13        encoder.write_i16(0);
14        encoder.write_array(Some(self.topics.as_slice()), |encoder, topic| {
15            encoder.write_string(topic)
16        })?;
17        encoder.write_nullable_bytes(self.user_data.as_deref())?;
18        Ok(encoder.into_bytes())
19    }
20
21    pub fn decode(bytes: &[u8]) -> Result<Self> {
22        let mut decoder = Decoder::new(bytes);
23        expect_version_v0("consumer protocol subscription", decoder.read_i16()?)?;
24        let subscription = Self {
25            topics: decoder
26                .read_array(
27                    "consumer protocol subscription topics",
28                    Decoder::read_string,
29                )?
30                .unwrap_or_default(),
31            user_data: decoder.read_nullable_bytes()?,
32        };
33        Ok(subscription)
34    }
35}
36
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct ConsumerProtocolAssignmentV0 {
39    pub assignments: Vec<ConsumerProtocolTopicAssignment>,
40    pub user_data: Option<Vec<u8>>,
41}
42
43impl ConsumerProtocolAssignmentV0 {
44    pub fn encode(&self) -> Result<Vec<u8>> {
45        let mut encoder = Encoder::new();
46        encoder.write_i16(0);
47        encoder.write_array(Some(self.assignments.as_slice()), |encoder, assignment| {
48            assignment.encode(encoder)
49        })?;
50        encoder.write_nullable_bytes(self.user_data.as_deref())?;
51        Ok(encoder.into_bytes())
52    }
53
54    pub fn decode(bytes: &[u8]) -> Result<Self> {
55        let mut decoder = Decoder::new(bytes);
56        expect_version_v0("consumer protocol assignment", decoder.read_i16()?)?;
57        let assignment = Self {
58            assignments: decoder
59                .read_array(
60                    "consumer protocol assignment topics",
61                    ConsumerProtocolTopicAssignment::decode,
62                )?
63                .unwrap_or_default(),
64            user_data: decoder.read_nullable_bytes()?,
65        };
66        Ok(assignment)
67    }
68}
69
70fn expect_version_v0(kind: &'static str, version: i16) -> Result<()> {
71    if version == 0 {
72        return Ok(());
73    }
74
75    Err(Error::UnsupportedVersion { kind, version })
76}
77
78#[derive(Debug, Clone, PartialEq, Eq)]
79pub struct ConsumerProtocolTopicAssignment {
80    pub topic: String,
81    pub partitions: Vec<i32>,
82}
83
84impl ConsumerProtocolTopicAssignment {
85    fn encode(&self, encoder: &mut Encoder) -> Result<()> {
86        encoder.write_string(&self.topic)?;
87        encoder.write_array(Some(self.partitions.as_slice()), |encoder, partition| {
88            encoder.write_i32(*partition);
89            Ok(())
90        })
91    }
92
93    fn decode(decoder: &mut Decoder<'_>) -> Result<Self> {
94        Ok(Self {
95            topic: decoder.read_string()?,
96            partitions: decoder
97                .read_array("consumer protocol assignment partitions", |decoder| {
98                    decoder.read_i32()
99                })?
100                .unwrap_or_default(),
101        })
102    }
103}
104
105#[cfg(test)]
106#[allow(clippy::unwrap_used)]
107mod tests {
108    use super::{
109        ConsumerProtocolAssignmentV0, ConsumerProtocolSubscriptionV0,
110        ConsumerProtocolTopicAssignment,
111    };
112
113    #[test]
114    fn encodes_consumer_protocol_subscription_v0() {
115        let subscription = ConsumerProtocolSubscriptionV0 {
116            topics: vec!["orders".to_owned(), "payments".to_owned()],
117            user_data: None,
118        };
119
120        assert_eq!(
121            subscription.encode().unwrap(),
122            [
123                0, 0, // version
124                0, 0, 0, 2, // topic count
125                0, 6, b'o', b'r', b'd', b'e', b'r', b's', // orders
126                0, 8, b'p', b'a', b'y', b'm', b'e', b'n', b't', b's', // payments
127                0xff, 0xff, 0xff, 0xff, // null user data
128            ]
129        );
130    }
131
132    #[test]
133    fn decodes_consumer_protocol_subscription_v0() {
134        let encoded = ConsumerProtocolSubscriptionV0 {
135            topics: vec!["orders".to_owned()],
136            user_data: Some(vec![1, 2, 3]),
137        }
138        .encode()
139        .unwrap();
140
141        let decoded = ConsumerProtocolSubscriptionV0::decode(&encoded).unwrap();
142
143        assert_eq!(
144            decoded,
145            ConsumerProtocolSubscriptionV0 {
146                topics: vec!["orders".to_owned()],
147                user_data: Some(vec![1, 2, 3]),
148            }
149        );
150    }
151
152    #[test]
153    fn rejects_unsupported_subscription_version() {
154        let error = ConsumerProtocolSubscriptionV0::decode(&[0, 1]).unwrap_err();
155
156        assert_eq!(
157            error,
158            crate::Error::UnsupportedVersion {
159                kind: "consumer protocol subscription",
160                version: 1,
161            }
162        );
163    }
164
165    #[test]
166    fn encodes_consumer_protocol_assignment_v0() {
167        let assignment = ConsumerProtocolAssignmentV0 {
168            assignments: vec![ConsumerProtocolTopicAssignment {
169                topic: "orders".to_owned(),
170                partitions: vec![0, 2],
171            }],
172            user_data: Some(vec![9]),
173        };
174
175        assert_eq!(
176            assignment.encode().unwrap(),
177            [
178                0, 0, // version
179                0, 0, 0, 1, // topic count
180                0, 6, b'o', b'r', b'd', b'e', b'r', b's', // topic
181                0, 0, 0, 2, // partition count
182                0, 0, 0, 0, // partition 0
183                0, 0, 0, 2, // partition 2
184                0, 0, 0, 1, 9, // user data
185            ]
186        );
187    }
188
189    #[test]
190    fn decodes_consumer_protocol_assignment_v0() {
191        let encoded = ConsumerProtocolAssignmentV0 {
192            assignments: vec![ConsumerProtocolTopicAssignment {
193                topic: "orders".to_owned(),
194                partitions: vec![0, 2],
195            }],
196            user_data: None,
197        }
198        .encode()
199        .unwrap();
200
201        let decoded = ConsumerProtocolAssignmentV0::decode(&encoded).unwrap();
202
203        assert_eq!(
204            decoded,
205            ConsumerProtocolAssignmentV0 {
206                assignments: vec![ConsumerProtocolTopicAssignment {
207                    topic: "orders".to_owned(),
208                    partitions: vec![0, 2],
209                }],
210                user_data: None,
211            }
212        );
213    }
214}