1use std::collections::{BTreeMap, BTreeSet};
23
24#[derive(Debug, Clone, PartialEq, Eq)]
26pub struct Subscription {
27 pub member_id: String,
28 pub topics: Vec<String>,
30 pub owned: Vec<TopicPartition>,
33 pub generation: i32,
42}
43
44#[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
60pub type Assignment = BTreeMap<String, Vec<TopicPartition>>;
66
67pub trait Assignor {
72 fn name(&self) -> &'static str;
74
75 fn assign(
81 &self,
82 members: &[Subscription],
83 partitions_per_topic: &BTreeMap<String, i32>,
84 ) -> Assignment;
85}
86
87#[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 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#[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 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 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
213fn 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#[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 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 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#[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 let target = StickyAssignor.assign(members, partitions_per_topic);
372
373 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 target
392 .into_iter()
393 .map(|(member_id, partitions)| {
394 let kept = partitions
395 .into_iter()
396 .filter(|tp| match owner.get(tp) {
397 Some(current) => *current == member_id,
400 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[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 #[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 #[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 #[test]
672 fn cooperative_withholds_a_partition_that_must_move() {
673 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 #[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 #[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 #[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 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 #[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}