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

krabka_client_consumer/assignor/
mod.rs
