1use std::{
12 collections::{BTreeMap, HashMap, HashSet},
13 sync::atomic,
14 time::Duration,
15};
16
17use krabka_protocol::{
18 owned::{
19 consumer_group_heartbeat_request,
20 consumer_group_heartbeat_request::{ConsumerGroupHeartbeatRequest, TopicPartitions},
21 consumer_group_heartbeat_response::ConsumerGroupHeartbeatResponse,
22 metadata_request::{MetadataRequest, MetadataRequestTopic},
23 },
24 primitives::uuid::Uuid as WireUuid,
25};
26use krabka_units::convert::TimeExt as _;
27use tokio::time::{Instant, sleep, sleep_until, timeout};
28use tokio_util::sync::CancellationToken;
29
30use crate::{
31 consumer, coordinator, error, offset_wire, rebalance_listener,
32 validate::refresh_tracked_topic_ids,
33};
34
35#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
37pub enum GroupProtocol {
38 #[default]
40 Classic,
41 Consumer,
43}
44
45pub const JOIN_GROUP_MEMBER_EPOCH: i32 = 0;
47pub const LEAVE_GROUP_MEMBER_EPOCH: i32 = -1;
49pub const LEAVE_GROUP_STATIC_MEMBER_EPOCH: i32 = -2;
51
52pub const FENCED_MEMBER_EPOCH: i16 = 110;
54const UNKNOWN_MEMBER_ID: i16 = 25;
56#[cfg(test)]
58const UNSUPPORTED_ASSIGNOR: i16 = 112;
59
60pub const RE2J_NEEDS_CONSUMER_PROTOCOL: &str = "Subscribe to RE2/J pattern is not supported when using the CLASSIC protocol defined in config group.protocol";
63
64pub type ResolvedAssignment = (Vec<(String, i32)>, HashSet<WireUuid>);
67
68pub type OwnedByTopicId = BTreeMap<[u8; 16], Vec<i32>>;
71
72#[derive(Clone, Debug, Default, PartialEq, Eq)]
76pub struct SentFields {
77 rebalance_timeout_ms: Option<i32>,
78 subscribed_topic_names: Option<Vec<String>>,
79 regex: Option<String>,
80 server_assignor: Option<String>,
81 topic_partitions: Option<OwnedByTopicId>,
82}
83
84#[derive(Clone, Debug)]
86pub struct HeartbeatFields<'a> {
87 pub group_id: &'a str,
88 pub member_id: &'a str,
89 pub member_epoch: i32,
90 pub instance_id: Option<&'a str>,
91 pub rack_id: Option<&'a str>,
92 pub rebalance_timeout_ms: i32,
93 pub subscribed_topic_names: Vec<String>,
96 pub regex: Option<String>,
98 pub server_assignor: Option<&'a str>,
99 pub owned: OwnedByTopicId,
101 pub joining: bool,
103}
104
105pub fn build_heartbeat(
108 fields: &HeartbeatFields<'_>,
109 sent: &mut SentFields,
110) -> ConsumerGroupHeartbeatRequest {
111 let mut request = ConsumerGroupHeartbeatRequest {
112 group_id: fields.group_id.to_owned(),
113 member_id: fields.member_id.to_owned(),
114 member_epoch: fields.member_epoch,
115 instance_id: fields.instance_id.map(str::to_owned),
116 ..Default::default()
117 };
118 let all = fields.joining;
119 if all {
120 request.rack_id = fields.rack_id.map(str::to_owned);
122 *sent = SentFields::default();
123 }
124 if all || sent.rebalance_timeout_ms != Some(fields.rebalance_timeout_ms) {
125 request.rebalance_timeout_ms = fields.rebalance_timeout_ms;
126 sent.rebalance_timeout_ms = Some(fields.rebalance_timeout_ms);
127 } else {
128 request.rebalance_timeout_ms = -1;
129 }
130 if all || sent.subscribed_topic_names.as_ref() != Some(&fields.subscribed_topic_names) {
131 request.subscribed_topic_names = Some(fields.subscribed_topic_names.clone());
132 sent.subscribed_topic_names = Some(fields.subscribed_topic_names.clone());
133 }
134 let pattern_changed = sent.regex != fields.regex;
138 if (all && fields.regex.is_some()) || pattern_changed {
139 request.subscribed_topic_regex = Some(fields.regex.clone().unwrap_or_default());
140 sent.regex.clone_from(&fields.regex);
141 }
142 if let Some(assignor) = fields.server_assignor
143 && (all || sent.server_assignor.as_deref() != Some(assignor))
144 {
145 request.server_assignor = Some(assignor.to_owned());
146 sent.server_assignor = Some(assignor.to_owned());
147 }
148 if all || sent.topic_partitions.as_ref() != Some(&fields.owned) {
149 request.topic_partitions = Some(
150 fields
151 .owned
152 .iter()
153 .map(|(topic_id, partitions)| TopicPartitions {
154 topic_id: WireUuid(*topic_id),
155 partitions: partitions.clone(),
156 ..Default::default()
157 })
158 .collect(),
159 );
160 sent.topic_partitions = Some(fields.owned.clone());
161 }
162 request
163}
164
165#[derive(Clone, Copy, Debug, Eq, PartialEq)]
168pub enum HeartbeatAction {
169 Ok,
171 Retry,
173 FindCoordinator,
175 Fenced,
178 Fatal,
180}
181
182pub fn heartbeat_action(error_code: i16) -> HeartbeatAction {
193 match error_code {
194 0 => HeartbeatAction::Ok,
195 coordinator::COORDINATOR_NOT_AVAILABLE | coordinator::NOT_COORDINATOR => {
196 HeartbeatAction::FindCoordinator
197 }
198 coordinator::COORDINATOR_LOAD_IN_PROGRESS => HeartbeatAction::Retry,
199 FENCED_MEMBER_EPOCH | UNKNOWN_MEMBER_ID => HeartbeatAction::Fenced,
200 _ => HeartbeatAction::Fatal,
205 }
206}
207
208pub fn leave_group_epoch(
212 group_instance_id: Option<&str>,
213 operation: consumer::GroupMembershipOperation,
214) -> i32 {
215 if operation == consumer::GroupMembershipOperation::LeaveGroup {
216 return LEAVE_GROUP_MEMBER_EPOCH;
217 }
218 if group_instance_id.is_some() {
219 LEAVE_GROUP_STATIC_MEMBER_EPOCH
220 } else {
221 LEAVE_GROUP_MEMBER_EPOCH
222 }
223}
224
225pub fn should_send_leave_heartbeat(
230 member_id: &str,
231 group_instance_id: Option<&str>,
232 operation: consumer::GroupMembershipOperation,
233) -> bool {
234 if member_id.is_empty() {
235 return false;
236 }
237 !(group_instance_id.is_none() && operation == consumer::GroupMembershipOperation::RemainInGroup)
238}
239
240pub fn assignment_partitions(
246 response: &ConsumerGroupHeartbeatResponse,
247 names: &HashMap<WireUuid, String>,
248) -> Option<ResolvedAssignment> {
249 let assignment = response.assignment.as_ref()?;
250 let mut partitions = Vec::new();
251 let mut unresolved = HashSet::new();
252 for topic in &assignment.topic_partitions {
253 match names.get(&topic.topic_id) {
254 Some(name) => partitions.extend(
255 topic
256 .partitions
257 .iter()
258 .map(|partition| (name.clone(), *partition)),
259 ),
260 None => {
261 unresolved.insert(topic.topic_id);
262 }
263 }
264 }
265 partitions.sort();
266 Some((partitions, unresolved))
267}
268
269pub fn owned_by_topic_id(
271 owned: &[(String, i32)],
272 topic_ids: &HashMap<String, WireUuid>,
273) -> OwnedByTopicId {
274 let mut by_id: OwnedByTopicId = BTreeMap::new();
275 for (topic, partition) in owned {
276 if let Some(topic_id) = topic_ids.get(topic) {
277 by_id.entry(topic_id.0).or_default().push(*partition);
278 }
279 }
280 for partitions in by_id.values_mut() {
281 partitions.sort_unstable();
282 }
283 by_id
284}
285
286pub async fn run(mut state: coordinator::CoordinatorState, shutdown: CancellationToken) {
294 if let Err(error) = state.client.refresh_metadata().await {
295 tracing::warn!(%error, "coordinator client metadata refresh failed at startup");
296 }
297 state.poll_timer = coordinator::PollTimer::new(state.max_poll_interval);
300 let mut interval = state.heartbeat_interval.to_std();
301 let mut send_now = true;
302 let mut polls_open = true;
303 let mut left_on_poll_timeout = false;
306 let mut unsubscribed = false;
309 let mut fatal = false;
312 let mut last_pattern_check: Option<Instant> = None;
313 loop {
314 if !send_now {
315 let event = tokio::select! {
316 biased;
317 () = shutdown.cancelled() => break,
318 () = sleep(interval), if !left_on_poll_timeout && !unsubscribed => {
319 TaskEvent::Heartbeat
320 }
321 changed = state.polls.changed(), if polls_open => {
322 if changed.is_err() {
323 polls_open = false;
324 continue;
325 }
326 TaskEvent::Poll
327 }
328 changed = state.subscription_changes.changed() => {
329 if changed.is_err() {
330 break;
331 }
332 TaskEvent::Subscription
333 }
334 () = sleep_until(state.poll_timer.deadline()),
335 if !left_on_poll_timeout && !unsubscribed =>
336 {
337 TaskEvent::PollTimeout
338 }
339 };
340 match event {
341 TaskEvent::Poll => {
342 state.poll_timer.reset();
343 if !left_on_poll_timeout {
344 continue;
345 }
346 left_on_poll_timeout = false;
350 state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
351 state.sent_heartbeat_fields = SentFields::default();
352 }
353 TaskEvent::PollTimeout => {
354 leave_after_poll_timeout(&mut state).await;
355 left_on_poll_timeout = true;
356 continue;
357 }
358 TaskEvent::Subscription => {
359 let subscription = state.subscription.borrow().clone();
360 if subscription.unsubscribes != state.seen_unsubscribes {
361 state.seen_unsubscribes = subscription.unsubscribes;
362 leave_after_unsubscribe(&mut state).await;
363 unsubscribed = true;
364 continue;
365 }
366 if subscription.is_none() {
367 continue;
368 }
369 if unsubscribed {
370 unsubscribed = false;
372 state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
373 state.sent_heartbeat_fields = SentFields::default();
374 }
375 if !subscription.topics.is_empty() {
379 state.client.metadata_topics().set(subscription.topics);
380 }
381 }
382 TaskEvent::Heartbeat => {}
383 }
384 }
385 if state.subscription.borrow().pattern.is_some()
389 && last_pattern_check.is_none_or(|last| {
390 coordinator::subscription_metadata_refresh_due(
391 last,
392 state.subscription_metadata_refresh_interval,
393 )
394 })
395 {
396 last_pattern_check = Some(Instant::now());
397 coordinator::refresh_pattern_topics(&mut state).await;
398 }
399 send_now = false;
400 let outcome = tokio::select! {
401 biased;
402 () = shutdown.cancelled() => break,
403 outcome = heartbeat_once(&mut state) => outcome,
404 };
405 match outcome {
406 HeartbeatOutcome::Interval(next) => interval = next,
407 HeartbeatOutcome::Reconciled(next) => {
408 interval = next;
409 send_now = true;
412 }
413 HeartbeatOutcome::Backoff => {
414 tokio::select! {
415 () = shutdown.cancelled() => break,
416 () = sleep(state.retry_policy.initial_backoff) => {}
417 }
418 send_now = true;
419 }
420 HeartbeatOutcome::Stop => {
421 fatal = true;
422 break;
423 }
424 }
425 }
426 let close = *state.close_operation.borrow();
427 if !left_on_poll_timeout
432 && !unsubscribed
433 && !fatal
434 && should_send_leave_heartbeat(
435 &state.member_id,
436 state.group_instance_id.as_deref(),
437 close.operation,
438 )
439 {
440 let epoch = leave_group_epoch(state.group_instance_id.as_deref(), close.operation);
441 state.member_epoch = epoch;
442 let leave_timeout = state.leave_group_timeout.to_std();
443 let leave = send_heartbeat(&mut state, HeartbeatKind::Leave).await;
444 let _ = timeout(leave_timeout, leave.send()).await;
445 }
446}
447
448#[derive(Clone, Copy, Debug, Eq, PartialEq)]
450enum TaskEvent {
451 Heartbeat,
453 Poll,
455 Subscription,
457 PollTimeout,
459}
460
461async fn leave_after_poll_timeout(state: &mut coordinator::CoordinatorState) {
468 coordinator::warn_poll_timeout(state);
469 set_leave_epoch(state);
470 let leave_timeout = state.leave_group_timeout.to_std();
471 let leave = send_heartbeat(state, HeartbeatKind::Leave).await;
472 let _ = timeout(leave_timeout, leave.send()).await;
473 coordinator::release_partitions_after_poll_timeout(state).await;
474 if let Err(error) = coordinator::call_lost_listener(state).await {
475 tracing::warn!(%error, group = %state.group_id, "the lost callback failed");
476 }
477 state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
478 state.sent_heartbeat_fields = SentFields::default();
479}
480
481async fn leave_after_unsubscribe(state: &mut coordinator::CoordinatorState) {
486 set_leave_epoch(state);
487 let leave_timeout = state.leave_group_timeout.to_std();
488 let leave = send_heartbeat(state, HeartbeatKind::Leave).await;
489 let _ = timeout(leave_timeout, leave.send()).await;
490 coordinator::forget_member_keeping_id(state).await;
491 state.rebalance_pending.send_replace(false);
492 state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
493 state.sent_heartbeat_fields = SentFields::default();
494}
495
496fn set_leave_epoch(state: &mut coordinator::CoordinatorState) {
497 state.member_epoch = leave_group_epoch(
498 state.group_instance_id.as_deref(),
499 consumer::GroupMembershipOperation::Default,
500 );
501}
502
503pub async fn heartbeat_during_callback(
511 state: &coordinator::CoordinatorState,
512 done: &CancellationToken,
513) {
514 let mut ticker = coordinator::heartbeat_ticker(state.heartbeat_interval);
515 loop {
516 tokio::select! {
517 biased;
518 () = done.cancelled() => return,
519 _ = ticker.tick() => {}
520 }
521 let mut sent = state.sent_heartbeat_fields.clone();
522 let fields = heartbeat_fields(state, HeartbeatKind::Interval).await;
523 let request = build_heartbeat(&fields, &mut sent);
524 let broker = state
525 .client
526 .broker(state.coordinator_id.load(atomic::Ordering::Relaxed));
527 let call = HeartbeatCall {
528 broker,
529 request,
530 needs_regex: state.subscription.borrow().regex.is_some(),
531 };
532 let result = tokio::select! {
533 biased;
534 () = done.cancelled() => return,
535 result = call.send() => result,
536 };
537 if let Err(error) = result {
538 tracing::warn!(%error, group = %state.group_id, "the heartbeat during a rebalance callback failed");
539 }
540 }
541}
542
543#[derive(Clone, Copy, Debug, Eq, PartialEq)]
545enum HeartbeatOutcome {
546 Interval(Duration),
548 Reconciled(Duration),
551 Backoff,
553 Stop,
555}
556
557async fn heartbeat_once(state: &mut coordinator::CoordinatorState) -> HeartbeatOutcome {
559 let kind = if state.member_epoch == JOIN_GROUP_MEMBER_EPOCH {
560 HeartbeatKind::Join
561 } else {
562 HeartbeatKind::Interval
563 };
564 let response = match send_heartbeat(state, kind).await.send().await {
565 Ok(response) => response,
566 Err(error) => {
567 tracing::warn!(%error, group = %state.group_id, "consumer group heartbeat failed");
568 state.sent_heartbeat_fields = SentFields::default();
572 if !coordinator::is_retriable_transport_error(&error) {
573 tracing::error!(
577 %error,
578 group = %state.group_id,
579 "the consumer group heartbeat failed for good; the member stops"
580 );
581 coordinator::report_fatal_error(
582 &state.poll_error,
583 error::ConsumerError::Client(error),
584 );
585 return HeartbeatOutcome::Stop;
586 }
587 state
588 .client
589 .evict_broker(state.coordinator_id.load(atomic::Ordering::Relaxed));
590 return HeartbeatOutcome::Backoff;
591 }
592 };
593 let interval = heartbeat_interval(response.heartbeat_interval_ms, state.heartbeat_interval);
594 if response.error_code != 0 {
595 state.sent_heartbeat_fields = SentFields::default();
598 }
599 match heartbeat_action(response.error_code) {
600 HeartbeatAction::Ok => {}
601 HeartbeatAction::Retry => return HeartbeatOutcome::Backoff,
602 HeartbeatAction::FindCoordinator => {
603 coordinator::find_coordinator_again(state).await;
604 return HeartbeatOutcome::Backoff;
605 }
606 HeartbeatAction::Fenced => {
607 tracing::warn!(
608 group = %state.group_id,
609 error_code = response.error_code,
610 "the coordinator fenced the member; giving up the partitions and joining again"
611 );
612 state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
613 coordinator::forget_member_keeping_id(state).await;
614 if let Err(error) = coordinator::call_lost_listener(state).await {
617 tracing::warn!(%error, group = %state.group_id, "the lost callback failed");
618 }
619 return HeartbeatOutcome::Backoff;
620 }
621 HeartbeatAction::Fatal => {
622 tracing::error!(
623 group = %state.group_id,
624 error_code = response.error_code,
625 error = ?response.error_message,
626 "the consumer group heartbeat failed with a fatal error"
627 );
628 coordinator::report_fatal_error(
629 &state.poll_error,
630 error::ConsumerError::Server(response.error_code),
631 );
632 return HeartbeatOutcome::Stop;
633 }
634 }
635 if let Some(member_id) = response.member_id.clone().filter(|id| !id.is_empty()) {
636 state.member_id = member_id;
637 }
638 state.member_epoch = response.member_epoch;
639 coordinator::publish_member_epoch(state).await;
640 let names = offset_wire::id_to_name(&state.topic_ids.lock().await.clone());
641 let Some((mut target, mut unresolved)) = assignment_partitions(&response, &names) else {
642 state.rebalance_pending.send_replace(false);
643 return HeartbeatOutcome::Interval(interval);
644 };
645 if !unresolved.is_empty() {
646 refresh_topic_ids(state, &unresolved).await;
651 let names = offset_wire::id_to_name(&state.topic_ids.lock().await.clone());
652 if let Some(resolved) = assignment_partitions(&response, &names) {
653 (target, unresolved) = resolved;
654 }
655 }
656 if !unresolved.is_empty() {
657 tracing::info!(
658 group = %state.group_id,
659 topics = unresolved.len(),
660 "the assignment names topic ids that the metadata does not name yet"
661 );
662 }
663 let topics: Vec<String> = target.iter().map(|(topic, _)| topic.clone()).collect();
667 state
668 .subscription
669 .send_if_modified(|subscription| subscription.store_assigned_topics(&topics));
670 match reconcile(state, &target).await {
671 Ok(true) => HeartbeatOutcome::Reconciled(interval),
675 Ok(false) => HeartbeatOutcome::Interval(interval),
676 Err(error) => {
677 tracing::warn!(%error, group = %state.group_id, "the reconciliation of the assignment failed");
678 HeartbeatOutcome::Backoff
679 }
680 }
681}
682
683#[derive(Clone, Copy, Debug, Eq, PartialEq)]
690enum HeartbeatKind {
691 Join,
693 Interval,
695 Leave,
697}
698
699async fn send_heartbeat(
706 state: &mut coordinator::CoordinatorState,
707 kind: HeartbeatKind,
708) -> HeartbeatCall<'_> {
709 let mut sent = state.sent_heartbeat_fields.clone();
710 let needs_regex = state.subscription.borrow().regex.is_some();
711 let request = build_heartbeat(&heartbeat_fields(state, kind).await, &mut sent);
712 state.sent_heartbeat_fields = sent;
713 let broker = state
714 .client
715 .broker(state.coordinator_id.load(atomic::Ordering::Relaxed));
716 HeartbeatCall {
717 broker,
718 request,
719 needs_regex,
720 }
721}
722
723async fn heartbeat_fields(
725 state: &coordinator::CoordinatorState,
726 kind: HeartbeatKind,
727) -> HeartbeatFields<'_> {
728 let topic_ids = state.topic_ids.lock().await.clone();
729 let owned = if kind == HeartbeatKind::Leave {
730 OwnedByTopicId::new()
731 } else {
732 owned_by_topic_id(&state.assigned.lock().await, &topic_ids)
733 };
734 let (subscribed_topic_names, regex) = {
735 let subscription = state.subscription.borrow();
736 let topics = if subscription.regex.is_some() {
740 Vec::new()
741 } else {
742 subscription.topics.clone()
743 };
744 (topics, subscription.regex.clone())
745 };
746 HeartbeatFields {
747 group_id: &state.group_id,
748 member_id: &state.member_id,
749 member_epoch: state.member_epoch,
750 instance_id: state.group_instance_id.as_deref(),
751 rack_id: state.client_rack.as_deref(),
752 rebalance_timeout_ms: consumer::protocol_millis_i32(state.max_poll_interval),
753 subscribed_topic_names,
754 regex,
755 server_assignor: state.server_assignor.as_deref(),
756 owned,
757 joining: kind == HeartbeatKind::Join,
758 }
759}
760
761struct RegexHeartbeat(ConsumerGroupHeartbeatRequest);
765
766impl krabka_protocol::Encode for RegexHeartbeat {
767 fn encode<B: bytes::BufMut>(
768 &self,
769 buf: &mut B,
770 version: i16,
771 ) -> Result<(), krabka_protocol::ProtocolError> {
772 self.0.encode(buf, version)
773 }
774
775 fn encoded_len(&self, version: i16) -> usize {
776 self.0.encoded_len(version)
777 }
778}
779
780impl krabka_protocol::ProtocolRequest for RegexHeartbeat {
781 const API_KEY: i16 = consumer_group_heartbeat_request::API_KEY;
782 const MIN_VERSION: i16 = 1;
783 const MAX_VERSION: i16 = consumer_group_heartbeat_request::MAX_VERSION;
784 const LATEST_STABLE_VERSION: i16 = consumer_group_heartbeat_request::LATEST_STABLE_VERSION;
785 const FLEXIBLE_MIN: i16 = consumer_group_heartbeat_request::FLEXIBLE_MIN;
786 type Response = ConsumerGroupHeartbeatResponse;
787}
788
789struct HeartbeatCall<'a> {
791 broker: krabka_client_core::BrokerHandle<'a>,
792 request: ConsumerGroupHeartbeatRequest,
793 needs_regex: bool,
796}
797
798impl HeartbeatCall<'_> {
799 async fn send(self) -> Result<ConsumerGroupHeartbeatResponse, krabka_client_core::ClientError> {
800 if self.needs_regex {
801 return self.broker.send(RegexHeartbeat(self.request)).await;
802 }
803 self.broker.send(self.request).await
804 }
805}
806
807async fn refresh_topic_ids(state: &coordinator::CoordinatorState, unresolved: &HashSet<WireUuid>) {
815 let by_id = state.subscription.borrow().regex.is_some();
816 let request = if by_id {
817 MetadataRequest {
818 topics: Some(
819 unresolved
820 .iter()
821 .map(|topic_id| MetadataRequestTopic {
822 topic_id: *topic_id,
823 name: None,
824 ..Default::default()
825 })
826 .collect(),
827 ),
828 allow_auto_topic_creation: false,
829 ..Default::default()
830 }
831 } else {
832 state.client.metadata_topics().request()
833 };
834 match state.client.refresh_metadata_with(request).await {
835 Ok(metadata) => {
836 refresh_tracked_topic_ids(&mut *state.topic_ids.lock().await, &metadata);
837 }
838 Err(error) => {
839 tracing::warn!(%error, group = %state.group_id, "metadata refresh failed");
840 }
841 }
842}
843
844fn heartbeat_interval(heartbeat_interval_ms: i32, configured: krabka_units::Time) -> Duration {
847 if heartbeat_interval_ms > 0 {
848 Duration::from_millis(u64::try_from(heartbeat_interval_ms).unwrap_or(0))
849 } else {
850 configured.to_std()
851 }
852}
853
854async fn reconcile(
861 state: &mut coordinator::CoordinatorState,
862 target: &[(String, i32)],
863) -> Result<bool, error::ConsumerError> {
864 let owned = state.assigned.lock().await.clone();
865 let target_set: HashSet<(String, i32)> = target.iter().cloned().collect();
866 let owned_set: HashSet<(String, i32)> = owned.iter().cloned().collect();
867 let mut revoked: Vec<(String, i32)> = owned
868 .iter()
869 .filter(|partition| !target_set.contains(*partition))
870 .cloned()
871 .collect();
872 let mut added: Vec<(String, i32)> = target
873 .iter()
874 .filter(|partition| !owned_set.contains(*partition))
875 .cloned()
876 .collect();
877 revoked.sort();
878 added.sort();
879 if revoked.is_empty() && added.is_empty() {
880 state.rebalance_pending.send_replace(false);
881 return Ok(false);
882 }
883 state.rebalance_pending.send_replace(true);
884 let epoch = state.member_epoch;
885 if !revoked.is_empty() {
886 coordinator::call_listener_kind(
887 state,
888 rebalance_listener::ListenerCallKind::Revoked,
889 revoked.clone(),
890 )
891 .await?;
892 let kept: Vec<(String, i32)> = owned
893 .iter()
894 .filter(|partition| target_set.contains(*partition))
895 .cloned()
896 .collect();
897 coordinator::publish_group_assignment(state, &kept, epoch).await;
898 }
899 coordinator::prime_added_offsets(state, &added).await?;
900 let gate = coordinator::assigned_callback_gate(state, &added);
901 coordinator::publish_group_assignment(state, target, epoch).await;
902 let result = coordinator::call_listener_kind(
903 state,
904 rebalance_listener::ListenerCallKind::Assigned,
905 added,
906 )
907 .await;
908 drop(gate);
909 state.rebalance_pending.send_replace(false);
910 result.map(|()| true)
911}
912
913#[cfg(test)]
914mod group_protocol_tests {
915 use std::sync::{Arc, Mutex};
916
917 use bytes::Buf as _;
918 use krabka_client_core::{MockBroker, mock_test_support::api_versions_response};
919 use krabka_protocol::{
920 Decode as _,
921 owned::{
922 api_versions_request,
923 api_versions_response::ApiVersionsResponse,
924 common::consumer_group_heartbeat_response::topic_partitions::TopicPartitions as AssignedPartitions,
925 consumer_group_heartbeat_request,
926 consumer_group_heartbeat_response::Assignment,
927 fetch_request,
928 fetch_response::FetchResponse,
929 find_coordinator_request,
930 find_coordinator_response::FindCoordinatorResponse,
931 list_offsets_request,
932 list_offsets_response::{
933 ListOffsetsPartitionResponse, ListOffsetsResponse, ListOffsetsTopicResponse,
934 },
935 metadata_request,
936 metadata_response::{MetadataResponse, MetadataResponseTopic},
937 offset_fetch_request,
938 offset_fetch_response::{
939 OffsetFetchResponse, OffsetFetchResponsePartition, OffsetFetchResponseTopic,
940 },
941 },
942 };
943 use krabka_units::{Time, millis, secs};
944
945 use super::*;
946 use crate::{
947 CloseOptions, Consumer,
948 consumer::GroupMembershipOperation,
949 test_support::{ListenerCall, ListenerCalls, RecordingListener, metadata_partitions},
950 };
951
952 const TOPIC: &str = "orders";
953 const TOPIC_ID: WireUuid = WireUuid([3; 16]);
954
955 #[derive(Clone, Debug)]
957 struct Answer {
958 error_code: i16,
959 member_epoch: i32,
960 assignment: Option<Vec<i32>>,
963 member_id: Option<String>,
966 }
967
968 fn assigns(member_epoch: i32, partitions: &[i32]) -> Answer {
970 Answer {
971 error_code: 0,
972 member_epoch,
973 assignment: Some(partitions.to_vec()),
974 member_id: None,
975 }
976 }
977
978 fn named(answer: Answer, member_id: &str) -> Answer {
980 Answer {
981 member_id: Some(member_id.to_owned()),
982 ..answer
983 }
984 }
985
986 fn fails(error_code: i16) -> Answer {
988 Answer {
989 error_code,
990 member_epoch: 0,
991 assignment: None,
992 member_id: None,
993 }
994 }
995
996 #[derive(Clone, Default)]
998 struct Recorder {
999 heartbeats: Arc<Mutex<Vec<ConsumerGroupHeartbeatRequest>>>,
1001 metadata_topics: Arc<Mutex<Vec<String>>>,
1003 asked_by_id: Arc<Mutex<bool>>,
1005 classic_heartbeats: Arc<Mutex<usize>>,
1008 }
1009
1010 impl Recorder {
1011 fn heartbeats(&self) -> Vec<ConsumerGroupHeartbeatRequest> {
1012 self.heartbeats.lock().expect("heartbeats lock").clone()
1013 }
1014
1015 fn acknowledged(&self) -> Vec<Vec<i32>> {
1017 acknowledged(&self.heartbeats())
1018 }
1019
1020 fn epochs(&self) -> Vec<i32> {
1021 self.heartbeats()
1022 .iter()
1023 .map(|request| request.member_epoch)
1024 .collect()
1025 }
1026
1027 fn first_heartbeat(&self) -> ConsumerGroupHeartbeatRequest {
1028 self.heartbeats()
1029 .into_iter()
1030 .next()
1031 .expect("a first heartbeat")
1032 }
1033 }
1034
1035 fn encode(response: &impl krabka_protocol::Encode, version: i16) -> Vec<u8> {
1036 let mut buffer = bytes::BytesMut::new();
1037 response.encode(&mut buffer, version).expect("encode");
1038 buffer.to_vec()
1039 }
1040
1041 fn flexible(response: &impl krabka_protocol::Encode, version: i16) -> Vec<u8> {
1044 let mut frame = vec![0];
1045 frame.extend(encode(response, version));
1046 frame
1047 }
1048
1049 fn api_versions() -> ApiVersionsResponse {
1051 api_versions_response([
1052 (api_versions_request::API_KEY, 3),
1053 (metadata_request::API_KEY, 12),
1054 (find_coordinator_request::API_KEY, 2),
1055 (offset_fetch_request::API_KEY, 5),
1056 (list_offsets_request::API_KEY, 7),
1057 (fetch_request::API_KEY, 12),
1058 (consumer_group_heartbeat_request::API_KEY, 1),
1059 ])
1060 }
1061
1062 fn metadata() -> MetadataResponse {
1064 MetadataResponse {
1065 topics: vec![MetadataResponseTopic {
1066 name: Some(TOPIC.into()),
1067 topic_id: TOPIC_ID,
1068 partitions: metadata_partitions(2),
1069 ..Default::default()
1070 }],
1071 ..Default::default()
1072 }
1073 }
1074
1075 fn empty_list_offset_partitions() -> Vec<ListOffsetsPartitionResponse> {
1077 (0..2)
1078 .map(|partition_index| ListOffsetsPartitionResponse {
1079 partition_index,
1080 timestamp: -1,
1081 offset: 0,
1082 leader_epoch: -1,
1083 ..Default::default()
1084 })
1085 .collect()
1086 }
1087
1088 fn list_offsets() -> ListOffsetsResponse {
1089 ListOffsetsResponse {
1090 topics: vec![ListOffsetsTopicResponse {
1091 name: TOPIC.into(),
1092 partitions: empty_list_offset_partitions(),
1093 ..Default::default()
1094 }],
1095 ..Default::default()
1096 }
1097 }
1098
1099 fn missing_committed_partitions() -> Vec<OffsetFetchResponsePartition> {
1102 (0..2)
1103 .map(|partition_index| OffsetFetchResponsePartition {
1104 partition_index,
1105 committed_offset: -1,
1106 committed_leader_epoch: -1,
1107 ..Default::default()
1108 })
1109 .collect()
1110 }
1111
1112 fn no_committed_offsets() -> OffsetFetchResponse {
1113 OffsetFetchResponse {
1114 topics: vec![OffsetFetchResponseTopic {
1115 name: TOPIC.into(),
1116 partitions: missing_committed_partitions(),
1117 ..Default::default()
1118 }],
1119 ..Default::default()
1120 }
1121 }
1122
1123 fn metadata_request_of(mut body: &[u8], version: i16) -> MetadataRequest {
1125 let client_id_len = usize::try_from(body.get_i16()).expect("client id length");
1126 body.advance(client_id_len);
1127 if version >= 9 {
1128 body.advance(1);
1129 }
1130 MetadataRequest::decode(&mut body, version).expect("decode")
1131 }
1132
1133 fn heartbeat_request(mut body: &[u8], version: i16) -> ConsumerGroupHeartbeatRequest {
1135 let client_id_len = usize::try_from(body.get_i16()).expect("client id length");
1136 body.advance(client_id_len);
1137 body.advance(1);
1139 ConsumerGroupHeartbeatRequest::decode(&mut body, version).expect("decode")
1140 }
1141
1142 async fn heartbeat_coordinator(answers: Vec<Answer>, recorder: &Recorder) -> MockBroker {
1146 heartbeat_coordinator_with(answers, recorder, false).await
1147 }
1148
1149 async fn heartbeat_coordinator_with(
1153 answers: Vec<Answer>,
1154 recorder: &Recorder,
1155 by_id_only: bool,
1156 ) -> MockBroker {
1157 let recorder = recorder.clone();
1158 let answers = Mutex::new(answers.into_iter());
1159 let last = Mutex::new(assigns(1, &[]));
1160 MockBroker::start(move |api_key, version, _corr_id, body| {
1161 if api_key == api_versions_request::API_KEY {
1162 return Some(encode(&api_versions(), 0));
1163 }
1164 if api_key == find_coordinator_request::API_KEY {
1165 return Some(encode(&FindCoordinatorResponse::default(), version));
1166 }
1167 if api_key == CLASSIC_HEARTBEAT_API_KEY {
1168 *recorder
1169 .classic_heartbeats
1170 .lock()
1171 .expect("classic heartbeats lock") += 1;
1172 return None;
1173 }
1174 if api_key == metadata_request::API_KEY {
1175 let request = metadata_request_of(body, version);
1176 let topics = request.topics.unwrap_or_default();
1177 let by_id = topics
1178 .iter()
1179 .any(|topic| topic.topic_id != WireUuid::default());
1180 if by_id {
1181 *recorder.asked_by_id.lock().expect("asked lock") = true;
1182 }
1183 recorder
1184 .metadata_topics
1185 .lock()
1186 .expect("metadata topics lock")
1187 .extend(topics.iter().filter_map(|topic| topic.name.clone()));
1188 if by_id_only && !by_id {
1189 return Some(flexible(&MetadataResponse::default(), version));
1190 }
1191 return Some(flexible(&metadata(), version));
1192 }
1193 if api_key == offset_fetch_request::API_KEY {
1194 return Some(encode(&no_committed_offsets(), version));
1195 }
1196 if api_key == list_offsets_request::API_KEY {
1197 return Some(flexible(&list_offsets(), version));
1198 }
1199 if api_key == fetch_request::API_KEY {
1200 return Some(flexible(&FetchResponse::default(), version));
1201 }
1202 if api_key != consumer_group_heartbeat_request::API_KEY {
1203 return None;
1204 }
1205 recorder
1206 .heartbeats
1207 .lock()
1208 .expect("heartbeats lock")
1209 .push(heartbeat_request(body, version));
1210 let next = answers.lock().expect("answers lock").next();
1211 let mut last = last.lock().expect("last lock");
1212 if let Some(next) = next {
1213 *last = next;
1214 }
1215 let response = ConsumerGroupHeartbeatResponse {
1216 error_code: last.error_code,
1217 member_epoch: last.member_epoch,
1218 member_id: last.member_id.clone(),
1219 heartbeat_interval_ms: 20,
1220 assignment: last.assignment.clone().map(|partitions| Assignment {
1221 topic_partitions: vec![AssignedPartitions {
1222 topic_id: TOPIC_ID,
1223 partitions,
1224 ..Default::default()
1225 }],
1226 ..Default::default()
1227 }),
1228 ..Default::default()
1229 };
1230 drop(last);
1231 Some(flexible(&response, version))
1232 })
1233 .await
1234 .unwrap()
1235 }
1236
1237 async fn assigned_coordinator(recorder: &Recorder) -> MockBroker {
1238 heartbeat_coordinator(vec![assigns(1, &[0, 1])], recorder).await
1239 }
1240
1241 struct AssignedGroupFixture {
1242 recorder: Recorder,
1243 mock: MockBroker,
1244 }
1245
1246 async fn assigned_group_fixture() -> AssignedGroupFixture {
1247 let recorder = Recorder::default();
1248 let mock = assigned_coordinator(&recorder).await;
1249 AssignedGroupFixture { recorder, mock }
1250 }
1251
1252 const CLASSIC_HEARTBEAT_API_KEY: i16 = 12;
1254
1255 #[derive(Debug, Eq, PartialEq)]
1257 struct Observed {
1258 acknowledged: Vec<Vec<i32>>,
1262 owned: Vec<i32>,
1264 last_epoch: i32,
1266 rejoins: usize,
1268 }
1269
1270 fn acknowledged(requests: &[ConsumerGroupHeartbeatRequest]) -> Vec<Vec<i32>> {
1272 requests
1273 .iter()
1274 .filter_map(|request| request.topic_partitions.as_ref())
1275 .map(|topics| {
1276 topics
1277 .iter()
1278 .flat_map(|topic| topic.partitions.clone())
1279 .collect()
1280 })
1281 .collect()
1282 }
1283
1284 enum GroupPollTarget<'a> {
1285 Acknowledged {
1286 recorder: &'a Recorder,
1287 count: usize,
1288 },
1289 Assignment,
1290 }
1291
1292 #[derive(bon::Builder)]
1293 struct GroupPollWait<'a> {
1294 consumer: &'a mut Consumer,
1295 target: GroupPollTarget<'a>,
1296 }
1297
1298 impl GroupPollWait<'_> {
1299 async fn wait(self) {
1300 let deadline = Instant::now() + Duration::from_secs(20);
1301 loop {
1302 let ready = match self.target {
1303 GroupPollTarget::Acknowledged { recorder, count } => {
1304 recorder.acknowledged().len() >= count
1305 }
1306 GroupPollTarget::Assignment => !self.consumer.assignment().await.is_empty(),
1307 };
1308 if ready || Instant::now() >= deadline {
1309 break;
1310 }
1311 let _ = self.consumer.poll(millis(20)).await;
1312 }
1313 }
1314 }
1315
1316 async fn run_member(answers: Vec<Answer>, instance_id: Option<&str>, acks: usize) -> Observed {
1319 let recorder = Recorder::default();
1320 let mock = heartbeat_coordinator(answers, &recorder).await;
1321 let mut consumer = TestConsumerOptions::builder()
1322 .mock(&mock)
1323 .server_assignor("uniform".into())
1324 .maybe_group_instance_id(instance_id.map(str::to_owned))
1325 .build()
1326 .start()
1327 .await;
1328 GroupPollWait::builder()
1329 .consumer(&mut consumer)
1330 .target(GroupPollTarget::Acknowledged {
1331 recorder: &recorder,
1332 count: acks,
1333 })
1334 .build()
1335 .wait()
1336 .await;
1337 let owned: Vec<i32> = consumer
1338 .assignment()
1339 .await
1340 .into_iter()
1341 .map(|(_, partition)| partition)
1342 .collect();
1343 consumer
1344 .close_with(CloseOptions::group_membership_operation(
1345 GroupMembershipOperation::Default,
1346 ))
1347 .await
1348 .expect("close");
1349 mock.stop();
1350 let epochs = recorder.epochs();
1351 Observed {
1352 acknowledged: recorder.acknowledged(),
1353 owned,
1354 last_epoch: *epochs.last().expect("a last heartbeat"),
1355 rejoins: epochs
1356 .iter()
1357 .skip(1)
1358 .filter(|epoch| **epoch == JOIN_GROUP_MEMBER_EPOCH)
1359 .count(),
1360 }
1361 }
1362
1363 enum SubscriptionMode {
1364 Topics(Vec<String>),
1365 Regex(String),
1366 Pattern(crate::TopicPattern),
1367 }
1368
1369 #[derive(bon::Builder)]
1370 struct TestConsumerOptions<'a> {
1371 mock: &'a MockBroker,
1372 #[builder(default = SubscriptionMode::Topics(vec![TOPIC.to_owned()]))]
1373 subscription: SubscriptionMode,
1374 server_assignor: Option<String>,
1375 #[builder(default = millis(20))]
1376 heartbeat_interval: Time,
1377 #[builder(default = secs(5))]
1378 request_timeout: Time,
1379 max_poll_interval: Option<Time>,
1380 group_instance_id: Option<String>,
1381 rebalance_listener: Option<Box<dyn crate::ConsumerRebalanceListener>>,
1382 }
1383
1384 impl TestConsumerOptions<'_> {
1385 async fn start(self) -> Consumer {
1386 let builder = Consumer::builder()
1387 .bootstrap(self.mock.addr.to_string())
1388 .group_id("group-a")
1389 .group_protocol(GroupProtocol::Consumer)
1390 .heartbeat_interval(self.heartbeat_interval)
1391 .request_timeout(self.request_timeout)
1392 .maybe_group_remote_assignor(self.server_assignor)
1393 .maybe_group_instance_id(self.group_instance_id)
1394 .maybe_rebalance_listener(self.rebalance_listener)
1395 .maybe_max_poll_interval(self.max_poll_interval);
1396 match self.subscription {
1397 SubscriptionMode::Topics(topics) => {
1398 builder.subscribe(topics).build().await.expect("build")
1399 }
1400 SubscriptionMode::Regex(regex) => {
1401 builder.subscribe_regex(regex).build().await.expect("build")
1402 }
1403 SubscriptionMode::Pattern(pattern) => builder
1404 .subscribe_pattern(pattern)
1405 .build()
1406 .await
1407 .expect("build"),
1408 }
1409 }
1410
1411 async fn start_and_assign(self) -> Consumer {
1412 let mut consumer = self.start().await;
1413 poll_until_assigned(&mut consumer).await;
1414 consumer
1415 }
1416 }
1417
1418 async fn poll_until_assigned(consumer: &mut Consumer) {
1419 let deadline = Instant::now() + Duration::from_secs(20);
1420 while consumer.assignment().await.is_empty() && Instant::now() < deadline {
1421 let _ = consumer.poll(millis(20)).await;
1422 }
1423 }
1424
1425 async fn poll_until_error(consumer: &mut Consumer) -> Option<String> {
1426 let deadline = Instant::now() + Duration::from_secs(20);
1427 let mut error = None;
1428 while error.is_none() && Instant::now() < deadline {
1429 error = consumer
1430 .poll(millis(20))
1431 .await
1432 .err()
1433 .map(|error| error.to_string());
1434 }
1435 error
1436 }
1437
1438 async fn started_consumer_group(mock: &MockBroker) -> Consumer {
1439 TestConsumerOptions::builder()
1440 .mock(mock)
1441 .build()
1442 .start_and_assign()
1443 .await
1444 }
1445
1446 async fn close_and_first_heartbeat(
1447 consumer: Consumer,
1448 mock: MockBroker,
1449 recorder: &Recorder,
1450 ) -> ConsumerGroupHeartbeatRequest {
1451 consumer.close().await.expect("close");
1452 mock.stop();
1453 recorder.first_heartbeat()
1454 }
1455
1456 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1459 async fn the_first_heartbeat_joins_the_group_with_every_field() {
1460 let AssignedGroupFixture { recorder, mock } = assigned_group_fixture().await;
1461 let consumer = TestConsumerOptions::builder()
1462 .mock(&mock)
1463 .server_assignor("uniform".into())
1464 .build()
1465 .start_and_assign()
1466 .await;
1467 let first = close_and_first_heartbeat(consumer, mock, &recorder).await;
1468 assert2::assert!(
1469 ConsumerGroupHeartbeatRequest {
1470 member_id: String::new(),
1471 ..first.clone()
1472 } == ConsumerGroupHeartbeatRequest {
1473 group_id: "group-a".to_owned(),
1474 member_id: String::new(),
1475 member_epoch: JOIN_GROUP_MEMBER_EPOCH,
1476 rebalance_timeout_ms: 300_000,
1477 subscribed_topic_names: Some(vec![TOPIC.to_owned()]),
1478 server_assignor: Some("uniform".to_owned()),
1479 topic_partitions: Some(Vec::new()),
1480 ..Default::default()
1481 }
1482 );
1483 assert2::assert!(uuid::Uuid::parse_str(&first.member_id).is_ok());
1485 }
1486
1487 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1493 async fn a_regular_expression_subscription_resolves_its_topics_by_id() {
1494 let recorder = Recorder::default();
1495 let mock = heartbeat_coordinator_with(vec![assigns(1, &[0, 1])], &recorder, true).await;
1496 let consumer = TestConsumerOptions::builder()
1497 .mock(&mock)
1498 .subscription(SubscriptionMode::Regex("orders-.*".into()))
1499 .build()
1500 .start_and_assign()
1501 .await;
1502 let assignment = consumer.assignment().await;
1503 let subscription = consumer.subscription();
1504 let first = close_and_first_heartbeat(consumer, mock, &recorder).await;
1505 assert2::assert!(
1506 (
1507 first.subscribed_topic_regex,
1508 first.subscribed_topic_names,
1509 *recorder.asked_by_id.lock().expect("asked lock"),
1510 assignment,
1511 subscription,
1512 ) == (
1513 Some("orders-.*".to_owned()),
1514 Some(Vec::new()),
1515 true,
1516 vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
1517 vec![TOPIC.to_owned()],
1518 )
1519 );
1520 }
1521
1522 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1526 async fn a_client_side_pattern_sends_the_topics_that_matched() {
1527 let AssignedGroupFixture { recorder, mock } = assigned_group_fixture().await;
1528 let consumer = TestConsumerOptions::builder()
1529 .mock(&mock)
1530 .subscription(SubscriptionMode::Pattern(crate::TopicPattern::new(
1531 |topic| topic.starts_with("ord"),
1532 )))
1533 .build()
1534 .start_and_assign()
1535 .await;
1536 let assignment = consumer.assignment().await;
1537 let first = close_and_first_heartbeat(consumer, mock, &recorder).await;
1538 assert2::assert!(
1539 (
1540 first.subscribed_topic_names,
1541 first.subscribed_topic_regex,
1542 assignment,
1543 ) == (
1544 Some(vec![TOPIC.to_owned()]),
1545 None,
1546 vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
1547 )
1548 );
1549 }
1550
1551 fn first_lost(calls: &ListenerCalls) -> Option<ListenerCall> {
1553 calls
1554 .lock()
1555 .expect("calls lock")
1556 .iter()
1557 .find(|(kind, _)| *kind == rebalance_listener::ListenerCallKind::Lost)
1558 .cloned()
1559 }
1560
1561 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1566 async fn a_fatal_heartbeat_error_reaches_poll() {
1567 let mut actual = Vec::new();
1568 for (name, answers, speaks_heartbeat) in [
1569 (
1570 "an unsupported assignor",
1571 vec![fails(UNSUPPORTED_ASSIGNOR)],
1572 true,
1573 ),
1574 ("a broker without the api", Vec::new(), false),
1575 ] {
1576 let recorder = Recorder::default();
1577 let mock = if speaks_heartbeat {
1578 heartbeat_coordinator(answers, &recorder).await
1579 } else {
1580 classic_only_coordinator().await
1581 };
1582 let mut consumer = TestConsumerOptions::builder()
1583 .mock(&mock)
1584 .build()
1585 .start()
1586 .await;
1587 let error = poll_until_error(&mut consumer).await;
1588 consumer.close().await.expect("close");
1589 mock.stop();
1590 actual.push((name, error));
1591 }
1592 assert2::assert!(
1593 actual
1594 == vec![
1595 (
1596 "an unsupported assignor",
1597 Some("broker error_code 112".to_owned()),
1598 ),
1599 (
1600 "a broker without the api",
1601 Some(
1602 "client: incompatible version: broker supports 2..=2, client wants 0..=1 for api_key 68"
1603 .to_owned()
1604 ),
1605 ),
1606 ]
1607 );
1608 }
1609
1610 async fn classic_only_coordinator() -> MockBroker {
1613 coordinator_with_heartbeat_versions(2, 2).await
1614 }
1615
1616 async fn coordinator_with_heartbeat_versions(min_version: i16, max_version: i16) -> MockBroker {
1619 MockBroker::start(move |api_key, version, _corr_id, _body| {
1620 if api_key == api_versions_request::API_KEY {
1621 let mut versions = api_versions();
1622 for api in &mut versions.api_keys {
1623 if api.api_key == consumer_group_heartbeat_request::API_KEY {
1624 api.min_version = min_version;
1625 api.max_version = max_version;
1626 }
1627 }
1628 return Some(encode(&versions, 0));
1629 }
1630 if api_key == find_coordinator_request::API_KEY {
1631 return Some(encode(&FindCoordinatorResponse::default(), version));
1632 }
1633 if api_key == metadata_request::API_KEY {
1634 return Some(flexible(&metadata(), version));
1635 }
1636 None
1637 })
1638 .await
1639 .unwrap()
1640 }
1641
1642 async fn regex_consumer(mock: &MockBroker) -> Consumer {
1643 TestConsumerOptions::builder()
1644 .mock(mock)
1645 .subscription(SubscriptionMode::Regex("orders-.*".into()))
1646 .build()
1647 .start()
1648 .await
1649 }
1650
1651 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1655 async fn a_regular_expression_needs_heartbeat_version_1() {
1656 let mock = coordinator_with_heartbeat_versions(0, 0).await;
1657 let mut consumer = regex_consumer(&mock).await;
1658 let error = poll_until_error(&mut consumer).await;
1659 consumer.close().await.expect("close");
1660 mock.stop();
1661 assert2::assert!(
1662 error
1663 == Some(
1664 "client: incompatible version: broker supports 0..=0, client wants 1..=1 for api_key 68"
1665 .to_owned()
1666 )
1667 );
1668 }
1669
1670 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1674 async fn a_regular_expression_subscription_keeps_fetching_after_a_resubscribe() {
1675 let recorder = Recorder::default();
1676 let mock = heartbeat_coordinator_with(vec![assigns(1, &[0, 1])], &recorder, true).await;
1677 let mut consumer = regex_consumer(&mock).await;
1678 let deadline = Instant::now() + Duration::from_secs(20);
1681 while Instant::now() < deadline {
1682 let _ = consumer.poll(millis(20)).await;
1683 let assigned = consumer.assignment().await;
1684 if !assigned.is_empty() && !consumer.group_fetches(&assigned).await.is_empty() {
1685 break;
1686 }
1687 }
1688 consumer.subscribe_regex("orders-.*").expect("subscribe");
1689 let assigned = consumer.assignment().await;
1690 let mut fetched: Vec<(String, i32)> = consumer
1691 .group_fetches(&assigned)
1692 .await
1693 .into_values()
1694 .flatten()
1695 .flat_map(|(topic, partitions)| {
1696 partitions
1697 .into_iter()
1698 .map(move |(partition, ..)| (topic.clone(), partition))
1699 })
1700 .collect();
1701 fetched.sort();
1702 consumer.close().await.expect("close");
1703 mock.stop();
1704 assert2::assert!(
1705 (assigned, fetched)
1706 == (
1707 vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
1708 vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
1709 )
1710 );
1711 }
1712
1713 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1717 async fn unsubscribe_leaves_the_group_and_a_new_subscription_joins_again() {
1718 let recorder = Recorder::default();
1719 let mock = assigned_coordinator(&recorder).await;
1720 let mut consumer = started_consumer_group(&mock).await;
1721 consumer.unsubscribe().await.expect("unsubscribe");
1722 let deadline = Instant::now() + Duration::from_secs(20);
1723 while !recorder.epochs().contains(&LEAVE_GROUP_MEMBER_EPOCH) && Instant::now() < deadline {
1724 sleep(Duration::from_millis(20)).await;
1725 }
1726 let after_unsubscribe = (consumer.assignment().await, recorder.epochs());
1727 consumer
1728 .subscribe([TOPIC.to_owned()])
1729 .await
1730 .expect("subscribe");
1731 poll_until_assigned(&mut consumer).await;
1732 let joined_again = consumer.assignment().await;
1733 consumer.close().await.expect("close");
1734 mock.stop();
1735 let (left, epochs) = after_unsubscribe;
1736 let joins_after_the_leave = epochs
1737 .iter()
1738 .skip_while(|epoch| **epoch != LEAVE_GROUP_MEMBER_EPOCH)
1739 .filter(|epoch| **epoch == JOIN_GROUP_MEMBER_EPOCH)
1740 .count();
1741 assert2::assert!(
1742 (
1743 left,
1744 epochs.contains(&LEAVE_GROUP_MEMBER_EPOCH),
1745 joins_after_the_leave,
1746 joined_again
1747 ) == (
1748 Vec::new(),
1749 true,
1750 0,
1751 vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
1752 )
1753 );
1754 }
1755
1756 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1761 async fn a_member_that_stops_polling_leaves_and_joins_again() {
1762 let recorder = Recorder::default();
1763 let calls: ListenerCalls = Arc::default();
1764 let mock = assigned_coordinator(&recorder).await;
1765 let mut consumer = TestConsumerOptions::builder()
1766 .mock(&mock)
1767 .max_poll_interval(millis(300))
1768 .rebalance_listener(Box::new(
1769 RecordingListener::builder()
1770 .calls(Arc::clone(&calls))
1771 .build(),
1772 ))
1773 .build()
1774 .start()
1775 .await;
1776 GroupPollWait::builder()
1777 .consumer(&mut consumer)
1778 .target(GroupPollTarget::Assignment)
1779 .build()
1780 .wait()
1781 .await;
1782 sleep(Duration::from_millis(700)).await;
1784 let left = consumer.assignment().await;
1785 let epochs_after_leave = recorder.epochs();
1786 GroupPollWait::builder()
1787 .consumer(&mut consumer)
1788 .target(GroupPollTarget::Assignment)
1789 .build()
1790 .wait()
1791 .await;
1792 let joined_again = consumer.assignment().await;
1793 consumer.close().await.expect("close");
1794 mock.stop();
1795 let lost = first_lost(&calls);
1798 assert2::assert!(
1799 (
1800 left,
1801 epochs_after_leave.contains(&LEAVE_GROUP_MEMBER_EPOCH),
1802 joined_again,
1803 lost,
1804 ) == (
1805 Vec::new(),
1806 true,
1807 vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
1808 Some((rebalance_listener::ListenerCallKind::Lost, vec![0, 1])),
1809 )
1810 );
1811 }
1812
1813 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1816 async fn a_fenced_member_gives_its_partitions_to_the_lost_callback() {
1817 let recorder = Recorder::default();
1818 let calls: ListenerCalls = Arc::default();
1819 let mock = heartbeat_coordinator(
1820 vec![
1821 assigns(1, &[0, 1]),
1822 fails(FENCED_MEMBER_EPOCH),
1823 assigns(1, &[0, 1]),
1824 ],
1825 &recorder,
1826 )
1827 .await;
1828 let mut consumer = TestConsumerOptions::builder()
1829 .mock(&mock)
1830 .rebalance_listener(Box::new(
1831 RecordingListener::builder()
1832 .calls(Arc::clone(&calls))
1833 .build(),
1834 ))
1835 .build()
1836 .start()
1837 .await;
1838 let member_id = consumer.member_id();
1839 let deadline = Instant::now() + Duration::from_secs(20);
1840 while first_lost(&calls).is_none() && Instant::now() < deadline {
1841 let _ = consumer.poll(millis(20)).await;
1842 }
1843 let kept_id = consumer.member_id();
1844 consumer.close().await.expect("close");
1845 mock.stop();
1846 assert2::assert!(
1847 (first_lost(&calls), kept_id == member_id)
1848 == (
1849 Some((rebalance_listener::ListenerCallKind::Lost, vec![0, 1])),
1850 true,
1851 )
1852 );
1853 }
1854
1855 const HEARTBEATS_IN_CALLBACK: usize = 3;
1858
1859 struct HeartbeatWaitingListener {
1863 recorder: Recorder,
1864 in_callback: Arc<Mutex<Option<usize>>>,
1867 }
1868
1869 #[async_trait::async_trait]
1870 impl crate::ConsumerRebalanceListener for HeartbeatWaitingListener {
1871 async fn on_partitions_revoked(
1872 &mut self,
1873 _consumer: &Consumer,
1874 _partitions: &[(String, i32)],
1875 ) -> Result<(), crate::RebalanceListenerError> {
1876 Ok(())
1877 }
1878
1879 async fn on_partitions_assigned(
1880 &mut self,
1881 _consumer: &Consumer,
1882 _partitions: &[(String, i32)],
1883 ) -> Result<(), crate::RebalanceListenerError> {
1884 let before = self.recorder.heartbeats().len();
1885 let bound = Instant::now() + Duration::from_secs(20);
1886 let mut arrived = 0;
1887 while arrived < HEARTBEATS_IN_CALLBACK && Instant::now() < bound {
1888 sleep(Duration::from_millis(5)).await;
1889 arrived = self.recorder.heartbeats().len() - before;
1890 }
1891 *self.in_callback.lock().expect("in callback lock") = Some(arrived);
1892 Ok(())
1893 }
1894 }
1895
1896 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1904 async fn heartbeats_go_on_during_a_rebalance_callback() {
1905 let recorder = Recorder::default();
1906 let in_callback: Arc<Mutex<Option<usize>>> = Arc::default();
1907 let mock = assigned_coordinator(&recorder).await;
1908 let mut consumer = TestConsumerOptions::builder()
1909 .mock(&mock)
1910 .rebalance_listener(Box::new(HeartbeatWaitingListener {
1911 recorder: recorder.clone(),
1912 in_callback: Arc::clone(&in_callback),
1913 }))
1914 .build()
1915 .start()
1916 .await;
1917 let deadline = Instant::now() + Duration::from_secs(30);
1918 while in_callback.lock().expect("in callback lock").is_none() && Instant::now() < deadline {
1919 let _ = consumer.poll(millis(200)).await;
1920 }
1921 let during_callback = *in_callback.lock().expect("in callback lock");
1922 consumer.close().await.expect("close");
1923 mock.stop();
1924 let classic = *recorder
1925 .classic_heartbeats
1926 .lock()
1927 .expect("classic heartbeats lock");
1928 assert2::assert!(
1929 (
1930 during_callback.map(|arrived| arrived >= HEARTBEATS_IN_CALLBACK),
1931 classic
1932 ) == (Some(true), 0)
1933 );
1934 }
1935
1936 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1939 async fn the_member_id_and_a_new_subscription_reach_the_consumer() {
1940 let recorder = Recorder::default();
1941 let mock =
1942 heartbeat_coordinator(vec![named(assigns(1, &[0, 1]), "server-id")], &recorder).await;
1943 let mut consumer = started_consumer_group(&mock).await;
1944 let member_id = consumer.member_id();
1945 consumer
1946 .subscribe([TOPIC.to_owned(), "payments".to_owned()])
1947 .await
1948 .expect("subscribe");
1949 let deadline = Instant::now() + Duration::from_secs(20);
1950 while !recorder
1951 .metadata_topics
1952 .lock()
1953 .expect("metadata topics lock")
1954 .iter()
1955 .any(|topic| topic == "payments")
1956 && Instant::now() < deadline
1957 {
1958 let _ = consumer.poll(millis(20)).await;
1959 }
1960 let asked_for_payments = recorder
1961 .metadata_topics
1962 .lock()
1963 .expect("metadata topics lock")
1964 .iter()
1965 .any(|topic| topic == "payments");
1966 consumer.close().await.expect("close");
1967 mock.stop();
1968 assert2::assert!((member_id, asked_for_payments) == ("server-id".to_owned(), true));
1969 }
1970
1971 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1976 async fn the_member_reconciles_and_leaves_as_kafka_does() {
1977 struct Case {
1978 name: &'static str,
1979 answers: Vec<Answer>,
1980 instance_id: Option<&'static str>,
1981 acks: usize,
1983 }
1984 let cases = [
1985 Case {
1986 name: "an assignment",
1987 answers: vec![assigns(1, &[0, 1])],
1988 instance_id: None,
1989 acks: 2,
1990 },
1991 Case {
1992 name: "an assignment that revokes a partition",
1993 answers: vec![assigns(1, &[0, 1]), assigns(2, &[0])],
1994 instance_id: None,
1995 acks: 3,
1996 },
1997 Case {
1998 name: "a fenced member",
1999 answers: vec![
2000 assigns(1, &[0, 1]),
2001 fails(FENCED_MEMBER_EPOCH),
2002 assigns(1, &[0, 1]),
2003 ],
2004 instance_id: None,
2005 acks: 4,
2006 },
2007 Case {
2008 name: "a static member",
2009 answers: vec![assigns(1, &[0, 1])],
2010 instance_id: Some("member-a"),
2011 acks: 2,
2012 },
2013 ];
2014 let mut actual = Vec::new();
2015 for case in cases {
2016 actual.push((
2017 case.name,
2018 run_member(case.answers, case.instance_id, case.acks).await,
2019 ));
2020 }
2021 assert2::assert!(
2022 actual
2023 == vec![
2024 (
2025 "an assignment",
2026 Observed {
2027 acknowledged: vec![vec![], vec![0, 1], vec![]],
2031 owned: vec![0, 1],
2032 last_epoch: LEAVE_GROUP_MEMBER_EPOCH,
2033 rejoins: 0,
2034 },
2035 ),
2036 (
2037 "an assignment that revokes a partition",
2038 Observed {
2039 acknowledged: vec![vec![], vec![0, 1], vec![0], vec![]],
2040 owned: vec![0],
2041 last_epoch: LEAVE_GROUP_MEMBER_EPOCH,
2042 rejoins: 0,
2043 },
2044 ),
2045 (
2046 "a fenced member",
2047 Observed {
2048 acknowledged: vec![vec![], vec![0, 1], vec![], vec![0, 1], vec![]],
2051 owned: vec![0, 1],
2052 last_epoch: LEAVE_GROUP_MEMBER_EPOCH,
2053 rejoins: 1,
2054 },
2055 ),
2056 (
2057 "a static member",
2058 Observed {
2059 acknowledged: vec![vec![], vec![0, 1], vec![]],
2060 owned: vec![0, 1],
2061 last_epoch: LEAVE_GROUP_STATIC_MEMBER_EPOCH,
2064 rejoins: 0,
2065 },
2066 ),
2067 ]
2068 );
2069 }
2070}
2071
2072#[cfg(test)]
2073mod tests {
2074 use consumer::GroupMembershipOperation;
2075 use krabka_protocol::owned::{
2076 common::consumer_group_heartbeat_response::topic_partitions::TopicPartitions as ResponseTopicPartitions,
2077 consumer_group_heartbeat_response::Assignment,
2078 };
2079
2080 use super::*;
2081
2082 const TOPIC_ID: WireUuid = WireUuid([7; 16]);
2083
2084 fn fields<'a>(member_epoch: i32, joining: bool, owned: OwnedByTopicId) -> HeartbeatFields<'a> {
2085 HeartbeatFields {
2086 group_id: "group-a",
2087 member_id: "member-a",
2088 member_epoch,
2089 instance_id: None,
2090 rack_id: Some("az-1"),
2091 rebalance_timeout_ms: 300_000,
2092 subscribed_topic_names: vec!["orders".to_owned()],
2093 regex: None,
2094 server_assignor: Some("uniform"),
2095 owned,
2096 joining,
2097 }
2098 }
2099
2100 #[test]
2103 fn a_heartbeat_repeats_only_the_fields_that_changed() {
2104 let mut sent = SentFields::default();
2105 let join = build_heartbeat(&fields(0, true, BTreeMap::new()), &mut sent);
2106 let steady = build_heartbeat(&fields(1, false, BTreeMap::new()), &mut sent);
2107 let owned = build_heartbeat(
2108 &fields(1, false, maplit::btreemap! {TOPIC_ID.0 => vec![0, 1]}),
2109 &mut sent,
2110 );
2111 assert2::assert!(
2112 (
2113 join.rack_id.clone(),
2114 join.rebalance_timeout_ms,
2115 join.subscribed_topic_names.clone(),
2116 join.server_assignor.clone(),
2117 join.topic_partitions,
2118 ) == (
2119 Some("az-1".to_owned()),
2120 300_000,
2121 Some(vec!["orders".to_owned()]),
2122 Some("uniform".to_owned()),
2123 Some(Vec::new()),
2124 )
2125 );
2126 assert2::assert!(
2127 (
2128 steady.rack_id,
2129 steady.rebalance_timeout_ms,
2130 steady.subscribed_topic_names,
2131 steady.server_assignor,
2132 steady.topic_partitions,
2133 ) == (None, -1, None, None, None)
2134 );
2135 assert2::assert!(
2136 owned.topic_partitions
2137 == Some(vec![TopicPartitions {
2138 topic_id: TOPIC_ID,
2139 partitions: vec![0, 1],
2140 ..Default::default()
2141 }])
2142 );
2143 }
2144
2145 #[test]
2149 fn a_heartbeat_carries_the_regular_expression_of_the_subscription() {
2150 let with_regex = |regex: Option<&'static str>, joining: bool| HeartbeatFields {
2151 subscribed_topic_names: Vec::new(),
2152 regex: regex.map(str::to_owned),
2153 ..fields(i32::from(!joining), joining, BTreeMap::new())
2154 };
2155 let mut sent = SentFields::default();
2156 let join = build_heartbeat(&with_regex(Some("orders-.*"), true), &mut sent);
2157 let steady = build_heartbeat(&with_regex(Some("orders-.*"), false), &mut sent);
2158 let changed = build_heartbeat(&with_regex(Some("orders-eu-.*"), false), &mut sent);
2159 let removed = build_heartbeat(&with_regex(None, false), &mut sent);
2160 let actual = [join, steady, changed, removed].map(|request| {
2161 (
2162 request.subscribed_topic_regex,
2163 request.subscribed_topic_names,
2164 )
2165 });
2166 assert2::assert!(
2167 actual
2168 == [
2169 (Some("orders-.*".to_owned()), Some(Vec::new())),
2170 (None, None),
2171 (Some("orders-eu-.*".to_owned()), None),
2172 (Some(String::new()), None),
2173 ]
2174 );
2175 }
2176
2177 #[test]
2179 fn heartbeat_errors_follow_kafkas_classes() {
2180 let actual = [0, 14, 15, 16, 25, 110, 111, 112, 81, 42, 35, 30].map(heartbeat_action);
2181 assert2::assert!(
2182 actual
2183 == [
2184 HeartbeatAction::Ok,
2185 HeartbeatAction::Retry,
2186 HeartbeatAction::FindCoordinator,
2187 HeartbeatAction::FindCoordinator,
2188 HeartbeatAction::Fenced,
2189 HeartbeatAction::Fenced,
2190 HeartbeatAction::Fatal,
2191 HeartbeatAction::Fatal,
2192 HeartbeatAction::Fatal,
2193 HeartbeatAction::Fatal,
2194 HeartbeatAction::Fatal,
2195 HeartbeatAction::Fatal,
2196 ]
2197 );
2198 }
2199
2200 #[test]
2202 fn a_leaving_member_follows_kafkas_epochs() {
2203 let mut actual = Vec::new();
2204 let mut wanted = Vec::new();
2205 for (name, instance_id, operation, expected) in [
2206 (
2207 "dynamic, default",
2208 None,
2209 GroupMembershipOperation::Default,
2210 (-1, true),
2211 ),
2212 (
2213 "static, default",
2214 Some("i-1"),
2215 GroupMembershipOperation::Default,
2216 (-2, true),
2217 ),
2218 (
2219 "static, leave group",
2220 Some("i-1"),
2221 GroupMembershipOperation::LeaveGroup,
2222 (-1, true),
2223 ),
2224 (
2225 "dynamic, remain in group",
2226 None,
2227 GroupMembershipOperation::RemainInGroup,
2228 (-1, false),
2229 ),
2230 (
2231 "static, remain in group",
2232 Some("i-1"),
2233 GroupMembershipOperation::RemainInGroup,
2234 (-2, true),
2235 ),
2236 ] {
2237 actual.push((
2238 name,
2239 (
2240 leave_group_epoch(instance_id, operation),
2241 should_send_leave_heartbeat("member-a", instance_id, operation),
2242 ),
2243 ));
2244 wanted.push((name, expected));
2245 }
2246 assert2::assert!(actual == wanted);
2247 }
2248
2249 #[test]
2252 fn an_assignment_resolves_the_topic_ids_of_the_metadata() {
2253 let other = WireUuid([9; 16]);
2254 let response = ConsumerGroupHeartbeatResponse {
2255 assignment: Some(Assignment {
2256 topic_partitions: vec![
2257 ResponseTopicPartitions {
2258 topic_id: TOPIC_ID,
2259 partitions: vec![1, 0],
2260 ..Default::default()
2261 },
2262 ResponseTopicPartitions {
2263 topic_id: other,
2264 partitions: vec![0],
2265 ..Default::default()
2266 },
2267 ],
2268 ..Default::default()
2269 }),
2270 ..Default::default()
2271 };
2272 let names = maplit::hashmap! {TOPIC_ID => "orders".to_owned()};
2273 assert2::assert!(
2274 assignment_partitions(&response, &names)
2275 == Some((
2276 vec![("orders".to_owned(), 0), ("orders".to_owned(), 1)],
2277 maplit::hashset! {other}
2278 ))
2279 );
2280 assert2::assert!(
2281 assignment_partitions(&ConsumerGroupHeartbeatResponse::default(), &names) == None
2282 );
2283 }
2284}