krabka_client_consumer/assignor/
mod.rs

1//! Partition assignors: the eager `range`, `roundrobin` and `sticky`, and the
2//! incremental `cooperative-sticky` from KIP-429.
3
4pub(crate) mod cooperative_sticky;
5pub(crate) mod range;
6pub(crate) mod round_robin;
7
8use std::collections::HashMap;
9
10use bytes::{Buf, BufMut, Bytes, BytesMut};
11
12use crate::builder::DecodedSubscription;
13
14/// Kafka's `ConsumerPartitionAssignor.RebalanceProtocol`. The order is the
15/// order of preference: a later variant is newer.
16#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
17pub(crate) enum RebalanceProtocol {
18    Eager,
19    Cooperative,
20}
21
22/// A built-in partition assignor, Kafka's `partition.assignment.strategy`.
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub enum Assignor {
25    /// `range`, Kafka's `RangeAssignor`.
26    Range,
27    /// `roundrobin`, Kafka's `RoundRobinAssignor`.
28    RoundRobin,
29    /// `sticky`, Kafka's eager `StickyAssignor`.
30    Sticky,
31    /// `cooperative-sticky`, Kafka's `CooperativeStickyAssignor` (KIP-429).
32    CooperativeSticky,
33}
34
35impl std::str::FromStr for Assignor {
36    type Err = String;
37
38    fn from_str(value: &str) -> Result<Self, Self::Err> {
39        match value {
40            "range" => Ok(Self::Range),
41            "roundrobin" => Ok(Self::RoundRobin),
42            "sticky" => Ok(Self::Sticky),
43            "cooperative-sticky" => Ok(Self::CooperativeSticky),
44            _ => Err(format!("invalid assignor: {value}")),
45        }
46    }
47}
48
49impl Assignor {
50    /// Kafka's default `partition.assignment.strategy`.
51    pub const DEFAULT_LIST: [Self; 2] = [Self::Range, Self::CooperativeSticky];
52
53    /// The protocol name in `JoinGroup` and `SyncGroup`.
54    pub(crate) fn protocol_name(self) -> &'static str {
55        match self {
56            Assignor::Range => "range",
57            Assignor::RoundRobin => "roundrobin",
58            Assignor::Sticky => "sticky",
59            Assignor::CooperativeSticky => "cooperative-sticky",
60        }
61    }
62
63    /// Kafka's `ConsumerPartitionAssignor.supportedProtocols`.
64    fn supported_protocols(self) -> &'static [RebalanceProtocol] {
65        match self {
66            Assignor::Range | Assignor::RoundRobin | Assignor::Sticky => {
67                &[RebalanceProtocol::Eager]
68            }
69            Assignor::CooperativeSticky => {
70                &[RebalanceProtocol::Cooperative, RebalanceProtocol::Eager]
71            }
72        }
73    }
74
75    /// The assignor with `name` in `assignors`.
76    pub(crate) fn by_protocol_name(assignors: &[Self], name: &str) -> Option<Self> {
77        assignors
78            .iter()
79            .copied()
80            .find(|assignor| assignor.protocol_name() == name)
81    }
82}
83
84/// Validate an assignor list and return its rebalance protocol.
85///
86/// Kafka's `ConsumerConfig` rejects a duplicate strategy, and
87/// `ConsumerCoordinator` uses the newest rebalance protocol that every assignor
88/// supports. A subscribing consumer needs at least one assignor
89/// (`ConsumerCoordinator.poll`).
90pub(crate) fn rebalance_protocol_of(assignors: &[Assignor]) -> Result<RebalanceProtocol, String> {
91    for (index, assignor) in assignors.iter().enumerate() {
92        if assignors[..index].contains(assignor) {
93            return Err(format!(
94                "consumer assignors: duplicate assignor {}",
95                assignor.protocol_name()
96            ));
97        }
98    }
99    let first = assignors
100        .first()
101        .ok_or_else(|| "consumer assignors must not be empty".to_owned())?;
102    first
103        .supported_protocols()
104        .iter()
105        .copied()
106        .filter(|protocol| {
107            assignors
108                .iter()
109                .all(|assignor| assignor.supported_protocols().contains(protocol))
110        })
111        .max()
112        .ok_or_else(|| "consumer assignors have no common rebalance protocol".to_owned())
113}
114
115/// The identity by which Kafka's `AbstractPartitionAssignor.MemberInfo` sorts
116/// members: static members first by `group.instance.id`, then dynamic members
117/// by member id.
118#[derive(Debug, Clone, PartialEq, Eq)]
119pub struct MemberKey {
120    pub member_id: String,
121    pub group_instance_id: Option<String>,
122}
123
124impl Ord for MemberKey {
125    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
126        match (&self.group_instance_id, &other.group_instance_id) {
127            (Some(a), Some(b)) => a.cmp(b),
128            (Some(_), None) => std::cmp::Ordering::Less,
129            (None, Some(_)) => std::cmp::Ordering::Greater,
130            (None, None) => self.member_id.cmp(&other.member_id),
131        }
132    }
133}
134
135impl PartialOrd for MemberKey {
136    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
137        Some(self.cmp(other))
138    }
139}
140
141/// One member of a `JoinGroup` response, with its decoded subscription.
142pub(crate) struct GroupMember {
143    pub key: MemberKey,
144    pub subscription: DecodedSubscription,
145}
146
147/// Kafka's `AbstractStickyAssignor.DEFAULT_GENERATION`.
148const DEFAULT_GENERATION: i32 = -1;
149
150/// The `userData` of the subscription for `assignor`.
151///
152/// `sticky` sends its last assignment and generation (`StickyAssignor`
153/// user data v1), and nothing before its first assignment.
154/// `cooperative-sticky` sends its generation. The other assignors send
155/// nothing.
156pub(crate) fn subscription_user_data(
157    assignor: Assignor,
158    last_assignment: Option<&[(String, i32)]>,
159    generation: i32,
160) -> Option<Bytes> {
161    match assignor {
162        Assignor::Range | Assignor::RoundRobin => None,
163        Assignor::Sticky => last_assignment.map(|partitions| {
164            let mut by_topic: std::collections::BTreeMap<&str, Vec<i32>> =
165                std::collections::BTreeMap::new();
166            for (topic, partition) in partitions {
167                by_topic.entry(topic).or_default().push(*partition);
168            }
169            let mut buf = BytesMut::new();
170            put_array_len(&mut buf, by_topic.len());
171            for (topic, mut partitions) in by_topic {
172                partitions.sort_unstable();
173                buf.put_i16(i16::try_from(topic.len()).unwrap_or(i16::MAX));
174                buf.put_slice(topic.as_bytes());
175                put_array_len(&mut buf, partitions.len());
176                for partition in partitions {
177                    buf.put_i32(partition);
178                }
179            }
180            buf.put_i32(generation);
181            buf.freeze()
182        }),
183        Assignor::CooperativeSticky => {
184            let mut buf = BytesMut::with_capacity(4);
185            buf.put_i32(generation);
186            Some(buf.freeze())
187        }
188    }
189}
190
191fn put_array_len(buf: &mut BytesMut, len: usize) {
192    buf.put_i32(i32::try_from(len).unwrap_or(i32::MAX));
193}
194
195/// The owned partitions and generation in `sticky` user data.
196///
197/// Kafka's `StickyAssignor.deserializeTopicPartitionAssignment` reads v1, then
198/// v0 without a generation, and ignores user data that neither parses.
199fn decode_sticky_user_data(user_data: Option<&Bytes>) -> (Vec<(String, i32)>, i32) {
200    let Some(user_data) = user_data.filter(|data| !data.is_empty()) else {
201        return (Vec::new(), DEFAULT_GENERATION);
202    };
203    let mut buf = user_data.clone();
204    let Some(partitions) = read_topic_partitions(&mut buf) else {
205        return (Vec::new(), DEFAULT_GENERATION);
206    };
207    let generation = if buf.remaining() >= 4 {
208        buf.get_i32()
209    } else {
210        DEFAULT_GENERATION
211    };
212    (partitions, generation)
213}
214
215fn read_topic_partitions(buf: &mut Bytes) -> Option<Vec<(String, i32)>> {
216    let mut partitions = Vec::new();
217    if buf.remaining() < 4 {
218        return None;
219    }
220    let topics = usize::try_from(buf.get_i32()).ok()?;
221    for _ in 0..topics {
222        if buf.remaining() < 2 {
223            return None;
224        }
225        let len = usize::try_from(buf.get_i16()).ok()?;
226        if buf.remaining() < len + 4 {
227            return None;
228        }
229        let topic = String::from_utf8(buf.split_to(len).to_vec()).ok()?;
230        let count = usize::try_from(buf.get_i32()).ok()?;
231        if buf.remaining() < count.checked_mul(4)? {
232            return None;
233        }
234        for _ in 0..count {
235            partitions.push((topic.clone(), buf.get_i32()));
236        }
237    }
238    Some(partitions)
239}
240
241/// Run `assignor` over the members of a `JoinGroup` response, as the group
242/// leader does in Kafka's `ConsumerCoordinator.onLeaderElected`.
243pub(crate) fn assign(
244    assignor: Assignor,
245    members: &[GroupMember],
246    topic_partitions: &HashMap<String, i32>,
247) -> HashMap<String, Vec<(String, i32)>> {
248    match assignor {
249        Assignor::Range => range::assign_members(
250            members
251                .iter()
252                .map(|member| (member.key.clone(), member.subscription.topics.clone()))
253                .collect(),
254            topic_partitions,
255        ),
256        Assignor::RoundRobin => round_robin::assign(
257            members
258                .iter()
259                .map(|member| (member.key.clone(), member.subscription.topics.clone()))
260                .collect(),
261            topic_partitions,
262        ),
263        Assignor::Sticky => {
264            // Kafka's `StickyAssignor.memberData` reads the owned partitions
265            // and the generation from the user data only.
266            let inputs: Vec<cooperative_sticky::MemberInput> = members
267                .iter()
268                .map(|member| {
269                    let (owned, generation) =
270                        decode_sticky_user_data(member.subscription.user_data.as_ref());
271                    (
272                        member.key.member_id.clone(),
273                        member.subscription.topics.clone(),
274                        owned,
275                        generation,
276                    )
277                })
278                .collect();
279            cooperative_sticky::assign_eager(&inputs, topic_partitions)
280        }
281        Assignor::CooperativeSticky => {
282            let inputs: Vec<cooperative_sticky::MemberInput> = members
283                .iter()
284                .map(|member| {
285                    (
286                        member.key.member_id.clone(),
287                        member.subscription.topics.clone(),
288                        member.subscription.owned.clone(),
289                        member.subscription.generation_id,
290                    )
291                })
292                .collect();
293            cooperative_sticky::assign(&inputs, topic_partitions)
294        }
295    }
296}
297
298#[cfg(test)]
299mod tests {
300
301    use super::*;
302
303    #[test]
304    fn assignor_values_parse_kafka_names() {
305        let actual = [
306            "range",
307            "roundrobin",
308            "sticky",
309            "cooperative-sticky",
310            "cooperative_sticky",
311            "round-robin",
312            "unknown",
313        ]
314        .map(|name| name.parse::<Assignor>().ok());
315        assert2::assert!(
316            actual
317                == [
318                    Some(Assignor::Range),
319                    Some(Assignor::RoundRobin),
320                    Some(Assignor::Sticky),
321                    Some(Assignor::CooperativeSticky),
322                    None,
323                    None,
324                    None,
325                ]
326        );
327    }
328
329    #[test]
330    fn assignor_protocol_names_match_kafka_protocols() {
331        let actual = [
332            Assignor::Range,
333            Assignor::RoundRobin,
334            Assignor::Sticky,
335            Assignor::CooperativeSticky,
336        ]
337        .map(Assignor::protocol_name);
338        assert2::assert!(actual == ["range", "roundrobin", "sticky", "cooperative-sticky"]);
339    }
340
341    /// Kafka's `ConsumerCoordinator` constructor keeps the rebalance protocols
342    /// that every assignor supports and uses the newest one.
343    #[test]
344    fn rebalance_protocol_is_the_newest_common_protocol() {
345        use Assignor::{CooperativeSticky, Range, RoundRobin, Sticky};
346        let mut actual = Vec::new();
347        let mut wanted = Vec::new();
348        for (name, assignors, expected) in [
349            (
350                "default list",
351                &Assignor::DEFAULT_LIST[..],
352                Ok(RebalanceProtocol::Eager),
353            ),
354            (
355                "cooperative sticky only",
356                &[CooperativeSticky][..],
357                Ok(RebalanceProtocol::Cooperative),
358            ),
359            (
360                "round robin",
361                &[RoundRobin][..],
362                Ok(RebalanceProtocol::Eager),
363            ),
364            (
365                "sticky and cooperative sticky",
366                &[CooperativeSticky, Sticky][..],
367                Ok(RebalanceProtocol::Eager),
368            ),
369            (
370                "empty",
371                &[][..],
372                Err("consumer assignors must not be empty".to_owned()),
373            ),
374            (
375                "duplicate",
376                &[Range, Range][..],
377                Err("consumer assignors: duplicate assignor range".to_owned()),
378            ),
379        ] {
380            actual.push((name, rebalance_protocol_of(assignors)));
381            wanted.push((name, expected));
382        }
383        assert2::assert!(actual == wanted);
384    }
385
386    /// The byte layout of Kafka's `StickyAssignor` user data v1 and
387    /// `CooperativeStickyAssignor` user data v0.
388    #[test]
389    fn subscription_user_data_matches_kafka_layouts() {
390        let owned = vec![
391            ("b".to_owned(), 1),
392            ("a".to_owned(), 2),
393            ("b".to_owned(), 0),
394        ];
395        let actual = [
396            subscription_user_data(Assignor::Range, Some(&owned), 7),
397            subscription_user_data(Assignor::RoundRobin, Some(&owned), 7),
398            subscription_user_data(Assignor::Sticky, None, 7),
399            subscription_user_data(Assignor::Sticky, Some(&owned), 7),
400            subscription_user_data(Assignor::CooperativeSticky, Some(&owned), 7),
401        ];
402        let sticky = Bytes::from_static(&[
403            0, 0, 0, 2, // two topics
404            0, 1, b'a', 0, 0, 0, 1, 0, 0, 0, 2, // a: [2]
405            0, 1, b'b', 0, 0, 0, 2, 0, 0, 0, 0, 0, 0, 0, 1, // b: [0, 1]
406            0, 0, 0, 7, // generation
407        ]);
408        assert2::assert!(
409            actual
410                == [
411                    None,
412                    None,
413                    None,
414                    Some(sticky.clone()),
415                    Some(Bytes::from_static(&[0, 0, 0, 7])),
416                ]
417        );
418        let v0 = sticky.slice(..sticky.len() - 4);
419        let decoded = [
420            decode_sticky_user_data(Some(&sticky)),
421            decode_sticky_user_data(Some(&v0)),
422            decode_sticky_user_data(Some(&Bytes::from_static(&[0, 0]))),
423            decode_sticky_user_data(None),
424        ];
425        let sorted = vec![
426            ("a".to_owned(), 2),
427            ("b".to_owned(), 0),
428            ("b".to_owned(), 1),
429        ];
430        assert2::assert!(
431            decoded
432                == [
433                    (sorted.clone(), 7),
434                    (sorted, -1),
435                    (Vec::new(), -1),
436                    (Vec::new(), -1),
437                ]
438        );
439    }
440
441    /// The eager `sticky` assignor moves partitions to a new member at once. The
442    /// `cooperative-sticky` assignor holds a moving partition back for the
443    /// next round.
444    #[test]
445    fn sticky_moves_partitions_at_once_and_cooperative_sticky_holds_them_back() {
446        let counts = maplit::hashmap! {"t".to_owned() => 4};
447        let owned = vec![
448            ("t".to_owned(), 0),
449            ("t".to_owned(), 1),
450            ("t".to_owned(), 2),
451            ("t".to_owned(), 3),
452        ];
453        let member = |id: &str, user_data: Option<Bytes>, owned: Vec<(String, i32)>, generation| {
454            GroupMember {
455                key: MemberKey {
456                    member_id: id.to_owned(),
457                    group_instance_id: None,
458                },
459                subscription: DecodedSubscription {
460                    topics: vec!["t".to_owned()],
461                    owned,
462                    generation_id: generation,
463                    rack_id: None,
464                    user_data,
465                },
466            }
467        };
468        let sticky_members = [
469            member(
470                "a",
471                subscription_user_data(Assignor::Sticky, Some(&owned), 3),
472                Vec::new(),
473                3,
474            ),
475            member("b", None, Vec::new(), -1),
476        ];
477        let cooperative_members = [
478            member(
479                "a",
480                subscription_user_data(Assignor::CooperativeSticky, Some(&owned), 3),
481                owned.clone(),
482                3,
483            ),
484            member(
485                "b",
486                subscription_user_data(Assignor::CooperativeSticky, None, -1),
487                Vec::new(),
488                -1,
489            ),
490        ];
491        let actual = (
492            assign(Assignor::Sticky, &sticky_members, &counts),
493            assign(Assignor::CooperativeSticky, &cooperative_members, &counts),
494        );
495        assert2::assert!(
496            actual
497                == (
498                    maplit::hashmap! {"a".to_owned() => vec![("t".to_owned(), 0), ("t".to_owned(), 1)], "b".to_owned() => vec![("t".to_owned(), 2), ("t".to_owned(), 3)]},
499                    maplit::hashmap! {"a".to_owned() => vec![("t".to_owned(), 0), ("t".to_owned(), 1)], "b".to_owned() => Vec::new()},
500                )
501        );
502    }
503
504    /// Kafka's `MemberInfo.compareTo`: static members first, by instance id.
505    #[test]
506    fn range_orders_static_members_by_group_instance_id() {
507        let counts = maplit::hashmap! {"t".to_owned() => 3};
508        let member = |id: &str, instance: Option<&str>| GroupMember {
509            key: MemberKey {
510                member_id: id.to_owned(),
511                group_instance_id: instance.map(str::to_owned),
512            },
513            subscription: DecodedSubscription {
514                topics: vec!["t".to_owned()],
515                owned: Vec::new(),
516                generation_id: -1,
517                rack_id: None,
518                user_data: None,
519            },
520        };
521        let members = [member("a", None), member("z", Some("i-1"))];
522        assert2::assert!(
523            assign(Assignor::Range, &members, &counts)
524                == maplit::hashmap! {"z".to_owned() => vec![("t".to_owned(), 0), ("t".to_owned(), 1)], "a".to_owned() => vec![("t".to_owned(), 2)]}
525        );
526    }
527}
krabka_client_consumer/assignor/mod.rs

krabka_client_consumer/assignor/
mod.rs
