Skip to main content

barnabas_core/
group.rs

1//! Consumer group assignment: who gets which partitions.
2//!
3//! Under the classic protocol the **group leader computes the assignment** —
4//! the coordinator only relays it — so this is client-side logic and belongs
5//! here, with no IO in sight. KIP-848 moves it to the coordinator, which is
6//! precisely why it lives behind the seam described in
7//! `docs/completing-the-client.md` rather than in the shared layer.
8//!
9//! # Matching the Java client matters here
10//!
11//! An assignor is not free to be clever. A group is usually mixed — a Rust
12//! consumer beside Java ones — and every member runs the *same* strategy on the
13//! *same* inputs and must reach the *same* answer, because only the leader's
14//! result is distributed and the others have to agree it is sane. Worse, a
15//! group whose members disagree about what `range` means will thrash: each
16//! rebalance elects a different leader and produces a different assignment.
17//!
18//! So these follow Java's algorithms, including the parts that look arbitrary
19//! — the sort orders especially, which is what makes the result deterministic
20//! across implementations.
21
22use std::collections::{BTreeMap, BTreeSet};
23
24/// One member's declared interest, as it arrives in `JoinGroup`.
25#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct Subscription {
27    pub member_id: String,
28    /// Topics this member wants. Order is not significant; assignors sort.
29    pub topics: Vec<String>,
30    /// What this member holds now, for assignors that try to keep it.
31    /// Empty for the stateless strategies.
32    pub owned: Vec<TopicPartition>,
33    /// The generation `owned` was true in, or -1 if the member did not say.
34    ///
35    /// **This is what makes an ownership claim trustworthy.** A member that was
36    /// away still advertises the partitions it used to hold, and a leader that
37    /// believes it will hand the same partition to two members at once. Kafka
38    /// added the field to `ConsumerProtocolSubscription` v2 for exactly this,
39    /// and the rule every client applies is the same: a claim from below the
40    /// highest generation in the group is stale and is ignored.
41    pub generation: i32,
42}
43
44/// A topic and one of its partitions.
45#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
46pub struct TopicPartition {
47    pub topic: String,
48    pub partition: i32,
49}
50
51impl TopicPartition {
52    pub fn new(topic: impl Into<String>, partition: i32) -> Self {
53        Self {
54            topic: topic.into(),
55            partition,
56        }
57    }
58}
59
60/// What the leader decided: member id → the partitions it owns.
61///
62/// A `BTreeMap` so the result is ordered and therefore diffable in a test
63/// failure, which matters more here than lookup speed — an assignment is
64/// computed once per rebalance.
65pub type Assignment = BTreeMap<String, Vec<TopicPartition>>;
66
67/// How partitions are shared out.
68///
69/// The name is what goes on the wire in `JoinGroup`; members negotiate a
70/// strategy they all support, so it must match Java's exactly.
71pub trait Assignor {
72    /// The protocol name Kafka knows this by.
73    fn name(&self) -> &'static str;
74
75    /// Divide `partitions_per_topic` among `members`.
76    ///
77    /// Every member appears in the result, with an empty vector if it got
78    /// nothing — a member that is absent from the assignment cannot tell
79    /// "assigned nothing" from "the leader forgot me".
80    fn assign(
81        &self,
82        members: &[Subscription],
83        partitions_per_topic: &BTreeMap<String, i32>,
84    ) -> Assignment;
85}
86
87/// Contiguous ranges per topic, Java's `RangeAssignor`.
88///
89/// For each topic the subscribed members are sorted by id and the partitions
90/// split into contiguous blocks. With 7 partitions and 3 members the split is
91/// 3/2/2, the remainder going to the earliest members.
92///
93/// **It is deliberately not balanced across topics.** With several topics of
94/// few partitions the first member accumulates the leftovers of each — the
95/// known wart of this strategy, and the reason `RoundRobin` exists. Reproduced
96/// rather than fixed, because a member that "fixes" it disagrees with the Java
97/// members in the same group.
98#[derive(Debug, Default, Clone, Copy)]
99pub struct RangeAssignor;
100
101impl Assignor for RangeAssignor {
102    fn name(&self) -> &'static str {
103        "range"
104    }
105
106    fn assign(
107        &self,
108        members: &[Subscription],
109        partitions_per_topic: &BTreeMap<String, i32>,
110    ) -> Assignment {
111        let mut assignment: Assignment = members
112            .iter()
113            .map(|m| (m.member_id.clone(), Vec::new()))
114            .collect();
115
116        for (topic, &count) in partitions_per_topic {
117            let mut subscribers: Vec<&str> = members
118                .iter()
119                .filter(|m| m.topics.iter().any(|t| t == topic))
120                .map(|m| m.member_id.as_str())
121                .collect();
122            if subscribers.is_empty() || count <= 0 {
123                continue;
124            }
125            subscribers.sort_unstable();
126
127            let members_count = i32::try_from(subscribers.len()).unwrap_or(i32::MAX);
128            let per_member = count / members_count;
129            let with_extra = count % members_count;
130
131            for (index, member) in subscribers.iter().enumerate() {
132                let index = i32::try_from(index).unwrap_or(i32::MAX);
133                // Java's arithmetic, kept in its shape so it is checkable
134                // against `RangeAssignor.java` line by line.
135                let start = per_member * index + index.min(with_extra);
136                let length = per_member + i32::from(index < with_extra);
137                let entry = assignment
138                    .get_mut(*member)
139                    .expect("every member is seeded above");
140                for partition in start..start + length {
141                    entry.push(TopicPartition::new(topic.clone(), partition));
142                }
143            }
144        }
145        assignment
146    }
147}
148
149/// One partition at a time around the members, Java's `RoundRobinAssignor`.
150///
151/// Every `(topic, partition)` in the group is laid out in order and dealt to
152/// the members in turn, skipping a member that did not subscribe to that topic.
153/// Balances across topics, which `Range` does not; in exchange it moves almost
154/// every partition when membership changes, which is what `sticky` addresses.
155#[derive(Debug, Default, Clone, Copy)]
156pub struct RoundRobinAssignor;
157
158impl Assignor for RoundRobinAssignor {
159    fn name(&self) -> &'static str {
160        "roundrobin"
161    }
162
163    fn assign(
164        &self,
165        members: &[Subscription],
166        partitions_per_topic: &BTreeMap<String, i32>,
167    ) -> Assignment {
168        let mut assignment: Assignment = members
169            .iter()
170            .map(|m| (m.member_id.clone(), Vec::new()))
171            .collect();
172
173        let mut sorted: Vec<&Subscription> = members.iter().collect();
174        sorted.sort_unstable_by(|a, b| a.member_id.cmp(&b.member_id));
175        if sorted.is_empty() {
176            return assignment;
177        }
178
179        // Every partition in the group, topic-major then partition — the order
180        // Java produces, and what makes two members agree.
181        let all: Vec<TopicPartition> = partitions_per_topic
182            .iter()
183            .flat_map(|(topic, &count)| {
184                (0..count.max(0)).map(move |p| TopicPartition::new(topic.clone(), p))
185            })
186            .collect();
187
188        let mut next = 0usize;
189        for tp in all {
190            // Advance to a member that wants this topic. If none does, the
191            // partition is unassigned, which is correct: nobody subscribed.
192            let mut looked_at = 0;
193            while looked_at < sorted.len()
194                && !sorted[next % sorted.len()].topics.contains(&tp.topic)
195            {
196                next += 1;
197                looked_at += 1;
198            }
199            if looked_at == sorted.len() {
200                continue;
201            }
202            let member = &sorted[next % sorted.len()].member_id;
203            assignment
204                .get_mut(member)
205                .expect("every member is seeded above")
206                .push(tp);
207            next += 1;
208        }
209        assignment
210    }
211}
212
213/// The ownership claims worth believing.
214///
215/// A member reporting a generation below the highest in the group has been away
216/// and its `owned` list describes a world that has moved on. Believing it is how
217/// a leader assigns one partition to two members.
218fn live_claims(members: &[Subscription]) -> BTreeMap<&str, &[TopicPartition]> {
219    let highest = members.iter().map(|m| m.generation).max().unwrap_or(-1);
220    members
221        .iter()
222        .map(|m| {
223            let owned: &[TopicPartition] = if m.generation >= highest && highest >= 0 {
224                &m.owned
225            } else {
226                &[]
227            };
228            (m.member_id.as_str(), owned)
229        })
230        .collect()
231}
232
233/// Keep what you had where possible, Java's `StickyAssignor` in spirit.
234///
235/// **Not Java's algorithm.** Java's does a balanced-then-shuffle pass with
236/// several refinement rounds; this keeps each member's still-valid partitions
237/// and deals the rest to whoever holds fewest. The results agree on the cases
238/// tested here and will not agree on every case.
239///
240/// That difference is only safe because **the leader's assignment is the one
241/// used** — a mixed group is consistent whoever leads, it just gets a slightly
242/// different balance depending. Where exact agreement is required, use
243/// `range` or `roundrobin`, which are reproduced exactly.
244#[derive(Debug, Default, Clone, Copy)]
245pub struct StickyAssignor;
246
247impl Assignor for StickyAssignor {
248    fn name(&self) -> &'static str {
249        "sticky"
250    }
251
252    fn assign(
253        &self,
254        members: &[Subscription],
255        partitions_per_topic: &BTreeMap<String, i32>,
256    ) -> Assignment {
257        let mut assignment: Assignment = members
258            .iter()
259            .map(|m| (m.member_id.clone(), Vec::new()))
260            .collect();
261
262        let mut sorted: Vec<&Subscription> = members.iter().collect();
263        sorted.sort_unstable_by(|a, b| a.member_id.cmp(&b.member_id));
264        if sorted.is_empty() {
265            return assignment;
266        }
267
268        let valid: BTreeSet<TopicPartition> = partitions_per_topic
269            .iter()
270            .flat_map(|(topic, &count)| {
271                (0..count.max(0)).map(move |p| TopicPartition::new(topic.clone(), p))
272            })
273            .collect();
274
275        // **Balance first, stickiness second.** Keeping every partition a member
276        // already holds is not sticky, it is inert: a member that joins an
277        // established group would never be given anything, because nothing is
278        // ever taken from anyone. So each member keeps at most its fair share,
279        // and the surplus is redealt.
280        let subscribers: Vec<&&Subscription> = sorted.iter().collect();
281        let quota_of = |member: &Subscription| -> usize {
282            let wanted: usize = valid
283                .iter()
284                .filter(|tp| member.topics.contains(&tp.topic))
285                .count();
286            let sharers = subscribers
287                .iter()
288                .filter(|m| m.topics.iter().any(|t| member.topics.contains(t)))
289                .count()
290                .max(1);
291            wanted.div_ceil(sharers)
292        };
293
294        let claims = live_claims(members);
295        let mut taken: BTreeSet<TopicPartition> = BTreeSet::new();
296        for member in &sorted {
297            let quota = quota_of(member);
298            for tp in claims[member.member_id.as_str()] {
299                if assignment[&member.member_id].len() >= quota {
300                    break;
301                }
302                if valid.contains(tp) && !taken.contains(tp) && member.topics.contains(&tp.topic) {
303                    taken.insert(tp.clone());
304                    assignment
305                        .get_mut(&member.member_id)
306                        .expect("seeded")
307                        .push(tp.clone());
308                }
309            }
310        }
311
312        // Whatever is left — never owned, owned by a member over quota, or
313        // owned by nobody subscribed — goes to whoever holds fewest.
314        let remaining: Vec<TopicPartition> = valid.difference(&taken).cloned().collect();
315        for tp in remaining {
316            let candidate = sorted
317                .iter()
318                .filter(|m| m.topics.contains(&tp.topic))
319                .min_by_key(|m| {
320                    (
321                        assignment.get(&m.member_id).map_or(0, Vec::len),
322                        m.member_id.clone(),
323                    )
324                });
325            if let Some(member) = candidate {
326                assignment
327                    .get_mut(&member.member_id)
328                    .expect("seeded")
329                    .push(tp.clone());
330            }
331        }
332
333        for partitions in assignment.values_mut() {
334            partitions.sort();
335        }
336        assignment
337    }
338}
339
340/// Sticky, but handing partitions over **one rebalance at a time**.
341///
342/// Java's `CooperativeStickyAssignor`, and the protocol change in KIP-429. The
343/// difference from every assignor above is not how the target is computed but
344/// what is *published*: a partition that must move from one member to another
345/// is given to **neither** in this round. Its current owner sees it missing
346/// from its assignment, revokes it, and rejoins; the next round hands it to its
347/// new owner.
348///
349/// That is what makes the rebalance incremental. Under the eager protocol every
350/// member gives up everything and stops consuming until the new assignment
351/// arrives — a stop-the-world pause proportional to the slowest member. Under
352/// this one a member keeps everything it is not losing and never stops reading
353/// it.
354///
355/// The cost is an extra rebalance round, which is why the assignment this
356/// returns is deliberately *incomplete* and must not be mistaken for a bug.
357#[derive(Debug, Default, Clone, Copy)]
358pub struct CooperativeStickyAssignor;
359
360impl Assignor for CooperativeStickyAssignor {
361    fn name(&self) -> &'static str {
362        "cooperative-sticky"
363    }
364
365    fn assign(
366        &self,
367        members: &[Subscription],
368        partitions_per_topic: &BTreeMap<String, i32>,
369    ) -> Assignment {
370        // The target: where each partition should end up.
371        let target = StickyAssignor.assign(members, partitions_per_topic);
372
373        // Who holds what right now — from the **raw** claims, deliberately.
374        //
375        // The generation filter belongs to computing the target, where an
376        // out-of-date claim would make a stale layout sticky. It must not be
377        // applied here: withholding is about who is *still reading* a
378        // partition, and a member whose claim is judged stale is reading it
379        // regardless. Filtering here hands its partitions to someone else while
380        // it still has them, which is the overlap this protocol exists to
381        // prevent. Java draws the same line — `memberData` filters by
382        // generation, `computePartitionsTransferringOwnership` does not.
383        let mut owner: BTreeMap<&TopicPartition, &str> = BTreeMap::new();
384        for member in members {
385            for tp in &member.owned {
386                owner.insert(tp, member.member_id.as_str());
387            }
388        }
389
390        // Publish only what is not being taken from someone else.
391        target
392            .into_iter()
393            .map(|(member_id, partitions)| {
394                let kept = partitions
395                    .into_iter()
396                    .filter(|tp| match owner.get(tp) {
397                        // Held by another member: it must revoke first, so this
398                        // round assigns it to nobody.
399                        Some(current) => *current == member_id,
400                        // Unowned — free to hand out now.
401                        None => true,
402                    })
403                    .collect();
404                (member_id, kept)
405            })
406            .collect()
407    }
408}
409
410#[cfg(test)]
411mod tests {
412    use super::*;
413
414    fn member(id: &str, topics: &[&str]) -> Subscription {
415        Subscription {
416            member_id: id.to_owned(),
417            topics: topics.iter().map(|t| (*t).to_owned()).collect(),
418            owned: Vec::new(),
419            generation: 1,
420        }
421    }
422
423    fn holding(id: &str, topics: &[&str], owned: &[(&str, i32)]) -> Subscription {
424        Subscription {
425            member_id: id.to_owned(),
426            topics: topics.iter().map(|t| (*t).to_owned()).collect(),
427            owned: owned
428                .iter()
429                .map(|(t, p)| TopicPartition::new(*t, *p))
430                .collect(),
431            generation: 1,
432        }
433    }
434
435    fn topics(entries: &[(&str, i32)]) -> BTreeMap<String, i32> {
436        entries.iter().map(|(t, c)| ((*t).to_owned(), *c)).collect()
437    }
438
439    fn partitions_of(assignment: &Assignment, member: &str) -> Vec<i32> {
440        assignment
441            .get(member)
442            .expect("member present")
443            .iter()
444            .map(|tp| tp.partition)
445            .collect()
446    }
447
448    /// **Parity with the Java client, checked against the Java client.**
449    ///
450    /// Every member of a group runs the same strategy and must reach the same
451    /// answer; a Rust member that disagrees makes the group thrash, electing a
452    /// different leader and producing a different assignment each rebalance.
453    /// So these expectations are not derived from reading `RangeAssignor.java`
454    /// — they are the output of running `kafka-clients` 3.9.0 itself:
455    ///
456    /// ```text
457    /// range       c0=[t:0,t:1]           c1=[t:2]
458    /// roundrobin  c0=[t:0,t:2]           c1=[t:1]
459    /// range       c0=[a:0,a:1,b:0,b:1]   c1=[a:2,b:2]
460    /// roundrobin  c0=[a:0,a:2,b:1]       c1=[a:1,b:0,b:2]
461    /// roundrobin  c0=[a:0]               c1=[a:1,b:0,b:1]
462    /// range       c0=[a:0,a:1,b:0,b:1]   c1=[a:2,a:3,b:2]   c2=[a:4]
463    /// roundrobin  c0=[a:0,a:3,b:0,b:2]   c1=[a:1,a:4,b:1]   c2=[a:2]
464    /// ```
465    ///
466    /// Reproduce with `scratchpad/oracle/Oracle.java` against a JDK image and
467    /// the `kafka-clients` jar from `apache/kafka:3.9.0`.
468    #[test]
469    fn range_and_roundrobin_match_the_java_client() {
470        fn flat(assignment: &Assignment, member: &str) -> Vec<String> {
471            assignment[member]
472                .iter()
473                .map(|tp| format!("{}:{}", tp.topic, tp.partition))
474                .collect()
475        }
476
477        let two = [member("c0", &["t"]), member("c1", &["t"])];
478        let a = RangeAssignor.assign(&two, &topics(&[("t", 3)]));
479        assert_eq!(flat(&a, "c0"), ["t:0", "t:1"]);
480        assert_eq!(flat(&a, "c1"), ["t:2"]);
481
482        let a = RoundRobinAssignor.assign(&two, &topics(&[("t", 3)]));
483        assert_eq!(flat(&a, "c0"), ["t:0", "t:2"]);
484        assert_eq!(flat(&a, "c1"), ["t:1"]);
485
486        let both = [member("c0", &["a", "b"]), member("c1", &["a", "b"])];
487        let spec = topics(&[("a", 3), ("b", 3)]);
488        let a = RangeAssignor.assign(&both, &spec);
489        assert_eq!(flat(&a, "c0"), ["a:0", "a:1", "b:0", "b:1"]);
490        assert_eq!(flat(&a, "c1"), ["a:2", "b:2"]);
491
492        let a = RoundRobinAssignor.assign(&both, &spec);
493        assert_eq!(flat(&a, "c0"), ["a:0", "a:2", "b:1"]);
494        assert_eq!(flat(&a, "c1"), ["a:1", "b:0", "b:2"]);
495
496        let uneven = [member("c0", &["a"]), member("c1", &["a", "b"])];
497        let a = RoundRobinAssignor.assign(&uneven, &topics(&[("a", 2), ("b", 2)]));
498        assert_eq!(flat(&a, "c0"), ["a:0"]);
499        assert_eq!(flat(&a, "c1"), ["a:1", "b:0", "b:1"]);
500
501        let three = [
502            member("c0", &["a", "b"]),
503            member("c1", &["a", "b"]),
504            member("c2", &["a"]),
505        ];
506        let spec = topics(&[("a", 5), ("b", 3)]);
507        let a = RangeAssignor.assign(&three, &spec);
508        assert_eq!(flat(&a, "c0"), ["a:0", "a:1", "b:0", "b:1"]);
509        assert_eq!(flat(&a, "c1"), ["a:2", "a:3", "b:2"]);
510        assert_eq!(flat(&a, "c2"), ["a:4"]);
511
512        let a = RoundRobinAssignor.assign(&three, &spec);
513        assert_eq!(flat(&a, "c0"), ["a:0", "a:3", "b:0", "b:2"]);
514        assert_eq!(flat(&a, "c1"), ["a:1", "a:4", "b:1"]);
515        assert_eq!(flat(&a, "c2"), ["a:2"]);
516    }
517
518    /// Java's documented example: 3 partitions, 2 members, the earlier member
519    /// takes the extra one.
520    #[test]
521    fn range_gives_the_remainder_to_the_earliest_members() {
522        let assignment = RangeAssignor.assign(
523            &[member("c0", &["t"]), member("c1", &["t"])],
524            &topics(&[("t", 3)]),
525        );
526        assert_eq!(partitions_of(&assignment, "c0"), vec![0, 1]);
527        assert_eq!(partitions_of(&assignment, "c1"), vec![2]);
528    }
529
530    #[test]
531    fn range_splits_evenly_when_it_divides() {
532        let assignment = RangeAssignor.assign(
533            &[member("c0", &["t"]), member("c1", &["t"])],
534            &topics(&[("t", 4)]),
535        );
536        assert_eq!(partitions_of(&assignment, "c0"), vec![0, 1]);
537        assert_eq!(partitions_of(&assignment, "c1"), vec![2, 3]);
538    }
539
540    /// **The wart, pinned deliberately.** Range works topic by topic, so with
541    /// two 3-partition topics `c0` takes the extra from *each*: 4 against 2.
542    /// A "fix" here would disagree with every Java member in the group.
543    #[test]
544    fn range_is_lopsided_across_topics_and_that_is_correct() {
545        let assignment = RangeAssignor.assign(
546            &[member("c0", &["a", "b"]), member("c1", &["a", "b"])],
547            &topics(&[("a", 3), ("b", 3)]),
548        );
549        assert_eq!(assignment["c0"].len(), 4);
550        assert_eq!(assignment["c1"].len(), 2);
551    }
552
553    /// The same input round-robin balances, which is the reason to choose it.
554    #[test]
555    fn roundrobin_balances_across_topics() {
556        let assignment = RoundRobinAssignor.assign(
557            &[member("c0", &["a", "b"]), member("c1", &["a", "b"])],
558            &topics(&[("a", 3), ("b", 3)]),
559        );
560        assert_eq!(assignment["c0"].len(), 3);
561        assert_eq!(assignment["c1"].len(), 3);
562    }
563
564    /// A member that did not subscribe is skipped rather than given the
565    /// partition, and the deal continues from the next member.
566    #[test]
567    fn roundrobin_skips_members_that_did_not_subscribe() {
568        let assignment = RoundRobinAssignor.assign(
569            &[member("c0", &["a"]), member("c1", &["a", "b"])],
570            &topics(&[("a", 2), ("b", 2)]),
571        );
572        assert!(assignment["c0"].iter().all(|tp| tp.topic == "a"));
573        assert_eq!(
574            assignment["c1"].iter().filter(|tp| tp.topic == "b").count(),
575            2
576        );
577    }
578
579    /// Nobody subscribed, so the partition goes unassigned rather than to an
580    /// uninterested member.
581    #[test]
582    fn a_topic_nobody_wants_is_left_alone() {
583        let assignment =
584            RoundRobinAssignor.assign(&[member("c0", &["a"])], &topics(&[("a", 1), ("z", 4)]));
585        assert_eq!(assignment["c0"].len(), 1);
586        assert!(assignment["c0"].iter().all(|tp| tp.topic == "a"));
587    }
588
589    /// Every member appears, even with nothing — "assigned nothing" must be
590    /// distinguishable from "the leader forgot me".
591    #[test]
592    fn every_member_appears_in_the_result() {
593        let assignment = RangeAssignor.assign(
594            &[member("c0", &["a"]), member("idle", &["nonexistent"])],
595            &topics(&[("a", 1)]),
596        );
597        assert!(assignment.contains_key("idle"));
598        assert!(assignment["idle"].is_empty());
599    }
600
601    /// The point of sticky: a member that keeps its subscription keeps its
602    /// partitions when someone else leaves.
603    #[test]
604    fn sticky_keeps_what_is_still_valid() {
605        let assignment = StickyAssignor.assign(
606            &[
607                holding("c0", &["t"], &[("t", 0), ("t", 1)]),
608                holding("c1", &["t"], &[("t", 2)]),
609            ],
610            &topics(&[("t", 4)]),
611        );
612        // 0 and 1 stay put; 3 is new and goes to whoever holds fewest.
613        assert!(assignment["c0"].contains(&TopicPartition::new("t", 0)));
614        assert!(assignment["c0"].contains(&TopicPartition::new("t", 1)));
615        assert!(assignment["c1"].contains(&TopicPartition::new("t", 2)));
616        assert!(assignment["c1"].contains(&TopicPartition::new("t", 3)));
617    }
618
619    /// **Sticky must balance, not merely preserve.** A member joining an
620    /// established group has to be given something, which means taking it from
621    /// someone. An earlier version kept every owned partition and only dealt
622    /// out unowned ones, so a new member sat idle forever and
623    /// `cooperative-sticky` had nothing to hand over.
624    #[test]
625    fn sticky_takes_from_the_over_provisioned_to_feed_a_new_member() {
626        let assignment = StickyAssignor.assign(
627            &[
628                holding("c0", &["t"], &[("t", 0), ("t", 1), ("t", 2), ("t", 3)]),
629                holding("c1", &["t"], &[]),
630            ],
631            &topics(&[("t", 4)]),
632        );
633        assert_eq!(assignment["c0"].len(), 2, "the incumbent gives up half");
634        assert_eq!(assignment["c1"].len(), 2, "the newcomer is fed");
635    }
636
637    /// A partition two members both claim to own is given to exactly one.
638    /// After a rebalance both may believe they hold it, and handing it to both
639    /// is duplicate consumption.
640    #[test]
641    fn sticky_never_assigns_a_partition_twice() {
642        let assignment = StickyAssignor.assign(
643            &[
644                holding("c0", &["t"], &[("t", 0)]),
645                holding("c1", &["t"], &[("t", 0)]),
646            ],
647            &topics(&[("t", 2)]),
648        );
649        let mut all: Vec<&TopicPartition> = assignment.values().flatten().collect();
650        let before = all.len();
651        all.sort();
652        all.dedup();
653        assert_eq!(all.len(), before, "a partition was assigned twice");
654        assert_eq!(before, 2, "both partitions must be assigned");
655    }
656
657    /// A partition the member owns but that no longer exists is dropped rather
658    /// than carried forward — topics shrink when they are recreated smaller.
659    #[test]
660    fn sticky_drops_partitions_that_no_longer_exist() {
661        let assignment = StickyAssignor.assign(
662            &[holding("c0", &["t"], &[("t", 0), ("t", 99)])],
663            &topics(&[("t", 1)]),
664        );
665        assert_eq!(assignment["c0"], vec![TopicPartition::new("t", 0)]);
666    }
667
668    /// **A partition moving between members is assigned to neither**, this
669    /// round. Its owner will revoke it and the next round hands it over — which
670    /// is the whole of KIP-429.
671    #[test]
672    fn cooperative_withholds_a_partition_that_must_move() {
673        // c0 holds both; c1 is new, so one has to move.
674        let members = [
675            holding("c0", &["t"], &[("t", 0), ("t", 1)]),
676            holding("c1", &["t"], &[]),
677        ];
678        let assignment = CooperativeStickyAssignor.assign(&members, &topics(&[("t", 2)]));
679
680        let total: usize = assignment.values().map(Vec::len).sum();
681        assert_eq!(
682            total, 1,
683            "the moving partition must be withheld: {assignment:?}"
684        );
685        assert_eq!(
686            assignment["c0"].len(),
687            1,
688            "c0 keeps the one it is not losing"
689        );
690        assert!(
691            assignment["c1"].is_empty(),
692            "c1 waits a round for its share"
693        );
694    }
695
696    /// Nothing is moving, so nothing is withheld and the round is complete.
697    #[test]
698    fn cooperative_publishes_everything_when_nothing_moves() {
699        let members = [holding("c0", &["t"], &[("t", 0), ("t", 1)])];
700        let assignment = CooperativeStickyAssignor.assign(&members, &topics(&[("t", 2)]));
701        assert_eq!(assignment["c0"].len(), 2);
702    }
703
704    /// An unowned partition needs no handover, so it is granted immediately.
705    #[test]
706    fn cooperative_grants_unowned_partitions_at_once() {
707        let members = [holding("c0", &["t"], &[("t", 0)])];
708        let assignment = CooperativeStickyAssignor.assign(&members, &topics(&[("t", 2)]));
709        assert_eq!(
710            assignment["c0"].len(),
711            2,
712            "the new partition needs no handover"
713        );
714    }
715
716    /// **A claim from an older generation is ignored.**
717    ///
718    /// A member that was away still advertises what it used to hold. Believing
719    /// it is how one partition ends up assigned to two members — silently,
720    /// since Kafka relays whatever the leader computed without checking.
721    #[test]
722    fn a_stale_ownership_claim_is_not_believed() {
723        let stale = Subscription {
724            member_id: "old".to_owned(),
725            topics: vec!["t".to_owned()],
726            owned: vec![TopicPartition::new("t", 0), TopicPartition::new("t", 1)],
727            generation: 1,
728        };
729        let current = Subscription {
730            member_id: "new".to_owned(),
731            topics: vec!["t".to_owned()],
732            owned: vec![],
733            generation: 5,
734        };
735
736        let assignment = CooperativeStickyAssignor.assign(&[stale, current], &topics(&[("t", 2)]));
737        // The stale claim is ignored for *stickiness* — the target does not
738        // preserve it — but the partitions are still withheld, because that
739        // member is still reading them until it says otherwise.
740        let total: usize = assignment.values().map(Vec::len).sum();
741        assert!(
742            total <= 2,
743            "no partition may be handed out twice: {assignment:?}"
744        );
745        let mut all: Vec<&TopicPartition> = assignment.values().flatten().collect();
746        let before = all.len();
747        all.sort();
748        all.dedup();
749        assert_eq!(all.len(), before, "a partition was assigned twice");
750    }
751
752    /// Every strategy must place every partition of a subscribed topic exactly
753    /// once. A gap is a partition nobody reads; an overlap is duplicate
754    /// consumption. Both are silent.
755    #[test]
756    fn every_partition_is_assigned_exactly_once() {
757        let members = [
758            member("c0", &["a", "b"]),
759            member("c1", &["a", "b"]),
760            member("c2", &["a"]),
761        ];
762        let spec = topics(&[("a", 5), ("b", 3)]);
763
764        for (name, assignment) in [
765            ("range", RangeAssignor.assign(&members, &spec)),
766            ("roundrobin", RoundRobinAssignor.assign(&members, &spec)),
767            ("sticky", StickyAssignor.assign(&members, &spec)),
768        ] {
769            let mut all: Vec<TopicPartition> = assignment.values().flatten().cloned().collect();
770            let count = all.len();
771            all.sort();
772            all.dedup();
773            assert_eq!(all.len(), count, "{name} assigned a partition twice");
774            assert_eq!(all.len(), 8, "{name} did not assign every partition");
775        }
776    }
777}