Skip to main content

krabka_client_consumer/
consumer_group.rs

1//! The consumer group protocol of KIP-848: `ConsumerGroupHeartbeat`.
2//!
3//! With `group_protocol(GroupProtocol::Consumer)` the member sends
4//! `ConsumerGroupHeartbeat` instead of `JoinGroup`, `SyncGroup` and
5//! `Heartbeat`. The group coordinator assigns the partitions and the member
6//! reconciles them: it gives up the partitions that the assignment no longer
7//! holds, takes the new ones, and acknowledges the assignment with the next
8//! heartbeat. This is Kafka's `ConsumerHeartbeatRequestManager` and
9//! `ConsumerMembershipManager`.
10
11use 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/// The group protocol of a consumer. Kafka's `group.protocol`.
36#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
37pub enum GroupProtocol {
38    /// `JoinGroup`, `SyncGroup` and `Heartbeat`, with client-side assignment.
39    #[default]
40    Classic,
41    /// `ConsumerGroupHeartbeat` with server-side assignment (KIP-848).
42    Consumer,
43}
44
45/// Kafka's `ConsumerGroupHeartbeatRequest.JOIN_GROUP_MEMBER_EPOCH`.
46pub const JOIN_GROUP_MEMBER_EPOCH: i32 = 0;
47/// Kafka's `ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH`.
48pub const LEAVE_GROUP_MEMBER_EPOCH: i32 = -1;
49/// Kafka's `ConsumerGroupHeartbeatRequest.LEAVE_GROUP_STATIC_MEMBER_EPOCH`.
50pub const LEAVE_GROUP_STATIC_MEMBER_EPOCH: i32 = -2;
51
52/// `FENCED_MEMBER_EPOCH`.
53pub const FENCED_MEMBER_EPOCH: i16 = 110;
54/// `UNKNOWN_MEMBER_ID`.
55const UNKNOWN_MEMBER_ID: i16 = 25;
56/// `UNSUPPORTED_ASSIGNOR`.
57#[cfg(test)]
58const UNSUPPORTED_ASSIGNOR: i16 = 112;
59
60/// Kafka's message for a regular expression subscription of a classic member:
61/// `ClassicKafkaConsumer.subscribe(SubscriptionPattern)`.
62pub 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
64/// The partitions of an assignment, and the topic ids that the metadata does
65/// not name yet.
66pub type ResolvedAssignment = (Vec<(String, i32)>, HashSet<WireUuid>);
67
68/// The partitions that a member owns, by topic id and sorted. The wire uuid
69/// has no order, so the map is keyed by its bytes.
70pub type OwnedByTopicId = BTreeMap<[u8; 16], Vec<i32>>;
71
72/// The fields of the last `ConsumerGroupHeartbeat`. Kafka's
73/// `ConsumerHeartbeatRequestManager.HeartbeatState.SentFields`: a field goes
74/// out again only when it changed, or when the member joins.
75#[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/// What the member sends in its next heartbeat.
85#[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    /// The subscribed topics, sorted. A regular expression subscription sends
94    /// none, because the coordinator matches the expression.
95    pub subscribed_topic_names: Vec<String>,
96    /// The regular expression of a broker-side pattern subscription.
97    pub regex: Option<String>,
98    pub server_assignor: Option<&'a str>,
99    /// The partitions that the member owns, by topic id.
100    pub owned: OwnedByTopicId,
101    /// `true` for the first heartbeat of a join, which sends every field.
102    pub joining: bool,
103}
104
105/// Build the next `ConsumerGroupHeartbeat` and record the fields that it
106/// carries.
107pub 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        // Kafka sends the rack only when the member joins.
121        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    // Kafka sends an empty expression to remove a pattern subscription, and
135    // sends the expression again only when it changed or when the member
136    // joins with one.
137    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/// What the member does with the error code of a heartbeat response. Kafka's
166/// `HeartbeatRequestManager.onErrorResponse`.
167#[derive(Clone, Copy, Debug, Eq, PartialEq)]
168pub enum HeartbeatAction {
169    /// The response carries the member state.
170    Ok,
171    /// Send the heartbeat again after the retry backoff.
172    Retry,
173    /// Find the coordinator again, then send the heartbeat again.
174    FindCoordinator,
175    /// The coordinator does not know the member any more: give up the
176    /// partitions and join again with epoch 0.
177    Fenced,
178    /// The application must see the error.
179    Fatal,
180}
181
182/// Map the error code of a `ConsumerGroupHeartbeat` response, as Kafka's
183/// `ConsumerHeartbeatRequestManager` does.
184///
185/// `COORDINATOR_NOT_AVAILABLE (15)` and `NOT_COORDINATOR (16)` make the member
186/// find the coordinator again. `COORDINATOR_LOAD_IN_PROGRESS (14)` retries.
187/// `FENCED_MEMBER_EPOCH (110)` and `UNKNOWN_MEMBER_ID (25)` make the member
188/// join again from epoch 0. `UNSUPPORTED_ASSIGNOR (112)`,
189/// `UNRELEASED_INSTANCE_ID (111)`, `GROUP_MAX_SIZE_REACHED (81)`,
190/// `INVALID_REQUEST (42)`, `UNSUPPORTED_VERSION (35)` and the authorization
191/// codes are fatal.
192pub 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        // `UNSUPPORTED_ASSIGNOR (112)`, `UNRELEASED_INSTANCE_ID (111)`,
201        // `GROUP_MAX_SIZE_REACHED (81)`, `INVALID_REQUEST (42)`,
202        // `UNSUPPORTED_VERSION (35)`, the authorization codes and every other
203        // code reach the application.
204        _ => HeartbeatAction::Fatal,
205    }
206}
207
208/// The epoch of the heartbeat that leaves the group. Kafka's
209/// `ConsumerMembershipManager.leaveGroupEpoch`: a static member that closes
210/// with the default operation keeps its place in the group with `-2`.
211pub 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
225/// Whether a closing member sends the heartbeat that leaves the group. Kafka's
226/// `ConsumerHeartbeatRequestManager.shouldSendLeaveHeartbeatNow`: a dynamic
227/// member that closes with `REMAIN_IN_GROUP` sends none and waits for the
228/// session timeout.
229pub 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
240/// The partitions of an assignment, by topic name.
241///
242/// A topic id that the metadata does not name yet stays unresolved, and the
243/// member reconciles it after the next metadata refresh, as Kafka's
244/// `AbstractMembershipManager` keeps unresolved assignments.
245pub 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
269/// The owned partitions by topic id, for the next heartbeat.
270pub 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
286/// Run the member of a consumer group with the KIP-848 protocol.
287///
288/// The task sends a `ConsumerGroupHeartbeat` at once and then each heartbeat
289/// interval that the coordinator returns. Each response can carry a new
290/// assignment, which the task reconciles before the next heartbeat
291/// acknowledges it. Kafka's `ConsumerHeartbeatRequestManager` and
292/// `ConsumerMembershipManager` do the same.
293pub 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    // Kafka's `Heartbeat.pollTimer`: `poll` resets it, and the member leaves
298    // the group when it expires.
299    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    // `true` after the poll timer expired: the member left the group and waits
304    // for a `poll` before it joins again. Kafka's `MemberState.STALE`.
305    let mut left_on_poll_timeout = false;
306    // `true` after `unsubscribe`: the member left the group and waits for a
307    // new subscription. Kafka's `AbstractMembershipManager.leaveGroup`.
308    let mut unsubscribed = false;
309    // `true` after a fatal error: Kafka's `transitionToSendingLeaveGroup`
310    // sends no leave heartbeat in the fatal state.
311    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                    // The `poll` came back: join the group again, as Kafka's
347                    // member leaves the stale state when the poll timer is
348                    // reset.
349                    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                        // A new subscription joins the group again.
371                        unsubscribed = false;
372                        state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
373                        state.sent_heartbeat_fields = SentFields::default();
374                    }
375                    // The task has its own client. Its metadata requests name
376                    // the topics of the new subscription, so it can resolve
377                    // the topic ids of the next assignment.
378                    if !subscription.topics.is_empty() {
379                        state.client.metadata_topics().set(subscription.topics);
380                    }
381                }
382                TaskEvent::Heartbeat => {}
383            }
384        }
385        // Kafka's `AsyncKafkaConsumer.updateAssignmentMetadataIfNeeded` matches
386        // a client-side pattern against the metadata in each poll. The member
387        // sends the topics that matched, and the coordinator assigns them.
388        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                // Kafka acknowledges a new assignment with the next heartbeat,
410                // which it sends at once.
411                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    // A member that already left the group, or that stopped for a fatal
428    // error, sends no leave heartbeat. Kafka's
429    // `AbstractMembershipManager.leaveGroup` returns at once when the member
430    // is not in the group.
431    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/// What wakes the heartbeat task.
449#[derive(Clone, Copy, Debug, Eq, PartialEq)]
450enum TaskEvent {
451    /// The heartbeat interval elapsed.
452    Heartbeat,
453    /// A `poll` came: reset the poll timer.
454    Poll,
455    /// The subscription changed.
456    Subscription,
457    /// The poll timer expired: leave the group.
458    PollTimeout,
459}
460
461/// Leave the group because the time between two `poll` calls was longer than
462/// `max_poll_interval`, and give up every partition.
463///
464/// Kafka's `AbstractHeartbeatRequestManager.poll` calls
465/// `transitionToSendingLeaveGroup(true)` for an expired poll timer, sends the
466/// leave heartbeat and then makes the member stale until the next `poll`.
467async 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
481/// Leave the group because the application unsubscribed. Kafka's
482/// `AbstractMembershipManager.leaveGroup`: the member sends a heartbeat with
483/// its leave epoch and gives up its partitions. The `Consumer` already ran the
484/// revoke callback and cleared the assignment.
485async 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
503/// Keep the member in the group while `poll` runs a rebalance listener
504/// callback. Kafka's heartbeat manager keeps sending `ConsumerGroupHeartbeat`
505/// while the application thread runs the callback.
506///
507/// The request carries only the fields that changed, so it repeats the member
508/// state and asks for nothing new. The heartbeat after the callback acts on
509/// the response of the member state.
510pub 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/// What the task does after one heartbeat.
544#[derive(Clone, Copy, Debug, Eq, PartialEq)]
545enum HeartbeatOutcome {
546    /// Wait this interval for the next heartbeat.
547    Interval(Duration),
548    /// A new assignment was reconciled: acknowledge it at once, then wait this
549    /// interval.
550    Reconciled(Duration),
551    /// Wait the retry backoff, then send again.
552    Backoff,
553    /// Stop the task.
554    Stop,
555}
556
557/// Send one heartbeat and act on the response.
558async 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            // The member does not know what reached the coordinator, so the
569            // next heartbeat sends every field again. Kafka's
570            // `ConsumerHeartbeatRequestManager.resetHeartbeatState`.
571            state.sent_heartbeat_fields = SentFields::default();
572            if !coordinator::is_retriable_transport_error(&error) {
573                // Kafka's `AbstractHeartbeatRequestManager.onFailure` retries
574                // only a `RetriableException`, and hands every other failure
575                // to the application. A broker without api key 68 fails here.
576                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        // Kafka's `AbstractHeartbeatRequestManager.onErrorResponse` resets the
596        // sent fields, so the heartbeat after an error carries them again.
597        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            // Kafka's `transitionToFenced` gives the partitions to
615            // `onPartitionsLost` before the member joins again.
616            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        // The assignment names its topics by id, so the member needs the ids
647        // of the metadata. Kafka's `Metadata` holds them for the consumer, and
648        // `ConsumerMembershipManager` keeps an unresolved assignment until a
649        // metadata update names its topics.
650        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    // The topics of a regular expression subscription come from the
664    // assignment, because the coordinator matches the expression. They give
665    // `poll` the topics of its metadata and its fetch filter.
666    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        // Kafka acknowledges only a reconciliation that moved partitions. An
672        // assignment that repeats what the member owns goes out again with the
673        // next interval heartbeat.
674        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/// Which heartbeat the member sends. Kafka sends every field only while the
684/// member joins (`MemberState.JOINING`), and a member that leaves gives up its
685/// assignment before the request
686/// (`AbstractMembershipManager.transitionToSendingLeaveGroup` sets
687/// `currentAssignment` to `NONE`), so the leave heartbeat carries an empty
688/// partition list.
689#[derive(Clone, Copy, Debug, Eq, PartialEq)]
690enum HeartbeatKind {
691    /// The first heartbeat of a member that joins the group.
692    Join,
693    /// A heartbeat of a member of the group.
694    Interval,
695    /// The heartbeat that leaves the group.
696    Leave,
697}
698
699/// Send the heartbeat of the current member state.
700///
701/// The fields that the request carried stay in `sent_heartbeat_fields`, so the
702/// next heartbeat sends only what changed. Kafka's
703/// `ConsumerHeartbeatRequestManager.HeartbeatState.buildRequestData` keeps them
704/// in the same way.
705async 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
723/// The heartbeat fields of the current member state.
724async 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        // A regular expression subscription names no topics: Kafka's
737        // `SubscriptionState.subscription()` is empty for
738        // `AUTO_PATTERN_RE2J`, and the coordinator matches the expression.
739        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
761/// A `ConsumerGroupHeartbeat` that carries `subscribed_topic_regex`. The
762/// field exists from version 1, so a broker that speaks only version 0 must
763/// fail the call in place of dropping the expression.
764struct 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
789/// One `ConsumerGroupHeartbeat` that is ready to go out.
790struct HeartbeatCall<'a> {
791    broker: krabka_client_core::BrokerHandle<'a>,
792    request: ConsumerGroupHeartbeatRequest,
793    /// `true` when the member subscribes to a regular expression, which needs
794    /// version 1.
795    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
807/// Store the topic ids of the metadata, so the member can resolve the topics
808/// of its assignment.
809///
810/// A regular expression subscription knows no topic names, so the request
811/// names the `unresolved` topic ids. Kafka's
812/// `ConsumerMetadata.newMetadataRequestBuilder` asks for the assigned topic
813/// ids of such a subscription, and for the topic names of every other one.
814async 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
844/// The heartbeat interval of a response, or the configured one when the
845/// response names none.
846fn 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
854/// Reconcile `target` with the partitions that the member owns. Returns
855/// `true` when the member gave up or took a partition.
856///
857/// Kafka's `AbstractMembershipManager.reconcile` gives up the partitions that
858/// the assignment no longer holds, with `onPartitionsRevoked`, and then takes
859/// the new ones with `onPartitionsAssigned`.
860async 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    /// One `ConsumerGroupHeartbeat` answer of the mock coordinator.
956    #[derive(Clone, Debug)]
957    struct Answer {
958        error_code: i16,
959        member_epoch: i32,
960        /// The assigned partitions of `orders`, or `None` for a response
961        /// without an assignment.
962        assignment: Option<Vec<i32>>,
963        /// The member id of the response. The coordinator can name another one
964        /// than the request.
965        member_id: Option<String>,
966    }
967
968    /// An answer that assigns `partitions` at `member_epoch`.
969    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    /// `answer` with the member id that the coordinator names.
979    fn named(answer: Answer, member_id: &str) -> Answer {
980        Answer {
981            member_id: Some(member_id.to_owned()),
982            ..answer
983        }
984    }
985
986    /// An answer that fails with `error_code`.
987    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    /// What the mock coordinator recorded.
997    #[derive(Clone, Default)]
998    struct Recorder {
999        /// Each `ConsumerGroupHeartbeat` that the mock decoded.
1000        heartbeats: Arc<Mutex<Vec<ConsumerGroupHeartbeatRequest>>>,
1001        /// The topic names of each `Metadata` request.
1002        metadata_topics: Arc<Mutex<Vec<String>>>,
1003        /// Whether a `Metadata` request named a topic by id.
1004        asked_by_id: Arc<Mutex<bool>>,
1005        /// The classic `Heartbeat` requests. A member of the consumer group
1006        /// protocol sends none.
1007        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        /// The owned partitions that each heartbeat carried.
1016        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    /// The response of a flexible api, after the tagged fields of its header.
1042    /// `ApiVersions` keeps the header of version 0, as Kafka does.
1043    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    /// The `ApiVersions` of a coordinator that speaks `ConsumerGroupHeartbeat`.
1050    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    /// The metadata of `orders`, which names the topic id of the assignment.
1063    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    /// The log start and end of `orders`: both partitions are empty.
1076    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    /// The committed offsets of the group: `orders` has none, so the member
1100    /// starts its partitions at the `auto_offset_reset`.
1101    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    /// The `Metadata` request of a frame, without its header.
1124    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    /// The `ConsumerGroupHeartbeat` request of a frame, without its header.
1134    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        // The flexible request header ends with its tagged fields.
1138        body.advance(1);
1139        ConsumerGroupHeartbeatRequest::decode(&mut body, version).expect("decode")
1140    }
1141
1142    /// A coordinator that answers each `ConsumerGroupHeartbeat` with the next
1143    /// entry of `answers`, repeats the last entry after that, and records every
1144    /// request that it decoded.
1145    async fn heartbeat_coordinator(answers: Vec<Answer>, recorder: &Recorder) -> MockBroker {
1146        heartbeat_coordinator_with(answers, recorder, false).await
1147    }
1148
1149    /// The same coordinator. With `by_id_only` it names the topic only for a
1150    /// `Metadata` request that asks for its id, as a broker does for a group
1151    /// with a broker-side pattern subscription.
1152    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    /// Kafka's `Heartbeat` api key, which the classic protocol uses.
1253    const CLASSIC_HEARTBEAT_API_KEY: i16 = 12;
1254
1255    /// What one case observed of the member.
1256    #[derive(Debug, Eq, PartialEq)]
1257    struct Observed {
1258        /// The owned partitions that the heartbeats carried. A heartbeat that
1259        /// repeats the field of the heartbeat before it carries none, so the
1260        /// list holds one entry for each change.
1261        acknowledged: Vec<Vec<i32>>,
1262        /// The partitions of `orders` that the member owned before it closed.
1263        owned: Vec<i32>,
1264        /// The member epoch of the last heartbeat.
1265        last_epoch: i32,
1266        /// How many heartbeats after the first one joined the group again.
1267        rejoins: usize,
1268    }
1269
1270    /// The owned partitions that each heartbeat of `requests` carried.
1271    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    /// Run one member against `answers` until its heartbeats carried `acks`
1317    /// partition lists, then close it.
1318    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    /// The first heartbeat of a member joins with epoch 0 and carries every
1457    /// field: Kafka's `ConsumerHeartbeatRequestManager.HeartbeatState`.
1458    #[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        // Kafka's `AsyncKafkaConsumer` generates a v4 uuid for the member id.
1484        assert2::assert!(uuid::Uuid::parse_str(&first.member_id).is_ok());
1485    }
1486
1487    /// A broker-side pattern subscription sends the regular expression and no
1488    /// topic names, and the member learns the topic of its assignment from a
1489    /// metadata request that names the topic id. Kafka's
1490    /// `subscribe(SubscriptionPattern)` and
1491    /// `ConsumerMetadata.newMetadataRequestBuilder`.
1492    #[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    /// A client-side pattern subscription matches the metadata itself, and the
1523    /// member sends the topics that matched. Kafka's
1524    /// `AsyncKafkaConsumer.updateAssignmentMetadataIfNeeded`.
1525    #[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    /// The first lost callback of `calls`.
1552    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    /// A fatal heartbeat response and a heartbeat that cannot go out both end
1562    /// the member and reach the application through `poll`. Kafka's
1563    /// `AbstractHeartbeatRequestManager.onErrorResponse` and `onFailure` hand
1564    /// such an error to the application thread.
1565    #[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    /// A broker whose `ConsumerGroupHeartbeat` versions do not overlap the
1611    /// ones of the client.
1612    async fn classic_only_coordinator() -> MockBroker {
1613        coordinator_with_heartbeat_versions(2, 2).await
1614    }
1615
1616    /// A coordinator that speaks `ConsumerGroupHeartbeat` in this version
1617    /// range and answers no heartbeat.
1618    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    /// A regular expression subscription needs `ConsumerGroupHeartbeat` v1,
1652    /// which added `subscribed_topic_regex`. A coordinator that speaks only
1653    /// v0 fails the call, because a v0 request would name no topic at all.
1654    #[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    /// The member fetches the partitions of a regular expression assignment
1671    /// also after the application subscribes to the same expression again,
1672    /// because the topics of such a subscription come from the assignment.
1673    #[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        // Poll until the member owns the partitions and their positions are
1679        // ready, so that the fetch of the next poll would go out.
1680        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    /// `unsubscribe` makes the member leave the group with its leave epoch and
1714    /// give up its partitions, and a new `subscribe` joins the group again.
1715    /// Kafka's `AbstractMembershipManager.leaveGroup`.
1716    #[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    /// The member leaves the group when the time between two `poll` calls is
1757    /// longer than `max_poll_interval`, gives its partitions to the lost
1758    /// callback, and joins again with the next `poll`. Kafka's
1759    /// `AbstractHeartbeatRequestManager.poll` and `transitionToStale`.
1760    #[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        // The application stops polling for longer than `max_poll_interval`.
1783        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        // A `poll` that ends before its callback ran makes the next `poll`
1796        // run the callback again, so the list can hold more than one call.
1797        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    /// A fenced member gives its partitions to the lost callback, and the
1814    /// member keeps its id. Kafka's `transitionToFenced`.
1815    #[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    /// The heartbeats that reach the coordinator while the assign callback of
1856    /// [`HeartbeatWaitingListener`] runs.
1857    const HEARTBEATS_IN_CALLBACK: usize = 3;
1858
1859    /// A rebalance listener whose assign callback returns only after
1860    /// [`HEARTBEATS_IN_CALLBACK`] more heartbeats reached the coordinator, or
1861    /// after a bound. It records how many arrived.
1862    struct HeartbeatWaitingListener {
1863        recorder: Recorder,
1864        /// The heartbeats that arrived during the assign callback, once it
1865        /// returned.
1866        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    /// The member keeps its membership with `ConsumerGroupHeartbeat` while a
1897    /// rebalance listener callback runs, and never sends the classic
1898    /// `Heartbeat`.
1899    ///
1900    /// The coordinator task publishes the assignment before it queues the
1901    /// assign callback, so a non-empty `assignment()` does not mean that the
1902    /// callback ran. The test waits for the callback itself.
1903    #[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    /// The member id of the coordinator reaches the consumer, and a new
1937    /// subscription reaches the metadata of the heartbeat task.
1938    #[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    /// The member reconciles each assignment of the coordinator, acknowledges
1972    /// it with the partitions that it owns, gives up every partition when the
1973    /// coordinator fences it, and leaves the group with the epoch of its
1974    /// membership.
1975    #[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            /// How many partition lists the heartbeats carry before the close.
1982            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                            // The join carries no partitions, the next
2028                            // heartbeat acknowledges the assignment, and the
2029                            // heartbeat that leaves gives the partitions up.
2030                            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                            // The fence releases the partitions, and the member
2049                            // joins again with epoch 0 and no partitions.
2050                            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                            // A static member that closes with the default
2062                            // operation keeps its place in the group.
2063                            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    /// Kafka's `HeartbeatState.buildRequestData` sends every field when the
2101    /// member joins, and after that only the fields that changed.
2102    #[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    /// Kafka's `HeartbeatState.buildRequestData` sends the regular expression
2146    /// of a broker-side pattern subscription when the member joins with one
2147    /// and when it changed, and sends an empty expression to remove it.
2148    #[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    /// Kafka's heartbeat error classes.
2178    #[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    /// Kafka's `leaveGroupEpoch` and `shouldSendLeaveHeartbeatNow`.
2201    #[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    /// The assignment names its topics by id. An id that the metadata does not
2250    /// name yet stays unresolved.
2251    #[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}