krafka 0.19.1

//! Fetch session management (KIP-227).
//!
//! Implements incremental fetch sessions to reduce fetch request payload
//! sizes. Instead of sending the full partition list on every poll, the
//! broker tracks session state and the client only sends partition changes.
//!
//! A per-broker `FetchSessionState` tracks:
//! - `session_id` (returned by the broker) and `session_epoch` (maintained by
//!   the client)
//! - The set of topics/partitions (with their fetch offsets and parameters)
//!   that are currently registered in the session
//!
//! On each fetch cycle the consumer computes a diff against the previous
//! session state and sends only the new/changed partitions in the `topics`
//! field plus any removed partitions in the `forgotten_topics` field.
//!
//! ## Topic keying (KIP-227 + KIP-516)
//!
//! When the broker supports Fetch v13+, each topic carries a 128-bit UUID
//! (`topic_id`). When a non-zero UUID is present the session uses it as the
//! primary key for the internal state map, reducing per-poll diff lookups
//! from variable-length string comparisons to fixed 16-byte equality tests.
//! Fallback to name-based keys is automatic when UUIDs are unavailable
//! (Fetch ≤ v12 or zero UUID).
//!
//! ## Session lifecycle
//!
//! - Epoch `0` ([`INITIAL_EPOCH`]): full fetch, opens or re-opens a session
//! - Epoch `1..=i32::MAX`: incremental fetches carrying only the diff
//! - Epoch `-1` ([`FINAL_EPOCH`]): closes the session
//!
//! Brokers have a bounded number of session slots. When the client gives up a
//! live session — rebalance, unsubscribe, close — it should send the final
//! epoch via [`FetchSessionCache::close_all`] rather than merely dropping
//! local state, otherwise the slot stays occupied until the broker evicts it
//! and other clients may meanwhile be rejected with
//! `FETCH_SESSION_ID_NOT_FOUND`.

use ahash::{AHashMap as HashMap, AHashSet as HashSet};

use crate::protocol::{FetchForgottenTopic, FetchPartitionRequest, FetchTopicRequest};
use crate::{BrokerId, PartitionId};

/// Epoch value indicating the initial (full) fetch.
pub const INITIAL_EPOCH: i32 = 0;

/// Epoch value that asks the broker to close the fetch session.
///
/// Sent in the `session_epoch` field of a `Fetch` request (v7+) together with
/// the established `session_id`. A broker keeps a bounded number of session
/// slots; without an explicit close the abandoned session lingers until LRU
/// eviction reclaims it, and in the meantime other clients can be turned away
/// with `FETCH_SESSION_ID_NOT_FOUND`.
pub(crate) const FINAL_EPOCH: i32 = -1;

/// All-zeroes UUID sentinel — indicates "no UUID" in Kafka wire protocol.
const ZERO_UUID: [u8; 16] = [0u8; 16];

/// Internal key used to track a topic inside a fetch session.
///
/// When the broker provides a non-zero `topic_id` (Fetch v13+), the UUID is
/// used directly, enabling O(1) diff lookups with 16-byte fixed-size keys.
/// When no UUID is available the topic name is used as a fallback.
#[derive(Debug, Clone, PartialEq, Eq, Hash)]
enum SessionKey {
    /// Topic name (Fetch v12 and below, or zero UUID).
    Name(String),
    /// Topic UUID (Fetch v13+).
    Uuid([u8; 16]),
}

impl SessionKey {
    fn from_request(topic: &FetchTopicRequest) -> Self {
        match topic.topic_id {
            Some(id) if id != ZERO_UUID => Self::Uuid(id),
            _ => Self::Name(topic.topic.clone()),
        }
    }
}

/// Tracked state for one topic inside a fetch session.
#[derive(Debug)]
struct TopicSession {
    /// Topic name — always stored so wire requests can be built correctly
    /// even when the primary session key is a UUID.
    name: String,
    /// Non-zero UUID when the session is keyed by UUID (Fetch v13+, KIP-516).
    ///
    /// Stored here so that `FetchForgottenTopic` entries for removed topics
    /// can carry the correct `topic_id` even after the topic is no longer in
    /// the `desired` list.
    uuid: Option<[u8; 16]>,
    /// Per-partition fetch state.
    partitions: HashMap<PartitionId, PartitionState>,
}

/// Snapshot of partition fetch parameters, used to detect changes between
/// consecutive fetches so that incremental requests only carry deltas.
///
/// Every field that is serialised on the wire for a partition is captured
/// here, and the diff compares the whole struct for equality. Tracking a
/// subset would let a changed field be silently dropped from an incremental
/// request, leaving the broker's session with the stale value for the rest of
/// the session's lifetime.
///
/// `current_leader_epoch` in particular must be tracked: it is serialised from
/// Fetch v9 onwards and is what lets the broker fence a consumer that is still
/// reading from a deposed leader (KIP-320). If an epoch bump were treated as
/// "unchanged", the broker would keep validating fetches against the old
/// epoch and the fencing check would never fire.
#[derive(Debug, Clone, PartialEq, Eq)]
struct PartitionState {
    fetch_offset: i64,
    log_start_offset: i64,
    partition_max_bytes: i32,
    current_leader_epoch: i32,
    last_fetched_epoch: i32,
}

impl PartitionState {
    /// Capture the wire-relevant fields of a partition request.
    fn from_request(part: &FetchPartitionRequest) -> Self {
        Self {
            fetch_offset: part.fetch_offset,
            log_start_offset: part.log_start_offset,
            partition_max_bytes: part.partition_max_bytes,
            current_leader_epoch: part.current_leader_epoch,
            last_fetched_epoch: part.last_fetched_epoch,
        }
    }
}

/// Per-broker fetch session state.
#[derive(Debug)]
pub struct FetchSessionState {
    /// Broker ID this session belongs to.
    ///
    /// Not read in production code but included in `Debug` output for
    /// diagnostics (the compiler intentionally ignores derive usage for
    /// dead-code analysis).
    #[allow(dead_code)]
    broker_id: BrokerId,
    /// Session ID returned by the broker (0 = no session).
    session_id: i32,
    /// Current epoch.
    epoch: i32,
    /// Topics currently registered in the session.
    ///
    /// Keyed by [`SessionKey::Uuid`] when a non-zero `topic_id` is available
    /// (Fetch v13+), otherwise by [`SessionKey::Name`].  Both variants map to
    /// a [`TopicSession`] which retains the human-readable topic name so that
    /// wire requests can be constructed correctly regardless of key type.
    topics: HashMap<SessionKey, TopicSession>,
}

/// The result of computing a fetch request from session state.
#[derive(Debug)]
pub struct FetchSessionRequest {
    /// Session ID to send.
    pub session_id: i32,
    /// Session epoch to send.
    pub session_epoch: i32,
    /// Topics/partitions to include in the request (new or changed).
    pub topics: Vec<FetchTopicRequest>,
    /// Topics/partitions to remove from the session.
    pub forgotten_topics: Vec<FetchForgottenTopic>,
    /// Whether this is a full fetch (epoch 0) vs incremental.
    pub is_full_fetch: bool,
}

/// Instruction to tear down one established fetch session on its broker.
///
/// Produced by [`FetchSessionState::close`] / [`FetchSessionCache::close_all`].
/// The caller sends a `Fetch` request (v7+) carrying these `session_id` and
/// `session_epoch` values with empty `topics` and `forgotten_topics`; the
/// broker then frees the session slot instead of holding it until eviction.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct FetchSessionClose {
    /// Broker holding the session to be closed.
    pub broker_id: BrokerId,
    /// Session ID previously assigned by that broker.
    pub session_id: i32,
    /// Always [`FINAL_EPOCH`]; carried explicitly so the caller can put it
    /// straight into the request.
    pub session_epoch: i32,
}

impl FetchSessionState {
    /// Create a new fetch session for a broker. Starts with no session.
    pub fn new(broker_id: BrokerId) -> Self {
        Self {
            broker_id,
            session_id: 0,
            epoch: INITIAL_EPOCH,
            topics: HashMap::new(),
        }
    }

    /// Returns true if the session is established (session_id != 0).
    pub fn has_session(&self) -> bool {
        self.session_id != 0
    }

    #[cfg(test)]
    pub fn broker_id(&self) -> BrokerId {
        self.broker_id
    }

    #[cfg(test)]
    pub fn session_id(&self) -> i32 {
        self.session_id
    }

    #[cfg(test)]
    pub fn epoch(&self) -> i32 {
        self.epoch
    }

    #[cfg(test)]
    pub fn partition_count(&self) -> usize {
        self.topics.values().map(|t| t.partitions.len()).sum()
    }

    /// Build the fetch request parameters by computing the diff between the
    /// desired partition set and the current session state.
    ///
    /// When the supplied topics carry non-zero `topic_id` values (Fetch v13+),
    /// the diff uses 16-byte UUID keys for O(1) lookup.  For Fetch ≤ v12 the
    /// fallback is topic-name keys, identical to the prior behaviour.
    ///
    /// Returns a `FetchSessionRequest` containing:
    /// - `session_id` / `session_epoch` to set on the wire
    /// - `topics` with only new/changed partitions (or all if full fetch)
    /// - `forgotten_topics` with removed partitions
    pub fn build_request(&self, desired: &[FetchTopicRequest]) -> FetchSessionRequest {
        if !self.has_session() {
            // No established session → full fetch (epoch 0).
            return FetchSessionRequest {
                session_id: 0,
                session_epoch: INITIAL_EPOCH,
                topics: desired.to_vec(),
                forgotten_topics: Vec::new(),
                is_full_fetch: true,
            };
        }

        // Incremental fetch: compute diff.
        // self.epoch already represents the next epoch to send;
        // it is bumped by update_from_response() after a successful response.
        let epoch = self.epoch;

        // Build a keyed map of the desired partitions for O(1) diff lookups
        // and removed-partition detection.  Keyed the same way session state
        // is stored (UUID when available, topic name otherwise) to ensure
        // correct matching.  This single map replaces both the `changed`
        // HashMap and the `desired_keys` HashSet used in a previous version.
        let mut desired_map: HashMap<SessionKey, HashMap<PartitionId, &FetchPartitionRequest>> =
            HashMap::with_capacity(desired.len());
        for topic in desired {
            let key = SessionKey::from_request(topic);
            let part_map = desired_map.entry(key).or_default();
            part_map.reserve(topic.partitions.len());
            for part in &topic.partitions {
                part_map.insert(part.partition, part);
            }
        }

        // 1. Find new or changed partitions.
        // Builds the output Vec<FetchTopicRequest> directly — no intermediate
        // HashMap allocation.
        let mut topics: Vec<FetchTopicRequest> = Vec::new();
        for topic in desired {
            let key = SessionKey::from_request(topic);
            let session_topic = self.topics.get(&key);
            let changed_partitions: Vec<FetchPartitionRequest> = topic
                .partitions
                .iter()
                .filter(|part| {
                    // A partition is re-sent when it is new to the session or
                    // when ANY tracked field differs from what the broker
                    // currently holds — comparing only a subset would leave the
                    // broker with stale values it can never learn about.
                    match session_topic.and_then(|t| t.partitions.get(&part.partition)) {
                        None => true,
                        Some(prev) => *prev != PartitionState::from_request(part),
                    }
                })
                .cloned()
                .collect();
            if !changed_partitions.is_empty() {
                topics.push(FetchTopicRequest {
                    topic: topic.topic.clone(),
                    topic_id: topic.topic_id,
                    partitions: changed_partitions,
                });
            }
        }

        // 2. Find removed topics/partitions.
        // Uses desired_map for O(1) lookup — eliminates a separate HashSet
        // allocation.  Uses session_topic.uuid for the forgotten-topic entry
        // so that UUID-keyed sessions (Fetch v13+, KIP-516) send the correct
        // topic_id to the broker even when the topic is no longer in `desired`.
        let mut forgotten_topics: Vec<FetchForgottenTopic> = Vec::new();
        for (key, session_topic) in &self.topics {
            if let Some(desired_parts) = desired_map.get(key) {
                // Topic still present — check for removed partitions.
                let removed: Vec<i32> = session_topic
                    .partitions
                    .keys()
                    .copied()
                    .filter(|p| !desired_parts.contains_key(p))
                    .collect();
                if !removed.is_empty() {
                    forgotten_topics.push(FetchForgottenTopic {
                        topic: session_topic.name.clone(),
                        topic_id: session_topic.uuid,
                        partitions: removed,
                    });
                }
            } else {
                // Entire topic removed.
                let parts: Vec<i32> = session_topic.partitions.keys().copied().collect();
                if !parts.is_empty() {
                    forgotten_topics.push(FetchForgottenTopic {
                        topic: session_topic.name.clone(),
                        topic_id: session_topic.uuid,
                        partitions: parts,
                    });
                }
            }
        }

        FetchSessionRequest {
            session_id: self.session_id,
            session_epoch: epoch,
            topics,
            forgotten_topics,
            is_full_fetch: false,
        }
    }

    /// Update session state after a successful fetch response.
    ///
    /// - If the broker returned a non-zero `session_id`, establish/continue
    ///   the session and bump the epoch.
    /// - If the broker returned `session_id == 0`, the session was closed.
    ///
    /// `desired` is the full partition list the consumer requested this cycle,
    /// used to rebuild the tracked topic/partition set.
    pub fn update_from_response(
        &mut self,
        response_session_id: i32,
        desired: &[FetchTopicRequest],
    ) {
        if response_session_id == 0 {
            // Broker doesn't support sessions or chose to close ours.
            self.reset();
            return;
        }

        // Establish or continue the session.
        self.session_id = response_session_id;
        self.epoch = self.next_epoch();

        // Rebuild tracked topic/partition set from the full desired list.
        // This ensures our state matches what the broker tracked.
        self.topics.clear();
        for topic in desired {
            let key = SessionKey::from_request(topic);
            let mut partitions = HashMap::with_capacity(topic.partitions.len());
            for part in &topic.partitions {
                partitions.insert(part.partition, PartitionState::from_request(part));
            }
            self.topics.insert(
                key,
                TopicSession {
                    name: topic.topic.clone(),
                    uuid: topic.topic_id.filter(|id| *id != ZERO_UUID),
                    partitions,
                },
            );
        }
    }

    /// Drop local session state without telling the broker.
    /// The next fetch will be a full fetch.
    ///
    /// Use this when the session is already gone or unusable on the broker
    /// side — a `FETCH_SESSION_ID_NOT_FOUND` / `INVALID_FETCH_SESSION_EPOCH`
    /// error, or a connection failure — where a close request would be
    /// pointless or impossible to deliver. When the session is still live and
    /// the client is voluntarily giving it up (rebalance, unsubscribe, close),
    /// prefer [`FetchSessionState::close`] so the broker can free the slot
    /// immediately.
    pub fn reset(&mut self) {
        self.session_id = 0;
        self.epoch = INITIAL_EPOCH;
        self.topics.clear();
    }

    /// Give up an established session, returning the request the caller must
    /// send so the broker releases its session slot.
    ///
    /// Local state is cleared unconditionally, so the next fetch is a full
    /// fetch whether or not the returned request is delivered successfully.
    /// Returns `None` when no session is established and there is therefore
    /// nothing for the broker to release.
    ///
    /// Sending is best-effort: the broker reclaims the slot via LRU eviction
    /// if the request never arrives, so callers may ignore transport errors.
    pub fn close(&mut self) -> Option<FetchSessionClose> {
        if !self.has_session() {
            self.reset();
            return None;
        }
        let close = FetchSessionClose {
            broker_id: self.broker_id,
            session_id: self.session_id,
            session_epoch: FINAL_EPOCH,
        };
        self.reset();
        Some(close)
    }

    fn next_epoch(&self) -> i32 {
        if self.epoch == i32::MAX {
            // Wrap around (matching the Java client).
            1
        } else {
            self.epoch + 1
        }
    }
}

/// Cache of fetch sessions, one per broker.
#[derive(Debug, Default)]
pub struct FetchSessionCache {
    sessions: HashMap<BrokerId, FetchSessionState>,
}

impl FetchSessionCache {
    /// Create a new empty cache.
    pub fn new() -> Self {
        Self {
            sessions: HashMap::new(),
        }
    }

    /// Get or create the session state for a broker.
    pub fn get_or_create(&mut self, broker_id: BrokerId) -> &mut FetchSessionState {
        self.sessions
            .entry(broker_id)
            .or_insert_with(|| FetchSessionState::new(broker_id))
    }

    /// Reset the session for a specific broker (e.g., on error).
    pub fn reset_broker(&mut self, broker_id: BrokerId) {
        if let Some(session) = self.sessions.get_mut(&broker_id) {
            session.reset();
        }
    }

    /// Clear every session and return the close requests the caller must send.
    ///
    /// All local state is dropped up front, so the cache is usable again
    /// immediately and the next fetch to each broker is a full fetch. Only
    /// brokers with an established session appear in the result.
    ///
    /// Sending the returned requests is best-effort — an undelivered close
    /// merely leaves the broker to evict the session on its own schedule.
    pub fn close_all(&mut self) -> Vec<FetchSessionClose> {
        self.sessions
            .values_mut()
            .filter_map(FetchSessionState::close)
            .collect()
    }

    /// Remove sessions for brokers not in the given set.
    ///
    /// Call this during metadata refresh or after a rebalance when the set of
    /// live brokers changes, so sessions for departed brokers do not accumulate
    /// indefinitely. The next fetch to a re-joined broker will start a fresh
    /// full fetch, which is correct behavior.
    pub(crate) fn retain_brokers(&mut self, broker_ids: &[BrokerId]) {
        let broker_set: HashSet<_> = broker_ids.iter().copied().collect();
        self.sessions.retain(|id, _| broker_set.contains(id));
    }
}

#[cfg(test)]
#[allow(clippy::unwrap_used, clippy::expect_used, clippy::panic)]
mod tests {
    use super::*;

    // ── Zero-UUID handling and epoch progression ──────────────────────
    //
    // Added after `cargo mutants` showed the suite could not distinguish the
    // real implementations from five different corruptions. All three sites
    // below fail *quietly* when they are wrong — the client keeps working, it
    // just stops getting KIP-227's benefit or starts keying sessions wrongly —
    // which is precisely why nothing noticed.

    const REAL_UUID: [u8; 16] = [7u8; 16];

    /// A zero UUID means "no topic ID" (Fetch v12 and below), not "the topic
    /// whose ID is zero".
    ///
    /// The guard picks between keying a session by UUID and keying it by name.
    /// Treating a zero UUID as real would collapse *every* pre-v13 topic onto
    /// the single key `Uuid([0; 16])`, so one topic's partition state would
    /// overwrite another's and the incremental diff would describe the wrong
    /// topic entirely.
    #[test]
    fn session_key_treats_a_zero_uuid_as_absent() {
        let mut with_real = make_topic_request("orders", &[(0, 0, 1024)]);
        with_real.topic_id = Some(REAL_UUID);
        assert_eq!(
            SessionKey::from_request(&with_real),
            SessionKey::Uuid(REAL_UUID),
            "a real UUID must key the session"
        );

        let mut with_zero = make_topic_request("orders", &[(0, 0, 1024)]);
        with_zero.topic_id = Some(ZERO_UUID);
        assert_eq!(
            SessionKey::from_request(&with_zero),
            SessionKey::Name("orders".to_string()),
            "a zero UUID is not an identifier — fall back to the name"
        );

        let absent = make_topic_request("orders", &[(0, 0, 1024)]);
        assert_eq!(
            SessionKey::from_request(&absent),
            SessionKey::Name("orders".to_string()),
            "no UUID at all falls back to the name"
        );

        // The two fallbacks must agree, or a broker that starts sending zero
        // UUIDs would silently re-key every existing session.
        assert_eq!(
            SessionKey::from_request(&with_zero),
            SessionKey::from_request(&absent)
        );
    }

    /// The stored UUID follows the same rule as the key.
    ///
    /// It is what `FetchForgottenTopic` entries carry, so storing a zero UUID
    /// as if it were real would send `topic_id: 0` for a removed topic and the
    /// broker would not know which topic to drop.
    #[test]
    fn update_from_response_stores_only_a_real_uuid() {
        let mut state = FetchSessionState::new(1);

        let mut real = make_topic_request("orders", &[(0, 0, 1024)]);
        real.topic_id = Some(REAL_UUID);
        state.update_from_response(99, std::slice::from_ref(&real));
        assert_eq!(
            state.topics.get(&SessionKey::Uuid(REAL_UUID)).unwrap().uuid,
            Some(REAL_UUID)
        );

        let mut zero = make_topic_request("payments", &[(0, 0, 1024)]);
        zero.topic_id = Some(ZERO_UUID);
        state.update_from_response(99, std::slice::from_ref(&zero));
        assert_eq!(
            state
                .topics
                .get(&SessionKey::Name("payments".to_string()))
                .unwrap()
                .uuid,
            None,
            "a zero UUID must not be stored as an identifier"
        );
    }

    /// The epoch advances by one and wraps to 1, never to 0.
    ///
    /// Epoch 0 means "start a new session" in KIP-227, so wrapping through it
    /// would silently turn an incremental fetch into a full one. A constant
    /// epoch is worse and just as quiet: the broker rejects it with
    /// `INVALID_FETCH_SESSION_EPOCH`, the client resets, and every fetch
    /// degrades to a full one — KIP-227's entire benefit lost with no error
    /// surfaced anywhere.
    #[test]
    fn next_epoch_advances_and_wraps_past_zero() {
        let mut state = FetchSessionState::new(1);
        assert_eq!(state.epoch, 0, "a fresh session starts at 0");
        assert_eq!(state.next_epoch(), 1);

        state.epoch = 1;
        assert_eq!(state.next_epoch(), 2, "the epoch must advance, not repeat");

        state.epoch = 41;
        assert_eq!(state.next_epoch(), 42);

        state.epoch = i32::MAX;
        assert_eq!(
            state.next_epoch(),
            1,
            "wrapping must skip 0, which means \"new session\""
        );
    }

    fn make_topic_request(topic: &str, partitions: &[(i32, i64, i32)]) -> FetchTopicRequest {
        FetchTopicRequest {
            topic: topic.to_string(),
            topic_id: None,
            partitions: partitions
                .iter()
                .map(|&(partition, offset, max_bytes)| FetchPartitionRequest {
                    partition,
                    current_leader_epoch: -1,
                    fetch_offset: offset,
                    last_fetched_epoch: -1,
                    log_start_offset: -1,
                    partition_max_bytes: max_bytes,
                    replica_directory_id: None,
                    high_watermark: None,
                })
                .collect(),
        }
    }

    fn make_topic_request_with_epoch(
        topic: &str,
        partitions: &[(i32, i64, i32, i32)],
    ) -> FetchTopicRequest {
        FetchTopicRequest {
            topic: topic.to_string(),
            topic_id: None,
            partitions: partitions
                .iter()
                .map(
                    |&(partition, offset, max_bytes, epoch)| FetchPartitionRequest {
                        partition,
                        current_leader_epoch: epoch,
                        fetch_offset: offset,
                        last_fetched_epoch: -1,
                        log_start_offset: -1,
                        partition_max_bytes: max_bytes,
                        replica_directory_id: None,
                        high_watermark: None,
                    },
                )
                .collect(),
        }
    }

    #[test]
    fn test_new_session_starts_with_no_session() {
        let state = FetchSessionState::new(1);
        assert_eq!(state.session_id(), 0);
        assert_eq!(state.epoch(), INITIAL_EPOCH);
        assert!(!state.has_session());
        assert_eq!(state.partition_count(), 0);
    }

    #[test]
    fn test_first_fetch_is_full() {
        let state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];

        let req = state.build_request(&desired);
        assert!(req.is_full_fetch);
        assert_eq!(req.session_id, 0);
        assert_eq!(req.session_epoch, INITIAL_EPOCH);
        assert_eq!(req.topics.len(), 1);
        assert!(req.forgotten_topics.is_empty());
    }

    #[test]
    fn test_session_established_after_response() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];

        // Simulate broker returning session_id=42
        state.update_from_response(42, &desired);

        assert!(state.has_session());
        assert_eq!(state.session_id(), 42);
        assert_eq!(state.epoch(), 1);
        assert_eq!(state.partition_count(), 1);
    }

    #[test]
    fn test_incremental_fetch_no_changes() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];

        state.update_from_response(42, &desired);

        // Same request again — no changes.
        let req = state.build_request(&desired);
        assert!(!req.is_full_fetch);
        assert_eq!(req.session_id, 42);
        assert_eq!(req.session_epoch, 1); // epoch maintained by client
        assert!(req.topics.is_empty()); // no changes
        assert!(req.forgotten_topics.is_empty()); // no removals
    }

    #[test]
    fn test_incremental_fetch_offset_changed() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);

        // Offset advanced.
        let desired2 = vec![make_topic_request("topic-a", &[(0, 200, 1048576)])];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);
        assert_eq!(req.topics.len(), 1);
        assert_eq!(req.topics[0].partitions[0].fetch_offset, 200);
        assert!(req.forgotten_topics.is_empty());
    }

    #[test]
    fn test_incremental_fetch_partition_added() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);

        // Add partition 1.
        let desired2 = vec![make_topic_request(
            "topic-a",
            &[(0, 100, 1048576), (1, 0, 1048576)],
        )];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);
        // Only the new partition 1 should appear.
        assert_eq!(req.topics.len(), 1);
        assert_eq!(req.topics[0].partitions.len(), 1);
        assert_eq!(req.topics[0].partitions[0].partition, 1);
        assert!(req.forgotten_topics.is_empty());
    }

    #[test]
    fn test_incremental_fetch_partition_removed() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request(
            "topic-a",
            &[(0, 100, 1048576), (1, 50, 1048576)],
        )];
        state.update_from_response(42, &desired);

        // Remove partition 1.
        let desired2 = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);
        assert!(req.topics.is_empty()); // no changes to p0
        assert_eq!(req.forgotten_topics.len(), 1);
        assert_eq!(req.forgotten_topics[0].topic, "topic-a");
        assert_eq!(req.forgotten_topics[0].partitions, vec![1]);
    }

    #[test]
    fn test_session_reset() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);
        assert!(state.has_session());

        state.reset();
        assert!(!state.has_session());
        assert_eq!(state.session_id(), 0);
        assert_eq!(state.epoch(), INITIAL_EPOCH);
        assert_eq!(state.partition_count(), 0);
    }

    #[test]
    fn test_broker_returns_zero_session_id_closes_session() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);
        assert!(state.has_session());

        // Broker closes session.
        state.update_from_response(0, &desired);
        assert!(!state.has_session());
    }

    #[test]
    fn test_epoch_wraps_at_max() {
        let mut state = FetchSessionState::new(1);
        state.session_id = 42;
        state.epoch = i32::MAX;

        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        let req = state.build_request(&desired);
        // build_request uses self.epoch directly.
        assert_eq!(req.session_epoch, i32::MAX);

        // update_from_response bumps via next_epoch, which wraps to 1.
        state.update_from_response(42, &desired);
        assert_eq!(state.epoch(), 1);
    }

    #[test]
    fn test_mixed_changes_and_removals() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![
            make_topic_request("topic-a", &[(0, 100, 1048576), (1, 50, 1048576)]),
            make_topic_request("topic-b", &[(0, 200, 1048576)]),
        ];
        state.update_from_response(42, &desired);

        // Change: topic-a p0 offset advanced.
        // Remove: topic-b p0.
        // Add: topic-a p2 (new).
        let desired2 = vec![make_topic_request(
            "topic-a",
            &[(0, 300, 1048576), (1, 50, 1048576), (2, 0, 1048576)],
        )];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);

        // Changed/new: topic-a p0 (offset changed) and p2 (new).
        let changed_partitions: Vec<i32> = req
            .topics
            .iter()
            .flat_map(|t| t.partitions.iter().map(|p| p.partition))
            .collect();
        assert!(changed_partitions.contains(&0));
        assert!(changed_partitions.contains(&2));
        assert!(!changed_partitions.contains(&1)); // p1 unchanged

        // Forgotten: topic-b p0.
        assert_eq!(req.forgotten_topics.len(), 1);
        assert_eq!(req.forgotten_topics[0].topic, "topic-b");
    }

    #[test]
    fn test_leader_epoch_change_detected() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request_with_epoch(
            "topic-a",
            &[(0, 100, 1048576, 5)],
        )];
        state.update_from_response(42, &desired);

        // Same offset and max_bytes, but the leader epoch advanced after a
        // leader election. current_leader_epoch is serialised from Fetch v9
        // on and drives KIP-320 fencing, so the partition must be re-sent —
        // otherwise the broker keeps fencing against epoch 5 forever.
        let desired2 = vec![make_topic_request_with_epoch(
            "topic-a",
            &[(0, 100, 1048576, 6)],
        )];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);
        assert_eq!(req.topics.len(), 1);
        assert_eq!(req.topics[0].partitions.len(), 1);
        assert_eq!(req.topics[0].partitions[0].partition, 0);
        assert_eq!(req.topics[0].partitions[0].current_leader_epoch, 6);
        assert!(req.forgotten_topics.is_empty());
    }

    #[test]
    fn test_leader_epoch_unchanged_is_not_resent() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request_with_epoch(
            "topic-a",
            &[(0, 100, 1048576, 5)],
        )];
        state.update_from_response(42, &desired);

        // Identical request including the epoch — nothing to send.
        let req = state.build_request(&desired);
        assert!(!req.is_full_fetch);
        assert!(req.topics.is_empty());
        assert!(req.forgotten_topics.is_empty());
    }

    #[test]
    fn test_leader_epoch_change_tracked_after_response() {
        // After the changed epoch is acknowledged, the new value becomes the
        // baseline and the same request no longer counts as a diff.
        let mut state = FetchSessionState::new(1);
        state.update_from_response(
            42,
            &[make_topic_request_with_epoch(
                "topic-a",
                &[(0, 100, 1048576, 5)],
            )],
        );

        let desired2 = vec![make_topic_request_with_epoch(
            "topic-a",
            &[(0, 100, 1048576, 6)],
        )];
        assert_eq!(state.build_request(&desired2).topics.len(), 1);

        state.update_from_response(42, &desired2);
        assert!(state.build_request(&desired2).topics.is_empty());
    }

    #[test]
    fn test_log_start_offset_change_detected() {
        let mut state = FetchSessionState::new(1);
        let mut desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);

        desired[0].partitions[0].log_start_offset = 50;
        let req = state.build_request(&desired);

        assert_eq!(req.topics.len(), 1);
        assert_eq!(req.topics[0].partitions[0].log_start_offset, 50);
    }

    #[test]
    fn test_last_fetched_epoch_change_detected() {
        let mut state = FetchSessionState::new(1);
        let mut desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);

        desired[0].partitions[0].last_fetched_epoch = 3;
        let req = state.build_request(&desired);

        assert_eq!(req.topics.len(), 1);
        assert_eq!(req.topics[0].partitions[0].last_fetched_epoch, 3);
    }

    #[test]
    fn test_close_returns_final_epoch_request() {
        let mut state = FetchSessionState::new(7);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);

        let close = state.close().expect("established session must be closed");
        assert_eq!(close.broker_id, 7);
        assert_eq!(close.session_id, 42);
        assert_eq!(close.session_epoch, FINAL_EPOCH);

        // Local state is cleared regardless of whether the request lands.
        assert!(!state.has_session());
        assert_eq!(state.epoch(), INITIAL_EPOCH);
        assert_eq!(state.partition_count(), 0);
        assert!(state.build_request(&desired).is_full_fetch);
    }

    #[test]
    fn test_close_without_session_is_noop() {
        let mut state = FetchSessionState::new(1);
        assert!(state.close().is_none());
        assert!(!state.has_session());
    }

    #[test]
    fn test_close_is_not_repeatable() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);

        assert!(state.close().is_some());
        // Session already relinquished — nothing left to tell the broker.
        assert!(state.close().is_none());
    }

    #[test]
    fn test_final_epoch_is_minus_one() {
        assert_eq!(FINAL_EPOCH, -1);
    }

    #[test]
    fn test_cache_close_all_returns_established_sessions() {
        let mut cache = FetchSessionCache::new();
        let desired = vec![make_topic_request("t", &[(0, 0, 1048576)])];
        cache.get_or_create(1).update_from_response(42, &desired);
        cache.get_or_create(2).update_from_response(43, &desired);
        // Broker 3 never established a session.
        cache.get_or_create(3);

        let mut closes = cache.close_all();
        closes.sort_unstable_by_key(|c| c.broker_id);

        assert_eq!(closes.len(), 2);
        assert_eq!(closes[0].broker_id, 1);
        assert_eq!(closes[0].session_id, 42);
        assert_eq!(closes[1].broker_id, 2);
        assert_eq!(closes[1].session_id, 43);
        assert!(closes.iter().all(|c| c.session_epoch == FINAL_EPOCH));

        // Every session is cleared, so a second close has nothing to send.
        assert!(!cache.get_or_create(1).has_session());
        assert!(!cache.get_or_create(2).has_session());
        assert!(cache.close_all().is_empty());
    }

    #[test]
    fn test_cache_get_or_create() {
        let mut cache = FetchSessionCache::new();
        {
            let session = cache.get_or_create(1);
            assert_eq!(session.broker_id(), 1);
            assert!(!session.has_session());
        }

        // Same broker returns same session.
        {
            let session = cache.get_or_create(1);
            session.session_id = 42;
        }
        {
            let session = cache.get_or_create(1);
            assert_eq!(session.session_id(), 42);
        }
    }

    #[test]
    fn test_cache_reset_broker() {
        let mut cache = FetchSessionCache::new();
        let desired = vec![make_topic_request("t", &[(0, 0, 1048576)])];
        cache.get_or_create(1).update_from_response(42, &desired);
        cache.get_or_create(2).update_from_response(43, &desired);

        cache.reset_broker(1);
        assert!(!cache.get_or_create(1).has_session());
        assert!(cache.get_or_create(2).has_session());
    }

    #[test]
    fn test_cache_retain_brokers() {
        let mut cache = FetchSessionCache::new();
        let desired = vec![make_topic_request("t", &[(0, 0, 1048576)])];
        cache.get_or_create(1).update_from_response(42, &desired);
        cache.get_or_create(2).update_from_response(43, &desired);
        cache.get_or_create(3).update_from_response(44, &desired);

        cache.retain_brokers(&[1, 3]);
        assert!(cache.get_or_create(1).has_session());
        assert!(!cache.get_or_create(2).has_session()); // was removed, recreated
        assert!(cache.get_or_create(3).has_session());
    }

    #[test]
    fn test_full_fetch_after_reset() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);

        state.reset();

        // After reset, next fetch should be full.
        let req = state.build_request(&desired);
        assert!(req.is_full_fetch);
        assert_eq!(req.session_id, 0);
        assert_eq!(req.session_epoch, INITIAL_EPOCH);
    }

    fn make_uuid_topic_request(
        topic: &str,
        uuid: [u8; 16],
        partitions: &[(i32, i64, i32)],
    ) -> FetchTopicRequest {
        FetchTopicRequest {
            topic: topic.to_string(),
            topic_id: Some(uuid),
            partitions: partitions
                .iter()
                .map(|&(partition, offset, max_bytes)| FetchPartitionRequest {
                    partition,
                    current_leader_epoch: -1,
                    fetch_offset: offset,
                    last_fetched_epoch: -1,
                    log_start_offset: -1,
                    partition_max_bytes: max_bytes,
                    replica_directory_id: None,
                    high_watermark: None,
                })
                .collect(),
        }
    }

    #[test]
    fn test_uuid_keyed_incremental_no_changes() {
        let uuid: [u8; 16] = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16];
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_uuid_topic_request(
            "topic-a",
            uuid,
            &[(0, 100, 1048576)],
        )];
        state.update_from_response(42, &desired);

        assert_eq!(state.partition_count(), 1);

        // Same request — no changes.
        let req = state.build_request(&desired);
        assert!(!req.is_full_fetch);
        assert!(req.topics.is_empty());
        assert!(req.forgotten_topics.is_empty());
    }

    #[test]
    fn test_uuid_keyed_offset_change() {
        let uuid: [u8; 16] = [0xAA; 16];
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_uuid_topic_request(
            "topic-a",
            uuid,
            &[(0, 100, 1048576)],
        )];
        state.update_from_response(42, &desired);

        let desired2 = vec![make_uuid_topic_request(
            "topic-a",
            uuid,
            &[(0, 200, 1048576)],
        )];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);
        assert_eq!(req.topics.len(), 1);
        assert_eq!(req.topics[0].topic, "topic-a");
        assert_eq!(req.topics[0].partitions[0].fetch_offset, 200);
    }

    #[test]
    fn test_uuid_keyed_partition_removed() {
        let uuid: [u8; 16] = [0xBB; 16];
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_uuid_topic_request(
            "topic-a",
            uuid,
            &[(0, 100, 1048576), (1, 50, 1048576)],
        )];
        state.update_from_response(42, &desired);

        // Remove partition 1.
        let desired2 = vec![make_uuid_topic_request(
            "topic-a",
            uuid,
            &[(0, 100, 1048576)],
        )];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);
        assert!(req.topics.is_empty()); // partition 0 unchanged
        assert_eq!(req.forgotten_topics.len(), 1);
        assert_eq!(req.forgotten_topics[0].topic, "topic-a");
        assert_eq!(req.forgotten_topics[0].partitions, vec![1]);
    }

    #[test]
    fn test_zero_uuid_falls_back_to_name_key() {
        // topic_id = [0; 16] (zero UUID sentinel) must fall back to name keying.
        let zero_uuid = [0u8; 16];
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_uuid_topic_request(
            "topic-a",
            zero_uuid,
            &[(0, 100, 1048576)],
        )];
        state.update_from_response(42, &desired);

        // Same desired → no diff.
        let req = state.build_request(&desired);
        assert!(!req.is_full_fetch);
        assert!(req.topics.is_empty());
        assert!(req.forgotten_topics.is_empty());
    }

    #[test]
    fn test_mixed_uuid_and_name_keying() {
        // Two topics: one with UUID, one without.
        let uuid: [u8; 16] = [0xCC; 16];
        let mut state = FetchSessionState::new(1);
        let desired = vec![
            make_uuid_topic_request("topic-uuid", uuid, &[(0, 100, 1048576)]),
            make_topic_request("topic-name", &[(0, 200, 1048576)]),
        ];
        state.update_from_response(42, &desired);
        assert_eq!(state.partition_count(), 2);

        // No changes.
        let req = state.build_request(&desired);
        assert!(!req.is_full_fetch);
        assert!(req.topics.is_empty());
        assert!(req.forgotten_topics.is_empty());
    }

    #[test]
    fn test_empty_desired_produces_all_forgotten() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request(
            "topic-a",
            &[(0, 100, 1048576), (1, 50, 1048576)],
        )];
        state.update_from_response(42, &desired);

        // All partitions removed.
        let desired2: Vec<FetchTopicRequest> = vec![];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);
        assert!(req.topics.is_empty());
        assert_eq!(req.forgotten_topics.len(), 1);
        assert_eq!(req.forgotten_topics[0].partitions.len(), 2);
    }

    #[test]
    fn test_max_bytes_change_detected() {
        let mut state = FetchSessionState::new(1);
        let desired = vec![make_topic_request("topic-a", &[(0, 100, 1048576)])];
        state.update_from_response(42, &desired);

        // max_bytes changed.
        let desired2 = vec![make_topic_request("topic-a", &[(0, 100, 2097152)])];
        let req = state.build_request(&desired2);

        assert!(!req.is_full_fetch);
        assert_eq!(req.topics.len(), 1);
        assert_eq!(req.topics[0].partitions[0].partition_max_bytes, 2097152);
    }
}