krabka-client-consumer 0.7.1

//! The consumer group protocol of KIP-848: `ConsumerGroupHeartbeat`.
//!
//! With `group_protocol(GroupProtocol::Consumer)` the member sends
//! `ConsumerGroupHeartbeat` instead of `JoinGroup`, `SyncGroup` and
//! `Heartbeat`. The group coordinator assigns the partitions and the member
//! reconciles them: it gives up the partitions that the assignment no longer
//! holds, takes the new ones, and acknowledges the assignment with the next
//! heartbeat. This is Kafka's `ConsumerHeartbeatRequestManager` and
//! `ConsumerMembershipManager`.

use std::collections::{BTreeMap, HashMap, HashSet};

use krabka_protocol::{
    owned::{
        consumer_group_heartbeat_request,
        consumer_group_heartbeat_request::{ConsumerGroupHeartbeatRequest, TopicPartitions},
        consumer_group_heartbeat_response::ConsumerGroupHeartbeatResponse,
        metadata_request::{MetadataRequest, MetadataRequestTopic},
    },
    primitives::uuid::Uuid as WireUuid,
};
use krabka_units::convert::TimeExt as _;

use crate::{consumer, coordinator, error, offset_wire, rebalance_listener};

/// The group protocol of a consumer. Kafka's `group.protocol`.
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum GroupProtocol {
    /// `JoinGroup`, `SyncGroup` and `Heartbeat`, with client-side assignment.
    #[default]
    Classic,
    /// `ConsumerGroupHeartbeat` with server-side assignment (KIP-848).
    Consumer,
}

/// Kafka's `ConsumerGroupHeartbeatRequest.JOIN_GROUP_MEMBER_EPOCH`.
pub(crate) const JOIN_GROUP_MEMBER_EPOCH: i32 = 0;
/// Kafka's `ConsumerGroupHeartbeatRequest.LEAVE_GROUP_MEMBER_EPOCH`.
pub(crate) const LEAVE_GROUP_MEMBER_EPOCH: i32 = -1;
/// Kafka's `ConsumerGroupHeartbeatRequest.LEAVE_GROUP_STATIC_MEMBER_EPOCH`.
pub(crate) const LEAVE_GROUP_STATIC_MEMBER_EPOCH: i32 = -2;

/// `FENCED_MEMBER_EPOCH`.
pub(crate) const FENCED_MEMBER_EPOCH: i16 = 110;
/// `UNKNOWN_MEMBER_ID`.
const UNKNOWN_MEMBER_ID: i16 = 25;
/// `UNSUPPORTED_ASSIGNOR`.
#[cfg(test)]
const UNSUPPORTED_ASSIGNOR: i16 = 112;

/// Kafka's message for a regular expression subscription of a classic member:
/// `ClassicKafkaConsumer.subscribe(SubscriptionPattern)`.
pub(crate) const RE2J_NEEDS_CONSUMER_PROTOCOL: &str = "Subscribe to RE2/J pattern is not supported when using the CLASSIC protocol defined in config group.protocol";

/// The partitions of an assignment, and the topic ids that the metadata does
/// not name yet.
pub(crate) type ResolvedAssignment = (Vec<(String, i32)>, HashSet<WireUuid>);

/// The partitions that a member owns, by topic id and sorted. The wire uuid
/// has no order, so the map is keyed by its bytes.
pub(crate) type OwnedByTopicId = BTreeMap<[u8; 16], Vec<i32>>;

/// The fields of the last `ConsumerGroupHeartbeat`. Kafka's
/// `ConsumerHeartbeatRequestManager.HeartbeatState.SentFields`: a field goes
/// out again only when it changed, or when the member joins.
#[derive(Clone, Debug, Default, PartialEq, Eq)]
pub(crate) struct SentFields {
    rebalance_timeout_ms: Option<i32>,
    subscribed_topic_names: Option<Vec<String>>,
    regex: Option<String>,
    server_assignor: Option<String>,
    topic_partitions: Option<OwnedByTopicId>,
}

/// What the member sends in its next heartbeat.
#[derive(Clone, Debug)]
pub(crate) struct HeartbeatFields<'a> {
    pub group_id: &'a str,
    pub member_id: &'a str,
    pub member_epoch: i32,
    pub instance_id: Option<&'a str>,
    pub rack_id: Option<&'a str>,
    pub rebalance_timeout_ms: i32,
    /// The subscribed topics, sorted. A regular expression subscription sends
    /// none, because the coordinator matches the expression.
    pub subscribed_topic_names: Vec<String>,
    /// The regular expression of a broker-side pattern subscription.
    pub regex: Option<String>,
    pub server_assignor: Option<&'a str>,
    /// The partitions that the member owns, by topic id.
    pub owned: OwnedByTopicId,
    /// `true` for the first heartbeat of a join, which sends every field.
    pub joining: bool,
}

/// Build the next `ConsumerGroupHeartbeat` and record the fields that it
/// carries.
pub(crate) fn build_heartbeat(
    fields: &HeartbeatFields<'_>,
    sent: &mut SentFields,
) -> ConsumerGroupHeartbeatRequest {
    let mut request = ConsumerGroupHeartbeatRequest {
        group_id: fields.group_id.to_owned(),
        member_id: fields.member_id.to_owned(),
        member_epoch: fields.member_epoch,
        instance_id: fields.instance_id.map(str::to_owned),
        ..Default::default()
    };
    let all = fields.joining;
    if all {
        // Kafka sends the rack only when the member joins.
        request.rack_id = fields.rack_id.map(str::to_owned);
        *sent = SentFields::default();
    }
    if all || sent.rebalance_timeout_ms != Some(fields.rebalance_timeout_ms) {
        request.rebalance_timeout_ms = fields.rebalance_timeout_ms;
        sent.rebalance_timeout_ms = Some(fields.rebalance_timeout_ms);
    } else {
        request.rebalance_timeout_ms = -1;
    }
    if all || sent.subscribed_topic_names.as_ref() != Some(&fields.subscribed_topic_names) {
        request.subscribed_topic_names = Some(fields.subscribed_topic_names.clone());
        sent.subscribed_topic_names = Some(fields.subscribed_topic_names.clone());
    }
    // Kafka sends an empty expression to remove a pattern subscription, and
    // sends the expression again only when it changed or when the member
    // joins with one.
    let pattern_changed = sent.regex != fields.regex;
    if (all && fields.regex.is_some()) || pattern_changed {
        request.subscribed_topic_regex = Some(fields.regex.clone().unwrap_or_default());
        sent.regex.clone_from(&fields.regex);
    }
    if let Some(assignor) = fields.server_assignor
        && (all || sent.server_assignor.as_deref() != Some(assignor))
    {
        request.server_assignor = Some(assignor.to_owned());
        sent.server_assignor = Some(assignor.to_owned());
    }
    if all || sent.topic_partitions.as_ref() != Some(&fields.owned) {
        request.topic_partitions = Some(
            fields
                .owned
                .iter()
                .map(|(topic_id, partitions)| TopicPartitions {
                    topic_id: WireUuid(*topic_id),
                    partitions: partitions.clone(),
                    ..Default::default()
                })
                .collect(),
        );
        sent.topic_partitions = Some(fields.owned.clone());
    }
    request
}

/// What the member does with the error code of a heartbeat response. Kafka's
/// `HeartbeatRequestManager.onErrorResponse`.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) enum HeartbeatAction {
    /// The response carries the member state.
    Ok,
    /// Send the heartbeat again after the retry backoff.
    Retry,
    /// Find the coordinator again, then send the heartbeat again.
    FindCoordinator,
    /// The coordinator does not know the member any more: give up the
    /// partitions and join again with epoch 0.
    Fenced,
    /// The application must see the error.
    Fatal,
}

/// Map the error code of a `ConsumerGroupHeartbeat` response, as Kafka's
/// `ConsumerHeartbeatRequestManager` does.
///
/// `COORDINATOR_NOT_AVAILABLE (15)` and `NOT_COORDINATOR (16)` make the member
/// find the coordinator again. `COORDINATOR_LOAD_IN_PROGRESS (14)` retries.
/// `FENCED_MEMBER_EPOCH (110)` and `UNKNOWN_MEMBER_ID (25)` make the member
/// join again from epoch 0. `UNSUPPORTED_ASSIGNOR (112)`,
/// `UNRELEASED_INSTANCE_ID (111)`, `GROUP_MAX_SIZE_REACHED (81)`,
/// `INVALID_REQUEST (42)`, `UNSUPPORTED_VERSION (35)` and the authorization
/// codes are fatal.
pub(crate) fn heartbeat_action(error_code: i16) -> HeartbeatAction {
    match error_code {
        0 => HeartbeatAction::Ok,
        coordinator::COORDINATOR_NOT_AVAILABLE | coordinator::NOT_COORDINATOR => {
            HeartbeatAction::FindCoordinator
        }
        coordinator::COORDINATOR_LOAD_IN_PROGRESS => HeartbeatAction::Retry,
        FENCED_MEMBER_EPOCH | UNKNOWN_MEMBER_ID => HeartbeatAction::Fenced,
        // `UNSUPPORTED_ASSIGNOR (112)`, `UNRELEASED_INSTANCE_ID (111)`,
        // `GROUP_MAX_SIZE_REACHED (81)`, `INVALID_REQUEST (42)`,
        // `UNSUPPORTED_VERSION (35)`, the authorization codes and every other
        // code reach the application.
        _ => HeartbeatAction::Fatal,
    }
}

/// The epoch of the heartbeat that leaves the group. Kafka's
/// `ConsumerMembershipManager.leaveGroupEpoch`: a static member that closes
/// with the default operation keeps its place in the group with `-2`.
pub(crate) fn leave_group_epoch(
    group_instance_id: Option<&str>,
    operation: consumer::GroupMembershipOperation,
) -> i32 {
    if operation == consumer::GroupMembershipOperation::LeaveGroup {
        return LEAVE_GROUP_MEMBER_EPOCH;
    }
    if group_instance_id.is_some() {
        LEAVE_GROUP_STATIC_MEMBER_EPOCH
    } else {
        LEAVE_GROUP_MEMBER_EPOCH
    }
}

/// Whether a closing member sends the heartbeat that leaves the group. Kafka's
/// `ConsumerHeartbeatRequestManager.shouldSendLeaveHeartbeatNow`: a dynamic
/// member that closes with `REMAIN_IN_GROUP` sends none and waits for the
/// session timeout.
pub(crate) fn should_send_leave_heartbeat(
    member_id: &str,
    group_instance_id: Option<&str>,
    operation: consumer::GroupMembershipOperation,
) -> bool {
    if member_id.is_empty() {
        return false;
    }
    !(group_instance_id.is_none() && operation == consumer::GroupMembershipOperation::RemainInGroup)
}

/// The partitions of an assignment, by topic name.
///
/// A topic id that the metadata does not name yet stays unresolved, and the
/// member reconciles it after the next metadata refresh, as Kafka's
/// `AbstractMembershipManager` keeps unresolved assignments.
pub(crate) fn assignment_partitions(
    response: &ConsumerGroupHeartbeatResponse,
    names: &HashMap<WireUuid, String>,
) -> Option<ResolvedAssignment> {
    let assignment = response.assignment.as_ref()?;
    let mut partitions = Vec::new();
    let mut unresolved = HashSet::new();
    for topic in &assignment.topic_partitions {
        match names.get(&topic.topic_id) {
            Some(name) => partitions.extend(
                topic
                    .partitions
                    .iter()
                    .map(|partition| (name.clone(), *partition)),
            ),
            None => {
                unresolved.insert(topic.topic_id);
            }
        }
    }
    partitions.sort();
    Some((partitions, unresolved))
}

/// The owned partitions by topic id, for the next heartbeat.
pub(crate) fn owned_by_topic_id(
    owned: &[(String, i32)],
    topic_ids: &HashMap<String, WireUuid>,
) -> OwnedByTopicId {
    let mut by_id: OwnedByTopicId = BTreeMap::new();
    for (topic, partition) in owned {
        if let Some(topic_id) = topic_ids.get(topic) {
            by_id.entry(topic_id.0).or_default().push(*partition);
        }
    }
    for partitions in by_id.values_mut() {
        partitions.sort_unstable();
    }
    by_id
}

/// Run the member of a consumer group with the KIP-848 protocol.
///
/// The task sends a `ConsumerGroupHeartbeat` at once and then each heartbeat
/// interval that the coordinator returns. Each response can carry a new
/// assignment, which the task reconciles before the next heartbeat
/// acknowledges it. Kafka's `ConsumerHeartbeatRequestManager` and
/// `ConsumerMembershipManager` do the same.
pub(crate) async fn run(
    mut state: coordinator::CoordinatorState,
    shutdown: tokio_util::sync::CancellationToken,
) {
    if let Err(error) = state.client.refresh_metadata().await {
        tracing::warn!(%error, "coordinator client metadata refresh failed at startup");
    }
    // Kafka's `Heartbeat.pollTimer`: `poll` resets it, and the member leaves
    // the group when it expires.
    state.poll_timer = coordinator::PollTimer::new(state.max_poll_interval);
    let mut interval = state.heartbeat_interval.to_std();
    let mut send_now = true;
    let mut polls_open = true;
    // `true` after the poll timer expired: the member left the group and waits
    // for a `poll` before it joins again. Kafka's `MemberState.STALE`.
    let mut left_on_poll_timeout = false;
    // `true` after `unsubscribe`: the member left the group and waits for a
    // new subscription. Kafka's `AbstractMembershipManager.leaveGroup`.
    let mut unsubscribed = false;
    // `true` after a fatal error: Kafka's `transitionToSendingLeaveGroup`
    // sends no leave heartbeat in the fatal state.
    let mut fatal = false;
    let mut last_pattern_check: Option<tokio::time::Instant> = None;
    loop {
        if !send_now {
            let event = tokio::select! {
                biased;
                () = shutdown.cancelled() => break,
                () = tokio::time::sleep(interval), if !left_on_poll_timeout && !unsubscribed => {
                    TaskEvent::Heartbeat
                }
                changed = state.polls.changed(), if polls_open => {
                    if changed.is_err() {
                        polls_open = false;
                        continue;
                    }
                    TaskEvent::Poll
                }
                changed = state.subscription_changes.changed() => {
                    if changed.is_err() {
                        break;
                    }
                    TaskEvent::Subscription
                }
                () = tokio::time::sleep_until(state.poll_timer.deadline()),
                    if !left_on_poll_timeout && !unsubscribed =>
                {
                    TaskEvent::PollTimeout
                }
            };
            match event {
                TaskEvent::Poll => {
                    state.poll_timer.reset();
                    if !left_on_poll_timeout {
                        continue;
                    }
                    // The `poll` came back: join the group again, as Kafka's
                    // member leaves the stale state when the poll timer is
                    // reset.
                    left_on_poll_timeout = false;
                    state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
                    state.sent_heartbeat_fields = SentFields::default();
                }
                TaskEvent::PollTimeout => {
                    leave_after_poll_timeout(&mut state).await;
                    left_on_poll_timeout = true;
                    continue;
                }
                TaskEvent::Subscription => {
                    let subscription = state.subscription.borrow().clone();
                    if subscription.unsubscribes != state.seen_unsubscribes {
                        state.seen_unsubscribes = subscription.unsubscribes;
                        leave_after_unsubscribe(&mut state).await;
                        unsubscribed = true;
                        continue;
                    }
                    if subscription.is_none() {
                        continue;
                    }
                    if unsubscribed {
                        // A new subscription joins the group again.
                        unsubscribed = false;
                        state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
                        state.sent_heartbeat_fields = SentFields::default();
                    }
                    // The task has its own client. Its metadata requests name
                    // the topics of the new subscription, so it can resolve
                    // the topic ids of the next assignment.
                    if !subscription.topics.is_empty() {
                        state.client.metadata_topics().set(subscription.topics);
                    }
                }
                TaskEvent::Heartbeat => {}
            }
        }
        // Kafka's `AsyncKafkaConsumer.updateAssignmentMetadataIfNeeded` matches
        // a client-side pattern against the metadata in each poll. The member
        // sends the topics that matched, and the coordinator assigns them.
        if state.subscription.borrow().pattern.is_some()
            && last_pattern_check.is_none_or(|last| {
                coordinator::subscription_metadata_refresh_due(
                    last,
                    state.subscription_metadata_refresh_interval,
                )
            })
        {
            last_pattern_check = Some(tokio::time::Instant::now());
            coordinator::refresh_pattern_topics(&mut state).await;
        }
        send_now = false;
        let outcome = tokio::select! {
            biased;
            () = shutdown.cancelled() => break,
            outcome = heartbeat_once(&mut state) => outcome,
        };
        match outcome {
            HeartbeatOutcome::Interval(next) => interval = next,
            HeartbeatOutcome::Reconciled(next) => {
                interval = next;
                // Kafka acknowledges a new assignment with the next heartbeat,
                // which it sends at once.
                send_now = true;
            }
            HeartbeatOutcome::Backoff => {
                tokio::select! {
                    () = shutdown.cancelled() => break,
                    () = tokio::time::sleep(state.retry_policy.initial_backoff) => {}
                }
                send_now = true;
            }
            HeartbeatOutcome::Stop => {
                fatal = true;
                break;
            }
        }
    }
    let close = *state.close_operation.borrow();
    // A member that already left the group, or that stopped for a fatal
    // error, sends no leave heartbeat. Kafka's
    // `AbstractMembershipManager.leaveGroup` returns at once when the member
    // is not in the group.
    if !left_on_poll_timeout
        && !unsubscribed
        && !fatal
        && should_send_leave_heartbeat(
            &state.member_id,
            state.group_instance_id.as_deref(),
            close.operation,
        )
    {
        let epoch = leave_group_epoch(state.group_instance_id.as_deref(), close.operation);
        state.member_epoch = epoch;
        let leave_timeout = state.leave_group_timeout.to_std();
        let leave = send_heartbeat(&mut state, HeartbeatKind::Leave).await;
        let _ = tokio::time::timeout(leave_timeout, leave.send()).await;
    }
}

/// What wakes the heartbeat task.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum TaskEvent {
    /// The heartbeat interval elapsed.
    Heartbeat,
    /// A `poll` came: reset the poll timer.
    Poll,
    /// The subscription changed.
    Subscription,
    /// The poll timer expired: leave the group.
    PollTimeout,
}

/// Leave the group because the time between two `poll` calls was longer than
/// `max_poll_interval`, and give up every partition.
///
/// Kafka's `AbstractHeartbeatRequestManager.poll` calls
/// `transitionToSendingLeaveGroup(true)` for an expired poll timer, sends the
/// leave heartbeat and then makes the member stale until the next `poll`.
async fn leave_after_poll_timeout(state: &mut coordinator::CoordinatorState) {
    coordinator::warn_poll_timeout(state);
    set_leave_epoch(state);
    let leave_timeout = state.leave_group_timeout.to_std();
    let leave = send_heartbeat(state, HeartbeatKind::Leave).await;
    let _ = tokio::time::timeout(leave_timeout, leave.send()).await;
    coordinator::release_partitions_after_poll_timeout(state).await;
    if let Err(error) = coordinator::call_lost_listener(state).await {
        tracing::warn!(%error, group = %state.group_id, "the lost callback failed");
    }
    state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
    state.sent_heartbeat_fields = SentFields::default();
}

/// Leave the group because the application unsubscribed. Kafka's
/// `AbstractMembershipManager.leaveGroup`: the member sends a heartbeat with
/// its leave epoch and gives up its partitions. The `Consumer` already ran the
/// revoke callback and cleared the assignment.
async fn leave_after_unsubscribe(state: &mut coordinator::CoordinatorState) {
    set_leave_epoch(state);
    let leave_timeout = state.leave_group_timeout.to_std();
    let leave = send_heartbeat(state, HeartbeatKind::Leave).await;
    let _ = tokio::time::timeout(leave_timeout, leave.send()).await;
    coordinator::forget_member_keeping_id(state).await;
    state.rebalance_pending.send_replace(false);
    state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
    state.sent_heartbeat_fields = SentFields::default();
}

fn set_leave_epoch(state: &mut coordinator::CoordinatorState) {
    state.member_epoch = leave_group_epoch(
        state.group_instance_id.as_deref(),
        consumer::GroupMembershipOperation::Default,
    );
}

/// Keep the member in the group while `poll` runs a rebalance listener
/// callback. Kafka's heartbeat manager keeps sending `ConsumerGroupHeartbeat`
/// while the application thread runs the callback.
///
/// The request carries only the fields that changed, so it repeats the member
/// state and asks for nothing new. The heartbeat after the callback acts on
/// the response of the member state.
pub(crate) async fn heartbeat_during_callback(
    state: &coordinator::CoordinatorState,
    done: &tokio_util::sync::CancellationToken,
) {
    let mut ticker = coordinator::heartbeat_ticker(state.heartbeat_interval);
    loop {
        tokio::select! {
            biased;
            () = done.cancelled() => return,
            _ = ticker.tick() => {}
        }
        let mut sent = state.sent_heartbeat_fields.clone();
        let fields = heartbeat_fields(state, HeartbeatKind::Interval).await;
        let request = build_heartbeat(&fields, &mut sent);
        let broker = state.client.broker(
            state
                .coordinator_id
                .load(std::sync::atomic::Ordering::Relaxed),
        );
        let call = HeartbeatCall {
            broker,
            request,
            needs_regex: state.subscription.borrow().regex.is_some(),
        };
        let result = tokio::select! {
            biased;
            () = done.cancelled() => return,
            result = call.send() => result,
        };
        if let Err(error) = result {
            tracing::warn!(%error, group = %state.group_id, "the heartbeat during a rebalance callback failed");
        }
    }
}

/// What the task does after one heartbeat.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum HeartbeatOutcome {
    /// Wait this interval for the next heartbeat.
    Interval(std::time::Duration),
    /// A new assignment was reconciled: acknowledge it at once, then wait this
    /// interval.
    Reconciled(std::time::Duration),
    /// Wait the retry backoff, then send again.
    Backoff,
    /// Stop the task.
    Stop,
}

/// Send one heartbeat and act on the response.
async fn heartbeat_once(state: &mut coordinator::CoordinatorState) -> HeartbeatOutcome {
    let kind = if state.member_epoch == JOIN_GROUP_MEMBER_EPOCH {
        HeartbeatKind::Join
    } else {
        HeartbeatKind::Interval
    };
    let response = match send_heartbeat(state, kind).await.send().await {
        Ok(response) => response,
        Err(error) => {
            tracing::warn!(%error, group = %state.group_id, "consumer group heartbeat failed");
            // The member does not know what reached the coordinator, so the
            // next heartbeat sends every field again. Kafka's
            // `ConsumerHeartbeatRequestManager.resetHeartbeatState`.
            state.sent_heartbeat_fields = SentFields::default();
            if !coordinator::is_retriable_transport_error(&error) {
                // Kafka's `AbstractHeartbeatRequestManager.onFailure` retries
                // only a `RetriableException`, and hands every other failure
                // to the application. A broker without api key 68 fails here.
                tracing::error!(
                    %error,
                    group = %state.group_id,
                    "the consumer group heartbeat failed for good; the member stops"
                );
                coordinator::report_fatal_error(
                    &state.poll_error,
                    error::ConsumerError::Client(error),
                );
                return HeartbeatOutcome::Stop;
            }
            state.client.evict_broker(
                state
                    .coordinator_id
                    .load(std::sync::atomic::Ordering::Relaxed),
            );
            return HeartbeatOutcome::Backoff;
        }
    };
    let interval = heartbeat_interval(response.heartbeat_interval_ms, state.heartbeat_interval);
    if response.error_code != 0 {
        // Kafka's `AbstractHeartbeatRequestManager.onErrorResponse` resets the
        // sent fields, so the heartbeat after an error carries them again.
        state.sent_heartbeat_fields = SentFields::default();
    }
    match heartbeat_action(response.error_code) {
        HeartbeatAction::Ok => {}
        HeartbeatAction::Retry => return HeartbeatOutcome::Backoff,
        HeartbeatAction::FindCoordinator => {
            coordinator::find_coordinator_again(state).await;
            return HeartbeatOutcome::Backoff;
        }
        HeartbeatAction::Fenced => {
            tracing::warn!(
                group = %state.group_id,
                error_code = response.error_code,
                "the coordinator fenced the member; giving up the partitions and joining again"
            );
            state.member_epoch = JOIN_GROUP_MEMBER_EPOCH;
            coordinator::forget_member_keeping_id(state).await;
            // Kafka's `transitionToFenced` gives the partitions to
            // `onPartitionsLost` before the member joins again.
            if let Err(error) = coordinator::call_lost_listener(state).await {
                tracing::warn!(%error, group = %state.group_id, "the lost callback failed");
            }
            return HeartbeatOutcome::Backoff;
        }
        HeartbeatAction::Fatal => {
            tracing::error!(
                group = %state.group_id,
                error_code = response.error_code,
                error = ?response.error_message,
                "the consumer group heartbeat failed with a fatal error"
            );
            coordinator::report_fatal_error(
                &state.poll_error,
                error::ConsumerError::Server(response.error_code),
            );
            return HeartbeatOutcome::Stop;
        }
    }
    if let Some(member_id) = response.member_id.clone().filter(|id| !id.is_empty()) {
        state.member_id = member_id;
    }
    state.member_epoch = response.member_epoch;
    coordinator::publish_member_epoch(state).await;
    let names = offset_wire::id_to_name(&state.topic_ids.lock().await.clone());
    let Some((mut target, mut unresolved)) = assignment_partitions(&response, &names) else {
        state.rebalance_pending.send_replace(false);
        return HeartbeatOutcome::Interval(interval);
    };
    if !unresolved.is_empty() {
        // The assignment names its topics by id, so the member needs the ids
        // of the metadata. Kafka's `Metadata` holds them for the consumer, and
        // `ConsumerMembershipManager` keeps an unresolved assignment until a
        // metadata update names its topics.
        refresh_topic_ids(state, &unresolved).await;
        let names = offset_wire::id_to_name(&state.topic_ids.lock().await.clone());
        if let Some(resolved) = assignment_partitions(&response, &names) {
            (target, unresolved) = resolved;
        }
    }
    if !unresolved.is_empty() {
        tracing::info!(
            group = %state.group_id,
            topics = unresolved.len(),
            "the assignment names topic ids that the metadata does not name yet"
        );
    }
    // The topics of a regular expression subscription come from the
    // assignment, because the coordinator matches the expression. They give
    // `poll` the topics of its metadata and its fetch filter.
    let topics: Vec<String> = target.iter().map(|(topic, _)| topic.clone()).collect();
    state
        .subscription
        .send_if_modified(|subscription| subscription.store_assigned_topics(&topics));
    match reconcile(state, &target).await {
        // Kafka acknowledges only a reconciliation that moved partitions. An
        // assignment that repeats what the member owns goes out again with the
        // next interval heartbeat.
        Ok(true) => HeartbeatOutcome::Reconciled(interval),
        Ok(false) => HeartbeatOutcome::Interval(interval),
        Err(error) => {
            tracing::warn!(%error, group = %state.group_id, "the reconciliation of the assignment failed");
            HeartbeatOutcome::Backoff
        }
    }
}

/// Which heartbeat the member sends. Kafka sends every field only while the
/// member joins (`MemberState.JOINING`), and a member that leaves gives up its
/// assignment before the request
/// (`AbstractMembershipManager.transitionToSendingLeaveGroup` sets
/// `currentAssignment` to `NONE`), so the leave heartbeat carries an empty
/// partition list.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
enum HeartbeatKind {
    /// The first heartbeat of a member that joins the group.
    Join,
    /// A heartbeat of a member of the group.
    Interval,
    /// The heartbeat that leaves the group.
    Leave,
}

/// Send the heartbeat of the current member state.
///
/// The fields that the request carried stay in `sent_heartbeat_fields`, so the
/// next heartbeat sends only what changed. Kafka's
/// `ConsumerHeartbeatRequestManager.HeartbeatState.buildRequestData` keeps them
/// in the same way.
async fn send_heartbeat(
    state: &mut coordinator::CoordinatorState,
    kind: HeartbeatKind,
) -> HeartbeatCall<'_> {
    let mut sent = state.sent_heartbeat_fields.clone();
    let needs_regex = state.subscription.borrow().regex.is_some();
    let request = build_heartbeat(&heartbeat_fields(state, kind).await, &mut sent);
    state.sent_heartbeat_fields = sent;
    let broker = state.client.broker(
        state
            .coordinator_id
            .load(std::sync::atomic::Ordering::Relaxed),
    );
    HeartbeatCall {
        broker,
        request,
        needs_regex,
    }
}

/// The heartbeat fields of the current member state.
async fn heartbeat_fields(
    state: &coordinator::CoordinatorState,
    kind: HeartbeatKind,
) -> HeartbeatFields<'_> {
    let topic_ids = state.topic_ids.lock().await.clone();
    let owned = if kind == HeartbeatKind::Leave {
        OwnedByTopicId::new()
    } else {
        owned_by_topic_id(&state.assigned.lock().await, &topic_ids)
    };
    let (subscribed_topic_names, regex) = {
        let subscription = state.subscription.borrow();
        // A regular expression subscription names no topics: Kafka's
        // `SubscriptionState.subscription()` is empty for
        // `AUTO_PATTERN_RE2J`, and the coordinator matches the expression.
        let topics = if subscription.regex.is_some() {
            Vec::new()
        } else {
            subscription.topics.clone()
        };
        (topics, subscription.regex.clone())
    };
    HeartbeatFields {
        group_id: &state.group_id,
        member_id: &state.member_id,
        member_epoch: state.member_epoch,
        instance_id: state.group_instance_id.as_deref(),
        rack_id: state.client_rack.as_deref(),
        rebalance_timeout_ms: consumer::protocol_millis_i32(state.max_poll_interval),
        subscribed_topic_names,
        regex,
        server_assignor: state.server_assignor.as_deref(),
        owned,
        joining: kind == HeartbeatKind::Join,
    }
}

/// A `ConsumerGroupHeartbeat` that carries `subscribed_topic_regex`. The
/// field exists from version 1, so a broker that speaks only version 0 must
/// fail the call in place of dropping the expression.
struct RegexHeartbeat(ConsumerGroupHeartbeatRequest);

impl krabka_protocol::Encode for RegexHeartbeat {
    fn encode<B: bytes::BufMut>(
        &self,
        buf: &mut B,
        version: i16,
    ) -> Result<(), krabka_protocol::ProtocolError> {
        self.0.encode(buf, version)
    }

    fn encoded_len(&self, version: i16) -> usize {
        self.0.encoded_len(version)
    }
}

impl krabka_protocol::ProtocolRequest for RegexHeartbeat {
    const API_KEY: i16 = consumer_group_heartbeat_request::API_KEY;
    const MIN_VERSION: i16 = 1;
    const MAX_VERSION: i16 = consumer_group_heartbeat_request::MAX_VERSION;
    const LATEST_STABLE_VERSION: i16 = consumer_group_heartbeat_request::LATEST_STABLE_VERSION;
    const FLEXIBLE_MIN: i16 = consumer_group_heartbeat_request::FLEXIBLE_MIN;
    type Response = ConsumerGroupHeartbeatResponse;
}

/// One `ConsumerGroupHeartbeat` that is ready to go out.
struct HeartbeatCall<'a> {
    broker: krabka_client_core::BrokerHandle<'a>,
    request: ConsumerGroupHeartbeatRequest,
    /// `true` when the member subscribes to a regular expression, which needs
    /// version 1.
    needs_regex: bool,
}

impl HeartbeatCall<'_> {
    async fn send(self) -> Result<ConsumerGroupHeartbeatResponse, krabka_client_core::ClientError> {
        if self.needs_regex {
            return self.broker.send(RegexHeartbeat(self.request)).await;
        }
        self.broker.send(self.request).await
    }
}

/// Store the topic ids of the metadata, so the member can resolve the topics
/// of its assignment.
///
/// A regular expression subscription knows no topic names, so the request
/// names the `unresolved` topic ids. Kafka's
/// `ConsumerMetadata.newMetadataRequestBuilder` asks for the assigned topic
/// ids of such a subscription, and for the topic names of every other one.
async fn refresh_topic_ids(state: &coordinator::CoordinatorState, unresolved: &HashSet<WireUuid>) {
    let by_id = state.subscription.borrow().regex.is_some();
    let request = if by_id {
        MetadataRequest {
            topics: Some(
                unresolved
                    .iter()
                    .map(|topic_id| MetadataRequestTopic {
                        topic_id: *topic_id,
                        name: None,
                        ..Default::default()
                    })
                    .collect(),
            ),
            allow_auto_topic_creation: false,
            ..Default::default()
        }
    } else {
        state.client.metadata_topics().request()
    };
    match state.client.refresh_metadata_with(request).await {
        Ok(metadata) => {
            crate::validate::refresh_tracked_topic_ids(
                &mut *state.topic_ids.lock().await,
                &metadata,
            );
        }
        Err(error) => {
            tracing::warn!(%error, group = %state.group_id, "metadata refresh failed");
        }
    }
}

/// The heartbeat interval of a response, or the configured one when the
/// response names none.
fn heartbeat_interval(
    heartbeat_interval_ms: i32,
    configured: krabka_units::Time,
) -> std::time::Duration {
    if heartbeat_interval_ms > 0 {
        std::time::Duration::from_millis(u64::try_from(heartbeat_interval_ms).unwrap_or(0))
    } else {
        configured.to_std()
    }
}

/// Reconcile `target` with the partitions that the member owns. Returns
/// `true` when the member gave up or took a partition.
///
/// Kafka's `AbstractMembershipManager.reconcile` gives up the partitions that
/// the assignment no longer holds, with `onPartitionsRevoked`, and then takes
/// the new ones with `onPartitionsAssigned`.
async fn reconcile(
    state: &mut coordinator::CoordinatorState,
    target: &[(String, i32)],
) -> Result<bool, error::ConsumerError> {
    let owned = state.assigned.lock().await.clone();
    let target_set: HashSet<(String, i32)> = target.iter().cloned().collect();
    let owned_set: HashSet<(String, i32)> = owned.iter().cloned().collect();
    let mut revoked: Vec<(String, i32)> = owned
        .iter()
        .filter(|partition| !target_set.contains(*partition))
        .cloned()
        .collect();
    let mut added: Vec<(String, i32)> = target
        .iter()
        .filter(|partition| !owned_set.contains(*partition))
        .cloned()
        .collect();
    revoked.sort();
    added.sort();
    if revoked.is_empty() && added.is_empty() {
        state.rebalance_pending.send_replace(false);
        return Ok(false);
    }
    state.rebalance_pending.send_replace(true);
    let epoch = state.member_epoch;
    if !revoked.is_empty() {
        coordinator::call_listener_kind(
            state,
            rebalance_listener::ListenerCallKind::Revoked,
            revoked.clone(),
        )
        .await?;
        let kept: Vec<(String, i32)> = owned
            .iter()
            .filter(|partition| target_set.contains(*partition))
            .cloned()
            .collect();
        coordinator::publish_group_assignment(state, &kept, epoch).await;
    }
    coordinator::prime_added_offsets(state, &added).await?;
    let gate = coordinator::assigned_callback_gate(state, &added);
    coordinator::publish_group_assignment(state, target, epoch).await;
    let result = coordinator::call_listener_kind(
        state,
        rebalance_listener::ListenerCallKind::Assigned,
        added,
    )
    .await;
    drop(gate);
    state.rebalance_pending.send_replace(false);
    result.map(|()| true)
}

#[cfg(test)]
mod group_protocol_tests {
    use std::sync::{Arc, Mutex};

    use bytes::Buf as _;
    use krabka_client_core::MockBroker;
    use krabka_protocol::{
        Decode as _,
        owned::{
            api_versions_request,
            api_versions_response::ApiVersionsResponse,
            common::consumer_group_heartbeat_response::topic_partitions::TopicPartitions as AssignedPartitions,
            consumer_group_heartbeat_request,
            consumer_group_heartbeat_response::Assignment,
            fetch_request,
            fetch_response::FetchResponse,
            find_coordinator_request,
            find_coordinator_response::FindCoordinatorResponse,
            list_offsets_request,
            list_offsets_response::{
                ListOffsetsPartitionResponse, ListOffsetsResponse, ListOffsetsTopicResponse,
            },
            metadata_request,
            metadata_response::{MetadataResponse, MetadataResponseTopic},
            offset_fetch_request,
            offset_fetch_response::{
                OffsetFetchResponse, OffsetFetchResponsePartition, OffsetFetchResponseTopic,
            },
        },
    };
    use krabka_units::{Time, millis, secs};

    use super::*;
    use crate::{
        CloseOptions, Consumer,
        consumer::GroupMembershipOperation,
        test_support::{ListenerCall, ListenerCalls, RecordingListener},
    };

    const TOPIC: &str = "orders";
    const TOPIC_ID: WireUuid = WireUuid([3; 16]);

    /// One `ConsumerGroupHeartbeat` answer of the mock coordinator.
    #[derive(Clone, Debug)]
    struct Answer {
        error_code: i16,
        member_epoch: i32,
        /// The assigned partitions of `orders`, or `None` for a response
        /// without an assignment.
        assignment: Option<Vec<i32>>,
        /// The member id of the response. The coordinator can name another one
        /// than the request.
        member_id: Option<String>,
    }

    /// An answer that assigns `partitions` at `member_epoch`.
    fn assigns(member_epoch: i32, partitions: &[i32]) -> Answer {
        Answer {
            error_code: 0,
            member_epoch,
            assignment: Some(partitions.to_vec()),
            member_id: None,
        }
    }

    /// `answer` with the member id that the coordinator names.
    fn named(answer: Answer, member_id: &str) -> Answer {
        Answer {
            member_id: Some(member_id.to_owned()),
            ..answer
        }
    }

    /// An answer that fails with `error_code`.
    fn fails(error_code: i16) -> Answer {
        Answer {
            error_code,
            member_epoch: 0,
            assignment: None,
            member_id: None,
        }
    }

    /// What the mock coordinator recorded.
    #[derive(Clone, Default)]
    struct Recorder {
        /// Each `ConsumerGroupHeartbeat` that the mock decoded.
        heartbeats: Arc<Mutex<Vec<ConsumerGroupHeartbeatRequest>>>,
        /// The topic names of each `Metadata` request.
        metadata_topics: Arc<Mutex<Vec<String>>>,
        /// Whether a `Metadata` request named a topic by id.
        asked_by_id: Arc<Mutex<bool>>,
        /// The classic `Heartbeat` requests. A member of the consumer group
        /// protocol sends none.
        classic_heartbeats: Arc<Mutex<usize>>,
    }

    impl Recorder {
        fn heartbeats(&self) -> Vec<ConsumerGroupHeartbeatRequest> {
            self.heartbeats.lock().expect("heartbeats lock").clone()
        }

        /// The owned partitions that each heartbeat carried.
        fn acknowledged(&self) -> Vec<Vec<i32>> {
            acknowledged(&self.heartbeats())
        }

        fn epochs(&self) -> Vec<i32> {
            self.heartbeats()
                .iter()
                .map(|request| request.member_epoch)
                .collect()
        }

        fn first_heartbeat(&self) -> ConsumerGroupHeartbeatRequest {
            self.heartbeats()
                .into_iter()
                .next()
                .expect("a first heartbeat")
        }
    }

    fn encode(response: &impl krabka_protocol::Encode, version: i16) -> Vec<u8> {
        let mut buffer = bytes::BytesMut::new();
        response.encode(&mut buffer, version).expect("encode");
        buffer.to_vec()
    }

    /// The response of a flexible api, after the tagged fields of its header.
    /// `ApiVersions` keeps the header of version 0, as Kafka does.
    fn flexible(response: &impl krabka_protocol::Encode, version: i16) -> Vec<u8> {
        let mut frame = vec![0];
        frame.extend(encode(response, version));
        frame
    }

    /// The `ApiVersions` of a coordinator that speaks `ConsumerGroupHeartbeat`.
    fn api_versions() -> ApiVersionsResponse {
        krabka_client_core::mock_test_support::api_versions_response([
            (api_versions_request::API_KEY, 3),
            (metadata_request::API_KEY, 12),
            (find_coordinator_request::API_KEY, 2),
            (offset_fetch_request::API_KEY, 5),
            (list_offsets_request::API_KEY, 7),
            (fetch_request::API_KEY, 12),
            (consumer_group_heartbeat_request::API_KEY, 1),
        ])
    }

    /// The metadata of `orders`, which names the topic id of the assignment.
    fn metadata() -> MetadataResponse {
        MetadataResponse {
            topics: vec![MetadataResponseTopic {
                name: Some(TOPIC.into()),
                topic_id: TOPIC_ID,
                partitions: crate::test_support::metadata_partitions(2),
                ..Default::default()
            }],
            ..Default::default()
        }
    }

    /// The log start and end of `orders`: both partitions are empty.
    fn empty_list_offset_partitions() -> Vec<ListOffsetsPartitionResponse> {
        (0..2)
            .map(|partition_index| ListOffsetsPartitionResponse {
                partition_index,
                timestamp: -1,
                offset: 0,
                leader_epoch: -1,
                ..Default::default()
            })
            .collect()
    }

    fn list_offsets() -> ListOffsetsResponse {
        ListOffsetsResponse {
            topics: vec![ListOffsetsTopicResponse {
                name: TOPIC.into(),
                partitions: empty_list_offset_partitions(),
                ..Default::default()
            }],
            ..Default::default()
        }
    }

    /// The committed offsets of the group: `orders` has none, so the member
    /// starts its partitions at the `auto_offset_reset`.
    fn missing_committed_partitions() -> Vec<OffsetFetchResponsePartition> {
        (0..2)
            .map(|partition_index| OffsetFetchResponsePartition {
                partition_index,
                committed_offset: -1,
                committed_leader_epoch: -1,
                ..Default::default()
            })
            .collect()
    }

    fn no_committed_offsets() -> OffsetFetchResponse {
        OffsetFetchResponse {
            topics: vec![OffsetFetchResponseTopic {
                name: TOPIC.into(),
                partitions: missing_committed_partitions(),
                ..Default::default()
            }],
            ..Default::default()
        }
    }

    /// The `Metadata` request of a frame, without its header.
    fn metadata_request_of(mut body: &[u8], version: i16) -> MetadataRequest {
        let client_id_len = usize::try_from(body.get_i16()).expect("client id length");
        body.advance(client_id_len);
        if version >= 9 {
            body.advance(1);
        }
        MetadataRequest::decode(&mut body, version).expect("decode")
    }

    /// The `ConsumerGroupHeartbeat` request of a frame, without its header.
    fn heartbeat_request(mut body: &[u8], version: i16) -> ConsumerGroupHeartbeatRequest {
        let client_id_len = usize::try_from(body.get_i16()).expect("client id length");
        body.advance(client_id_len);
        // The flexible request header ends with its tagged fields.
        body.advance(1);
        ConsumerGroupHeartbeatRequest::decode(&mut body, version).expect("decode")
    }

    /// A coordinator that answers each `ConsumerGroupHeartbeat` with the next
    /// entry of `answers`, repeats the last entry after that, and records every
    /// request that it decoded.
    async fn heartbeat_coordinator(answers: Vec<Answer>, recorder: &Recorder) -> MockBroker {
        heartbeat_coordinator_with(answers, recorder, false).await
    }

    /// The same coordinator. With `by_id_only` it names the topic only for a
    /// `Metadata` request that asks for its id, as a broker does for a group
    /// with a broker-side pattern subscription.
    async fn heartbeat_coordinator_with(
        answers: Vec<Answer>,
        recorder: &Recorder,
        by_id_only: bool,
    ) -> MockBroker {
        let recorder = recorder.clone();
        let answers = Mutex::new(answers.into_iter());
        let last = Mutex::new(assigns(1, &[]));
        MockBroker::start(move |api_key, version, _corr_id, body| {
            if api_key == api_versions_request::API_KEY {
                return Some(encode(&api_versions(), 0));
            }
            if api_key == find_coordinator_request::API_KEY {
                return Some(encode(&FindCoordinatorResponse::default(), version));
            }
            if api_key == CLASSIC_HEARTBEAT_API_KEY {
                *recorder
                    .classic_heartbeats
                    .lock()
                    .expect("classic heartbeats lock") += 1;
                return None;
            }
            if api_key == metadata_request::API_KEY {
                let request = metadata_request_of(body, version);
                let topics = request.topics.unwrap_or_default();
                let by_id = topics
                    .iter()
                    .any(|topic| topic.topic_id != WireUuid::default());
                if by_id {
                    *recorder.asked_by_id.lock().expect("asked lock") = true;
                }
                recorder
                    .metadata_topics
                    .lock()
                    .expect("metadata topics lock")
                    .extend(topics.iter().filter_map(|topic| topic.name.clone()));
                if by_id_only && !by_id {
                    return Some(flexible(&MetadataResponse::default(), version));
                }
                return Some(flexible(&metadata(), version));
            }
            if api_key == offset_fetch_request::API_KEY {
                return Some(encode(&no_committed_offsets(), version));
            }
            if api_key == list_offsets_request::API_KEY {
                return Some(flexible(&list_offsets(), version));
            }
            if api_key == fetch_request::API_KEY {
                return Some(flexible(&FetchResponse::default(), version));
            }
            if api_key != consumer_group_heartbeat_request::API_KEY {
                return None;
            }
            recorder
                .heartbeats
                .lock()
                .expect("heartbeats lock")
                .push(heartbeat_request(body, version));
            let mut last = last.lock().expect("last lock");
            if let Some(next) = answers.lock().expect("answers lock").next() {
                *last = next;
            }
            let response = ConsumerGroupHeartbeatResponse {
                error_code: last.error_code,
                member_epoch: last.member_epoch,
                member_id: last.member_id.clone(),
                heartbeat_interval_ms: 20,
                assignment: last.assignment.clone().map(|partitions| Assignment {
                    topic_partitions: vec![AssignedPartitions {
                        topic_id: TOPIC_ID,
                        partitions,
                        ..Default::default()
                    }],
                    ..Default::default()
                }),
                ..Default::default()
            };
            Some(flexible(&response, version))
        })
        .await
    }

    async fn assigned_coordinator(recorder: &Recorder) -> MockBroker {
        heartbeat_coordinator(vec![assigns(1, &[0, 1])], recorder).await
    }

    struct AssignedGroupFixture {
        recorder: Recorder,
        mock: MockBroker,
    }

    async fn assigned_group_fixture() -> AssignedGroupFixture {
        let recorder = Recorder::default();
        let mock = assigned_coordinator(&recorder).await;
        AssignedGroupFixture { recorder, mock }
    }

    /// Kafka's `Heartbeat` api key, which the classic protocol uses.
    const CLASSIC_HEARTBEAT_API_KEY: i16 = 12;

    /// What one case observed of the member.
    #[derive(Debug, Eq, PartialEq)]
    struct Observed {
        /// The owned partitions that the heartbeats carried. A heartbeat that
        /// repeats the field of the heartbeat before it carries none, so the
        /// list holds one entry for each change.
        acknowledged: Vec<Vec<i32>>,
        /// The partitions of `orders` that the member owned before it closed.
        owned: Vec<i32>,
        /// The member epoch of the last heartbeat.
        last_epoch: i32,
        /// How many heartbeats after the first one joined the group again.
        rejoins: usize,
    }

    /// The owned partitions that each heartbeat of `requests` carried.
    fn acknowledged(requests: &[ConsumerGroupHeartbeatRequest]) -> Vec<Vec<i32>> {
        requests
            .iter()
            .filter_map(|request| request.topic_partitions.as_ref())
            .map(|topics| {
                topics
                    .iter()
                    .flat_map(|topic| topic.partitions.clone())
                    .collect()
            })
            .collect()
    }

    enum GroupPollTarget<'a> {
        Acknowledged {
            recorder: &'a Recorder,
            count: usize,
        },
        Assignment,
    }

    #[derive(bon::Builder)]
    struct GroupPollWait<'a> {
        consumer: &'a mut Consumer,
        target: GroupPollTarget<'a>,
    }

    impl GroupPollWait<'_> {
        async fn wait(self) {
            let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(20);
            loop {
                let ready = match self.target {
                    GroupPollTarget::Acknowledged { recorder, count } => {
                        recorder.acknowledged().len() >= count
                    }
                    GroupPollTarget::Assignment => !self.consumer.assignment().await.is_empty(),
                };
                if ready || tokio::time::Instant::now() >= deadline {
                    break;
                }
                let _ = self.consumer.poll(millis(20)).await;
            }
        }
    }

    /// Run one member against `answers` until its heartbeats carried `acks`
    /// partition lists, then close it.
    async fn run_member(answers: Vec<Answer>, instance_id: Option<&str>, acks: usize) -> Observed {
        let recorder = Recorder::default();
        let mock = heartbeat_coordinator(answers, &recorder).await;
        let mut consumer = TestConsumerOptions::builder()
            .mock(&mock)
            .server_assignor("uniform".into())
            .maybe_group_instance_id(instance_id.map(str::to_owned))
            .build()
            .start()
            .await;
        GroupPollWait::builder()
            .consumer(&mut consumer)
            .target(GroupPollTarget::Acknowledged {
                recorder: &recorder,
                count: acks,
            })
            .build()
            .wait()
            .await;
        let owned: Vec<i32> = consumer
            .assignment()
            .await
            .into_iter()
            .map(|(_, partition)| partition)
            .collect();
        consumer
            .close_with(CloseOptions::group_membership_operation(
                GroupMembershipOperation::Default,
            ))
            .await
            .expect("close");
        mock.stop();
        let epochs = recorder.epochs();
        Observed {
            acknowledged: recorder.acknowledged(),
            owned,
            last_epoch: *epochs.last().expect("a last heartbeat"),
            rejoins: epochs
                .iter()
                .skip(1)
                .filter(|epoch| **epoch == JOIN_GROUP_MEMBER_EPOCH)
                .count(),
        }
    }

    enum SubscriptionMode {
        Topics(Vec<String>),
        Regex(String),
        Pattern(crate::TopicPattern),
    }

    #[derive(bon::Builder)]
    struct TestConsumerOptions<'a> {
        mock: &'a MockBroker,
        #[builder(default = SubscriptionMode::Topics(vec![TOPIC.to_owned()]))]
        subscription: SubscriptionMode,
        server_assignor: Option<String>,
        #[builder(default = millis(20))]
        heartbeat_interval: Time,
        #[builder(default = secs(5))]
        request_timeout: Time,
        max_poll_interval: Option<Time>,
        group_instance_id: Option<String>,
        rebalance_listener: Option<Box<dyn crate::ConsumerRebalanceListener>>,
    }

    impl TestConsumerOptions<'_> {
        async fn start(self) -> Consumer {
            let builder = Consumer::builder()
                .bootstrap(self.mock.addr.to_string())
                .group_id("group-a")
                .group_protocol(GroupProtocol::Consumer)
                .heartbeat_interval(self.heartbeat_interval)
                .request_timeout(self.request_timeout)
                .maybe_group_remote_assignor(self.server_assignor)
                .maybe_group_instance_id(self.group_instance_id)
                .maybe_rebalance_listener(self.rebalance_listener)
                .maybe_max_poll_interval(self.max_poll_interval);
            match self.subscription {
                SubscriptionMode::Topics(topics) => {
                    builder.subscribe(topics).build().await.expect("build")
                }
                SubscriptionMode::Regex(regex) => {
                    builder.subscribe_regex(regex).build().await.expect("build")
                }
                SubscriptionMode::Pattern(pattern) => builder
                    .subscribe_pattern(pattern)
                    .build()
                    .await
                    .expect("build"),
            }
        }

        async fn start_and_assign(self) -> Consumer {
            let mut consumer = self.start().await;
            poll_until_assigned(&mut consumer).await;
            consumer
        }
    }

    async fn poll_until_assigned(consumer: &mut Consumer) {
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(20);
        while consumer.assignment().await.is_empty() && tokio::time::Instant::now() < deadline {
            let _ = consumer.poll(millis(20)).await;
        }
    }

    async fn poll_until_error(consumer: &mut Consumer) -> Option<String> {
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(20);
        let mut error = None;
        while error.is_none() && tokio::time::Instant::now() < deadline {
            error = consumer
                .poll(millis(20))
                .await
                .err()
                .map(|error| error.to_string());
        }
        error
    }

    async fn started_consumer_group(mock: &MockBroker) -> Consumer {
        TestConsumerOptions::builder()
            .mock(mock)
            .build()
            .start_and_assign()
            .await
    }

    async fn close_and_first_heartbeat(
        consumer: Consumer,
        mock: MockBroker,
        recorder: &Recorder,
    ) -> ConsumerGroupHeartbeatRequest {
        consumer.close().await.expect("close");
        mock.stop();
        recorder.first_heartbeat()
    }

    /// The first heartbeat of a member joins with epoch 0 and carries every
    /// field: Kafka's `ConsumerHeartbeatRequestManager.HeartbeatState`.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn the_first_heartbeat_joins_the_group_with_every_field() {
        let AssignedGroupFixture { recorder, mock } = assigned_group_fixture().await;
        let consumer = TestConsumerOptions::builder()
            .mock(&mock)
            .server_assignor("uniform".into())
            .build()
            .start_and_assign()
            .await;
        let first = close_and_first_heartbeat(consumer, mock, &recorder).await;
        assert2::assert!(
            ConsumerGroupHeartbeatRequest {
                member_id: String::new(),
                ..first.clone()
            } == ConsumerGroupHeartbeatRequest {
                group_id: "group-a".to_owned(),
                member_id: String::new(),
                member_epoch: JOIN_GROUP_MEMBER_EPOCH,
                rebalance_timeout_ms: 300_000,
                subscribed_topic_names: Some(vec![TOPIC.to_owned()]),
                server_assignor: Some("uniform".to_owned()),
                topic_partitions: Some(Vec::new()),
                ..Default::default()
            }
        );
        // Kafka's `AsyncKafkaConsumer` generates a v4 uuid for the member id.
        assert2::assert!(uuid::Uuid::parse_str(&first.member_id).is_ok());
    }

    /// A broker-side pattern subscription sends the regular expression and no
    /// topic names, and the member learns the topic of its assignment from a
    /// metadata request that names the topic id. Kafka's
    /// `subscribe(SubscriptionPattern)` and
    /// `ConsumerMetadata.newMetadataRequestBuilder`.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_regular_expression_subscription_resolves_its_topics_by_id() {
        let recorder = Recorder::default();
        let mock = heartbeat_coordinator_with(vec![assigns(1, &[0, 1])], &recorder, true).await;
        let consumer = TestConsumerOptions::builder()
            .mock(&mock)
            .subscription(SubscriptionMode::Regex("orders-.*".into()))
            .build()
            .start_and_assign()
            .await;
        let assignment = consumer.assignment().await;
        let subscription = consumer.subscription();
        let first = close_and_first_heartbeat(consumer, mock, &recorder).await;
        assert2::assert!(
            (
                first.subscribed_topic_regex,
                first.subscribed_topic_names,
                *recorder.asked_by_id.lock().expect("asked lock"),
                assignment,
                subscription,
            ) == (
                Some("orders-.*".to_owned()),
                Some(Vec::new()),
                true,
                vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
                vec![TOPIC.to_owned()],
            )
        );
    }

    /// A client-side pattern subscription matches the metadata itself, and the
    /// member sends the topics that matched. Kafka's
    /// `AsyncKafkaConsumer.updateAssignmentMetadataIfNeeded`.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_client_side_pattern_sends_the_topics_that_matched() {
        let AssignedGroupFixture { recorder, mock } = assigned_group_fixture().await;
        let consumer = TestConsumerOptions::builder()
            .mock(&mock)
            .subscription(SubscriptionMode::Pattern(crate::TopicPattern::new(
                |topic| topic.starts_with("ord"),
            )))
            .build()
            .start_and_assign()
            .await;
        let assignment = consumer.assignment().await;
        let first = close_and_first_heartbeat(consumer, mock, &recorder).await;
        assert2::assert!(
            (
                first.subscribed_topic_names,
                first.subscribed_topic_regex,
                assignment,
            ) == (
                Some(vec![TOPIC.to_owned()]),
                None,
                vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
            )
        );
    }

    /// The first lost callback of `calls`.
    fn first_lost(calls: &ListenerCalls) -> Option<ListenerCall> {
        calls
            .lock()
            .expect("calls lock")
            .iter()
            .find(|(kind, _)| *kind == rebalance_listener::ListenerCallKind::Lost)
            .cloned()
    }

    /// A fatal heartbeat response and a heartbeat that cannot go out both end
    /// the member and reach the application through `poll`. Kafka's
    /// `AbstractHeartbeatRequestManager.onErrorResponse` and `onFailure` hand
    /// such an error to the application thread.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_fatal_heartbeat_error_reaches_poll() {
        let mut actual = Vec::new();
        for (name, answers, speaks_heartbeat) in [
            (
                "an unsupported assignor",
                vec![fails(UNSUPPORTED_ASSIGNOR)],
                true,
            ),
            ("a broker without the api", Vec::new(), false),
        ] {
            let recorder = Recorder::default();
            let mock = if speaks_heartbeat {
                heartbeat_coordinator(answers, &recorder).await
            } else {
                classic_only_coordinator().await
            };
            let mut consumer = TestConsumerOptions::builder()
                .mock(&mock)
                .build()
                .start()
                .await;
            let error = poll_until_error(&mut consumer).await;
            consumer.close().await.expect("close");
            mock.stop();
            actual.push((name, error));
        }
        assert2::assert!(
            actual
                == vec![
                    (
                        "an unsupported assignor",
                        Some("broker error_code 112".to_owned()),
                    ),
                    (
                        "a broker without the api",
                        Some(
                            "client: incompatible version: broker supports 2..=2, client wants 0..=1 for api_key 68"
                                .to_owned()
                        ),
                    ),
                ]
        );
    }

    /// A broker whose `ConsumerGroupHeartbeat` versions do not overlap the
    /// ones of the client.
    async fn classic_only_coordinator() -> MockBroker {
        coordinator_with_heartbeat_versions(2, 2).await
    }

    /// A coordinator that speaks `ConsumerGroupHeartbeat` in this version
    /// range and answers no heartbeat.
    async fn coordinator_with_heartbeat_versions(min_version: i16, max_version: i16) -> MockBroker {
        MockBroker::start(move |api_key, version, _corr_id, _body| {
            if api_key == api_versions_request::API_KEY {
                let mut versions = api_versions();
                for api in &mut versions.api_keys {
                    if api.api_key == consumer_group_heartbeat_request::API_KEY {
                        api.min_version = min_version;
                        api.max_version = max_version;
                    }
                }
                return Some(encode(&versions, 0));
            }
            if api_key == find_coordinator_request::API_KEY {
                return Some(encode(&FindCoordinatorResponse::default(), version));
            }
            if api_key == metadata_request::API_KEY {
                return Some(flexible(&metadata(), version));
            }
            None
        })
        .await
    }

    async fn regex_consumer(mock: &MockBroker) -> Consumer {
        TestConsumerOptions::builder()
            .mock(mock)
            .subscription(SubscriptionMode::Regex("orders-.*".into()))
            .build()
            .start()
            .await
    }

    /// A regular expression subscription needs `ConsumerGroupHeartbeat` v1,
    /// which added `subscribed_topic_regex`. A coordinator that speaks only
    /// v0 fails the call, because a v0 request would name no topic at all.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_regular_expression_needs_heartbeat_version_1() {
        let mock = coordinator_with_heartbeat_versions(0, 0).await;
        let mut consumer = regex_consumer(&mock).await;
        let error = poll_until_error(&mut consumer).await;
        consumer.close().await.expect("close");
        mock.stop();
        assert2::assert!(
            error
                == Some(
                    "client: incompatible version: broker supports 0..=0, client wants 1..=1 for api_key 68"
                        .to_owned()
                )
        );
    }

    /// The member fetches the partitions of a regular expression assignment
    /// also after the application subscribes to the same expression again,
    /// because the topics of such a subscription come from the assignment.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_regular_expression_subscription_keeps_fetching_after_a_resubscribe() {
        let recorder = Recorder::default();
        let mock = heartbeat_coordinator_with(vec![assigns(1, &[0, 1])], &recorder, true).await;
        let mut consumer = regex_consumer(&mock).await;
        // Poll until the member owns the partitions and their positions are
        // ready, so that the fetch of the next poll would go out.
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(20);
        while tokio::time::Instant::now() < deadline {
            let _ = consumer.poll(millis(20)).await;
            let assigned = consumer.assignment().await;
            if !assigned.is_empty() && !consumer.group_fetches(&assigned).await.is_empty() {
                break;
            }
        }
        consumer.subscribe_regex("orders-.*").expect("subscribe");
        let assigned = consumer.assignment().await;
        let mut fetched: Vec<(String, i32)> = consumer
            .group_fetches(&assigned)
            .await
            .into_values()
            .flatten()
            .flat_map(|(topic, partitions)| {
                partitions
                    .into_iter()
                    .map(move |(partition, ..)| (topic.clone(), partition))
            })
            .collect();
        fetched.sort();
        consumer.close().await.expect("close");
        mock.stop();
        assert2::assert!(
            (assigned, fetched)
                == (
                    vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
                    vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
                )
        );
    }

    /// `unsubscribe` makes the member leave the group with its leave epoch and
    /// give up its partitions, and a new `subscribe` joins the group again.
    /// Kafka's `AbstractMembershipManager.leaveGroup`.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn unsubscribe_leaves_the_group_and_a_new_subscription_joins_again() {
        let recorder = Recorder::default();
        let mock = assigned_coordinator(&recorder).await;
        let mut consumer = started_consumer_group(&mock).await;
        consumer.unsubscribe().await.expect("unsubscribe");
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(20);
        while !recorder.epochs().contains(&LEAVE_GROUP_MEMBER_EPOCH)
            && tokio::time::Instant::now() < deadline
        {
            tokio::time::sleep(std::time::Duration::from_millis(20)).await;
        }
        let after_unsubscribe = (consumer.assignment().await, recorder.epochs());
        consumer
            .subscribe([TOPIC.to_owned()])
            .await
            .expect("subscribe");
        poll_until_assigned(&mut consumer).await;
        let joined_again = consumer.assignment().await;
        consumer.close().await.expect("close");
        mock.stop();
        let (left, epochs) = after_unsubscribe;
        let joins_after_the_leave = epochs
            .iter()
            .skip_while(|epoch| **epoch != LEAVE_GROUP_MEMBER_EPOCH)
            .filter(|epoch| **epoch == JOIN_GROUP_MEMBER_EPOCH)
            .count();
        assert2::assert!(
            (
                left,
                epochs.contains(&LEAVE_GROUP_MEMBER_EPOCH),
                joins_after_the_leave,
                joined_again
            ) == (
                Vec::new(),
                true,
                0,
                vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
            )
        );
    }

    /// The member leaves the group when the time between two `poll` calls is
    /// longer than `max_poll_interval`, gives its partitions to the lost
    /// callback, and joins again with the next `poll`. Kafka's
    /// `AbstractHeartbeatRequestManager.poll` and `transitionToStale`.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_member_that_stops_polling_leaves_and_joins_again() {
        let recorder = Recorder::default();
        let calls: ListenerCalls = Arc::default();
        let mock = assigned_coordinator(&recorder).await;
        let mut consumer = TestConsumerOptions::builder()
            .mock(&mock)
            .max_poll_interval(millis(300))
            .rebalance_listener(Box::new(
                RecordingListener::builder()
                    .calls(Arc::clone(&calls))
                    .build(),
            ))
            .build()
            .start()
            .await;
        GroupPollWait::builder()
            .consumer(&mut consumer)
            .target(GroupPollTarget::Assignment)
            .build()
            .wait()
            .await;
        // The application stops polling for longer than `max_poll_interval`.
        tokio::time::sleep(std::time::Duration::from_millis(700)).await;
        let left = consumer.assignment().await;
        let epochs_after_leave = recorder.epochs();
        GroupPollWait::builder()
            .consumer(&mut consumer)
            .target(GroupPollTarget::Assignment)
            .build()
            .wait()
            .await;
        let joined_again = consumer.assignment().await;
        consumer.close().await.expect("close");
        mock.stop();
        // A `poll` that ends before its callback ran makes the next `poll`
        // run the callback again, so the list can hold more than one call.
        let lost = first_lost(&calls);
        assert2::assert!(
            (
                left,
                epochs_after_leave.contains(&LEAVE_GROUP_MEMBER_EPOCH),
                joined_again,
                lost,
            ) == (
                Vec::new(),
                true,
                vec![(TOPIC.to_owned(), 0), (TOPIC.to_owned(), 1)],
                Some((rebalance_listener::ListenerCallKind::Lost, vec![0, 1])),
            )
        );
    }

    /// A fenced member gives its partitions to the lost callback, and the
    /// member keeps its id. Kafka's `transitionToFenced`.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn a_fenced_member_gives_its_partitions_to_the_lost_callback() {
        let recorder = Recorder::default();
        let calls: ListenerCalls = Arc::default();
        let mock = heartbeat_coordinator(
            vec![
                assigns(1, &[0, 1]),
                fails(FENCED_MEMBER_EPOCH),
                assigns(1, &[0, 1]),
            ],
            &recorder,
        )
        .await;
        let mut consumer = TestConsumerOptions::builder()
            .mock(&mock)
            .rebalance_listener(Box::new(
                RecordingListener::builder()
                    .calls(Arc::clone(&calls))
                    .build(),
            ))
            .build()
            .start()
            .await;
        let member_id = consumer.member_id();
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(20);
        while first_lost(&calls).is_none() && tokio::time::Instant::now() < deadline {
            let _ = consumer.poll(millis(20)).await;
        }
        let kept_id = consumer.member_id();
        consumer.close().await.expect("close");
        mock.stop();
        assert2::assert!(
            (first_lost(&calls), kept_id == member_id)
                == (
                    Some((rebalance_listener::ListenerCallKind::Lost, vec![0, 1])),
                    true,
                )
        );
    }

    /// The heartbeats that reach the coordinator while the assign callback of
    /// [`HeartbeatWaitingListener`] runs.
    const HEARTBEATS_IN_CALLBACK: usize = 3;

    /// A rebalance listener whose assign callback returns only after
    /// [`HEARTBEATS_IN_CALLBACK`] more heartbeats reached the coordinator, or
    /// after a bound. It records how many arrived.
    struct HeartbeatWaitingListener {
        recorder: Recorder,
        /// The heartbeats that arrived during the assign callback, once it
        /// returned.
        in_callback: Arc<Mutex<Option<usize>>>,
    }

    #[async_trait::async_trait]
    impl crate::ConsumerRebalanceListener for HeartbeatWaitingListener {
        async fn on_partitions_revoked(
            &mut self,
            _consumer: &Consumer,
            _partitions: &[(String, i32)],
        ) -> Result<(), crate::RebalanceListenerError> {
            Ok(())
        }

        async fn on_partitions_assigned(
            &mut self,
            _consumer: &Consumer,
            _partitions: &[(String, i32)],
        ) -> Result<(), crate::RebalanceListenerError> {
            let before = self.recorder.heartbeats().len();
            let bound = tokio::time::Instant::now() + std::time::Duration::from_secs(20);
            let mut arrived = 0;
            while arrived < HEARTBEATS_IN_CALLBACK && tokio::time::Instant::now() < bound {
                tokio::time::sleep(std::time::Duration::from_millis(5)).await;
                arrived = self.recorder.heartbeats().len() - before;
            }
            *self.in_callback.lock().expect("in callback lock") = Some(arrived);
            Ok(())
        }
    }

    /// The member keeps its membership with `ConsumerGroupHeartbeat` while a
    /// rebalance listener callback runs, and never sends the classic
    /// `Heartbeat`.
    ///
    /// The coordinator task publishes the assignment before it queues the
    /// assign callback, so a non-empty `assignment()` does not mean that the
    /// callback ran. The test waits for the callback itself.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn heartbeats_go_on_during_a_rebalance_callback() {
        let recorder = Recorder::default();
        let in_callback: Arc<Mutex<Option<usize>>> = Arc::default();
        let mock = assigned_coordinator(&recorder).await;
        let mut consumer = TestConsumerOptions::builder()
            .mock(&mock)
            .rebalance_listener(Box::new(HeartbeatWaitingListener {
                recorder: recorder.clone(),
                in_callback: Arc::clone(&in_callback),
            }))
            .build()
            .start()
            .await;
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(30);
        while in_callback.lock().expect("in callback lock").is_none()
            && tokio::time::Instant::now() < deadline
        {
            let _ = consumer.poll(millis(200)).await;
        }
        let during_callback = *in_callback.lock().expect("in callback lock");
        consumer.close().await.expect("close");
        mock.stop();
        let classic = *recorder
            .classic_heartbeats
            .lock()
            .expect("classic heartbeats lock");
        assert2::assert!(
            (
                during_callback.map(|arrived| arrived >= HEARTBEATS_IN_CALLBACK),
                classic
            ) == (Some(true), 0)
        );
    }

    /// The member id of the coordinator reaches the consumer, and a new
    /// subscription reaches the metadata of the heartbeat task.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn the_member_id_and_a_new_subscription_reach_the_consumer() {
        let recorder = Recorder::default();
        let mock =
            heartbeat_coordinator(vec![named(assigns(1, &[0, 1]), "server-id")], &recorder).await;
        let mut consumer = started_consumer_group(&mock).await;
        let member_id = consumer.member_id();
        consumer
            .subscribe([TOPIC.to_owned(), "payments".to_owned()])
            .await
            .expect("subscribe");
        let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(20);
        while !recorder
            .metadata_topics
            .lock()
            .expect("metadata topics lock")
            .iter()
            .any(|topic| topic == "payments")
            && tokio::time::Instant::now() < deadline
        {
            let _ = consumer.poll(millis(20)).await;
        }
        let asked_for_payments = recorder
            .metadata_topics
            .lock()
            .expect("metadata topics lock")
            .iter()
            .any(|topic| topic == "payments");
        consumer.close().await.expect("close");
        mock.stop();
        assert2::assert!((member_id, asked_for_payments) == ("server-id".to_owned(), true));
    }

    /// The member reconciles each assignment of the coordinator, acknowledges
    /// it with the partitions that it owns, gives up every partition when the
    /// coordinator fences it, and leaves the group with the epoch of its
    /// membership.
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn the_member_reconciles_and_leaves_as_kafka_does() {
        struct Case {
            name: &'static str,
            answers: Vec<Answer>,
            instance_id: Option<&'static str>,
            /// How many partition lists the heartbeats carry before the close.
            acks: usize,
        }
        let cases = [
            Case {
                name: "an assignment",
                answers: vec![assigns(1, &[0, 1])],
                instance_id: None,
                acks: 2,
            },
            Case {
                name: "an assignment that revokes a partition",
                answers: vec![assigns(1, &[0, 1]), assigns(2, &[0])],
                instance_id: None,
                acks: 3,
            },
            Case {
                name: "a fenced member",
                answers: vec![
                    assigns(1, &[0, 1]),
                    fails(FENCED_MEMBER_EPOCH),
                    assigns(1, &[0, 1]),
                ],
                instance_id: None,
                acks: 4,
            },
            Case {
                name: "a static member",
                answers: vec![assigns(1, &[0, 1])],
                instance_id: Some("member-a"),
                acks: 2,
            },
        ];
        let mut actual = Vec::new();
        for case in cases {
            actual.push((
                case.name,
                run_member(case.answers, case.instance_id, case.acks).await,
            ));
        }
        assert2::assert!(
            actual
                == vec![
                    (
                        "an assignment",
                        Observed {
                            // The join carries no partitions, the next
                            // heartbeat acknowledges the assignment, and the
                            // heartbeat that leaves gives the partitions up.
                            acknowledged: vec![vec![], vec![0, 1], vec![]],
                            owned: vec![0, 1],
                            last_epoch: LEAVE_GROUP_MEMBER_EPOCH,
                            rejoins: 0,
                        },
                    ),
                    (
                        "an assignment that revokes a partition",
                        Observed {
                            acknowledged: vec![vec![], vec![0, 1], vec![0], vec![]],
                            owned: vec![0],
                            last_epoch: LEAVE_GROUP_MEMBER_EPOCH,
                            rejoins: 0,
                        },
                    ),
                    (
                        "a fenced member",
                        Observed {
                            // The fence releases the partitions, and the member
                            // joins again with epoch 0 and no partitions.
                            acknowledged: vec![vec![], vec![0, 1], vec![], vec![0, 1], vec![]],
                            owned: vec![0, 1],
                            last_epoch: LEAVE_GROUP_MEMBER_EPOCH,
                            rejoins: 1,
                        },
                    ),
                    (
                        "a static member",
                        Observed {
                            acknowledged: vec![vec![], vec![0, 1], vec![]],
                            owned: vec![0, 1],
                            // A static member that closes with the default
                            // operation keeps its place in the group.
                            last_epoch: LEAVE_GROUP_STATIC_MEMBER_EPOCH,
                            rejoins: 0,
                        },
                    ),
                ]
        );
    }
}

#[cfg(test)]
mod tests {
    use consumer::GroupMembershipOperation;
    use krabka_protocol::owned::{
        common::consumer_group_heartbeat_response::topic_partitions::TopicPartitions as ResponseTopicPartitions,
        consumer_group_heartbeat_response::Assignment,
    };

    use super::*;

    const TOPIC_ID: WireUuid = WireUuid([7; 16]);

    fn fields<'a>(member_epoch: i32, joining: bool, owned: OwnedByTopicId) -> HeartbeatFields<'a> {
        HeartbeatFields {
            group_id: "group-a",
            member_id: "member-a",
            member_epoch,
            instance_id: None,
            rack_id: Some("az-1"),
            rebalance_timeout_ms: 300_000,
            subscribed_topic_names: vec!["orders".to_owned()],
            regex: None,
            server_assignor: Some("uniform"),
            owned,
            joining,
        }
    }

    /// Kafka's `HeartbeatState.buildRequestData` sends every field when the
    /// member joins, and after that only the fields that changed.
    #[test]
    fn a_heartbeat_repeats_only_the_fields_that_changed() {
        let mut sent = SentFields::default();
        let join = build_heartbeat(&fields(0, true, BTreeMap::new()), &mut sent);
        let steady = build_heartbeat(&fields(1, false, BTreeMap::new()), &mut sent);
        let owned = build_heartbeat(
            &fields(1, false, maplit::btreemap! {TOPIC_ID.0 => vec![0, 1]}),
            &mut sent,
        );
        assert2::assert!(
            (
                join.rack_id.clone(),
                join.rebalance_timeout_ms,
                join.subscribed_topic_names.clone(),
                join.server_assignor.clone(),
                join.topic_partitions.clone(),
            ) == (
                Some("az-1".to_owned()),
                300_000,
                Some(vec!["orders".to_owned()]),
                Some("uniform".to_owned()),
                Some(Vec::new()),
            )
        );
        assert2::assert!(
            (
                steady.rack_id,
                steady.rebalance_timeout_ms,
                steady.subscribed_topic_names,
                steady.server_assignor,
                steady.topic_partitions,
            ) == (None, -1, None, None, None)
        );
        assert2::assert!(
            owned.topic_partitions
                == Some(vec![TopicPartitions {
                    topic_id: TOPIC_ID,
                    partitions: vec![0, 1],
                    ..Default::default()
                }])
        );
    }

    /// Kafka's `HeartbeatState.buildRequestData` sends the regular expression
    /// of a broker-side pattern subscription when the member joins with one
    /// and when it changed, and sends an empty expression to remove it.
    #[test]
    fn a_heartbeat_carries_the_regular_expression_of_the_subscription() {
        let with_regex = |regex: Option<&'static str>, joining: bool| HeartbeatFields {
            subscribed_topic_names: Vec::new(),
            regex: regex.map(str::to_owned),
            ..fields(i32::from(!joining), joining, BTreeMap::new())
        };
        let mut sent = SentFields::default();
        let join = build_heartbeat(&with_regex(Some("orders-.*"), true), &mut sent);
        let steady = build_heartbeat(&with_regex(Some("orders-.*"), false), &mut sent);
        let changed = build_heartbeat(&with_regex(Some("orders-eu-.*"), false), &mut sent);
        let removed = build_heartbeat(&with_regex(None, false), &mut sent);
        let actual = [join, steady, changed, removed].map(|request| {
            (
                request.subscribed_topic_regex,
                request.subscribed_topic_names,
            )
        });
        assert2::assert!(
            actual
                == [
                    (Some("orders-.*".to_owned()), Some(Vec::new())),
                    (None, None),
                    (Some("orders-eu-.*".to_owned()), None),
                    (Some(String::new()), None),
                ]
        );
    }

    /// Kafka's heartbeat error classes.
    #[test]
    fn heartbeat_errors_follow_kafkas_classes() {
        let actual = [0, 14, 15, 16, 25, 110, 111, 112, 81, 42, 35, 30].map(heartbeat_action);
        assert2::assert!(
            actual
                == [
                    HeartbeatAction::Ok,
                    HeartbeatAction::Retry,
                    HeartbeatAction::FindCoordinator,
                    HeartbeatAction::FindCoordinator,
                    HeartbeatAction::Fenced,
                    HeartbeatAction::Fenced,
                    HeartbeatAction::Fatal,
                    HeartbeatAction::Fatal,
                    HeartbeatAction::Fatal,
                    HeartbeatAction::Fatal,
                    HeartbeatAction::Fatal,
                    HeartbeatAction::Fatal,
                ]
        );
    }

    /// Kafka's `leaveGroupEpoch` and `shouldSendLeaveHeartbeatNow`.
    #[test]
    fn a_leaving_member_follows_kafkas_epochs() {
        let mut actual = Vec::new();
        let mut wanted = Vec::new();
        for (name, instance_id, operation, expected) in [
            (
                "dynamic, default",
                None,
                GroupMembershipOperation::Default,
                (-1, true),
            ),
            (
                "static, default",
                Some("i-1"),
                GroupMembershipOperation::Default,
                (-2, true),
            ),
            (
                "static, leave group",
                Some("i-1"),
                GroupMembershipOperation::LeaveGroup,
                (-1, true),
            ),
            (
                "dynamic, remain in group",
                None,
                GroupMembershipOperation::RemainInGroup,
                (-1, false),
            ),
            (
                "static, remain in group",
                Some("i-1"),
                GroupMembershipOperation::RemainInGroup,
                (-2, true),
            ),
        ] {
            actual.push((
                name,
                (
                    leave_group_epoch(instance_id, operation),
                    should_send_leave_heartbeat("member-a", instance_id, operation),
                ),
            ));
            wanted.push((name, expected));
        }
        assert2::assert!(actual == wanted);
    }

    /// The assignment names its topics by id. An id that the metadata does not
    /// name yet stays unresolved.
    #[test]
    fn an_assignment_resolves_the_topic_ids_of_the_metadata() {
        let other = WireUuid([9; 16]);
        let response = ConsumerGroupHeartbeatResponse {
            assignment: Some(Assignment {
                topic_partitions: vec![
                    ResponseTopicPartitions {
                        topic_id: TOPIC_ID,
                        partitions: vec![1, 0],
                        ..Default::default()
                    },
                    ResponseTopicPartitions {
                        topic_id: other,
                        partitions: vec![0],
                        ..Default::default()
                    },
                ],
                ..Default::default()
            }),
            ..Default::default()
        };
        let names = maplit::hashmap! {TOPIC_ID => "orders".to_owned()};
        assert2::assert!(
            assignment_partitions(&response, &names)
                == Some((
                    vec![("orders".to_owned(), 0), ("orders".to_owned(), 1)],
                    maplit::hashset! {other}
                ))
        );
        assert2::assert!(
            assignment_partitions(&ConsumerGroupHeartbeatResponse::default(), &names) == None
        );
    }
}