krabka_client_consumer/share/

poll.rs

//! `ShareConsumer::poll` + acknowledgement (`ShareFetch` / `ShareAcknowledge`).
//!
//! `poll()` issues one `ShareFetch` over the live assignment. It pairs each
//! acquired record with the `acquired_records` range that the broker returns,
//! so each record carries the broker's `delivery_count`. It also keeps the
//! ranges for the next poll's implicit auto-`Accept`.
//!
//! ## Acknowledgement
//!
//! - [`Implicit`](super::types::ShareAckMode::Implicit) is the default. The
//!   *next* `poll()`, and `close()`, implicitly `Accept` every range that the
//!   previous `poll()` returned. The application does nothing.
//! - [`Explicit`](super::types::ShareAckMode::Explicit): the application calls
//!   [`acknowledge`](ShareConsumer::acknowledge) per record. The next `poll()`
//!   flushes the staged acks as a piggyback, or
//!   [`commit`](ShareConsumer::commit) flushes them in a standalone
//!   `ShareAcknowledge`.
//!
//! ## Session epoch
//!
//! The broker's share-session cache opens at epoch 0 and stores 1. It then
//! expects every later `ShareFetch` *or* `ShareAcknowledge` to carry the stored
//! epoch, and it increments the epoch on each accepted request. This module
//! mirrors that exactly: it sends `self.share_session_epoch` and advances it by
//! one after every successful `ShareFetch` / `ShareAcknowledge`, the sequence
//! 0 → 1 → 2 → …. A wrong epoch makes the broker drop the session with
//! `INVALID_SHARE_SESSION_EPOCH`.

use std::{collections::HashMap, mem, time::Duration};

use krabka_protocol::{
    owned::{
        share_acknowledge_request::{
            AcknowledgePartition, AcknowledgeTopic, AcknowledgementBatch as AckAckBatch,
            ShareAcknowledgeRequest,
        },
        share_acknowledge_response::ShareAcknowledgeResponse,
        share_fetch_request::{
            AcknowledgementBatch as FetchAckBatch, FetchPartition, FetchTopic, ShareFetchRequest,
        },
        share_fetch_response::PartitionData,
    },
    primitives::uuid::Uuid as WireUuid,
};
use krabka_units::{
    ByteSize, Time,
    convert::{ByteSizeExt as _, TimeExt as _},
};
use tokio::time::{Instant, sleep};
use tracing::field;

use super::{
    consumer::ShareConsumer,
    types::{ShareAckMode, ShareAckType, ShareAcquireMode, ShareConsumerRecord},
};
use crate::{coordinator, error::ConsumerError, poll::record_timestamp};

fn build_share_fetch_topics(
    assignment: &[(WireUuid, String, i32)],
    acks: &HashMap<(WireUuid, i32), Vec<FetchAckBatch>>,
) -> Vec<FetchTopic> {
    let mut by_topic: HashMap<WireUuid, Vec<(i32, Vec<FetchAckBatch>)>> = HashMap::new();
    for (tid, _name, partition) in assignment {
        let packs = acks.get(&(*tid, *partition)).cloned().unwrap_or_default();
        by_topic.entry(*tid).or_default().push((*partition, packs));
    }

    by_topic
        .into_iter()
        .map(|(topic_id, parts)| FetchTopic {
            topic_id,
            partitions: parts
                .into_iter()
                .map(
                    |(partition_index, acknowledgement_batches)| FetchPartition {
                        partition_index,
                        acknowledgement_batches,
                        ..Default::default()
                    },
                )
                .collect(),
            ..Default::default()
        })
        .collect()
}

#[bon::builder]
fn build_share_fetch_request(
    #[builder(into)] group_id: String,
    #[builder(into)] member_id: String,
    share_session_epoch: i32,
    timeout: Time,
    min: ByteSize,
    max: ByteSize,
    max_records: i32,
    acquire_mode: ShareAcquireMode,
    #[builder(default)] topics: Vec<FetchTopic>,
) -> ShareFetchRequest {
    ShareFetchRequest {
        group_id: Some(group_id),
        member_id: Some(member_id),
        share_session_epoch,
        // Truncate rather than round, and floor at zero: `max_wait_ms` is a
        // wire field, so a fractional millisecond must not round up past the
        // caller's budget and an elapsed budget must not go negative.
        max_wait_ms: i32::try_from(timeout.millis_i64_trunc().max(0)).unwrap_or(i32::MAX),
        min_bytes: min.bytes_i32(),
        max_bytes: max.bytes_i32(),
        max_records,
        batch_size: max_records,
        share_acquire_mode: acquire_mode.wire(),
        topics,
        ..Default::default()
    }
}

fn build_share_ack_request(
    group_id: String,
    member_id: String,
    share_session_epoch: i32,
    is_renew_ack: bool,
    topics: Vec<AcknowledgeTopic>,
) -> ShareAcknowledgeRequest {
    ShareAcknowledgeRequest {
        group_id: Some(group_id),
        member_id: Some(member_id),
        share_session_epoch,
        is_renew_ack,
        topics,
        ..Default::default()
    }
}

fn response_has_error(error_code: i16) -> bool {
    error_code != 0
}

fn first_share_ack_partition_error(response: &ShareAcknowledgeResponse) -> Option<i16> {
    response
        .responses
        .iter()
        .flat_map(|topic| &topic.partitions)
        .map(|partition| partition.error_code)
        .find(|code| response_has_error(*code))
}

fn range_len(first: i64, last: i64) -> usize {
    usize::try_from(last - first + 1).unwrap_or(0)
}

fn offset_in_range(first: i64, offset: i64, last: i64) -> bool {
    first <= offset && offset <= last
}

fn record_offset(base_offset: i64, offset_delta: i32) -> i64 {
    base_offset + i64::from(offset_delta)
}

/// The records of one `ShareFetch` partition row.
///
/// The function skips control batches. It pairs each record with the
/// `delivery_count` of the acquired range that holds its offset, or 0 when no
/// range holds it. Kafka's `ShareCompletedFetch` reads the timestamp and the
/// timestamp type from the record batch, as the classic consumer does.
fn share_partition_records(topic_name: &str, part: &PartitionData) -> Vec<ShareConsumerRecord> {
    let Some(batches) = part.records.as_ref().and_then(|payload| payload.as_v2()) else {
        return Vec::new();
    };
    let mut out = Vec::new();
    for batch in batches {
        if batch.attributes.is_control_batch() {
            continue;
        }
        for r in &batch.records {
            let offset = record_offset(batch.base_offset, r.offset_delta);
            // Pair the record with the acquired range that contains it to
            // read the broker's delivery_count for this offset.
            let delivery_count = part
                .acquired_records
                .iter()
                .find(|ar| offset_in_range(ar.first_offset, offset, ar.last_offset))
                .map_or(0, |ar| ar.delivery_count);
            let (timestamp, timestamp_type) = record_timestamp(batch, r);
            out.push(ShareConsumerRecord {
                topic: topic_name.to_owned(),
                partition: part.partition_index,
                offset,
                timestamp,
                timestamp_type,
                key: r.key.clone(),
                value: r.value.clone(),
                headers: r
                    .headers
                    .iter()
                    .map(|header| (header.key.clone(), header.value.clone()))
                    .collect(),
                delivery_count,
            });
        }
    }
    out
}

impl ShareConsumer {
    /// Acquire and return the next batch of records.
    ///
    /// This method piggybacks the acknowledgements for the previous `poll()` on
    /// the `ShareFetch`. Those are the implicit auto-`Accept` ranges, or the
    /// drained explicit `acknowledge()` calls. It then decodes the acquired
    /// records and pairs each one with the broker-reported `delivery_count`
    /// from its `acquired_records` range.
    ///
    /// With no assignment yet, this method waits for the heartbeat loop to
    /// assign partitions within `timeout`. The fetch uses the remaining wait
    /// budget. If the assignment stays empty, it returns empty at the deadline.
    #[tracing::instrument(
        name = "share_consumer.poll",
        level = "debug",
        skip_all,
        fields(
            group_id = %self.group_id,
            member_id = %self.member_id,
            session_epoch = self.share_session_epoch,
            timeout_ms = timeout.millis_i64_trunc(),
            assigned_partitions = field::Empty,
            records = field::Empty,
        ),
        err(level = "debug")
    )]
    /// # Errors
    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
    pub async fn poll(&mut self, timeout: Time) -> Result<Vec<ShareConsumerRecord>, ConsumerError> {
        let deadline = Instant::now() + timeout.to_std();
        let assignment = loop {
            // A heartbeat can assign partitions or report a fatal error while
            // this poll waits for the initial assignment or a rejoin.
            if let Some(error) = coordinator::take_poll_error(&self.poll_error) {
                return Err(error);
            }
            let assignment = self.assignment.lock().await.clone();
            if !assignment.is_empty() {
                break assignment;
            }
            let remaining = deadline.saturating_duration_since(Instant::now());
            if remaining.is_zero() {
                return Ok(Vec::new());
            }
            sleep(remaining.min(Duration::from_millis(20))).await;
        };
        tracing::Span::current().record("assigned_partitions", assignment.len());

        // Build the piggyback acknowledgement batches per (topic_id, partition)
        // from the previous poll, draining the source so each ack is sent once.
        let acks = self.take_piggyback_acks();

        // Group assigned partitions by topic id, attaching the (topic_id,
        // partition) acks to the matching partition entry.
        let topics = build_share_fetch_topics(&assignment, &acks);

        let resp = self
            .client
            .send(
                build_share_fetch_request()
                    .group_id(self.group_id.clone())
                    .member_id(self.member_id.clone())
                    .share_session_epoch(self.share_session_epoch)
                    .timeout(Time::from_std(
                        deadline.saturating_duration_since(Instant::now()),
                    ))
                    .min(self.fetch_min)
                    .max(self.fetch_max)
                    .max_records(self.fetch_max_records)
                    .acquire_mode(self.acquire_mode)
                    .topics(topics)
                    .call(),
            )
            .await?;

        if response_has_error(resp.error_code) {
            return Err(ConsumerError::Server(resp.error_code));
        }
        // A successful ShareFetch consumes one session epoch; advance to the
        // value the broker now expects (it stored ours + 1).
        self.share_session_epoch = self.share_session_epoch.wrapping_add(1);

        // Reverse-map topic id → name for the returned rows (the fetch response
        // carries only topic_id). Fall back to the cached topic_names.
        let name_for: HashMap<WireUuid, String> = assignment
            .iter()
            .map(|(tid, name, _)| (*tid, name.clone()))
            .collect();

        let mut out: Vec<ShareConsumerRecord> = Vec::new();
        let mut delivered: Vec<(WireUuid, i32, i64, i64)> = Vec::new();
        for topic in &resp.responses {
            let topic_name = name_for.get(&topic.topic_id).cloned().unwrap_or_default();
            for part in &topic.partitions {
                if response_has_error(part.acknowledge_error_code) {
                    tracing::warn!(
                        topic = %topic_name,
                        partition = part.partition_index,
                        acknowledge_error_code = part.acknowledge_error_code,
                        "share fetch piggyback acknowledge error"
                    );
                }
                if response_has_error(part.error_code) {
                    tracing::warn!(
                        topic = %topic_name,
                        partition = part.partition_index,
                        error_code = part.error_code,
                        "share fetch partition error"
                    );
                    continue;
                }

                // Remember the acquired ranges for the next implicit auto-Accept.
                for ar in &part.acquired_records {
                    delivered.push((
                        topic.topic_id,
                        part.partition_index,
                        ar.first_offset,
                        ar.last_offset,
                    ));
                }

                out.extend(share_partition_records(&topic_name, part));
            }
        }

        // Implicit mode auto-Accepts these ranges on the next poll/close; explicit
        // mode acknowledges per record, so the ranges are not auto-accepted.
        self.prev_delivered = delivered;
        tracing::Span::current().record("records", out.len());
        Ok(out)
    }

    /// Stage an explicit acknowledgement for `record`, in Explicit mode.
    ///
    /// The next [`poll`](ShareConsumer::poll) flushes the ack as a piggyback,
    /// or [`commit`](ShareConsumer::commit) flushes it in a standalone
    /// `ShareAcknowledge`.
    ///
    /// # Errors
    ///
    /// Returns [`ConsumerError::IllegalState`] in
    /// [`Implicit`](super::types::ShareAckMode::Implicit) mode. In that mode the
    /// next poll or close auto-`Accept`s every delivered record, so this method
    /// cannot honor an explicit `acknowledge()`. Staging it would silently leak
    /// into `pending_acks`, which the implicit path never flushes. This mirrors
    /// the JVM `KafkaShareConsumer`, which raises `IllegalStateException` if you
    /// explicitly acknowledge while in implicit acknowledgement mode.
    pub fn acknowledge(
        &mut self,
        record: &ShareConsumerRecord,
        ack: ShareAckType,
    ) -> Result<(), ConsumerError> {
        if self.ack_mode == ShareAckMode::Implicit {
            return Err(ConsumerError::IllegalState(
                "acknowledge() is not allowed in implicit ack mode; \
                 records are auto-accepted on the next poll/close"
                    .into(),
            ));
        }
        let topic_id = self.topic_id_for(&record.topic);
        self.pending_acks.push((
            topic_id,
            record.partition,
            record.offset,
            record.offset,
            ack.wire(),
        ));
        Ok(())
    }

    /// Renew the acquisition lock on a single delivered `record` (KIP-1222).
    ///
    /// This method sends a standalone `ShareAcknowledge` with
    /// `is_renew_ack = true` and `acknowledge_types = [4]` for the record's
    /// offset. The Renew wire type extends the broker-side lock deadline without
    /// changing the record's state. Like [`acknowledge`](ShareConsumer::acknowledge),
    /// this is only valid in explicit ack mode. It advances the session epoch on
    /// success.
    ///
    /// # Errors
    ///
    /// Returns [`ConsumerError::IllegalState`] in
    /// [`Implicit`](super::types::ShareAckMode::Implicit) mode, where the next
    /// poll or close auto-accepts the records, so a lock renewal is meaningless.
    /// Returns [`ConsumerError::Server`] if the broker rejects the renew.
    #[tracing::instrument(
        name = "share_consumer.renew",
        level = "debug",
        skip_all,
        fields(
            group_id = %self.group_id,
            partition = record.partition,
            offset = record.offset,
        ),
        err(level = "debug")
    )]
    pub async fn renew(&mut self, record: &ShareConsumerRecord) -> Result<(), ConsumerError> {
        if self.ack_mode == ShareAckMode::Implicit {
            return Err(ConsumerError::IllegalState(
                "renew() is not allowed in implicit ack mode; \
                 records are auto-accepted on the next poll/close"
                    .into(),
            ));
        }
        let topic_id = self.topic_id_for(&record.topic);
        let topics = build_ack_topics(vec![(
            topic_id,
            record.partition,
            record.offset,
            record.offset,
            4,
        )]);

        let resp = self
            .client
            .send(build_share_ack_request(
                self.group_id.clone(),
                self.member_id.clone(),
                self.share_session_epoch,
                true,
                topics,
            ))
            .await?;
        if response_has_error(resp.error_code) {
            return Err(ConsumerError::Server(resp.error_code));
        }
        self.share_session_epoch = self.share_session_epoch.wrapping_add(1);
        if let Some(error_code) = first_share_ack_partition_error(&resp) {
            return Err(ConsumerError::Server(error_code));
        }
        Ok(())
    }

    /// Flush the staged explicit acknowledgements in a standalone
    /// `ShareAcknowledge`. This is a no-op when nothing is staged.
    #[tracing::instrument(
        name = "share_consumer.commit",
        level = "debug",
        skip_all,
        fields(group_id = %self.group_id, pending_acks = self.pending_acks.len()),
        err(level = "debug")
    )]
    /// # Errors
    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
    pub async fn commit(&mut self) -> Result<(), ConsumerError> {
        self.flush_pending_acks().await
    }

    /// Drain `pending_acks` into a `ShareAcknowledge`.
    ///
    /// This method advances the session epoch on success, because an accepted
    /// `ShareAcknowledge` consumes one epoch exactly like a `ShareFetch`. It is
    /// a no-op, with no epoch advance, when nothing is staged.
    #[tracing::instrument(
        name = "share_consumer.flush_pending_acks",
        level = "debug",
        skip_all,
        fields(
            group_id = %self.group_id,
            session_epoch = self.share_session_epoch,
            pending_acks = self.pending_acks.len(),
        ),
        err(level = "debug")
    )]
    pub(crate) async fn flush_pending_acks(&mut self) -> Result<(), ConsumerError> {
        if self.pending_acks.is_empty() {
            return Ok(());
        }
        let drained = mem::take(&mut self.pending_acks);
        let topics = build_ack_topics(drained);

        let resp = self
            .client
            .send(build_share_ack_request(
                self.group_id.clone(),
                self.member_id.clone(),
                self.share_session_epoch,
                false,
                topics,
            ))
            .await?;
        if response_has_error(resp.error_code) {
            return Err(ConsumerError::Server(resp.error_code));
        }
        self.share_session_epoch = self.share_session_epoch.wrapping_add(1);
        if let Some(error_code) = first_share_ack_partition_error(&resp) {
            return Err(ConsumerError::Server(error_code));
        }
        Ok(())
    }

    /// Build the piggyback acknowledgement batches for the next `ShareFetch`,
    /// keyed by `(topic_id, partition)`. This method consumes the source state.
    ///
    /// - Implicit: one `Accept` batch per previously-delivered range.
    /// - Explicit: the drained `pending_acks`, grouped into per-offset batches.
    fn take_piggyback_acks(&mut self) -> HashMap<(WireUuid, i32), Vec<FetchAckBatch>> {
        let mut out: HashMap<(WireUuid, i32), Vec<FetchAckBatch>> = HashMap::new();
        match self.ack_mode {
            ShareAckMode::Implicit => {
                for (tid, partition, first, last) in mem::take(&mut self.prev_delivered) {
                    let count = range_len(first, last);
                    out.entry((tid, partition))
                        .or_default()
                        .push(FetchAckBatch {
                            first_offset: first,
                            last_offset: last,
                            acknowledge_types: vec![ShareAckType::Accept.wire(); count],
                            ..Default::default()
                        });
                }
            }
            ShareAckMode::Explicit => {
                for (tid, partition, first, last, ack) in mem::take(&mut self.pending_acks) {
                    let count = range_len(first, last);
                    out.entry((tid, partition))
                        .or_default()
                        .push(FetchAckBatch {
                            first_offset: first,
                            last_offset: last,
                            acknowledge_types: vec![ack; count],
                            ..Default::default()
                        });
                }
                // Explicit mode never auto-accepts; clear any stale ranges.
                self.prev_delivered.clear();
            }
        }
        out
    }

    /// Resolve a topic id from a topic name through the live assignment or the
    /// cached `topic_names`.
    ///
    /// This method returns the zero uuid if the name is unknown. The broker then
    /// rejects the ack, which surfaces the misuse instead of a silent mis-ack.
    fn topic_id_for(&self, name: &str) -> WireUuid {
        // The assignment carries (topic_id, name, partition); use it first since
        // it is the set the application is acking against. `try_lock` keeps this
        // sync (acknowledge() takes &mut self, not async).
        if let Ok(assignment) = self.assignment.try_lock()
            && let Some((tid, _, _)) = assignment.iter().find(|(_, n, _)| n == name)
        {
            return *tid;
        }
        if let Ok(names) = self.topic_names.try_lock()
            && let Some((tid, _)) = names.iter().find(|(_, n)| n.as_str() == name)
        {
            return *tid;
        }
        WireUuid::default()
    }
}

/// Group `(topic_id, partition, first, last, ack_wire)` acks into the
/// `ShareAcknowledge` topic/partition/batch shape. This function coalesces by
/// topic and by partition.
fn build_ack_topics(acks: Vec<(WireUuid, i32, i64, i64, i8)>) -> Vec<AcknowledgeTopic> {
    let mut by_topic: HashMap<WireUuid, HashMap<i32, Vec<AckAckBatch>>> = HashMap::new();
    for (tid, partition, first, last, ack) in acks {
        let count = range_len(first, last);
        by_topic
            .entry(tid)
            .or_default()
            .entry(partition)
            .or_default()
            .push(AckAckBatch {
                first_offset: first,
                last_offset: last,
                acknowledge_types: vec![ack; count],
                ..Default::default()
            });
    }
    by_topic
        .into_iter()
        .map(|(topic_id, parts)| AcknowledgeTopic {
            topic_id,
            partitions: parts
                .into_iter()
                .map(
                    |(partition_index, acknowledgement_batches)| AcknowledgePartition {
                        partition_index,
                        acknowledgement_batches,
                        ..Default::default()
                    },
                )
                .collect(),
            ..Default::default()
        })
        .collect()
}

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

    use assert2::check;
    use krabka_client_core::Client;
    use krabka_protocol::tagged_fields::UnknownTaggedFields;
    use tokio::sync::Mutex;
    use tokio_util::sync::CancellationToken;

    use super::*;
    use crate::{consumer::TimestampType, share};

    fn id(n: u8) -> WireUuid {
        let mut b = [0u8; 16];
        b[15] = n;
        WireUuid(b)
    }

    #[test]
    fn standalone_ack_surfaces_partition_error() {
        use krabka_protocol::owned::share_acknowledge_response::{
            PartitionData, ShareAcknowledgeTopicResponse,
        };

        let response = ShareAcknowledgeResponse {
            responses: vec![ShareAcknowledgeTopicResponse {
                partitions: vec![PartitionData {
                    error_code: 121,
                    ..Default::default()
                }],
                ..Default::default()
            }],
            ..Default::default()
        };

        assert2::assert!(first_share_ack_partition_error(&response) == Some(121));
    }

    async fn test_consumer(ack_mode: ShareAckMode) -> ShareConsumer {
        ShareConsumer {
            client: Client::builder()
                .bootstrap("127.0.0.1:1")
                .client_id("share-poll-test")
                .build()
                .await
                .unwrap(),
            group_id: "group-a".into(),
            member_id: "member-a".into(),
            member_epoch: Arc::new(Mutex::new(3)),
            assignment: Arc::new(Mutex::new(vec![(id(7), "topic-a".into(), 2)])),
            topic_names: Arc::new(Mutex::new(HashMap::new())),
            share_session_epoch: 4,
            fetch_min: share::DEFAULT_SHARE_CONSUMER_FETCH_MIN,
            fetch_max: share::DEFAULT_SHARE_CONSUMER_FETCH_MAX,
            fetch_max_records: share::DEFAULT_SHARE_CONSUMER_FETCH_MAX_RECORDS,
            acquire_mode: ShareAcquireMode::BatchOptimized,
            ack_mode,
            pending_acks: Vec::new(),
            prev_delivered: Vec::new(),
            shutdown: CancellationToken::new(),
            hb_handle: None,
            poll_error: coordinator::PollErrorSlot::default(),
        }
    }

    #[tokio::test]
    async fn poll_waits_for_assignment_and_fetches_with_the_remaining_budget() {
        use bytes::{Buf as _, Bytes, BytesMut};
        use krabka_client_core::MockBroker;
        use krabka_protocol::{
            Decode as _, Encode as _,
            owned::{
                api_versions_request,
                api_versions_response::{ApiVersion, ApiVersionsResponse},
                share_fetch_request,
                share_fetch_response::{
                    AcquiredRecords, ShareFetchResponse, ShareFetchableTopicResponse,
                },
            },
            records::{Record, RecordBatch, RecordsPayload},
        };

        let fetched = Arc::new(StdMutex::new(None));
        let observed = Arc::clone(&fetched);
        let broker = MockBroker::start(move |api_key, version, _corr, mut body| {
            let mut encoded = BytesMut::new();
            if api_key == api_versions_request::API_KEY {
                ApiVersionsResponse {
                    api_keys: vec![ApiVersion {
                        api_key: share_fetch_request::API_KEY,
                        min_version: share_fetch_request::MIN_VERSION,
                        max_version: share_fetch_request::MAX_VERSION,
                        ..Default::default()
                    }],
                    ..Default::default()
                }
                .encode(&mut encoded, 0)
                .unwrap();
            } else if api_key == share_fetch_request::API_KEY {
                let client_id_len = usize::try_from(body.get_i16()).unwrap_or(0);
                body.advance(client_id_len);
                assert2::assert!(body.get_u8() == 0); // Empty flexible request-header tags.
                *observed.lock().unwrap() =
                    Some(ShareFetchRequest::decode(&mut body, version).unwrap());
                encoded.extend_from_slice(&[0]); // Empty flexible response-header tags.
                ShareFetchResponse {
                    responses: vec![ShareFetchableTopicResponse {
                        topic_id: id(7),
                        partitions: vec![PartitionData {
                            partition_index: 2,
                            records: Some(RecordsPayload::V2(vec![RecordBatch {
                                base_offset: 7,
                                records: vec![Record {
                                    value: Some(Bytes::from_static(b"work")),
                                    ..Default::default()
                                }],
                                ..Default::default()
                            }])),
                            acquired_records: vec![AcquiredRecords {
                                first_offset: 7,
                                last_offset: 7,
                                delivery_count: 1,
                                ..Default::default()
                            }],
                            ..Default::default()
                        }],
                        ..Default::default()
                    }],
                    ..Default::default()
                }
                .encode(&mut encoded, version)
                .unwrap();
            } else {
                panic!("unexpected API key {api_key}");
            }
            Some(encoded.to_vec())
        })
        .await
        .unwrap();
        let mut consumer = test_consumer(ShareAckMode::Explicit).await;
        consumer.client = Client::builder()
            .bootstrap(broker.addr.to_string())
            .build()
            .await
            .unwrap();
        consumer.assignment.lock().await.clear();
        consumer.share_session_epoch = 0;
        let assignment = Arc::clone(&consumer.assignment);
        let assigned = tokio::spawn(async move {
            sleep(Duration::from_millis(100)).await;
            assignment.lock().await.push((id(7), "topic-a".into(), 2));
        });
        let started = Instant::now();
        let records = consumer.poll(krabka_units::secs(1)).await.unwrap();
        let elapsed = started.elapsed();
        assigned.await.unwrap();
        broker.stop();

        assert2::assert!(elapsed < Duration::from_secs(1));
        assert2::assert!(records.len() == 1);
        assert2::assert!(
            (records[0].offset, records[0].value.as_deref()) == (7, Some(&b"work"[..]))
        );
        assert2::assert!(consumer.share_session_epoch == 1);
        let request = fetched
            .lock()
            .unwrap()
            .clone()
            .expect("assignment triggers a ShareFetch in this poll");
        assert2::assert!((1..=900).contains(&request.max_wait_ms));
        assert2::assert!(request.topics[0].partitions[0].partition_index == 2);
    }

    #[tokio::test(start_paused = true)]
    async fn poll_without_assignment_respects_zero_and_finite_deadlines() {
        let mut consumer = test_consumer(ShareAckMode::Explicit).await;
        consumer.assignment.lock().await.clear();
        for timeout_ms in [0, 37] {
            let started = Instant::now();
            let records = consumer
                .poll(krabka_units::millis(timeout_ms))
                .await
                .unwrap();
            assert2::assert!(records.is_empty());
            assert2::assert!(started.elapsed() == Duration::from_millis(u64::from(timeout_ms)));
        }
        assert2::assert!(consumer.share_session_epoch == 4);
    }

    #[tokio::test(start_paused = true)]
    async fn poll_notices_fatal_errors_while_waiting_for_assignment() {
        let mut consumer = test_consumer(ShareAckMode::Explicit).await;
        consumer.assignment.lock().await.clear();
        let poll_error = Arc::clone(&consumer.poll_error);
        let rejected = tokio::spawn(async move {
            sleep(Duration::from_millis(5)).await;
            coordinator::report_rejoin_error(
                &poll_error,
                ConsumerError::GroupAuthorizationFailed("group-a".into()),
            );
        });
        let started = Instant::now();
        let result = consumer.poll(krabka_units::secs(1)).await;
        rejected.await.unwrap();
        assert2::assert!(
            matches!(result, Err(ConsumerError::GroupAuthorizationFailed(group)) if group == "group-a")
        );
        assert2::assert!(started.elapsed() <= Duration::from_millis(20));
    }

    fn only<T>(items: &[T]) -> &T {
        assert2::assert!(items.len() == 1);
        &items[0]
    }

    #[test]
    fn share_fetch_request_preserves_acquire_mode_limits_and_timeout_bounds() {
        let topic = FetchTopic {
            topic_id: id(7),
            partitions: vec![FetchPartition {
                partition_index: 2,
                partition_max_bytes: 123,
                acknowledgement_batches: Vec::new(),
                ..Default::default()
            }],
            ..Default::default()
        };

        let req = build_share_fetch_request()
            .group_id("group-a")
            .member_id("member-a")
            .share_session_epoch(4)
            .timeout(krabka_units::millis(250))
            .min(krabka_units::bytes(7))
            .max(krabka_units::bytes(65_536))
            .max_records(37)
            .acquire_mode(ShareAcquireMode::RecordLimit)
            .topics(vec![topic.clone()])
            .call();

        assert2::assert!(
            req == ShareFetchRequest {
                group_id: Some("group-a".into()),
                member_id: Some("member-a".into()),
                share_session_epoch: 4,
                max_wait_ms: 250,
                min_bytes: 7,
                max_bytes: 65_536,
                max_records: 37,
                batch_size: 37,
                share_acquire_mode: 1,
                is_renew_ack: false,
                topics: vec![topic],
                forgotten_topics_data: Vec::new(),
                unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
            }
        );

        let saturated = build_share_fetch_request()
            .group_id("group-a")
            .member_id("member-a")
            .share_session_epoch(4)
            .timeout(Time::from_millis(i64::from(u32::MAX)))
            .min(krabka_units::bytes(7))
            .max(krabka_units::bytes(65_536))
            .max_records(37)
            .acquire_mode(ShareAcquireMode::BatchOptimized)
            .call();
        assert2::assert!((saturated.max_wait_ms, saturated.share_acquire_mode) == (i32::MAX, 0));

        // A sub-millisecond remainder truncates: rounding 250.9 ms up to 251
        // would hold the broker's Fetch open past the caller's budget.
        let fractional = build_share_fetch_request()
            .group_id("group-a")
            .member_id("member-a")
            .share_session_epoch(4)
            .timeout(krabka_units::micros(250_900))
            .min(krabka_units::bytes(7))
            .max(krabka_units::bytes(65_536))
            .max_records(37)
            .acquire_mode(ShareAcquireMode::BatchOptimized)
            .call();
        assert2::assert!(fractional.max_wait_ms == 250);
    }

    #[test]
    fn share_fetch_request_encodes_default_sizes_as_their_exact_wire_integers() {
        let req = build_share_fetch_request()
            .group_id("group-a")
            .member_id("member-a")
            .share_session_epoch(0)
            .timeout(krabka_units::millis(500))
            .min(krabka_units::bytes(1))
            .max(krabka_units::mebibytes(50))
            .max_records(500)
            .acquire_mode(ShareAcquireMode::BatchOptimized)
            .call();

        assert2::assert!(
            req == ShareFetchRequest {
                group_id: Some("group-a".into()),
                member_id: Some("member-a".into()),
                share_session_epoch: 0,
                max_wait_ms: 500,
                min_bytes: 1,
                max_bytes: 52_428_800,
                max_records: 500,
                batch_size: 500,
                share_acquire_mode: 0,
                is_renew_ack: false,
                topics: Vec::new(),
                forgotten_topics_data: Vec::new(),
                unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
            }
        );
    }

    #[test]
    fn share_fetch_topics_group_assignment_and_attach_partition_acks() {
        let ack = FetchAckBatch {
            first_offset: 10,
            last_offset: 11,
            acknowledge_types: vec![ShareAckType::Accept.wire(); 2],
            ..Default::default()
        };
        let mut acks = HashMap::new();
        acks.insert((id(7), 2), vec![ack.clone()]);

        let topics = build_share_fetch_topics(
            &[
                (id(7), "topic-a".into(), 2),
                (id(7), "topic-a".into(), 3),
                (id(8), "topic-b".into(), 1),
            ],
            &acks,
        );

        let topic = topics.iter().find(|topic| topic.topic_id == id(7)).unwrap();
        assert2::assert!(topic.partitions.len() == 2);
        let part = topic
            .partitions
            .iter()
            .find(|part| part.partition_index == 2)
            .unwrap();
        assert2::assert!(
            (
                part.partition_max_bytes,
                part.acknowledgement_batches.as_slice()
            ) == (0, &[ack][..])
        );
        let empty = topic
            .partitions
            .iter()
            .find(|part| part.partition_index == 3)
            .unwrap();
        assert2::assert!(empty.acknowledgement_batches.is_empty());
    }

    #[test]
    fn share_ack_request_preserves_identity_epoch_renew_flag_and_topics() {
        let topics = build_ack_topics(vec![(id(7), 2, 10, 12, ShareAckType::Reject.wire())]);

        let req =
            build_share_ack_request("group-a".into(), "member-a".into(), 5, true, topics.clone());

        assert2::assert!(
            req == ShareAcknowledgeRequest {
                group_id: Some("group-a".into()),
                member_id: Some("member-a".into()),
                share_session_epoch: 5,
                is_renew_ack: true,
                topics,
                unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
            }
        );
    }

    #[test]
    fn response_and_record_helpers_preserve_boundaries() {
        for (name, code, expected) in [("success", 0, false), ("error", 17, true)] {
            check!(response_has_error(code) == expected, "case {name}");
        }
        for (name, first, last, expected) in [
            ("inclusive range", 10, 12, 3),
            ("reversed range", 12, 10, 0),
        ] {
            check!(range_len(first, last) == expected, "case {name}");
        }
        for (name, first, offset, last, expected) in [
            ("lower bound", 10, 10, 12, true),
            ("upper bound", 10, 12, 12, true),
            ("below range", 10, 9, 12, false),
            ("above range", 10, 13, 12, false),
        ] {
            check!(
                offset_in_range(first, offset, last) == expected,
                "case {name}"
            );
        }
        check!(record_offset(100, 7) == 107);
    }

    #[tokio::test]
    async fn acknowledge_rejects_implicit_mode_and_stages_explicit_record() {
        let record = ShareConsumerRecord {
            topic: "topic-a".into(),
            partition: 2,
            offset: 10,
            timestamp: 0,
            timestamp_type: TimestampType::CreateTime,
            key: None,
            value: None,
            headers: Vec::new(),
            delivery_count: 1,
        };

        let mut implicit = test_consumer(ShareAckMode::Implicit).await;
        assert2::assert!(
            implicit
                .acknowledge(&record, ShareAckType::Accept)
                .unwrap_err()
                .to_string()
                .contains("implicit ack mode")
        );

        let mut explicit = test_consumer(ShareAckMode::Explicit).await;
        explicit
            .acknowledge(&record, ShareAckType::Release)
            .unwrap();
        assert2::assert!(
            explicit.pending_acks == vec![(id(7), 2, 10, 10, ShareAckType::Release.wire())]
        );
    }

    #[tokio::test]
    async fn renew_rejects_implicit_mode_before_sending() {
        let record = ShareConsumerRecord {
            topic: "topic-a".into(),
            partition: 2,
            offset: 10,
            timestamp: 0,
            timestamp_type: TimestampType::CreateTime,
            key: None,
            value: None,
            headers: Vec::new(),
            delivery_count: 1,
        };
        let mut consumer = test_consumer(ShareAckMode::Implicit).await;

        assert2::assert!(
            consumer
                .renew(&record)
                .await
                .unwrap_err()
                .to_string()
                .contains("implicit ack mode")
        );
    }

    #[tokio::test]
    async fn renew_sends_the_wire_type_and_preserves_response_epoch_boundaries() {
        use bytes::{Buf as _, BytesMut};
        use krabka_client_core::MockBroker;
        use krabka_protocol::{
            Decode as _, Encode as _,
            owned::{
                api_versions_request,
                api_versions_response::{ApiVersion, ApiVersionsResponse},
                share_acknowledge_request,
                share_acknowledge_response::{PartitionData, ShareAcknowledgeTopicResponse},
            },
        };

        let requests = Arc::new(StdMutex::new(Vec::new()));
        let observed = Arc::clone(&requests);
        let mut responses = [(0, 0), (17, 0), (0, 121)].into_iter();
        let broker = MockBroker::start(move |api_key, version, _corr, mut body| {
            let mut encoded = BytesMut::new();
            if api_key == api_versions_request::API_KEY {
                ApiVersionsResponse {
                    api_keys: vec![ApiVersion {
                        api_key: share_acknowledge_request::API_KEY,
                        min_version: 2,
                        max_version: 2,
                        ..Default::default()
                    }],
                    ..Default::default()
                }
                .encode(&mut encoded, 0)
                .unwrap();
            } else if api_key == share_acknowledge_request::API_KEY {
                let client_id_len = usize::try_from(body.get_i16()).unwrap_or(0);
                body.advance(client_id_len);
                assert2::assert!(body.get_u8() == 0);
                observed
                    .lock()
                    .unwrap()
                    .push(ShareAcknowledgeRequest::decode(&mut body, version).unwrap());
                let (error_code, partition_error) =
                    responses.next().expect("one response per renew");
                encoded.extend_from_slice(&[0]);
                ShareAcknowledgeResponse {
                    error_code,
                    responses: vec![ShareAcknowledgeTopicResponse {
                        topic_id: id(7),
                        partitions: vec![PartitionData {
                            partition_index: 2,
                            error_code: partition_error,
                            ..Default::default()
                        }],
                        ..Default::default()
                    }],
                    ..Default::default()
                }
                .encode(&mut encoded, version)
                .unwrap();
            } else {
                panic!("unexpected API key {api_key}");
            }
            Some(encoded.to_vec())
        })
        .await
        .unwrap();
        let mut consumer = test_consumer(ShareAckMode::Explicit).await;
        consumer.client = Client::builder()
            .bootstrap(broker.addr.to_string())
            .build()
            .await
            .unwrap();
        let record = ShareConsumerRecord {
            topic: "topic-a".into(),
            partition: 2,
            offset: 10,
            timestamp: 0,
            timestamp_type: TimestampType::CreateTime,
            key: None,
            value: None,
            headers: Vec::new(),
            delivery_count: 1,
        };
        let mut outcomes = Vec::new();
        for _ in 0..3 {
            let error_code = match consumer.renew(&record).await {
                Ok(()) => 0,
                Err(ConsumerError::Server(code)) => code,
                Err(error) => panic!("unexpected renew error: {error}"),
            };
            outcomes.push((error_code, consumer.share_session_epoch));
        }
        broker.stop();
        assert2::assert!(outcomes == vec![(0, 5), (17, 5), (121, 6)]);
        let requests = requests.lock().unwrap();
        assert2::assert!(requests.len() == 3);
        for (request, epoch) in requests.iter().zip([4, 5, 5]) {
            assert2::assert!(request.group_id.as_deref() == Some("group-a"));
            assert2::assert!(request.member_id.as_deref() == Some("member-a"));
            assert2::assert!((request.share_session_epoch, request.is_renew_ack) == (epoch, true));
            let topic = only(&request.topics);
            assert2::assert!(topic.topic_id == id(7));
            let partition = only(&topic.partitions);
            assert2::assert!(partition.partition_index == 2);
            assert2::assert!(
                *only(&partition.acknowledgement_batches)
                    == AckAckBatch {
                        first_offset: 10,
                        last_offset: 10,
                        acknowledge_types: vec![4],
                        ..Default::default()
                    }
            );
        }
    }

    #[tokio::test]
    async fn take_piggyback_acks_drains_implicit_deliveries_as_accept_ranges() {
        let mut consumer = test_consumer(ShareAckMode::Implicit).await;
        consumer.prev_delivered = vec![(id(7), 2, 10, 12)];

        let acks = consumer.take_piggyback_acks();

        assert2::assert!(consumer.prev_delivered.is_empty());
        let batch = only(&acks[&(id(7), 2)]);
        assert2::assert!(
            *batch
                == FetchAckBatch {
                    first_offset: 10,
                    last_offset: 12,
                    acknowledge_types: vec![1, 1, 1],
                    unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
                }
        );
    }

    #[tokio::test]
    async fn take_piggyback_acks_drains_explicit_pending_and_clears_stale_deliveries() {
        let mut consumer = test_consumer(ShareAckMode::Explicit).await;
        consumer.prev_delivered = vec![(id(7), 2, 1, 1)];
        consumer.pending_acks = vec![(id(7), 2, 10, 11, ShareAckType::Reject.wire())];

        let acks = consumer.take_piggyback_acks();

        assert2::assert!(
            (
                consumer.prev_delivered.is_empty(),
                consumer.pending_acks.is_empty()
            ) == (true, true)
        );
        let batch = only(&acks[&(id(7), 2)]);
        assert2::assert!(
            *batch
                == FetchAckBatch {
                    first_offset: 10,
                    last_offset: 11,
                    acknowledge_types: vec![3, 3],
                    unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
                }
        );
    }

    #[tokio::test]
    async fn topic_id_for_prefers_assignment_then_names_and_defaults_unknown() {
        let consumer = test_consumer(ShareAckMode::Explicit).await;
        consumer
            .topic_names
            .lock()
            .await
            .insert(id(8), "topic-b".into());

        for (name, expected) in [
            ("topic-a", id(7)),
            ("topic-b", id(8)),
            ("missing", WireUuid::default()),
        ] {
            assert2::assert!(consumer.topic_id_for(name) == expected);
        }
    }

    #[test]
    fn build_ack_topics_groups_offsets_by_topic_and_partition() {
        let topics = build_ack_topics(vec![
            (id(7), 2, 10, 12, ShareAckType::Accept.wire()),
            (id(7), 3, 20, 20, ShareAckType::Release.wire()),
            (id(8), 1, 30, 31, ShareAckType::Reject.wire()),
        ]);

        let topic = topics.iter().find(|topic| topic.topic_id == id(7)).unwrap();
        assert2::assert!(topic.partitions.len() == 2);
        let part = topic
            .partitions
            .iter()
            .find(|part| part.partition_index == 2)
            .unwrap();
        let batch = only(&part.acknowledgement_batches);
        assert2::assert!(
            *batch
                == AckAckBatch {
                    first_offset: 10,
                    last_offset: 12,
                    acknowledge_types: vec![1, 1, 1],
                    unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
                }
        );

        let renew = build_ack_topics(vec![(id(7), 2, 10, 10, 4)]);
        let renew_batch = only(&only(&only(&renew).partitions).acknowledgement_batches);
        assert2::assert!(
            *renew_batch
                == AckAckBatch {
                    first_offset: 10,
                    last_offset: 10,
                    acknowledge_types: vec![4],
                    unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
                }
        );
    }

    /// Kafka's `ShareCompletedFetch` gives each record the timestamp and the
    /// timestamp type of its record batch, as the classic consumer does.
    #[test]
    fn share_records_take_the_batch_timestamp_type() {
        use krabka_protocol::owned::share_fetch_response::AcquiredRecords;

        use crate::poll::timestamp_cases;

        for case in &timestamp_cases::CASES {
            let part = PartitionData {
                partition_index: 2,
                records: Some(timestamp_cases::batch(case).into()),
                acquired_records: vec![AcquiredRecords {
                    first_offset: timestamp_cases::BASE_OFFSET,
                    last_offset: timestamp_cases::BASE_OFFSET + 1,
                    delivery_count: 1,
                    ..Default::default()
                }],
                ..Default::default()
            };

            let expected: Vec<ShareConsumerRecord> = case
                .expected
                .iter()
                .zip(timestamp_cases::BASE_OFFSET..)
                .zip(timestamp_cases::VALUES)
                .map(
                    |(((timestamp, timestamp_type), offset), value)| ShareConsumerRecord {
                        topic: "topic-a".into(),
                        partition: 2,
                        offset,
                        timestamp: *timestamp,
                        timestamp_type: *timestamp_type,
                        key: None,
                        value: Some(bytes::Bytes::from_static(value)),
                        headers: Vec::new(),
                        delivery_count: 1,
                    },
                )
                .collect();
            check!(
                share_partition_records("topic-a", &part) == expected,
                "case {}",
                case.name
            );
        }
    }
}