1use crate::group::{Assignment, Subscription, TopicPartition};
28
29pub mod codes {
31 pub const NONE: i16 = 0;
32 pub const COORDINATOR_LOAD_IN_PROGRESS: i16 = 14;
33 pub const COORDINATOR_NOT_AVAILABLE: i16 = 15;
34 pub const NOT_COORDINATOR: i16 = 16;
35 pub const ILLEGAL_GENERATION: i16 = 22;
36 pub const UNKNOWN_MEMBER_ID: i16 = 25;
37 pub const REBALANCE_IN_PROGRESS: i16 = 27;
38 pub const FENCED_INSTANCE_ID: i16 = 82;
39 pub const MEMBER_ID_REQUIRED: i16 = 79;
40}
41
42pub const NO_GENERATION: i32 = -1;
44
45#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
47pub enum RebalanceProtocol {
48 #[default]
53 Eager,
54 Cooperative,
61}
62
63#[derive(Debug, Clone, Copy, PartialEq, Eq)]
64pub enum MemberState {
65 Unjoined,
67 Joining,
69 Syncing,
71 Stable,
74}
75
76#[derive(Debug, Clone, PartialEq, Eq)]
78pub enum Step {
79 Join { member_id: String },
81 AssignAndSync { members: Vec<Subscription> },
84 Sync,
86 Heartbeat,
88 FindCoordinator,
90}
91
92#[derive(Debug, Clone)]
94pub struct GroupMember {
95 group_id: String,
96 member_id: String,
97 generation: i32,
98 state: MemberState,
99 topics: Vec<String>,
100 assignment: Vec<TopicPartition>,
101 leader: bool,
102 protocol: RebalanceProtocol,
103 lost: Vec<TopicPartition>,
106 pending_members: Vec<Subscription>,
115}
116
117impl GroupMember {
118 #[must_use]
119 pub fn new(group_id: impl Into<String>, topics: Vec<String>) -> Self {
120 Self {
121 group_id: group_id.into(),
122 member_id: String::new(),
123 generation: NO_GENERATION,
124 state: MemberState::Unjoined,
125 topics,
126 assignment: Vec::new(),
127 leader: false,
128 protocol: RebalanceProtocol::Eager,
129 lost: Vec::new(),
130 pending_members: Vec::new(),
131 }
132 }
133
134 #[must_use]
136 pub fn with_protocol(mut self, protocol: RebalanceProtocol) -> Self {
137 self.protocol = protocol;
138 self
139 }
140
141 #[must_use]
142 pub fn protocol(&self) -> RebalanceProtocol {
143 self.protocol
144 }
145
146 #[must_use]
150 pub fn lost(&self) -> &[TopicPartition] {
151 &self.lost
152 }
153
154 #[must_use]
155 pub fn group_id(&self) -> &str {
156 &self.group_id
157 }
158
159 #[must_use]
161 pub fn topics(&self) -> &[String] {
162 &self.topics
163 }
164
165 #[must_use]
166 pub fn state(&self) -> MemberState {
167 self.state
168 }
169
170 #[must_use]
171 pub fn generation(&self) -> i32 {
172 self.generation
173 }
174
175 #[must_use]
176 pub fn member_id(&self) -> &str {
177 &self.member_id
178 }
179
180 #[must_use]
181 pub fn is_leader(&self) -> bool {
182 self.leader
183 }
184
185 #[must_use]
187 pub fn assignment(&self) -> &[TopicPartition] {
188 &self.assignment
189 }
190
191 #[must_use]
199 pub fn can_commit(&self) -> bool {
200 self.state == MemberState::Stable && self.generation != NO_GENERATION
201 }
202
203 #[must_use]
205 pub fn step(&self) -> Step {
206 match self.state {
207 MemberState::Unjoined | MemberState::Joining => Step::Join {
208 member_id: self.member_id.clone(),
209 },
210 MemberState::Syncing if self.leader => Step::AssignAndSync {
211 members: self.pending_members.clone(),
212 },
213 MemberState::Syncing => Step::Sync,
214 MemberState::Stable => Step::Heartbeat,
215 }
216 }
217
218 #[must_use]
220 pub fn subscription(&self) -> Subscription {
221 Subscription {
222 member_id: self.member_id.clone(),
223 topics: self.topics.clone(),
224 owned: self.assignment.clone(),
228 generation: self.generation,
229 }
230 }
231
232 pub fn set_topics(&mut self, topics: Vec<String>) {
235 if topics != self.topics {
236 self.topics = topics;
237 self.revoke_and_rejoin();
238 }
239 }
240
241 pub fn request_rejoin(&mut self) {
251 if self.state != MemberState::Unjoined {
252 self.revoke_and_rejoin();
253 }
254 }
255
256 pub fn on_join(
260 &mut self,
261 error: i16,
262 generation: i32,
263 member_id: &str,
264 leader_id: &str,
265 members: Vec<Subscription>,
266 ) -> Step {
267 match error {
268 codes::NONE => {
269 self.member_id = member_id.to_owned();
270 self.generation = generation;
271 self.leader = leader_id == member_id;
272 self.state = MemberState::Syncing;
273 if self.leader {
274 self.pending_members = members.clone();
275 Step::AssignAndSync { members }
276 } else {
277 self.pending_members.clear();
278 Step::Sync
279 }
280 }
281 codes::MEMBER_ID_REQUIRED => {
284 self.member_id = member_id.to_owned();
285 self.state = MemberState::Unjoined;
286 Step::Join {
287 member_id: self.member_id.clone(),
288 }
289 }
290 codes::UNKNOWN_MEMBER_ID | codes::FENCED_INSTANCE_ID => {
291 self.member_id = String::new();
293 self.revoke_and_rejoin();
294 Step::Join {
295 member_id: String::new(),
296 }
297 }
298 codes::COORDINATOR_NOT_AVAILABLE
299 | codes::NOT_COORDINATOR
300 | codes::COORDINATOR_LOAD_IN_PROGRESS => {
301 self.state = MemberState::Unjoined;
302 Step::FindCoordinator
303 }
304 _ => {
305 self.revoke_and_rejoin();
306 Step::Join {
307 member_id: self.member_id.clone(),
308 }
309 }
310 }
311 }
312
313 pub fn on_sync(&mut self, error: i16, assigned: Vec<TopicPartition>) -> Step {
315 match error {
316 codes::NONE => {
317 self.lost.clear();
318 if self.protocol == RebalanceProtocol::Cooperative {
319 let mut assigned_sorted = assigned.clone();
324 assigned_sorted.sort();
325 self.lost = self
326 .assignment
327 .iter()
328 .filter(|tp| !assigned_sorted.contains(tp))
329 .cloned()
330 .collect();
331 }
332 self.assignment = assigned;
333 self.assignment.sort();
334 self.state = MemberState::Stable;
335
336 if !self.lost.is_empty() {
337 self.state = MemberState::Unjoined;
348 self.leader = false;
349 self.pending_members.clear();
350 return Step::Join {
351 member_id: self.member_id.clone(),
352 };
353 }
354 Step::Heartbeat
355 }
356 codes::REBALANCE_IN_PROGRESS | codes::ILLEGAL_GENERATION => {
358 self.revoke_and_rejoin();
359 Step::Join {
360 member_id: self.member_id.clone(),
361 }
362 }
363 codes::UNKNOWN_MEMBER_ID | codes::FENCED_INSTANCE_ID => {
364 self.member_id = String::new();
365 self.revoke_and_rejoin();
366 Step::Join {
367 member_id: String::new(),
368 }
369 }
370 codes::COORDINATOR_NOT_AVAILABLE | codes::NOT_COORDINATOR => {
371 self.revoke_and_rejoin();
372 Step::FindCoordinator
373 }
374 _ => {
375 self.revoke_and_rejoin();
376 Step::Join {
377 member_id: self.member_id.clone(),
378 }
379 }
380 }
381 }
382
383 pub fn on_heartbeat(&mut self, error: i16) -> Step {
385 match error {
386 codes::NONE => Step::Heartbeat,
387 codes::REBALANCE_IN_PROGRESS | codes::ILLEGAL_GENERATION => {
390 self.revoke_and_rejoin();
391 Step::Join {
392 member_id: self.member_id.clone(),
393 }
394 }
395 codes::UNKNOWN_MEMBER_ID | codes::FENCED_INSTANCE_ID => {
396 self.member_id = String::new();
397 self.revoke_and_rejoin();
398 Step::Join {
399 member_id: String::new(),
400 }
401 }
402 codes::COORDINATOR_NOT_AVAILABLE | codes::NOT_COORDINATOR => {
403 self.revoke_and_rejoin();
404 Step::FindCoordinator
405 }
406 _ => {
407 self.revoke_and_rejoin();
408 Step::Join {
409 member_id: self.member_id.clone(),
410 }
411 }
412 }
413 }
414
415 pub fn on_leave(&mut self) {
418 self.member_id = String::new();
419 self.revoke_and_rejoin();
420 }
421
422 #[must_use]
424 pub fn my_share(&self, assignment: &Assignment) -> Vec<TopicPartition> {
425 assignment.get(&self.member_id).cloned().unwrap_or_default()
426 }
427
428 fn revoke_and_rejoin(&mut self) {
434 if self.protocol == RebalanceProtocol::Eager {
438 self.lost = std::mem::take(&mut self.assignment);
443 } else {
444 self.lost.clear();
445 }
446 self.pending_members.clear();
447 self.leader = false;
455 self.state = MemberState::Unjoined;
456 }
457}
458
459#[cfg(test)]
460mod tests {
461 use super::*;
462
463 fn member() -> GroupMember {
464 GroupMember::new("g", vec!["t".to_owned()])
465 }
466
467 fn tp(partition: i32) -> TopicPartition {
468 TopicPartition::new("t", partition)
469 }
470
471 #[test]
473 fn a_first_join_is_refused_and_then_accepted() {
474 let mut m = member();
475 assert_eq!(
476 m.step(),
477 Step::Join {
478 member_id: String::new()
479 }
480 );
481
482 let step = m.on_join(codes::MEMBER_ID_REQUIRED, NO_GENERATION, "m-1", "", vec![]);
484 assert_eq!(
485 step,
486 Step::Join {
487 member_id: "m-1".to_owned()
488 }
489 );
490 assert_eq!(m.member_id(), "m-1");
491 assert_eq!(m.state(), MemberState::Unjoined);
492
493 let step = m.on_join(codes::NONE, 7, "m-1", "m-2", vec![]);
494 assert_eq!(step, Step::Sync, "a follower syncs with no assignment");
495 assert_eq!(m.state(), MemberState::Syncing);
496 assert!(!m.is_leader());
497
498 assert_eq!(m.on_sync(codes::NONE, vec![tp(0), tp(1)]), Step::Heartbeat);
499 assert_eq!(m.state(), MemberState::Stable);
500 assert_eq!(m.assignment(), &[tp(0), tp(1)]);
501 assert_eq!(m.generation(), 7);
502 }
503
504 #[test]
506 fn the_leader_is_asked_to_assign() {
507 let mut m = member();
508 let members = vec![Subscription {
509 member_id: "m-1".to_owned(),
510 topics: vec!["t".to_owned()],
511 owned: vec![],
512 generation: 1,
513 }];
514 let step = m.on_join(codes::NONE, 3, "m-1", "m-1", members.clone());
515 assert_eq!(step, Step::AssignAndSync { members });
516 assert!(m.is_leader());
517 }
518
519 #[test]
526 fn the_leader_is_still_the_leader_on_the_next_step() {
527 let mut m = member();
528 let members = vec![Subscription {
529 member_id: "m-1".to_owned(),
530 topics: vec!["t".to_owned()],
531 owned: vec![],
532 generation: 1,
533 }];
534 m.on_join(codes::NONE, 3, "m-1", "m-1", members.clone());
535 assert_eq!(m.step(), Step::AssignAndSync { members });
536 }
537
538 #[test]
540 fn a_follower_is_still_a_follower_on_the_next_step() {
541 let mut m = member();
542 m.on_join(codes::NONE, 3, "m-1", "m-2", vec![]);
543 assert_eq!(m.step(), Step::Sync);
544 }
545
546 #[test]
550 fn commits_are_refused_outside_stable() {
551 let mut m = member();
552 assert!(!m.can_commit(), "not a member yet");
553
554 m.on_join(codes::NONE, 1, "m-1", "m-1", vec![]);
555 assert!(!m.can_commit(), "syncing is not stable");
556
557 m.on_sync(codes::NONE, vec![tp(0)]);
558 assert!(m.can_commit());
559
560 m.on_heartbeat(codes::REBALANCE_IN_PROGRESS);
561 assert!(!m.can_commit(), "a rebalance suspends commits");
562 }
563
564 #[test]
567 fn a_rebalance_revokes_the_assignment_immediately() {
568 let mut m = member();
569 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
570 m.on_sync(codes::NONE, vec![tp(0), tp(1)]);
571 assert_eq!(m.assignment().len(), 2);
572
573 let step = m.on_heartbeat(codes::REBALANCE_IN_PROGRESS);
574 assert_eq!(
575 step,
576 Step::Join {
577 member_id: "m-1".to_owned()
578 }
579 );
580 assert!(m.assignment().is_empty(), "partitions must be given up");
581 assert_eq!(m.state(), MemberState::Unjoined);
582 assert_eq!(m.generation(), 1);
587 }
588
589 #[test]
591 fn an_illegal_generation_revokes_too() {
592 let mut m = member();
593 m.on_join(codes::NONE, 4, "m-1", "m-2", vec![]);
594 m.on_sync(codes::NONE, vec![tp(0)]);
595
596 m.on_heartbeat(codes::ILLEGAL_GENERATION);
597 assert!(m.assignment().is_empty());
598 assert!(!m.can_commit());
599 assert_eq!(
600 m.member_id(),
601 "m-1",
602 "the member id survives a generation bump"
603 );
604 }
605
606 #[test]
609 fn an_unknown_member_id_forgets_the_identity() {
610 let mut m = member();
611 m.on_join(codes::NONE, 4, "m-1", "m-2", vec![]);
612 m.on_sync(codes::NONE, vec![tp(0)]);
613
614 let step = m.on_heartbeat(codes::UNKNOWN_MEMBER_ID);
615 assert_eq!(
616 step,
617 Step::Join {
618 member_id: String::new()
619 }
620 );
621 assert_eq!(m.member_id(), "", "the id is no longer ours to use");
622 assert!(m.assignment().is_empty());
623 }
624
625 #[test]
627 fn a_lost_coordinator_is_rediscovered() {
628 let mut m = member();
629 m.on_join(codes::NONE, 2, "m-1", "m-2", vec![]);
630 m.on_sync(codes::NONE, vec![tp(0)]);
631
632 assert_eq!(
633 m.on_heartbeat(codes::NOT_COORDINATOR),
634 Step::FindCoordinator
635 );
636 assert!(
637 m.assignment().is_empty(),
638 "still revoked: we cannot heartbeat"
639 );
640 }
641
642 #[test]
644 fn a_rebalance_during_sync_restarts_the_join() {
645 let mut m = member();
646 m.on_join(codes::NONE, 5, "m-1", "m-1", vec![]);
647 assert_eq!(m.state(), MemberState::Syncing);
648
649 let step = m.on_sync(codes::REBALANCE_IN_PROGRESS, vec![]);
650 assert_eq!(
651 step,
652 Step::Join {
653 member_id: "m-1".to_owned()
654 }
655 );
656 assert!(
657 !m.is_leader(),
658 "leadership is not carried across a rebalance"
659 );
660 }
661
662 #[test]
665 fn changing_topics_forces_a_rejoin() {
666 let mut m = member();
667 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
668 m.on_sync(codes::NONE, vec![tp(0)]);
669
670 m.set_topics(vec!["t".to_owned(), "u".to_owned()]);
671 assert_eq!(m.state(), MemberState::Unjoined);
672 assert!(m.assignment().is_empty());
673
674 m.on_join(codes::NONE, 2, "m-1", "m-2", vec![]);
676 m.on_sync(codes::NONE, vec![tp(0)]);
677 m.set_topics(vec!["t".to_owned(), "u".to_owned()]);
678 assert_eq!(m.state(), MemberState::Stable);
679 }
680
681 #[test]
684 fn the_subscription_carries_what_is_owned() {
685 let mut m = member();
686 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
687 m.on_sync(codes::NONE, vec![tp(3)]);
688
689 let s = m.subscription();
690 assert_eq!(s.member_id, "m-1");
691 assert_eq!(s.topics, vec!["t".to_owned()]);
692 assert_eq!(s.owned, vec![tp(3)]);
693 }
694
695 #[test]
698 fn leaving_gives_everything_up() {
699 let mut m = member();
700 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
701 m.on_sync(codes::NONE, vec![tp(0)]);
702
703 m.on_leave();
704 assert_eq!(m.member_id(), "");
705 assert!(m.assignment().is_empty());
706 assert!(!m.can_commit());
707 }
708
709 fn cooperative() -> GroupMember {
710 GroupMember::new("g", vec!["t".to_owned()]).with_protocol(RebalanceProtocol::Cooperative)
711 }
712
713 #[test]
717 fn a_cooperative_rebalance_keeps_the_partitions() {
718 let mut m = cooperative();
719 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
720 m.on_sync(codes::NONE, vec![tp(0), tp(1)]);
721
722 m.on_heartbeat(codes::REBALANCE_IN_PROGRESS);
723 assert_eq!(
724 m.assignment(),
725 &[tp(0), tp(1)],
726 "cooperative members keep what nobody has taken yet"
727 );
728 assert!(
729 !m.can_commit(),
730 "but they are not stable, so they do not commit"
731 );
732 }
733
734 #[test]
736 fn an_eager_rebalance_gives_everything_up() {
737 let mut m = member();
738 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
739 m.on_sync(codes::NONE, vec![tp(0), tp(1)]);
740
741 m.on_heartbeat(codes::REBALANCE_IN_PROGRESS);
742 assert!(m.assignment().is_empty());
743 }
744
745 #[test]
749 fn a_smaller_assignment_is_a_handover() {
750 let mut m = cooperative();
751 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
752 m.on_sync(codes::NONE, vec![tp(0), tp(1), tp(2)]);
753 assert!(m.lost().is_empty());
754
755 m.on_heartbeat(codes::REBALANCE_IN_PROGRESS);
756 m.on_join(codes::NONE, 2, "m-1", "m-2", vec![]);
757 let step = m.on_sync(codes::NONE, vec![tp(0)]);
758
759 assert_eq!(m.lost(), &[tp(1), tp(2)], "the two that moved");
760 assert_eq!(m.assignment(), &[tp(0)], "and the one that did not");
761 assert_eq!(
762 step,
763 Step::Join {
764 member_id: "m-1".to_owned()
765 },
766 "rejoin at once: the released partitions have no owner until we do"
767 );
768 }
769
770 #[test]
773 fn a_handover_rejoin_keeps_its_generation() {
774 let mut m = cooperative();
775 m.on_join(codes::NONE, 4, "m-1", "m-2", vec![]);
776 m.on_sync(codes::NONE, vec![tp(0), tp(1)]);
777 m.on_heartbeat(codes::REBALANCE_IN_PROGRESS);
778 m.on_join(codes::NONE, 5, "m-1", "m-2", vec![]);
779 m.on_sync(codes::NONE, vec![tp(0)]);
780
781 assert_eq!(
782 m.generation(),
783 5,
784 "still a member of the generation just assigned"
785 );
786 assert_eq!(
787 m.subscription().generation,
788 5,
789 "and it says so, so its claim on tp(0) is believed"
790 );
791 }
792
793 #[test]
795 fn an_unchanged_assignment_does_not_rejoin() {
796 let mut m = cooperative();
797 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
798 assert_eq!(m.on_sync(codes::NONE, vec![tp(0)]), Step::Heartbeat);
799 assert!(m.lost().is_empty());
800 assert_eq!(m.state(), MemberState::Stable);
801 }
802
803 #[test]
807 fn outside_stable_it_owns_nothing_and_commits_nothing() {
808 let errors = [
809 codes::REBALANCE_IN_PROGRESS,
810 codes::ILLEGAL_GENERATION,
811 codes::UNKNOWN_MEMBER_ID,
812 codes::FENCED_INSTANCE_ID,
813 codes::NOT_COORDINATOR,
814 codes::COORDINATOR_NOT_AVAILABLE,
815 9_999, ];
817
818 for error in errors {
819 let mut m = member();
820 let _ = RebalanceProtocol::default();
821 m.on_join(codes::NONE, 1, "m-1", "m-2", vec![]);
822 m.on_sync(codes::NONE, vec![tp(0), tp(1)]);
823 assert!(m.can_commit());
824
825 m.on_heartbeat(error);
826 assert_ne!(m.state(), MemberState::Stable, "error {error}");
827 assert!(m.assignment().is_empty(), "error {error} kept partitions");
828 assert!(!m.can_commit(), "error {error} still allowed a commit");
829 }
830 }
831}