Skip to main content

krabka_client_consumer/
poll.rs

1//! `Consumer::poll` issues one `Fetch` that covers every assigned partition,
2//! advances next-offsets, and returns the decoded records.
3
4use std::{
5    collections::{BTreeMap, BTreeSet, HashMap, HashSet, VecDeque},
6    future,
7    future::Future,
8    mem,
9    sync::{Arc, PoisonError},
10    time::{Duration, SystemTime, UNIX_EPOCH},
11};
12
13use futures_util::future::{OptionFuture, join_all};
14use krabka_client_core::VersionBoundedRequest;
15#[cfg(test)]
16use krabka_client_core::mock_test_support;
17use krabka_ids::LeaderEpoch;
18#[cfg(test)]
19use krabka_protocol::owned::fetch_request;
20#[cfg(test)]
21use krabka_protocol::records::RecordsPayload;
22use krabka_protocol::{
23    ProtocolRequest,
24    owned::{
25        fetch_request::{FetchPartition, FetchRequest, FetchTopic, ForgottenTopic},
26        fetch_response,
27        fetch_response::FetchResponse,
28        list_offsets_request::{self, ListOffsetsPartition, ListOffsetsRequest, ListOffsetsTopic},
29        list_offsets_response::ListOffsetsResponse,
30    },
31    primitives::uuid,
32    records::{Record, RecordBatch},
33};
34use krabka_units::{
35    ByteSize, Time,
36    convert::{ByteSizeExt as _, TimeExt as _},
37    mebibytes,
38};
39use tokio::{
40    sync::{Notify, oneshot},
41    task::JoinHandle,
42    time::{self, Instant, sleep_until},
43};
44use tracing::field;
45
46#[cfg(test)]
47use crate::poll;
48#[cfg(test)]
49use crate::test_support;
50use crate::{
51    assignor,
52    builder::{AutoOffsetReset, IsolationLevel},
53    consumer,
54    consumer::{Consumer, ConsumerRecord, Header, TimestampType},
55    control, coordinator,
56    error::{ConsumerError, is_transport_error},
57    fetch_buffer::BufferedPartition,
58    fetch_session::{
59        FINAL_EPOCH, FetchSession, INVALID_SESSION_ID, SessionPartition, SessionRequest,
60    },
61    offset_wire,
62    position::{PartitionPosition, with_offsets_and_positions},
63    validate::should_await_validation,
64};
65
66/// Synthetic leader id that means "leader unknown → use the bootstrap
67/// connection".
68///
69/// It matches `BrokerPool`'s bootstrap slot, so a fallback Fetch goes out
70/// through `Client::send` and not `Client::broker(id)`.
71const BOOTSTRAP_LEADER: i32 = -1;
72const UNKNOWN_FETCH_OFFSET: i64 = -1;
73const UNKNOWN_LEADER_ID: i32 = -1;
74/// `OFFSET_OUT_OF_RANGE`.
75const OFFSET_OUT_OF_RANGE: i16 = 1;
76/// The offset the wire uses for "the broker did not answer this".
77const UNKNOWN_OFFSET: i64 = -1;
78pub const DEFAULT_FETCH_PARTITION_MAX: ByteSize = mebibytes(1);
79pub const DEFAULT_FETCH_MAX: ByteSize = mebibytes(50);
80
81/// One fetchable partition's request fields:
82/// `(partition, fetch_offset, current_leader_epoch, last_fetched_epoch)`.
83type FetchSpec = (i32, i64, LeaderEpoch, LeaderEpoch);
84
85/// Partitions to fetch, grouped first by leader id, then by topic.
86type FetchByLeader = HashMap<i32, HashMap<String, Vec<FetchSpec>>>;
87
88fn fetch_leader_id(leader_id: i32, knows_leader: bool) -> i32 {
89    if leader_id >= 0 && knows_leader {
90        leader_id
91    } else {
92        BOOTSTRAP_LEADER
93    }
94}
95
96fn should_use_bootstrap_leader(leader: i32) -> bool {
97    leader == BOOTSTRAP_LEADER
98}
99
100fn fetch_offset_or_unknown(offsets: &HashMap<(String, i32), i64>, key: &(String, i32)) -> i64 {
101    offsets.get(key).copied().unwrap_or(UNKNOWN_FETCH_OFFSET)
102}
103
104fn is_read_committed(isolation_level: IsolationLevel) -> bool {
105    isolation_level == IsolationLevel::ReadCommitted
106}
107
108pub fn is_transient_transport_error(e: &krabka_client_core::ClientError) -> bool {
109    is_transport_error(e, true)
110}
111
112fn is_transient_poll_error(e: &ConsumerError) -> bool {
113    matches!(e, ConsumerError::Client(client) if is_transient_transport_error(client))
114}
115
116/// What [`Consumer::poll`] does with one partition row of a `Fetch` response,
117/// chosen by the row's `error_code`.
118#[derive(Clone, Copy, Debug, Eq, PartialEq)]
119enum FetchPartitionAction {
120    /// `NONE`: decode the records.
121    Records,
122    /// `OFFSET_OUT_OF_RANGE`: resolve a new position with the reset policy.
123    ResetOffset,
124    /// `NOT_LEADER_OR_FOLLOWER`: adopt the leader hint, or refresh metadata.
125    Reroute,
126    /// `FENCED_LEADER_EPOCH` or `UNKNOWN_LEADER_EPOCH`: validate the position
127    /// against fresher metadata.
128    RevalidateEpoch,
129    /// `UNKNOWN_TOPIC_OR_PARTITION` or `UNKNOWN_TOPIC_ID`: skip the row and
130    /// refresh metadata. The next poll fetches the partition again.
131    RefreshMetadata,
132    /// Any other code: fail the poll with `ConsumerError::Server(code)`.
133    Fail(i16),
134}
135
136/// Map a `Fetch` partition `error_code` to its [`FetchPartitionAction`].
137///
138/// Kafka's `FetchCollector.handleInitializeErrors` also requests a metadata
139/// update for `REPLICA_NOT_AVAILABLE` (9), `KAFKA_STORAGE_ERROR` (56),
140/// `OFFSET_NOT_AVAILABLE` (78) and `INCONSISTENT_TOPIC_ID` (103), the errors of
141/// a follower that serves a fetch (KIP-392).
142///
143/// Kafka's `FetchCollector.handleInitializeErrors` requests a metadata update
144/// for `UNKNOWN_TOPIC_OR_PARTITION` (3) and for `UNKNOWN_TOPIC_ID` (100), and
145/// does not raise an error to the application. A Fetch v13 or later names the
146/// topic by id only. The broker answers 100 when its metadata image does not
147/// hold that id yet, or when the topic was deleted and created again.
148fn classify_fetch_partition_error(error_code: i16) -> FetchPartitionAction {
149    match error_code {
150        0 => FetchPartitionAction::Records,
151        1 /* OFFSET_OUT_OF_RANGE */ => FetchPartitionAction::ResetOffset,
152        6 /* NOT_LEADER_OR_FOLLOWER */ => FetchPartitionAction::Reroute,
153        74 /* FENCED_LEADER_EPOCH */ | 75 /* UNKNOWN_LEADER_EPOCH */ => {
154            FetchPartitionAction::RevalidateEpoch
155        }
156        3 /* UNKNOWN_TOPIC_OR_PARTITION */
157        | 9 /* REPLICA_NOT_AVAILABLE */
158        | 56 /* KAFKA_STORAGE_ERROR */
159        | 78 /* OFFSET_NOT_AVAILABLE */
160        | 100 /* UNKNOWN_TOPIC_ID */
161        | 103 /* INCONSISTENT_TOPIC_ID */ => FetchPartitionAction::RefreshMetadata,
162        other => FetchPartitionAction::Fail(other),
163    }
164}
165
166fn aborted_txn_started(first_offset: i64, batch_base_offset: i64) -> bool {
167    first_offset <= batch_base_offset
168}
169
170fn should_drop_aborted_batch(
171    read_committed: bool,
172    is_transactional: bool,
173    producer_is_aborted: bool,
174) -> bool {
175    read_committed && is_transactional && producer_is_aborted
176}
177
178fn record_offset(base_offset: i64, offset_delta: i32) -> i64 {
179    base_offset + i64::from(offset_delta)
180}
181
182/// The timestamp and timestamp type of `record` in `batch`.
183///
184/// Kafka's `DefaultRecordBatch` reads `max_timestamp` as the log append time
185/// when the batch attributes say `LogAppendTime`. `DefaultRecord.readFrom`
186/// then gives that time to every record and ignores the record delta.
187pub fn record_timestamp(batch: &RecordBatch, record: &Record) -> (i64, TimestampType) {
188    let timestamp_type = TimestampType::from(batch.attributes.timestamp_type());
189    let timestamp = match timestamp_type {
190        TimestampType::CreateTime => batch.base_timestamp + record.timestamp_delta,
191        TimestampType::LogAppendTime => batch.max_timestamp,
192    };
193    (timestamp, timestamp_type)
194}
195
196/// The `max_bytes` of one partition in a Fetch request.
197fn session_partition(
198    topic_id: uuid::Uuid,
199    (_, fetch_offset, leader_epoch, last_fetched_epoch): FetchSpec,
200    partition_max: ByteSize,
201) -> SessionPartition {
202    SessionPartition {
203        topic_id,
204        fetch_offset,
205        // Unwrap the leader epochs to raw wire `int32` at the FetchRequest
206        // encode boundary.
207        current_leader_epoch: leader_epoch.get(),
208        last_fetched_epoch: last_fetched_epoch.get(),
209        partition_max_bytes: partition_max.bytes_i32(),
210    }
211}
212
213/// Build the Fetch request of one broker from its session request.
214///
215/// Kafka's `AbstractFetch.createFetchRequest`: `max_wait_ms` is
216/// `fetch.max.wait.ms`, and the session id, epoch, partitions and forgotten
217/// partitions come from the fetch session (KIP-227).
218fn build_fetch_request(
219    max_wait_ms: i32,
220    isolation_level: IsolationLevel,
221    min: ByteSize,
222    max: ByteSize,
223    rack_id: &str,
224    session: SessionRequest,
225) -> FetchRequest {
226    let mut topics: Vec<FetchTopic> = Vec::new();
227    for ((topic, partition), data) in session.partitions {
228        if topics.last().is_none_or(|last| last.topic != topic) {
229            topics.push(FetchTopic {
230                topic: topic.clone(),
231                topic_id: data.topic_id,
232                ..Default::default()
233            });
234        }
235        if let Some(last) = topics.last_mut() {
236            last.partitions.push(FetchPartition {
237                partition,
238                fetch_offset: data.fetch_offset,
239                current_leader_epoch: data.current_leader_epoch,
240                last_fetched_epoch: data.last_fetched_epoch,
241                partition_max_bytes: data.partition_max_bytes,
242                ..Default::default()
243            });
244        }
245    }
246    let mut forgotten_topics_data: Vec<ForgottenTopic> = Vec::new();
247    for ((topic, partition), topic_id) in session.forgotten {
248        if forgotten_topics_data
249            .last()
250            .is_none_or(|last| last.topic != topic)
251        {
252            forgotten_topics_data.push(ForgottenTopic {
253                topic: topic.clone(),
254                topic_id,
255                ..Default::default()
256            });
257        }
258        if let Some(last) = forgotten_topics_data.last_mut() {
259            last.partitions.push(partition);
260        }
261    }
262    FetchRequest {
263        max_wait_ms,
264        min_bytes: min.bytes_i32(),
265        max_bytes: max.bytes_i32(),
266        isolation_level: isolation_level.wire(),
267        session_id: session.session_id,
268        session_epoch: session.session_epoch,
269        topics,
270        forgotten_topics_data,
271        // KIP-392: `AbstractFetch.createFetchRequest` sends `client.rack`.
272        rack_id: rack_id.to_owned(),
273        ..Default::default()
274    }
275}
276
277/// A Fetch that runs in a task, with the fetch offset of each partition that
278/// it asked for.
279struct InFlightFetch {
280    handle: JoinHandle<()>,
281    /// The result of the request. The task sends it before it notifies
282    /// `Fetches::completed`, so a waiter that wakes can always take it.
283    result: oneshot::Receiver<Result<FetchResponse, krabka_client_core::ClientError>>,
284    requested: HashMap<(String, i32), i64>,
285    /// The partitions that this Fetch asked from a preferred read replica.
286    from_replica: HashSet<(String, i32)>,
287}
288
289/// A Fetch whose task ended: its broker, the fetch offsets it asked for, and
290/// its result, or `None` when the task ended without a result.
291struct CompletedFetch {
292    leader: i32,
293    requested: HashMap<(String, i32), i64>,
294    from_replica: HashSet<(String, i32)>,
295    result: Option<Result<FetchResponse, krabka_client_core::ClientError>>,
296}
297
298/// The Fetch requests of a consumer: at most one in flight per broker, and the
299/// fetch session of each broker.
300///
301/// Kafka's `AbstractFetch` sends a Fetch to every broker that has no pending
302/// Fetch, and keeps a `FetchSessionHandler` per broker. A response that arrives
303/// after `poll` returned waits here for the next `poll`.
304#[derive(Default)]
305pub struct Fetches {
306    sessions: HashMap<i32, FetchSession>,
307    in_flight: HashMap<i32, InFlightFetch>,
308    /// The Fetch requests whose result the consumer took from `in_flight` and
309    /// did not process yet.
310    ready: Vec<CompletedFetch>,
311    completed: Arc<Notify>,
312    /// KIP-392: the replica that a broker named in `preferred_read_replica`,
313    /// with the time until which the consumer fetches from it. Kafka's
314    /// `SubscriptionState.TopicPartitionState.preferredReadReplica`.
315    preferred_read_replicas: HashMap<(String, i32), (i32, Instant)>,
316}
317
318impl Fetches {
319    /// Move each Fetch whose task sent its result, or ended, to `ready`.
320    fn collect_ready(&mut self) {
321        let leaders: Vec<i32> = self.in_flight.keys().copied().collect();
322        for leader in leaders {
323            let Some(fetch) = self.in_flight.get_mut(&leader) else {
324                continue;
325            };
326            let result = match fetch.result.try_recv() {
327                Ok(result) => Some(result),
328                Err(oneshot::error::TryRecvError::Closed) => None,
329                Err(oneshot::error::TryRecvError::Empty) => continue,
330            };
331            if let Some(fetch) = self.in_flight.remove(&leader) {
332                self.ready.push(CompletedFetch {
333                    leader,
334                    requested: fetch.requested,
335                    from_replica: fetch.from_replica,
336                    result,
337                });
338            }
339        }
340    }
341}
342
343impl Drop for Fetches {
344    fn drop(&mut self) {
345        for fetch in self.in_flight.values() {
346            fetch.handle.abort();
347        }
348    }
349}
350
351/// `ListOffsets` timestamp that asks for the log end offset.
352pub const LATEST_TIMESTAMP: i64 = -1;
353/// `ListOffsets` timestamp that asks for the log start offset.
354pub const EARLIEST_TIMESTAMP: i64 = -2;
355
356/// Placeholder for "reset with `ListOffsets`", resolved before the next Fetch:
357/// the log end, or the offset for the timestamp of `by_duration`.
358pub const LATEST_SENTINEL: i64 = i64::MAX;
359
360/// Placeholder for the reset of [`Consumer::seek_to_beginning`], resolved with
361/// `ListOffsets(-2)`: Kafka's `requestOffsetReset(partitions, EARLIEST)`.
362/// `auto_offset_reset = Earliest` plants the same sentinel at build time and
363/// on `OFFSET_OUT_OF_RANGE`, so `position()` never reports an offset before
364/// the round trip resolves the real log start.
365pub const BEGINNING_SENTINEL: i64 = i64::MAX - 1;
366
367/// Placeholder for the reset of [`Consumer::seek_to_end`]: Kafka's
368/// `requestOffsetReset(partitions, LATEST)`, resolved with `ListOffsets(-1)`.
369pub const END_SENTINEL: i64 = i64::MAX - 2;
370
371/// Placeholder for the position of a manually assigned partition: the
372/// committed offset, or a reset when the group has none. Kafka's
373/// `TopicPartitionState.shouldInitialize`.
374pub const COMMITTED_SENTINEL: i64 = i64::MAX - 3;
375
376/// Placeholder for a partition with no committed offset under
377/// `auto_offset_reset = None`. Kafka never auto-resets under `none`
378/// (`SubscriptionState.resetInitializingPositions`,
379/// `OffsetFetcherUtils.offsetResetStrategyWithValidTimestamp`); it raises
380/// `NoOffsetForPartitionException` from `poll()` and `position()` instead of
381/// planting a `ListOffsets` sentinel. This placeholder is therefore never
382/// resolved by [`Consumer::resolve_reset_sentinels`]: it stays until a commit
383/// or a seek gives the partition a real position.
384pub const NO_OFFSET_SENTINEL: i64 = i64::MAX - 4;
385
386/// Whether `next_offset` is a placeholder that a position update did not
387/// resolve yet: a reset (Kafka's `TopicPartitionState.awaitingReset`), the
388/// committed offset of a new manual assignment, or a partition with no offset
389/// under `auto_offset_reset = None`. Such a partition has no valid position,
390/// so the consumer does not fetch or commit it.
391pub fn is_reset_sentinel(next_offset: i64) -> bool {
392    next_offset >= NO_OFFSET_SENTINEL
393}
394
395/// The `ListOffsets` timestamp that resolves a [`LATEST_SENTINEL`] for the
396/// reset policy at `now_ms`.
397///
398/// Kafka's `AutoOffsetResetStrategy.timestamp` gives `-1` for `latest` and
399/// `now - duration` for `by_duration`.
400fn reset_timestamp(policy: AutoOffsetReset, now_ms: i64) -> i64 {
401    match policy {
402        AutoOffsetReset::ByDuration(duration) => {
403            now_ms.saturating_sub(i64::try_from(duration.as_millis()).unwrap_or(i64::MAX))
404        }
405        AutoOffsetReset::Earliest | AutoOffsetReset::Latest | AutoOffsetReset::None => {
406            LATEST_TIMESTAMP
407        }
408    }
409}
410
411/// Run `future`, or return [`ConsumerError::Wakeup`] when `wakeup` comes
412/// first.
413async fn until_woken<F: Future>(
414    wakeup: Option<&control::WakeupHandle>,
415    future: F,
416) -> Result<F::Output, ConsumerError> {
417    match wakeup {
418        None => Ok(future.await),
419        Some(wakeup) => tokio::select! {
420            biased;
421            () = wakeup.woken() => Err(ConsumerError::Wakeup),
422            output = future => Ok(output),
423        },
424    }
425}
426
427/// Milliseconds since the Unix epoch.
428fn unix_now_ms() -> i64 {
429    SystemTime::now()
430        .duration_since(UNIX_EPOCH)
431        .map_or(0, |elapsed| {
432            i64::try_from(elapsed.as_millis()).unwrap_or(i64::MAX)
433        })
434}
435
436/// `ListOffsets` partitions for one broker, by topic:
437/// `(partition, current_leader_epoch)`.
438type ListOffsetsSpecs = BTreeMap<String, Vec<(i32, LeaderEpoch)>>;
439
440/// Build the consumer's `ListOffsets` request for one broker.
441///
442/// Kafka's `ListOffsetsRequest.Builder.forConsumer` sends the consumer's own
443/// isolation level, so a `read_committed` consumer that asks for the latest
444/// offset gets the last stable offset and not the high watermark. The broker
445/// reads `isolation_level` from version 2, so a `read_committed` consumer
446/// sends [`ReadCommittedListOffsets`], which negotiates version 2 or higher.
447/// `current_leader_epoch` (KIP-320, version 4) comes from metadata and is `-1`
448/// when the epoch is unknown, as in `OffsetFetcher.groupListOffsetRequests`.
449fn build_offsets_request(
450    by_topic: ListOffsetsSpecs,
451    timestamp: i64,
452    isolation_level: IsolationLevel,
453) -> ListOffsetsRequest {
454    let topics: Vec<ListOffsetsTopic> = by_topic
455        .into_iter()
456        .map(|(name, partitions)| ListOffsetsTopic {
457            name,
458            partitions: partitions
459                .into_iter()
460                .map(|(partition_index, leader_epoch)| ListOffsetsPartition {
461                    partition_index,
462                    // Unwrap the leader epoch to the raw wire `int32`.
463                    current_leader_epoch: leader_epoch.get(),
464                    timestamp,
465                    ..Default::default()
466                })
467                .collect(),
468            ..Default::default()
469        })
470        .collect();
471    ListOffsetsRequest {
472        replica_id: -1,
473        isolation_level: isolation_level.wire(),
474        topics,
475        ..Default::default()
476    }
477}
478
479/// A `ListOffsets` request from a `read_committed` consumer, from version 2.
480///
481/// `isolation_level` exists from version 2. At version 0 or 1 the codec
482/// leaves it out, and the broker answers with `read_uncommitted` semantics:
483/// the high watermark and not the last stable offset. Kafka's
484/// `ListOffsetsRequest.Builder.forConsumer` sets the oldest allowed version
485/// to 2 for `READ_COMMITTED`. This type gives the same range to version
486/// negotiation. When the broker supports only version 0 or 1, the send fails
487/// with `ClientError::IncompatibleVersion` before the request goes out.
488/// Kafka fails with `UnsupportedVersionException` in that case.
489pub type ReadCommittedListOffsets =
490    VersionBoundedRequest<ListOffsetsRequest, 2, { list_offsets_request::MAX_VERSION }>;
491
492/// Group `ListOffsets` partitions by the broker that gets the request.
493///
494/// A partition goes to its leader when the leader id is known and the pool
495/// can dial it, and to the bootstrap connection otherwise. That is the rule
496/// that `Consumer::group_fetches` uses for Fetch. Each partition carries the
497/// leader epoch from its position, which is `-1` when metadata has not
498/// reported one.
499fn group_list_offsets(
500    keys: &[(String, i32)],
501    positions: &HashMap<(String, i32), PartitionPosition>,
502    knows_broker: impl Fn(i32) -> bool,
503) -> BTreeMap<i32, ListOffsetsSpecs> {
504    let mut grouped: BTreeMap<i32, ListOffsetsSpecs> = BTreeMap::new();
505    for key in keys {
506        let position = positions.get(key).copied().unwrap_or_default();
507        let leader = fetch_leader_id(position.leader_id, knows_broker(position.leader_id));
508        grouped
509            .entry(leader)
510            .or_default()
511            .entry(key.0.clone())
512            .or_default()
513            .push((key.1, position.leader_epoch));
514    }
515    grouped
516}
517
518/// What the consumer does with one partition row of a `ListOffsets` answer.
519#[derive(Clone, Copy, Debug, Eq, PartialEq)]
520enum ListOffsetsRowAction {
521    /// `NONE` with a known offset: set the position to the offset.
522    Position(i64),
523    /// `TOPIC_AUTHORIZATION_FAILED`: fail the poll.
524    Unauthorized,
525    /// Any other row: keep the position and try again after a metadata
526    /// refresh.
527    Retry,
528}
529
530/// Map a `ListOffsets` partition row to its [`ListOffsetsRowAction`].
531///
532/// Kafka's `OffsetFetcherUtils.handleListOffsetResponse` sets a position only
533/// for `NONE` with an offset other than `-1`. It throws
534/// `TopicAuthorizationException` for `TOPIC_AUTHORIZATION_FAILED`. It retries
535/// every other code, which includes codes that it does not name.
536fn classify_list_offsets_row(error_code: i16, offset: i64) -> ListOffsetsRowAction {
537    match error_code {
538        0 if offset != UNKNOWN_OFFSET => ListOffsetsRowAction::Position(offset),
539        29 /* TOPIC_AUTHORIZATION_FAILED */ => ListOffsetsRowAction::Unauthorized,
540        _ => ListOffsetsRowAction::Retry,
541    }
542}
543
544/// What the consumer takes from the `ListOffsets` answers for a set of
545/// partitions.
546#[derive(Clone, Debug, Default, Eq, PartialEq)]
547struct ListOffsetsResult {
548    /// The offset for each requested partition that got one.
549    offsets: HashMap<(String, i32), i64>,
550    /// `true` when a requested partition got no offset and must be tried
551    /// again after a metadata refresh.
552    retry: bool,
553    /// The topics that answered `TOPIC_AUTHORIZATION_FAILED`.
554    unauthorized: BTreeSet<String>,
555}
556
557impl ListOffsetsResult {
558    /// Collect the answers for `requested`.
559    ///
560    /// A requested partition that no answer names is retried. An answer with
561    /// an unauthorized row gives no offsets, because Kafka's
562    /// `handleListOffsetResponse` throws for the whole response.
563    fn collect(requested: &[(String, i32)], answers: &[ListOffsetsResponse]) -> Self {
564        let mut result = Self::default();
565        for answer in answers {
566            let mut offsets = HashMap::new();
567            let mut unauthorized = BTreeSet::new();
568            for topic in &answer.topics {
569                for row in &topic.partitions {
570                    match classify_list_offsets_row(row.error_code, row.offset) {
571                        ListOffsetsRowAction::Position(offset) => {
572                            offsets.insert((topic.name.clone(), row.partition_index), offset);
573                        }
574                        ListOffsetsRowAction::Unauthorized => {
575                            unauthorized.insert(topic.name.clone());
576                        }
577                        ListOffsetsRowAction::Retry => {}
578                    }
579                }
580            }
581            if unauthorized.is_empty() {
582                result.offsets.extend(offsets);
583            } else {
584                result.unauthorized.extend(unauthorized);
585            }
586        }
587        result.offsets.retain(|key, _| requested.contains(key));
588        result.retry = requested
589            .iter()
590            .any(|key| !result.offsets.contains_key(key));
591        result
592    }
593
594    /// `Err(TopicAuthorizationFailed)` when a topic was not authorized.
595    fn authorization(&self) -> Result<(), ConsumerError> {
596        if self.unauthorized.is_empty() {
597            Ok(())
598        } else {
599            Err(ConsumerError::TopicAuthorizationFailed(
600                self.unauthorized.clone(),
601            ))
602        }
603    }
604}
605
606impl Consumer {
607    /// Returns the records from every v2 batch that the broker returned per
608    /// assigned partition, or an empty vec on timeout.
609    ///
610    /// Under `read_committed` isolation, this method filters out control
611    /// batches and records that belong to aborted transactions on the client
612    /// side, with the response's `aborted_transactions` list. The broker returns
613    /// verbatim bytes.
614    ///
615    /// `poll` returns at most `max_poll_records` records. It keeps the rest of
616    /// a fetch for the next call, which then sends no Fetch.
617    ///
618    /// A rebalance that the coordinator asks for starts only when the
619    /// application calls `poll`. The internal coordinator task then runs the
620    /// join and mutates the live `assigned` snapshot in place. For the eager
621    /// protocol, `poll` waits for that join up to `timeout` and fetches nothing
622    /// before it completes. `poll` also resets the `max_poll_interval` timer.
623    #[tracing::instrument(
624        name = "consumer.poll",
625        level = "debug",
626        skip_all,
627        fields(
628            group_id = %self.group_id,
629            timeout_ms = timeout.millis_i64_trunc(),
630            assigned_partitions = field::Empty,
631            leaders = field::Empty,
632            records = field::Empty,
633        ),
634        err(level = "debug")
635    )]
636    /// # Errors
637    /// Returns an error when configuration is invalid, protocol encoding fails, the broker rejects the request, or transport I/O fails.
638    pub async fn poll(&mut self, timeout: Time) -> Result<Vec<ConsumerRecord>, ConsumerError> {
639        let deadline = Instant::now()
640            + Duration::from_millis(u64::try_from(timeout.millis_i64_trunc()).unwrap_or(0));
641        // Kafka's `ClassicKafkaConsumer.poll` throws `IllegalStateException`
642        // without a subscription.
643        if self.subscription.borrow().is_none() && !self.subscription.borrow().manual_assignment {
644            return Err(ConsumerError::NotSubscribed);
645        }
646        // Kafka's `ConsumerNetworkClient.maybeTriggerWakeup`: a pending
647        // wakeup fails the `poll` before it blocks.
648        if self.wakeup.take() {
649            return Err(ConsumerError::Wakeup);
650        }
651        if !self.prepare_poll(deadline).await? {
652            return Ok(Vec::new());
653        }
654        if !self.wait_for_rebalance(deadline).await? {
655            return Ok(Vec::new());
656        }
657        let wakeup = self.wakeup.clone();
658
659        // Records of an earlier fetch come first, without a new Fetch.
660        let buffered = self.drain_fetch_buffer().await;
661        if !buffered.is_empty() {
662            tracing::Span::current().record("records", buffered.len());
663            return Ok(buffered);
664        }
665
666        // 2. Build a FetchRequest covering every assigned partition.
667        let assigned = self.assigned.lock().await.clone();
668        tracing::Span::current().record("assigned_partitions", assigned.len());
669        if assigned.is_empty() {
670            tokio::select! {
671                () = sleep_until(deadline) => return Ok(Vec::new()),
672                () = wakeup.woken() => return Err(ConsumerError::Wakeup),
673            }
674        }
675
676        let by_leader = self.group_fetches(&assigned).await;
677        tracing::Span::current().record("leaders", by_leader.len());
678        let topic_ids = self.topic_ids.lock().await.clone();
679        self.send_fetches(by_leader, &topic_ids);
680        self.fetches.collect_ready();
681        if self.fetches.in_flight.is_empty() && self.fetches.ready.is_empty() {
682            tokio::select! {
683                () = self.wait_without_fetches(&assigned, deadline) => return Ok(Vec::new()),
684                () = wakeup.woken() => return Err(ConsumerError::Wakeup),
685            }
686        }
687        tokio::select! {
688            biased;
689            () = self.wait_for_fetches(deadline) => {}
690            // The Fetch requests stay in flight. The next `poll` takes their
691            // results.
692            () = wakeup.woken() => return Err(ConsumerError::Wakeup),
693        }
694        let responses = self.take_completed_fetches(&topic_ids).await?;
695
696        self.process_fetch_responses(responses, &topic_ids).await?;
697        let records = self.drain_fetch_buffer().await;
698        tracing::Span::current().record("records", records.len());
699        Ok(records)
700    }
701
702    /// Wait until the join that the coordinator task runs for this `poll`
703    /// completes, or until `deadline`. Return `false` when the join did not
704    /// complete in time.
705    ///
706    /// For the eager protocol, Kafka's `ConsumerCoordinator.onJoinPrepare`
707    /// revokes every partition before the `JoinGroup` inside `poll`, so that
708    /// `poll` fetches nothing until the join completes. The cooperative
709    /// protocol keeps the owned partitions, and `poll` fetches them while the
710    /// join runs. A member without partitions has nothing to fetch either.
711    ///
712    /// The rebalance listener calls that the join asks for run here, inside
713    /// `poll`.
714    ///
715    /// # Errors
716    ///
717    /// Returns [`ConsumerError::RebalanceListenerFailed`] when a listener
718    /// callback failed.
719    pub(crate) async fn wait_for_rebalance(
720        &mut self,
721        deadline: Instant,
722    ) -> Result<bool, ConsumerError> {
723        // Kafka's `onJoinComplete` keeps the first callback error and throws
724        // it after the rebalance step. A failed callback therefore does not
725        // end this `poll` while the join that queued it still runs.
726        let mut first_error = self.run_pending_listener_calls().await.err();
727        if !*self.rebalance_pending.borrow() {
728            return first_error.map_or(Ok(true), Err);
729        }
730        let eager = self.rebalance_protocol == assignor::RebalanceProtocol::Eager;
731        if first_error.is_none() && !eager && !self.assigned.lock().await.is_empty() {
732            return Ok(true);
733        }
734        // The coordinator task can request the rejoin after the start of this
735        // `poll` signalled it. Signal again, so that this `poll` starts the
736        // join that it waits for, as Kafka's `ensureActiveGroup` does.
737        coordinator::note_poll(&self.poll_signal);
738        let wakeup = self.wakeup.clone();
739        let joined = loop {
740            // A slow callback must not keep this `poll` past its timeout. The
741            // next `poll` runs the calls that come later.
742            if Instant::now() >= deadline {
743                break false;
744            }
745            let call = tokio::select! {
746                biased;
747                call = self.listener_calls.recv() => call,
748                joined = self.rebalance_pending.wait_for(|pending| !*pending) => {
749                    // A closed channel means that the coordinator task
750                    // stopped. `poll` then continues with the assignment that
751                    // it has.
752                    let _ = joined;
753                    break true;
754                }
755                () = sleep_until(deadline) => break false,
756                // A wakeup ends the wait between two callbacks, never inside
757                // one.
758                () = wakeup.woken() => {
759                    first_error.get_or_insert(ConsumerError::Wakeup);
760                    break false;
761                }
762            };
763            match call {
764                Some(call) => {
765                    if let Err(error) = self.complete_listener_call(call).await {
766                        first_error.get_or_insert(error);
767                    }
768                }
769                // No listener calls can come: wait for the join alone.
770                None => {
771                    break tokio::select! {
772                        biased;
773                        joined = self.rebalance_pending.wait_for(|pending| !*pending) => {
774                            let _ = joined;
775                            true
776                        }
777                        () = sleep_until(deadline) => false,
778                        () = wakeup.woken() => {
779                            first_error.get_or_insert(ConsumerError::Wakeup);
780                            false
781                        }
782                    };
783                }
784            }
785        };
786        // The join can end with calls that are still waiting, for example the
787        // assign callback.
788        if let Err(error) = self.run_pending_listener_calls().await {
789            first_error.get_or_insert(error);
790        }
791        first_error.map_or(Ok(joined), Err)
792    }
793
794    /// Return up to `max_poll_records` buffered records, and move the consumed
795    /// positions past them.
796    ///
797    /// The `assigned` guard lives until the positions move, so the coordinator
798    /// cannot revoke a partition between the ownership check and the drain.
799    /// The lock order is `assigned`, then `next_offsets`, then `positions`, as
800    /// in the coordinator task.
801    async fn drain_fetch_buffer(&mut self) -> Vec<ConsumerRecord> {
802        let ownership_ids = self.commit_identity.lock().await.ownership_ids.clone();
803        let paused = self.paused_of(&ownership_ids);
804        let assigned_guard = self.assigned.lock().await;
805        let assigned: HashSet<(String, i32)> = assigned_guard.iter().cloned().collect();
806        let mut offsets = self.next_offsets.lock().await;
807        let mut positions = self.positions.lock().await;
808        let records = self.fetch_buffer.drain(
809            self.max_poll_records,
810            &assigned,
811            &paused,
812            &mut offsets,
813            &mut positions,
814        );
815        drop(positions);
816        drop(offsets);
817        drop(assigned_guard);
818        records
819    }
820    /// Wait when this `poll` sends no Fetch and none is in flight, for example
821    /// when every partition is paused.
822    ///
823    /// Kafka's `ClassicKafkaConsumer.pollForFetches` waits for the poll timeout,
824    /// or for the retry backoff while a partition has no valid position. A
825    /// listener call that waits also ends the wait after the retry backoff, so
826    /// the next `poll` runs it.
827    async fn wait_without_fetches(&self, assigned: &[(String, i32)], deadline: Instant) {
828        let all_positions = {
829            let offsets = self.next_offsets.lock().await;
830            let positions = self.positions.lock().await;
831            assigned.iter().all(|key| {
832                offsets
833                    .get(key)
834                    .is_some_and(|offset| !is_reset_sentinel(*offset))
835                    && !positions
836                        .get(key)
837                        .is_some_and(|position| position.awaiting_validation)
838            })
839        };
840        let until = if all_positions && self.listener_calls.is_empty() {
841            deadline
842        } else {
843            deadline.min(Instant::now() + self.retry_policy.initial_backoff)
844        };
845        // A rebalance that starts during the wait queues listener calls that
846        // this consumer must run, so it ends the wait, as Kafka's
847        // `pollForFetches` waits at most `coordinator.timeToNextPoll`.
848        let mut rebalance_pending = self.rebalance_pending.clone();
849        rebalance_pending.borrow_and_update();
850        let rebalance = async {
851            // Without a coordinator task nothing changes the flag.
852            if rebalance_pending.changed().await.is_err() {
853                future::pending::<()>().await;
854            }
855        };
856        tokio::select! {
857            () = sleep_until(until) => {}
858            () = rebalance => {}
859        }
860    }
861
862    /// Decode the fetch responses into the fetch buffer, and act on the
863    /// partition errors.
864    async fn process_fetch_responses(
865        &mut self,
866        responses: Vec<fetch_response::FetchResponse>,
867        topic_ids: &HashMap<String, uuid::Uuid>,
868    ) -> Result<(), ConsumerError> {
869        // 3. Decode each partition's RecordBatches, advance next-offsets.
870        //
871        // The wire-level `records` field can carry multiple concatenated
872        // RecordBatches; we iterate every v2 batch, emit one ConsumerRecord
873        // per Record, and bump next_offsets to the highest seen offset + 1.
874        // Reverse-map topic_id → name. At Fetch v ≥ 13 the response carries
875        // only `topic_id`; `topic.topic` is empty.
876        let id_to_name = offset_wire::id_to_name(topic_ids);
877
878        // Re-snapshot the assignment: a cooperative rebalance may have
879        // revoked partitions while this Fetch was in flight. Records for
880        // partitions we no longer own must be dropped — the new owner will
881        // serve them from the offset we committed at revoke time. Snapshot
882        // before locking `next_offsets` to keep the coordinator's
883        // assigned→next_offsets lock order (avoids deadlock).
884        let still_owned: HashSet<(String, i32)> =
885            self.assigned.lock().await.iter().cloned().collect();
886
887        let mut fetched: Vec<BufferedPartition> = Vec::new();
888        let mut refresh_after_processing = false;
889        // The partitions that answered `OFFSET_OUT_OF_RANGE` under `Earliest`
890        // or `None`, with the offset each was fetched from. Both policies need
891        // the broker's real log start, and reading that is an RPC, so the loop
892        // records them and the code below the guard resolves them.
893        let mut out_of_range: Vec<((String, i32), i64)> = Vec::new();
894        // The fetch positions that this loop resets, for the auto commit
895        // before a `JoinGroup`.
896        let mut polled_resets: Vec<((String, i32), Option<i64>)> = Vec::new();
897        // KIP-392 preferred read replica changes: `Some(replica)` to fetch from
898        // it, `None` to go back to the leader.
899        let mut replica_updates: Vec<((String, i32), Option<i32>)> = Vec::new();
900        let mut offsets = self.next_offsets.lock().await;
901        for topic in responses.iter().flat_map(|resp| &resp.responses) {
902            let topic_name = if topic.topic.is_empty() {
903                id_to_name.get(&topic.topic_id).cloned().unwrap_or_default()
904            } else {
905                topic.topic.clone()
906            };
907            for part in &topic.partitions {
908                // Drop records for partitions revoked while this Fetch was
909                // in flight (cooperative rebalance transparency).
910                if !still_owned.contains(&(topic_name.clone(), part.partition_index)) {
911                    continue;
912                }
913
914                let key = (topic_name.clone(), part.partition_index);
915
916                // KIP-320 in-band truncation: leader served no records and told
917                // us where to truncate (diverging_epoch.end_offset >= 0).
918                if part.diverging_epoch.end_offset >= 0 {
919                    self.handle_truncation_in_poll(
920                        &mut offsets,
921                        &key,
922                        part.diverging_epoch.end_offset,
923                    )?;
924                    polled_resets.push((key.clone(), Some(part.diverging_epoch.end_offset)));
925                    continue;
926                }
927                // Error-first: inspect the partition error_code before decoding.
928                let action = classify_fetch_partition_error(part.error_code);
929                // Kafka's `FetchCollector.handleInitializeErrors` clears the
930                // preferred read replica with each metadata update request.
931                if matches!(
932                    action,
933                    FetchPartitionAction::Reroute
934                        | FetchPartitionAction::RevalidateEpoch
935                        | FetchPartitionAction::RefreshMetadata
936                ) {
937                    replica_updates.push((key.clone(), None));
938                }
939                match action {
940                    FetchPartitionAction::Records => {
941                        // `FetchCollector.updatePartitionState`.
942                        if part.preferred_read_replica >= 0 {
943                            replica_updates.push((key.clone(), Some(part.preferred_read_replica)));
944                        }
945                    }
946                    FetchPartitionAction::ResetOffset => {
947                        // The response cannot say where the log now starts. Apache
948                        // Kafka builds an errored partition with `log_start_offset`,
949                        // `high_watermark` and `last_stable_offset` all set to -1, so
950                        // `part.log_start_offset` is -1 here whatever the real log
951                        // start is. Reading it and fetching from it wedges the
952                        // partition: every following Fetch asks for -1 and gets
953                        // OFFSET_OUT_OF_RANGE again. A hardcoded 0 wedges it the same
954                        // way once retention has moved the log start past 0.
955                        //
956                        // So every policy resolves its position with a ListOffsets
957                        // instead. Latest and by_duration plant a sentinel that
958                        // `resolve_reset_sentinels` replaces before the next Fetch.
959                        // Earliest and None instead go to `out_of_range`, so the
960                        // true log start can be read once the offsets guard is
961                        // released; None then reports the error instead of using it.
962                        match self.auto_offset_reset {
963                            AutoOffsetReset::Latest | AutoOffsetReset::ByDuration(_) => {
964                                offsets.insert(key.clone(), LATEST_SENTINEL);
965                                polled_resets.push((key.clone(), None));
966                            }
967                            AutoOffsetReset::Earliest | AutoOffsetReset::None => {
968                                let fetch_offset = fetch_offset_or_unknown(&offsets, &key);
969                                out_of_range.push((key.clone(), fetch_offset));
970                            }
971                        }
972                        continue;
973                    }
974                    FetchPartitionAction::Reroute => {
975                        // A routing miss, NOT a truncation: we sent the Fetch to
976                        // a broker that no longer leads this partition (e.g. a
977                        // leadership change since the last metadata refresh).
978                        // Re-target the leader so the next poll routes correctly;
979                        // do NOT set awaiting_validation (nothing diverged).
980                        let mut positions = self.positions.lock().await;
981                        if part.current_leader.leader_id >= 0 {
982                            // The broker handed us the new leader inline (KIP-320
983                            // current_leader hint). Adopt it immediately.
984                            let p = positions.entry(key.clone()).or_default();
985                            p.leader_id = part.current_leader.leader_id;
986                            // Wrap the KIP-320 current-leader hint (raw wire
987                            // `int32`) at the Fetch-response decode boundary.
988                            p.leader_epoch = LeaderEpoch(part.current_leader.leader_epoch);
989                        } else {
990                            // No hint: force a metadata refresh after this loop
991                            // so the next poll learns the new leader. Reset the
992                            // stale leader id so the bootstrap fallback (and a
993                            // re-flag, if metadata advances the epoch) kicks in.
994                            if let Some(p) = positions.get_mut(&key) {
995                                p.leader_id = UNKNOWN_LEADER_ID;
996                            }
997                            drop(positions);
998                            refresh_after_processing = true;
999                        }
1000                        continue;
1001                    }
1002                    FetchPartitionAction::RevalidateEpoch => {
1003                        let mut positions = self.positions.lock().await;
1004                        if let Some(p) = positions.get_mut(&key) {
1005                            // Force refresh_leader_epochs to re-flag against
1006                            // fresher metadata next poll (any real epoch >= 0 > -1).
1007                            p.leader_epoch = LeaderEpoch(UNKNOWN_LEADER_ID);
1008                            self.client.metadata_topics().request_update();
1009                            // Only gate on validation when we have a consumed epoch
1010                            // to validate against. A never-consumed partition
1011                            // (offset_epoch < 0) has nothing to validate; flagging it
1012                            // would wedge it — validate_positions skips offset_epoch
1013                            // < 0, and the fetch builder skips awaiting_validation.
1014                            if p.offset_epoch.is_known() {
1015                                p.awaiting_validation = true;
1016                            }
1017                        }
1018                        drop(positions);
1019                        continue;
1020                    }
1021                    FetchPartitionAction::RefreshMetadata => {
1022                        // The broker does not know the topic or the topic id. Skip the
1023                        // row and refresh metadata, as Kafka's consumer does. The
1024                        // refresh at the start of the next poll also stores a new
1025                        // topic id for a topic that was created again.
1026                        tracing::warn!(
1027                            topic = %topic_name,
1028                            partition = part.partition_index,
1029                            error_code = part.error_code,
1030                            "fetch partition names a topic the broker does not know; refreshing metadata"
1031                        );
1032                        refresh_after_processing = true;
1033                        continue;
1034                    }
1035                    FetchPartitionAction::Fail(code) => {
1036                        return Err(ConsumerError::Server(code));
1037                    }
1038                }
1039
1040                record_readable_end(
1041                    &mut *self.end_offsets.lock().await,
1042                    key.clone(),
1043                    self.isolation_level,
1044                    part.high_watermark,
1045                    part.last_stable_offset,
1046                );
1047
1048                if let Some(partition) =
1049                    self.process_partition_records(&offsets, &key, &topic_name, part)
1050                {
1051                    fetched.push(partition);
1052                }
1053            }
1054        }
1055        if let Some(auto_commit) = &self.auto_commit {
1056            auto_commit.reset_polled(polled_resets).await;
1057        }
1058        // Drop the offsets guard before any `.await`: refreshing metadata is an
1059        // RPC, and we must never hold a Mutex guard across an await point.
1060        drop(offsets);
1061        let now = Instant::now();
1062        for (key, replica) in replica_updates {
1063            match replica {
1064                // `TopicPartitionState.updatePreferredReadReplica` restarts the
1065                // expiry only when the replica changes.
1066                Some(replica) => {
1067                    let expires = now + self.metadata_max_age.to_std();
1068                    let entry = self
1069                        .fetches
1070                        .preferred_read_replicas
1071                        .entry(key)
1072                        .or_insert((replica, expires));
1073                    if entry.0 != replica {
1074                        *entry = (replica, expires);
1075                    }
1076                }
1077                None => {
1078                    self.fetches.preferred_read_replicas.remove(&key);
1079                }
1080            }
1081        }
1082        for partition in fetched {
1083            self.fetch_buffer.push(partition);
1084        }
1085        if !out_of_range.is_empty() && self.recover_out_of_range(&out_of_range).await? {
1086            refresh_after_processing = true;
1087        }
1088        if refresh_after_processing {
1089            // A NOT_LEADER_OR_FOLLOWER without a current_leader hint means our
1090            // cached leader is stale, and UNKNOWN_TOPIC_OR_PARTITION and
1091            // UNKNOWN_TOPIC_ID mean our topic metadata is stale. Kafka's
1092            // fetcher calls `Metadata.requestUpdate`; the next poll refreshes
1093            // before it routes a fetch.
1094            self.client.metadata_topics().request_update();
1095        }
1096        Ok(())
1097    }
1098
1099    /// Send a Fetch to each leader in `by_leader` that has no Fetch in flight.
1100    /// Each request runs in its own task, so the leaders answer in parallel.
1101    fn send_fetches(&mut self, by_leader: FetchByLeader, topic_ids: &HashMap<String, uuid::Uuid>) {
1102        let max_wait_ms = consumer::protocol_millis_i32(self.fetch_max_wait);
1103        for (leader, by_topic) in by_leader {
1104            if self.fetches.in_flight.contains_key(&leader) {
1105                continue;
1106            }
1107            let wanted: BTreeMap<(String, i32), SessionPartition> = by_topic
1108                .into_iter()
1109                .flat_map(|(topic, specs)| {
1110                    let topic_id = topic_ids.get(&topic).copied().unwrap_or_default();
1111                    let partition_max = self.fetch_partition_max;
1112                    specs.into_iter().map(move |spec| {
1113                        (
1114                            (topic.clone(), spec.0),
1115                            session_partition(topic_id, spec, partition_max),
1116                        )
1117                    })
1118                })
1119                .collect();
1120            let requested = wanted
1121                .iter()
1122                .map(|(key, partition)| (key.clone(), partition.fetch_offset))
1123                .collect();
1124            let from_replica = wanted
1125                .keys()
1126                .filter(|key| {
1127                    self.fetches
1128                        .preferred_read_replicas
1129                        .get(*key)
1130                        .is_some_and(|(replica, _)| *replica == leader)
1131                })
1132                .cloned()
1133                .collect();
1134            let session = self
1135                .fetches
1136                .sessions
1137                .entry(leader)
1138                .or_default()
1139                .build(wanted);
1140            let request = build_fetch_request(
1141                max_wait_ms,
1142                self.isolation_level,
1143                self.fetch_min,
1144                self.fetch_max,
1145                self.client_rack.as_deref().unwrap_or_default(),
1146                session,
1147            );
1148            let client = self.client.clone();
1149            let completed = Arc::clone(&self.fetches.completed);
1150            let (sender, result) = oneshot::channel();
1151            let handle = tokio::spawn(async move {
1152                let response = if should_use_bootstrap_leader(leader) {
1153                    client.send(request).await
1154                } else {
1155                    client.broker(leader).send(request).await
1156                };
1157                let _ = sender.send(response);
1158                completed.notify_one();
1159            });
1160            self.fetches.in_flight.insert(
1161                leader,
1162                InFlightFetch {
1163                    handle,
1164                    result,
1165                    requested,
1166                    from_replica,
1167                },
1168            );
1169        }
1170    }
1171
1172    /// Wait until a Fetch in flight completes, or until `deadline`. Kafka's
1173    /// `ClassicKafkaConsumer.pollForFetches` bounds the network wait with the
1174    /// poll timer and returns as soon as a fetch is available.
1175    async fn wait_for_fetches(&mut self, deadline: Instant) {
1176        loop {
1177            self.fetches.collect_ready();
1178            if self.fetches.in_flight.is_empty() || !self.fetches.ready.is_empty() {
1179                return;
1180            }
1181            tokio::select! {
1182                () = self.fetches.completed.notified() => {}
1183                () = sleep_until(deadline) => return,
1184            }
1185        }
1186    }
1187
1188    /// Take the responses of the completed Fetch requests, and update the
1189    /// fetch session of each broker.
1190    ///
1191    /// The function drops the data of a partition whose fetch position changed
1192    /// after the request, as Kafka's `FetchCollector.initialize` discards a
1193    /// stale fetch. A response that the fetch session rejects gives no data, as
1194    /// in `AbstractFetch.handleFetchSuccess`.
1195    ///
1196    /// # Errors
1197    ///
1198    /// Returns a transport error that is not transient.
1199    async fn take_completed_fetches(
1200        &mut self,
1201        topic_ids: &HashMap<String, uuid::Uuid>,
1202    ) -> Result<Vec<FetchResponse>, ConsumerError> {
1203        self.fetches.collect_ready();
1204        if self.fetches.ready.is_empty() {
1205            return Ok(Vec::new());
1206        }
1207        let id_to_name = offset_wire::id_to_name(topic_ids);
1208        let offsets = self.next_offsets.lock().await.clone();
1209        let mut responses = Vec::new();
1210        let mut failure = None;
1211        for fetch in mem::take(&mut self.fetches.ready) {
1212            let leader = fetch.leader;
1213            let session = self.fetches.sessions.entry(leader).or_default();
1214            match fetch.result {
1215                Some(Ok(mut response)) => {
1216                    let partitions = response
1217                        .responses
1218                        .iter()
1219                        .map(|topic| topic.partitions.len())
1220                        .sum();
1221                    if !session.handle_response(
1222                        response.error_code,
1223                        response.session_id,
1224                        partitions,
1225                        response.throttle_time_ms,
1226                    ) {
1227                        tracing::debug!(
1228                            leader,
1229                            error_code = response.error_code,
1230                            "fetch session answered with an error; the next fetch is a full fetch"
1231                        );
1232                        continue;
1233                    }
1234                    let mut replica_out_of_range = Vec::new();
1235                    for topic in &mut response.responses {
1236                        let name = if topic.topic.is_empty() {
1237                            id_to_name.get(&topic.topic_id).cloned().unwrap_or_default()
1238                        } else {
1239                            topic.topic.clone()
1240                        };
1241                        topic.partitions.retain(|partition| {
1242                            let key = (name.clone(), partition.partition_index);
1243                            // A preferred replica that answers out of range
1244                            // gives no reset: Kafka's `handleInitializeErrors`
1245                            // clears the replica and fetches from the leader.
1246                            // The request decides this, as the preference can
1247                            // change while the Fetch is in flight.
1248                            if partition.error_code == OFFSET_OUT_OF_RANGE
1249                                && fetch.from_replica.contains(&key)
1250                            {
1251                                replica_out_of_range.push(key);
1252                                return false;
1253                            }
1254                            fetch.requested.get(&key).is_some_and(|requested| {
1255                                offsets.get(&key).copied().unwrap_or(0) == *requested
1256                            })
1257                        });
1258                    }
1259                    for key in replica_out_of_range {
1260                        self.fetches.preferred_read_replicas.remove(&key);
1261                    }
1262                    responses.push(response);
1263                }
1264                Some(Err(error)) => {
1265                    session.handle_error();
1266                    // `AbstractFetch.handleFetchFailure` clears the preferred read
1267                    // replica of each partition of the failed session.
1268                    for key in fetch.requested.keys() {
1269                        self.fetches.preferred_read_replicas.remove(key);
1270                    }
1271                    if is_transient_transport_error(&error) {
1272                        if should_use_bootstrap_leader(leader) {
1273                            self.client.reconnect_bootstrap().await;
1274                        } else {
1275                            self.client.evict_broker(leader);
1276                        }
1277                    } else if failure.is_none() {
1278                        failure = Some(error);
1279                    }
1280                }
1281                // The task ended without a result: it panicked or was aborted.
1282                None => session.handle_error(),
1283            }
1284        }
1285        if let Some(error) = failure {
1286            return Err(error.into());
1287        }
1288        Ok(responses)
1289    }
1290
1291    /// Close the fetch session of each broker. Kafka's `AbstractFetch.close`
1292    /// sends a Fetch with the session id, epoch `-1` and no partitions to each
1293    /// broker that holds a session.
1294    pub(crate) async fn close_fetch_sessions(&self, timeout: Duration) {
1295        let max_wait_ms = consumer::protocol_millis_i32(self.fetch_max_wait);
1296        let closes: Vec<_> = self
1297            .fetches
1298            .sessions
1299            .iter()
1300            .filter(|(_, session)| session.session_id() != INVALID_SESSION_ID)
1301            .map(|(leader, session)| {
1302                let request = build_fetch_request(
1303                    max_wait_ms,
1304                    self.isolation_level,
1305                    self.fetch_min,
1306                    self.fetch_max,
1307                    self.client_rack.as_deref().unwrap_or_default(),
1308                    SessionRequest {
1309                        session_id: session.session_id(),
1310                        session_epoch: FINAL_EPOCH,
1311                        partitions: BTreeMap::new(),
1312                        forgotten: Vec::new(),
1313                    },
1314                );
1315                let client = self.client.clone();
1316                let leader = *leader;
1317                async move {
1318                    let _ = if should_use_bootstrap_leader(leader) {
1319                        client.send(request).await
1320                    } else {
1321                        client.broker(leader).send(request).await
1322                    };
1323                }
1324            })
1325            .collect();
1326        let _ = time::timeout(timeout, join_all(closes)).await;
1327    }
1328
1329    pub(crate) async fn group_fetches(&mut self, assigned: &[(String, i32)]) -> FetchByLeader {
1330        let subscription = self.subscription.borrow().clone();
1331        let awaiting_callback = self
1332            .assigned_callback_pending
1333            .lock()
1334            .unwrap_or_else(PoisonError::into_inner)
1335            .clone();
1336        let ownership_ids = self.commit_identity.lock().await.ownership_ids.clone();
1337        let paused = self.paused_of(&ownership_ids);
1338        let buffered = self.fetch_buffer.buffered_partitions();
1339        let mut grouped: FetchByLeader = HashMap::new();
1340        let now = Instant::now();
1341        let mut refresh_metadata = false;
1342        with_offsets_and_positions(&self.next_offsets, &self.positions, |offsets, positions| {
1343            for (t, p) in assigned {
1344                if awaiting_callback.contains(&(t.clone(), *p)) {
1345                    continue;
1346                }
1347                // Kafka's `SubscriptionState.isFetchableAndSubscribed`: with a
1348                // topic subscription, a partition of a topic that the
1349                // consumer no longer subscribes to is not fetched. A pattern
1350                // subscription has no topic list of the application, and a
1351                // regular expression subscription has none at all, because the
1352                // group coordinator matches it.
1353                if !subscription.manual_assignment
1354                    && subscription.pattern.is_none()
1355                    && subscription.regex.is_none()
1356                    && !subscription.contains(t)
1357                {
1358                    continue;
1359                }
1360                // Skip partitions still awaiting validation — they must not be
1361                // fetched until proven consistent.
1362                if positions
1363                    .get(&(t.clone(), *p))
1364                    .is_some_and(|x| x.awaiting_validation)
1365                {
1366                    continue;
1367                }
1368                let next = offsets.get(&(t.clone(), *p)).copied().unwrap_or(0);
1369                // A reset sentinel is a reset that `ListOffsets` did not
1370                // resolve yet. Kafka does not fetch a partition that awaits a
1371                // reset, and the sentinel is no offset to fetch from.
1372                if is_reset_sentinel(next) {
1373                    continue;
1374                }
1375                // Kafka's `AbstractFetch.fetchablePartitions` leaves out a
1376                // paused partition and a partition with buffered records.
1377                if paused.contains(&(t.clone(), *p)) || buffered.contains(&(t.clone(), *p)) {
1378                    continue;
1379                }
1380                let pos = positions.get(&(t.clone(), *p)).copied().unwrap_or_default();
1381                // Route to the leader when its id is known AND the pool has a
1382                // dialable address for it; otherwise fall back to the bootstrap
1383                // connection. `knows_broker` is a synchronous registry lookup
1384                // (no await), so it's safe to call while the offsets/positions
1385                // guards are held. A leader whose advertised address is unusable
1386                // (e.g. port 0 from an in-process test broker) is treated as
1387                // unknown — the bootstrap broker is the leader in that
1388                // single-broker case anyway.
1389                let mut leader =
1390                    fetch_leader_id(pos.leader_id, self.client.knows_broker(pos.leader_id));
1391                // KIP-392: Kafka's `AbstractFetch.selectReadReplica` fetches from
1392                // the preferred read replica until it expires. A replica that is
1393                // not in the metadata clears the preference and refreshes it.
1394                let key = (t.clone(), *p);
1395                match self.fetches.preferred_read_replicas.get(&key).copied() {
1396                    Some((_, expires)) if now > expires => {
1397                        self.fetches.preferred_read_replicas.remove(&key);
1398                    }
1399                    Some((replica, _)) if self.client.knows_broker(replica) => leader = replica,
1400                    Some(_) => {
1401                        self.fetches.preferred_read_replicas.remove(&key);
1402                        refresh_metadata = true;
1403                    }
1404                    None => {}
1405                }
1406                // Kafka's `AbstractFetch.prepareFetchRequests` skips a broker
1407                // with a pending Fetch.
1408                if self.fetches.in_flight.contains_key(&leader) {
1409                    continue;
1410                }
1411                grouped
1412                    .entry(leader)
1413                    .or_default()
1414                    .entry(t.clone())
1415                    .or_default()
1416                    .push((*p, next, pos.leader_epoch, pos.offset_epoch));
1417            }
1418        })
1419        .await;
1420        if refresh_metadata {
1421            // The next poll refreshes, as for the other metadata updates of
1422            // the fetch path.
1423            self.client.metadata_topics().request_update();
1424        }
1425
1426        grouped
1427    }
1428
1429    async fn prepare_poll(&self, deadline: Instant) -> Result<bool, ConsumerError> {
1430        // Kafka's consumer raises a fatal `OffsetFetch` error from `poll()`.
1431        // A rejoin in the coordinator task leaves such an error here.
1432        if let Some(error) = coordinator::take_poll_error(&self.poll_error) {
1433            return Err(error);
1434        }
1435        // Kafka's `ConsumerCoordinator.poll` sends the interval auto commit
1436        // before `updateFetchPositions`.
1437        self.maybe_auto_commit_async().await;
1438        // This `poll` resets the poll timer of the coordinator task, and lets
1439        // it start a rejoin that waits for a `poll`. The positions for the
1440        // commit before that join are already recorded above.
1441        coordinator::note_poll(&self.poll_signal);
1442        // Metadata comes first. A partition that has no position yet (a new
1443        // assignment without a committed offset) gets its leader id and
1444        // epoch here, so its first `ListOffsets` goes to the leader and not
1445        // to the bootstrap broker. Kafka's `OffsetFetcher.groupListOffsetRequests`
1446        // routes with `metadata.currentLeader(tp)` in the same way.
1447        self.sync_metadata_topics();
1448        let wakeup = self.wakeup.clone();
1449        self.update_fetch_positions(Some(&wakeup), Some(deadline))
1450            .await
1451    }
1452
1453    /// Name the subscribed topics in the metadata requests of this consumer.
1454    /// The coordinator task changes the topics of a pattern subscription on
1455    /// its own client, so `poll` learns the leaders of a topic that a pattern
1456    /// adds. Kafka's `ConsumerMetadata` reads the subscription of each update.
1457    fn sync_metadata_topics(&self) {
1458        let topics = {
1459            let subscription = self.subscription.borrow();
1460            (!subscription.is_none()).then(|| subscription.topics.clone())
1461        };
1462        if let Some(topics) = topics
1463            && self.client.metadata_topics().names() != topics
1464        {
1465            self.client.metadata_topics().set(topics);
1466        }
1467    }
1468
1469    /// Whether the metadata must be refreshed before this consumer fetches.
1470    ///
1471    /// Kafka's consumer refreshes when `Metadata.updateRequested` holds: the
1472    /// metadata is older than `metadata.max.age.ms`, the topics changed, a
1473    /// fetch error asked for an update, or an assigned partition has no known
1474    /// leader. It does not refresh on every `poll`. A refresh on every poll
1475    /// pipelines a `Metadata` request behind a parked long-poll `Fetch` on the
1476    /// same connection, and the broker answers it only when the fetch
1477    /// returns, so the poll stalls for up to `fetch.max.wait.ms` and sends no
1478    /// fetch to the other leaders meanwhile (krabka-io/krabka-broker#1394).
1479    pub(crate) async fn metadata_due(&self) -> bool {
1480        let topics = self.client.metadata_topics();
1481        if topics.update_requested()
1482            || topics
1483                .last_refreshed()
1484                .is_none_or(|at| at.elapsed() >= self.metadata_max_age.to_std())
1485        {
1486            return true;
1487        }
1488        let assigned = self.assigned.lock().await.clone();
1489        let positions = self.positions.lock().await;
1490        assigned.iter().any(|key| {
1491            positions
1492                .get(key)
1493                .is_none_or(|position| position.leader_id < 0)
1494        })
1495    }
1496
1497    /// Refresh the leader epochs when [`metadata_due`](Self::metadata_due),
1498    /// resolve the offset resets and validate the positions. Kafka's
1499    /// `KafkaConsumer.updateFetchPositions`.
1500    ///
1501    /// Return `false` after a transient error, when the caller must try again
1502    /// later.
1503    ///
1504    /// With `wakeup`, a wakeup ends each request of the update, as Kafka's
1505    /// `ConsumerNetworkClient.poll` throws `WakeupException` while it waits.
1506    /// The truncation that the validation found is applied without a wakeup
1507    /// check.
1508    pub(crate) async fn update_fetch_positions(
1509        &self,
1510        wakeup: Option<&control::WakeupHandle>,
1511        deadline: Option<Instant>,
1512    ) -> Result<bool, ConsumerError> {
1513        if self.metadata_due().await
1514            && let Err(error) = until_woken(wakeup, self.refresh_leader_epochs()).await?
1515        {
1516            if is_transient_poll_error(&error) {
1517                self.client.reconnect_bootstrap().await;
1518                return Ok(false);
1519            }
1520            return Err(error);
1521        }
1522        if let Err(error) = until_woken(wakeup, self.resolve_committed_sentinels(deadline)).await? {
1523            if is_transient_poll_error(&error) {
1524                self.client.reconnect_bootstrap().await;
1525                return Ok(false);
1526            }
1527            return Err(error);
1528        }
1529        // Kafka's `resetInitializingPositions` throws `NoOffsetForPartitionException`
1530        // for every partition that has no committed offset and no reset
1531        // policy, before it ever sends a `ListOffsets`.
1532        self.fail_partitions_without_offset().await?;
1533        if let Err(error) = until_woken(wakeup, self.resolve_reset_sentinels()).await? {
1534            if is_transient_poll_error(&error) {
1535                self.client.reconnect_bootstrap().await;
1536                return Ok(false);
1537            }
1538            return Err(error);
1539        }
1540        let truncated = match until_woken(wakeup, self.validate_positions()).await? {
1541            Ok(truncated) => truncated,
1542            Err(error) if is_transient_poll_error(&error) => {
1543                self.client.reconnect_bootstrap().await;
1544                return Ok(false);
1545            }
1546            Err(error) => return Err(error),
1547        };
1548        if !truncated.is_empty() {
1549            self.apply_truncation(&truncated).await?;
1550        }
1551        Ok(true)
1552    }
1553
1554    /// Decode the records of one partition row into a [`BufferedPartition`].
1555    /// The consumed position does not move here. It moves when `poll` returns
1556    /// the records.
1557    fn process_partition_records(
1558        &self,
1559        offsets: &HashMap<(String, i32), i64>,
1560        key: &(String, i32),
1561        topic_name: &str,
1562        part: &fetch_response::PartitionData,
1563    ) -> Option<BufferedPartition> {
1564        // Legacy MessageSet payloads are skipped here; the consumer
1565        // only handles v2 batches.
1566        let batches = part.records.as_ref()?.as_v2()?;
1567        // The broker returns whole record batches whose last offset is
1568        // >= the requested fetch_offset, even when the batch starts
1569        // before it (e.g. after an OFFSET_OUT_OF_RANGE reset or when
1570        // a single large batch straddles log_start). Kafka's JVM
1571        // client skips any records below the position; we do the same.
1572        // Capture the position now — before `next_offset_after` updates
1573        // it — so the filter baseline matches the actual fetch offset.
1574        let fetch_floor = offsets.get(key).copied().unwrap_or(0);
1575        let mut records = VecDeque::new();
1576        // read_committed filtering happens entirely client-side: the
1577        // broker returns verbatim on-disk bytes (control batches,
1578        // aborted records and all) plus an `aborted_transactions`
1579        // list. We replay Kafka's algorithm — walk batches in offset
1580        // order, tracking which producer_ids have an open aborted
1581        // transaction, and drop transactional records from those.
1582        let read_committed = is_read_committed(self.isolation_level);
1583        // Aborted txns sorted by first_offset; consumed front-to-back
1584        // as batch offsets advance past each entry's start.
1585        let mut aborted: VecDeque<(i64, i64)> = if read_committed {
1586            let mut v: Vec<(i64, i64)> = part
1587                .aborted_transactions
1588                .as_deref()
1589                .unwrap_or(&[])
1590                .iter()
1591                .map(|a| (a.first_offset, a.producer_id))
1592                .collect();
1593            v.sort_unstable();
1594            v.into()
1595        } else {
1596            VecDeque::new()
1597        };
1598        // producer_ids with a currently-open aborted transaction.
1599        let mut aborted_pids: HashSet<i64> = HashSet::new();
1600        for batch in batches {
1601            // Move every aborted txn that starts at or before this
1602            // batch into the active set.
1603            if read_committed {
1604                while let Some(&(first_offset, pid)) = aborted.front() {
1605                    if aborted_txn_started(first_offset, batch.base_offset) {
1606                        aborted_pids.insert(pid);
1607                        aborted.pop_front();
1608                    } else {
1609                        break;
1610                    }
1611                }
1612            }
1613            // Control batches (commit/abort markers) carry no user
1614            // records. A control batch for a producer ends its aborted
1615            // transaction; drop the batch either way.
1616            if batch.attributes.is_control_batch() {
1617                if read_committed {
1618                    aborted_pids.remove(&batch.producer_id);
1619                }
1620                continue;
1621            }
1622            // Drop transactional records belonging to an aborted txn.
1623            if should_drop_aborted_batch(
1624                read_committed,
1625                batch.attributes.is_transactional(),
1626                aborted_pids.contains(&batch.producer_id),
1627            ) {
1628                continue;
1629            }
1630            for r in &batch.records {
1631                let offset = record_offset(batch.base_offset, r.offset_delta);
1632                // Skip records that precede the fetch floor: the broker
1633                // returned a whole batch whose base_offset < our
1634                // position (straddle case — see fetch_floor comment).
1635                if offset < fetch_floor {
1636                    continue;
1637                }
1638                let (timestamp, timestamp_type) = record_timestamp(batch, r);
1639                records.push_back(ConsumerRecord {
1640                    topic: topic_name.to_owned(),
1641                    partition: part.partition_index,
1642                    offset,
1643                    leader_epoch: batch.partition_leader_epoch,
1644                    timestamp,
1645                    timestamp_type,
1646                    key: r.key.clone(),
1647                    value: r.value.clone(),
1648                    headers: r
1649                        .headers
1650                        .iter()
1651                        .map(|h| Header {
1652                            key: h.key.clone(),
1653                            value: h.value.clone(),
1654                        })
1655                        .collect(),
1656                });
1657            }
1658        }
1659        // When `poll` returns the last record, the position moves past every
1660        // fetched batch, and the position's offset_epoch becomes the highest
1661        // batch leader epoch, so the next Fetch sends the correct
1662        // last_fetched_epoch (KIP-320).
1663        let next_offset = next_offset_after(batches)?;
1664        Some(BufferedPartition {
1665            key: key.clone(),
1666            position: fetch_floor,
1667            records,
1668            next_offset,
1669            // Wrap the batch's raw wire `partition_leader_epoch` (`int32`) at
1670            // the RecordBatch decode boundary.
1671            last_epoch: batches
1672                .iter()
1673                .map(|batch| batch.partition_leader_epoch)
1674                .max()
1675                .map(LeaderEpoch),
1676        })
1677    }
1678}
1679
1680const fn readable_end_offset(
1681    isolation: IsolationLevel,
1682    high_watermark: i64,
1683    last_stable_offset: i64,
1684) -> i64 {
1685    match isolation {
1686        IsolationLevel::ReadCommitted => last_stable_offset,
1687        IsolationLevel::ReadUncommitted => high_watermark,
1688    }
1689}
1690
1691fn record_readable_end(
1692    ends: &mut HashMap<(String, i32), i64>,
1693    key: (String, i32),
1694    isolation: IsolationLevel,
1695    high_watermark: i64,
1696    last_stable_offset: i64,
1697) {
1698    ends.insert(
1699        key,
1700        readable_end_offset(isolation, high_watermark, last_stable_offset),
1701    );
1702}
1703
1704/// The offset to fetch next after consuming `batches`: one past the highest
1705/// `base_offset + last_offset_delta` across all decoded batches.
1706///
1707/// This function returns `None` when there are no batches, which leaves the
1708/// offset unchanged. The consumer uses it to advance past control and aborted
1709/// batches that emit no records, instead of re-fetching them.
1710fn next_offset_after(batches: &[RecordBatch]) -> Option<i64> {
1711    batches
1712        .iter()
1713        .map(|b| b.base_offset + i64::from(b.last_offset_delta) + 1)
1714        .max()
1715}
1716
1717impl Consumer {
1718    /// Replace each reset sentinel in `next_offsets` with the offset from
1719    /// `ListOffsets`.
1720    ///
1721    /// A `LATEST_SENTINEL` resolves to the log end (`timestamp=-1`), or for
1722    /// `by_duration` to the first offset at or after now minus the duration.
1723    /// `auto_offset_reset = Latest` plants those sentinels at build time, and
1724    /// the `OFFSET_OUT_OF_RANGE` arm of the poll loop plants them again.
1725    /// [`Consumer::seek_to_beginning`] plants a `BEGINNING_SENTINEL`
1726    /// (`timestamp=-2`) and [`Consumer::seek_to_end`] an `END_SENTINEL`
1727    /// (`timestamp=-1`).
1728    ///
1729    /// This runs in `prepare_poll`, after the metadata refresh. A partition
1730    /// whose row has an error keeps its sentinel. `group_fetches` skips it, and
1731    /// the metadata refresh at the start of the next `prepare_poll` lets that
1732    /// poll try again.
1733    ///
1734    /// # Errors
1735    ///
1736    /// Returns `TopicAuthorizationFailed` when the broker does not authorize
1737    /// a topic, and a transport error that is not transient.
1738    #[tracing::instrument(
1739        name = "consumer.resolve_reset_sentinels",
1740        level = "debug",
1741        skip_all,
1742        fields(group_id = %self.group_id, sentinels = field::Empty),
1743        err(level = "debug")
1744    )]
1745    async fn resolve_reset_sentinels(&self) -> Result<(), ConsumerError> {
1746        let by_sentinel = {
1747            let offsets = self.next_offsets.lock().await;
1748            [
1749                (
1750                    LATEST_SENTINEL,
1751                    reset_timestamp(self.auto_offset_reset, unix_now_ms()),
1752                ),
1753                (BEGINNING_SENTINEL, EARLIEST_TIMESTAMP),
1754                (END_SENTINEL, LATEST_TIMESTAMP),
1755            ]
1756            .map(|(sentinel, timestamp)| (sentinel, timestamp, keys_at(&offsets, sentinel)))
1757        };
1758        let count: usize = by_sentinel.iter().map(|(_, _, keys)| keys.len()).sum();
1759        if count == 0 {
1760            return Ok(());
1761        }
1762        tracing::Span::current().record("sentinels", count);
1763        let mut results = Vec::new();
1764        for (sentinel, timestamp, keys) in by_sentinel {
1765            if !keys.is_empty() {
1766                results.push((sentinel, self.list_offsets(&keys, timestamp).await?));
1767            }
1768        }
1769        {
1770            // A seek or a rebalance can change a position while the request
1771            // is in flight. Replace only a sentinel that is still there.
1772            let mut offsets = self.next_offsets.lock().await;
1773            for (sentinel, result) in &results {
1774                for (key, offset) in &result.offsets {
1775                    if let Some(next) = offsets.get_mut(key)
1776                        && *next == *sentinel
1777                    {
1778                        *next = *offset;
1779                    }
1780                }
1781            }
1782        }
1783        // A row with a retriable error keeps its sentinel. Kafka's
1784        // `OffsetFetcher` requests a metadata update for it, so the next poll
1785        // sends the retry to the current leader. An unauthorized topic fails
1786        // the poll instead.
1787        if results
1788            .iter()
1789            .any(|(_, result)| result.retry && result.unauthorized.is_empty())
1790        {
1791            self.client.metadata_topics().request_update();
1792        }
1793        results
1794            .into_iter()
1795            .try_for_each(|(_, result)| result.authorization())
1796    }
1797
1798    /// Fail with [`ConsumerError::NoOffsetForPartition`] when a partition has
1799    /// no committed offset under `auto_offset_reset = None`.
1800    ///
1801    /// `auto_offset_reset = None` never plants a `ListOffsets` sentinel: a
1802    /// partition without a commit keeps [`NO_OFFSET_SENTINEL`] in
1803    /// `next_offsets` until a commit or a seek gives it a real position, so
1804    /// this check fires again on every `poll()` and `position()` until then,
1805    /// as Kafka's `resetInitializingPositions` does.
1806    ///
1807    /// # Errors
1808    ///
1809    /// Returns [`ConsumerError::NoOffsetForPartition`] naming every partition
1810    /// with no offset.
1811    async fn fail_partitions_without_offset(&self) -> Result<(), ConsumerError> {
1812        let missing: BTreeSet<(String, i32)> = self
1813            .next_offsets
1814            .lock()
1815            .await
1816            .iter()
1817            .filter(|&(_, &offset)| offset == NO_OFFSET_SENTINEL)
1818            .map(|(key, _)| key.clone())
1819            .collect();
1820        if missing.is_empty() {
1821            Ok(())
1822        } else {
1823            Err(ConsumerError::NoOffsetForPartition(missing))
1824        }
1825    }
1826
1827    /// Replace each `COMMITTED_SENTINEL` with the committed offset of the
1828    /// group, or with the reset of `auto_offset_reset` when the group has none
1829    /// or the consumer has no group id. Kafka's
1830    /// `refreshCommittedOffsetsIfNeeded` and `resetInitializingPositions`.
1831    ///
1832    /// # Errors
1833    ///
1834    /// Returns the error of the `OffsetFetch` request.
1835    pub(crate) async fn resolve_committed_sentinels(
1836        &self,
1837        deadline: Option<Instant>,
1838    ) -> Result<(), ConsumerError> {
1839        let keys = keys_at(&*self.next_offsets.lock().await, COMMITTED_SENTINEL);
1840        if keys.is_empty() {
1841            return Ok(());
1842        }
1843        // Kafka's `updateFetchPositions` gives the committed offsets the timer
1844        // of the `poll`.
1845        let deadline =
1846            deadline.unwrap_or_else(|| Instant::now() + self.default_api_timeout.to_std());
1847        let committed = if self.group_id.is_empty() {
1848            HashMap::new()
1849        } else {
1850            self.committed_until(&keys, deadline).await?
1851        };
1852        let reset = consumer::reset_starting_offset(self.auto_offset_reset);
1853        with_offsets_and_positions(&self.next_offsets, &self.positions, |offsets, positions| {
1854            for key in keys {
1855                let Some(next) = offsets.get_mut(&key) else {
1856                    continue;
1857                };
1858                if *next != COMMITTED_SENTINEL {
1859                    continue;
1860                }
1861                match committed.get(&key).cloned().flatten() {
1862                    Some(offset) => {
1863                        *next = offset.offset;
1864                        let position = positions.entry(key).or_default();
1865                        position.offset_epoch = LeaderEpoch(offset.leader_epoch.unwrap_or(-1));
1866                        // Kafka's `TopicPartitionState.seekUnvalidated` waits for
1867                        // `OffsetForLeaderEpoch` when the committed epoch is older
1868                        // than the leader epoch of the metadata.
1869                        position.awaiting_validation =
1870                            should_await_validation(position.leader_epoch, position.offset_epoch);
1871                    }
1872                    None => *next = reset,
1873                }
1874            }
1875        })
1876        .await;
1877        Ok(())
1878    }
1879
1880    /// Send `ListOffsets(timestamp)` for `keys` to each partition leader.
1881    ///
1882    /// This is Kafka's `OffsetFetcher.groupListOffsetRequests` and
1883    /// `sendListOffsetRequest`: one request per leader, with the consumer's
1884    /// isolation level and the leader epoch from metadata. No lock is held
1885    /// across a request. A transient transport error drops that connection,
1886    /// and its partitions are retried, as Kafka retries a
1887    /// `RetriableException`.
1888    ///
1889    /// # Errors
1890    ///
1891    /// Returns a transport error that is not transient.
1892    #[cfg_attr(test, mutants::skip)] // cargo-mutants: ListOffsets RPC orchestration, exercised by the mock-broker test
1893    async fn list_offsets(
1894        &self,
1895        keys: &[(String, i32)],
1896        timestamp: i64,
1897    ) -> Result<ListOffsetsResult, ConsumerError> {
1898        let by_leader = group_list_offsets(keys, &*self.positions.lock().await, |id| {
1899            self.client.knows_broker(id)
1900        });
1901        let mut answers = Vec::with_capacity(by_leader.len());
1902        for (leader, by_topic) in by_leader {
1903            let request = build_offsets_request(by_topic, timestamp, self.isolation_level);
1904            let answer = if is_read_committed(self.isolation_level) {
1905                self.send_list_offsets(leader, ReadCommittedListOffsets::new(request))
1906                    .await
1907            } else {
1908                self.send_list_offsets(leader, request).await
1909            };
1910            match answer {
1911                Ok(answer) => answers.push(answer),
1912                Err(e) if is_transient_transport_error(&e) => {
1913                    if should_use_bootstrap_leader(leader) {
1914                        self.client.reconnect_bootstrap().await;
1915                    } else {
1916                        self.client.evict_broker(leader);
1917                    }
1918                }
1919                Err(e) => return Err(e.into()),
1920            }
1921        }
1922        Ok(ListOffsetsResult::collect(keys, &answers))
1923    }
1924
1925    /// Send one `ListOffsets` request to `leader`, or to the bootstrap
1926    /// connection when the leader is not known.
1927    async fn send_list_offsets<R>(
1928        &self,
1929        leader: i32,
1930        request: R,
1931    ) -> Result<ListOffsetsResponse, krabka_client_core::ClientError>
1932    where
1933        R: ProtocolRequest<Response = ListOffsetsResponse>,
1934    {
1935        if should_use_bootstrap_leader(leader) {
1936            self.client.send(request).await
1937        } else {
1938            self.client.broker(leader).send(request).await
1939        }
1940    }
1941
1942    /// Recover the partitions that answered `OFFSET_OUT_OF_RANGE`.
1943    ///
1944    /// An errored Fetch partition carries no usable `log_start_offset`, so one
1945    /// `ListOffsets(timestamp=-2)` reads the real log start for the whole set.
1946    /// `Earliest` then fetches from there. `None` reports the first partition
1947    /// as a truncation, and the log start is the safe offset the caller seeks
1948    /// to.
1949    ///
1950    /// The caller has already released the `next_offsets` guard, so the RPC
1951    /// here holds no lock. A partition whose row has an error keeps its
1952    /// position, so the next Fetch gets `OFFSET_OUT_OF_RANGE` again and this
1953    /// path runs again. The return value is `true` when the caller must
1954    /// refresh metadata before that.
1955    ///
1956    /// # Errors
1957    ///
1958    /// Returns `TopicAuthorizationFailed` when the broker does not authorize
1959    /// a topic, and `LogTruncation` under `auto.offset.reset=none`.
1960    #[tracing::instrument(
1961        name = "consumer.recover_out_of_range",
1962        level = "debug",
1963        skip_all,
1964        fields(group_id = %self.group_id, partitions = out_of_range.len()),
1965        err(level = "debug")
1966    )]
1967    async fn recover_out_of_range(
1968        &self,
1969        out_of_range: &[((String, i32), i64)],
1970    ) -> Result<bool, ConsumerError> {
1971        let keys: Vec<(String, i32)> = out_of_range.iter().map(|(key, _)| key.clone()).collect();
1972        let result = self.list_offsets(&keys, EARLIEST_TIMESTAMP).await?;
1973        result.authorization()?;
1974
1975        if self.auto_offset_reset == AutoOffsetReset::None {
1976            if let Some(truncation) = out_of_range_truncation(out_of_range, &result.offsets) {
1977                return Err(truncation);
1978            }
1979            return Ok(false);
1980        }
1981
1982        let mut offsets = self.next_offsets.lock().await;
1983        let log_starts = out_of_range_positions(out_of_range, &result.offsets);
1984        for (key, log_start) in &log_starts {
1985            offsets.insert(key.clone(), *log_start);
1986        }
1987        let reset_polled = self.auto_commit.as_ref().map(|auto_commit| {
1988            auto_commit.reset_polled(
1989                log_starts
1990                    .into_iter()
1991                    .map(|(key, log_start)| (key, Some(log_start))),
1992            )
1993        });
1994        // `polled` follows the reset before `next_offsets` unlocks: the lock
1995        // order N→AP of the `lock_order_model` test.
1996        OptionFuture::from(reset_polled).await;
1997        drop(offsets);
1998        Ok(result.retry)
1999    }
2000}
2001
2002/// The fetch position to write for each out-of-range partition.
2003///
2004/// A partition the answer does not mention keeps the position it has. The next
2005/// poll fetches it again and gets `OFFSET_OUT_OF_RANGE` again, which retries
2006/// this path. That is better than inventing an offset for it.
2007fn out_of_range_positions(
2008    out_of_range: &[((String, i32), i64)],
2009    log_starts: &HashMap<(String, i32), i64>,
2010) -> Vec<((String, i32), i64)> {
2011    out_of_range
2012        .iter()
2013        .filter_map(|(key, _)| log_starts.get(key).map(|start| (key.clone(), *start)))
2014        .collect()
2015}
2016
2017/// The truncation that `auto.offset.reset=none` reports.
2018///
2019/// It names the first out-of-range partition, the offset the fetch asked for,
2020/// and the log start the caller seeks to. `UNKNOWN_OFFSET` as the safe offset
2021/// says the broker gave no answer for that partition.
2022fn out_of_range_truncation(
2023    out_of_range: &[((String, i32), i64)],
2024    log_starts: &HashMap<(String, i32), i64>,
2025) -> Option<ConsumerError> {
2026    let ((topic, partition), fetch_offset) = out_of_range.first()?;
2027    Some(ConsumerError::LogTruncation {
2028        topic: topic.clone(),
2029        partition: *partition,
2030        fetch_offset: *fetch_offset,
2031        safe_offset: log_starts
2032            .get(&(topic.clone(), *partition))
2033            .copied()
2034            .unwrap_or(UNKNOWN_OFFSET),
2035    })
2036}
2037
2038/// The keys whose next offset is exactly `sentinel`.
2039fn keys_at(offsets: &HashMap<(String, i32), i64>, sentinel: i64) -> Vec<(String, i32)> {
2040    offsets
2041        .iter()
2042        .filter(|(_, value)| **value == sentinel)
2043        .map(|(key, _)| key.clone())
2044        .collect()
2045}
2046
2047impl Consumer {
2048    /// Apply the truncations that the proactive validate pass detected to
2049    /// `next_offsets`.
2050    ///
2051    /// This method honors `auto.offset.reset`. With `None` it errors on the
2052    /// first truncated partition.
2053    #[tracing::instrument(
2054        name = "consumer.apply_truncation",
2055        level = "debug",
2056        skip_all,
2057        fields(group_id = %self.group_id, truncated = truncated.len()),
2058        err(level = "debug")
2059    )]
2060    async fn apply_truncation(
2061        &self,
2062        truncated: &HashMap<(String, i32), i64>,
2063    ) -> Result<(), ConsumerError> {
2064        let mut offsets = self.next_offsets.lock().await;
2065        for (key, safe_offset) in truncated {
2066            if self.auto_offset_reset == AutoOffsetReset::None {
2067                let fetch_offset = fetch_offset_or_unknown(&offsets, key);
2068                return Err(ConsumerError::LogTruncation {
2069                    topic: key.0.clone(),
2070                    partition: key.1,
2071                    fetch_offset,
2072                    safe_offset: *safe_offset,
2073                });
2074            }
2075            offsets.insert(key.clone(), *safe_offset);
2076        }
2077        let reset_polled = self.auto_commit.as_ref().map(|auto_commit| {
2078            auto_commit.reset_polled(
2079                truncated
2080                    .iter()
2081                    .map(|(key, safe_offset)| (key.clone(), Some(*safe_offset))),
2082            )
2083        });
2084        // `polled` follows the reset before `next_offsets` unlocks: the lock
2085        // order N→AP of the `lock_order_model` test.
2086        OptionFuture::from(reset_polled).await;
2087        drop(offsets);
2088        Ok(())
2089    }
2090
2091    /// In-band `diverging_epoch` handler for use inside the poll loop, while the
2092    /// `next_offsets` guard is already held.
2093    fn handle_truncation_in_poll(
2094        &self,
2095        offsets: &mut HashMap<(String, i32), i64>,
2096        key: &(String, i32),
2097        safe_offset: i64,
2098    ) -> Result<(), ConsumerError> {
2099        if self.auto_offset_reset == AutoOffsetReset::None {
2100            let fetch_offset = fetch_offset_or_unknown(offsets, key);
2101            return Err(ConsumerError::LogTruncation {
2102                topic: key.0.clone(),
2103                partition: key.1,
2104                fetch_offset,
2105                safe_offset,
2106            });
2107        }
2108        offsets.insert(key.clone(), safe_offset);
2109        Ok(())
2110    }
2111}
2112
2113#[cfg(test)]
2114mod offset_advance_tests {
2115    use std::collections::HashMap;
2116
2117    use assert2::check;
2118    use krabka_protocol::{
2119        owned::{
2120            fetch_request::ReplicaState,
2121            list_offsets_response::{ListOffsetsPartitionResponse, ListOffsetsTopicResponse},
2122        },
2123        primitives::uuid::Uuid as WireUuid,
2124        records::{RecordBatch, RecordsPayload},
2125        tagged_fields::UnknownTaggedFields,
2126    };
2127
2128    use super::*;
2129
2130    fn id(n: u8) -> WireUuid {
2131        let mut b = [0u8; 16];
2132        b[15] = n;
2133        WireUuid(b)
2134    }
2135
2136    #[test]
2137    fn fetch_partition_error_codes_map_to_kafka_consumer_actions() {
2138        for (name, error_code, expected) in [
2139            ("none", 0, FetchPartitionAction::Records),
2140            ("offset out of range", 1, FetchPartitionAction::ResetOffset),
2141            (
2142                "unknown topic or partition",
2143                3,
2144                FetchPartitionAction::RefreshMetadata,
2145            ),
2146            ("not leader or follower", 6, FetchPartitionAction::Reroute),
2147            (
2148                "fenced leader epoch",
2149                74,
2150                FetchPartitionAction::RevalidateEpoch,
2151            ),
2152            (
2153                "unknown leader epoch",
2154                75,
2155                FetchPartitionAction::RevalidateEpoch,
2156            ),
2157            (
2158                "unknown topic id",
2159                100,
2160                FetchPartitionAction::RefreshMetadata,
2161            ),
2162            (
2163                "replica not available",
2164                9,
2165                FetchPartitionAction::RefreshMetadata,
2166            ),
2167            (
2168                "kafka storage error",
2169                56,
2170                FetchPartitionAction::RefreshMetadata,
2171            ),
2172            (
2173                "offset not available",
2174                78,
2175                FetchPartitionAction::RefreshMetadata,
2176            ),
2177            (
2178                "inconsistent topic id",
2179                103,
2180                FetchPartitionAction::RefreshMetadata,
2181            ),
2182            (
2183                "topic authorization failed",
2184                29,
2185                FetchPartitionAction::Fail(29),
2186            ),
2187        ] {
2188            check!(
2189                classify_fetch_partition_error(error_code) == expected,
2190                "case {name}"
2191            );
2192        }
2193    }
2194
2195    #[test]
2196    fn fetch_leader_id_uses_known_non_negative_leader_or_bootstrap() {
2197        check!(BOOTSTRAP_LEADER == -1);
2198        check!(UNKNOWN_LEADER_ID == -1);
2199        check!(UNKNOWN_FETCH_OFFSET == -1);
2200        for (name, leader, known, expected) in [
2201            ("known non-negative", 3, true, 3),
2202            ("negative leader", -1, true, BOOTSTRAP_LEADER),
2203            ("unknown broker", 3, false, BOOTSTRAP_LEADER),
2204        ] {
2205            check!(fetch_leader_id(leader, known) == expected, "case {name}");
2206        }
2207        for (name, leader, expected) in [
2208            ("bootstrap sentinel", BOOTSTRAP_LEADER, true),
2209            ("broker leader", 3, false),
2210        ] {
2211            check!(
2212                should_use_bootstrap_leader(leader) == expected,
2213                "case {name}"
2214            );
2215        }
2216    }
2217
2218    #[test]
2219    fn readable_end_respects_consumer_isolation() {
2220        check!(readable_end_offset(IsolationLevel::ReadUncommitted, 12, 9) == 12);
2221        check!(readable_end_offset(IsolationLevel::ReadCommitted, 12, 9) == 9);
2222
2223        let mut ends = HashMap::new();
2224        record_readable_end(
2225            &mut ends,
2226            ("topic-a".into(), 0),
2227            IsolationLevel::ReadCommitted,
2228            12,
2229            9,
2230        );
2231        check!(ends.get(&("topic-a".into(), 0)) == Some(&9));
2232    }
2233
2234    #[test]
2235    fn poll_helpers_preserve_sentinel_filter_and_record_math_boundaries() {
2236        let key = ("topic-a".to_owned(), 2);
2237        let mut offsets = HashMap::new();
2238
2239        assert2::assert!(fetch_offset_or_unknown(&offsets, &key) == UNKNOWN_FETCH_OFFSET);
2240        offsets.insert(key.clone(), 42);
2241        check!(fetch_offset_or_unknown(&offsets, &key) == 42);
2242
2243        for (name, isolation, expected) in [
2244            ("read committed", IsolationLevel::ReadCommitted, true),
2245            ("read uncommitted", IsolationLevel::ReadUncommitted, false),
2246        ] {
2247            check!(is_read_committed(isolation) == expected, "case {name}");
2248        }
2249        for (name, first, offset, expected) in [
2250            ("at first offset", 10, 10, true),
2251            ("after first offset", 10, 11, true),
2252            ("before first offset", 11, 10, false),
2253        ] {
2254            check!(
2255                aborted_txn_started(first, offset) == expected,
2256                "case {name}"
2257            );
2258        }
2259        for (name, committed, aborted, started, expected) in [
2260            ("drop aborted", true, true, true, true),
2261            ("uncommitted isolation", false, true, true, false),
2262            ("not aborted", true, false, true, false),
2263            ("transaction not started", true, true, false, false),
2264        ] {
2265            check!(
2266                should_drop_aborted_batch(committed, aborted, started) == expected,
2267                "case {name}"
2268            );
2269        }
2270        check!(record_offset(100, 7) == 107);
2271    }
2272
2273    #[test]
2274    fn transient_transport_error_classification_is_narrow() {
2275        use std::io;
2276
2277        use krabka_client_core::ClientError;
2278
2279        let transport_cases = vec![
2280            ("disconnected", ClientError::Disconnected, true),
2281            (
2282                "timeout",
2283                ClientError::Timeout(krabka_units::millis(10)),
2284                true,
2285            ),
2286            (
2287                "connection reset",
2288                ClientError::Io(io::Error::new(io::ErrorKind::ConnectionReset, "reset")),
2289                true,
2290            ),
2291            (
2292                "connect refused",
2293                ClientError::Connect {
2294                    addr: "127.0.0.1:9092".parse().unwrap(),
2295                    source: io::Error::new(io::ErrorKind::ConnectionRefused, "refused"),
2296                },
2297                true,
2298            ),
2299            ("server error", ClientError::Server { error_code: 6 }, false),
2300            (
2301                "incompatible version",
2302                ClientError::IncompatibleVersion {
2303                    api_key: 1,
2304                    broker_min: 0,
2305                    broker_max: 1,
2306                    client_min: 2,
2307                    client_max: 3,
2308                },
2309                false,
2310            ),
2311        ];
2312        for (name, error, expected) in &transport_cases {
2313            check!(
2314                is_transient_transport_error(error) == *expected,
2315                "case {name}"
2316            );
2317        }
2318        let poll_cases = vec![
2319            (
2320                "client connect",
2321                ConsumerError::Client(ClientError::Connect {
2322                    addr: "127.0.0.1:9092".parse().unwrap(),
2323                    source: io::Error::new(io::ErrorKind::ConnectionRefused, "refused"),
2324                }),
2325                true,
2326            ),
2327            ("server", ConsumerError::Server(6), false),
2328        ];
2329        for (name, error, expected) in &poll_cases {
2330            check!(is_transient_poll_error(error) == *expected, "case {name}");
2331        }
2332    }
2333
2334    /// Kafka's `AbstractFetch.createFetchRequest` and `FetchRequest.Builder`:
2335    /// the limits, the isolation level, the session fields, the partitions by
2336    /// topic, and the forgotten partitions by topic.
2337    #[test]
2338    fn build_fetch_request_carries_limits_session_partitions_and_forgotten_topics() {
2339        let partition = |fetch_offset| SessionPartition {
2340            topic_id: id(7),
2341            fetch_offset,
2342            current_leader_epoch: 5,
2343            last_fetched_epoch: 4,
2344            partition_max_bytes: 128 * 1024,
2345        };
2346        let req = build_fetch_request(
2347            500,
2348            IsolationLevel::ReadCommitted,
2349            krabka_units::bytes(7),
2350            mebibytes(2),
2351            "az-1",
2352            SessionRequest {
2353                session_id: 77,
2354                session_epoch: 3,
2355                partitions: maplit::btreemap! {("topic-a".to_owned(), 2) => partition(42), ("topic-a".to_owned(), 3) => partition(43)},
2356                forgotten: vec![
2357                    (("topic-b".to_owned(), 0), id(8)),
2358                    (("topic-b".to_owned(), 1), id(8)),
2359                ],
2360            },
2361        );
2362        let fetch_partition = |partition, fetch_offset| FetchPartition {
2363            partition,
2364            current_leader_epoch: 5,
2365            fetch_offset,
2366            last_fetched_epoch: 4,
2367            log_start_offset: -1,
2368            partition_max_bytes: 128 * 1024,
2369            replica_directory_id: WireUuid::default(),
2370            high_watermark: i64::MAX,
2371            unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2372        };
2373        assert2::assert!(
2374            req == FetchRequest {
2375                replica_id: -1,
2376                max_wait_ms: 500,
2377                min_bytes: 7,
2378                max_bytes: 2 * 1024 * 1024,
2379                isolation_level: 1, // read_committed wire value
2380                session_id: 77,
2381                session_epoch: 3,
2382                topics: vec![FetchTopic {
2383                    topic: "topic-a".into(),
2384                    topic_id: id(7),
2385                    partitions: vec![fetch_partition(2, 42), fetch_partition(3, 43)],
2386                    unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2387                }],
2388                forgotten_topics_data: vec![ForgottenTopic {
2389                    topic: "topic-b".into(),
2390                    topic_id: id(8),
2391                    partitions: vec![0, 1],
2392                    unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2393                }],
2394                rack_id: "az-1".into(),
2395                cluster_id: None,
2396                replica_state: ReplicaState {
2397                    replica_id: -1,
2398                    replica_epoch: -1,
2399                    unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2400                },
2401                unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2402            }
2403        );
2404    }
2405
2406    /// The request carries the caller's timestamp verbatim, the consumer's
2407    /// isolation level and the leader epoch, and asks as a consumer, which
2408    /// is `replica_id = -1`.
2409    ///
2410    /// `-1` asks for the log end and `-2` for the log start. The
2411    /// `OFFSET_OUT_OF_RANGE` recovery path depends on the second one, because
2412    /// an errored Fetch partition reports no usable `log_start_offset`.
2413    #[test]
2414    fn offsets_request_carries_the_timestamp_isolation_and_leader_epoch() {
2415        for (timestamp, isolation, isolation_level, leader_epoch) in [
2416            (LATEST_TIMESTAMP, IsolationLevel::ReadCommitted, 1, 4),
2417            (LATEST_TIMESTAMP, IsolationLevel::ReadUncommitted, 0, -1),
2418            (EARLIEST_TIMESTAMP, IsolationLevel::ReadCommitted, 1, -1),
2419            (EARLIEST_TIMESTAMP, IsolationLevel::ReadUncommitted, 0, 9),
2420        ] {
2421            let by_topic =
2422                maplit::btreemap! {"topic-a".to_owned() => vec![(3, LeaderEpoch(leader_epoch))]};
2423            let req = build_offsets_request(by_topic, timestamp, isolation);
2424
2425            check!(
2426                req == ListOffsetsRequest {
2427                    replica_id: -1,
2428                    isolation_level,
2429                    topics: vec![ListOffsetsTopic {
2430                        name: "topic-a".into(),
2431                        partitions: vec![ListOffsetsPartition {
2432                            partition_index: 3,
2433                            current_leader_epoch: leader_epoch,
2434                            timestamp,
2435                            unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2436                        }],
2437                        unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2438                    }],
2439                    timeout_ms: 0,
2440                    unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2441                },
2442                "timestamp {timestamp}, isolation {isolation:?}"
2443            );
2444        }
2445    }
2446
2447    /// A partition goes to its known leader with the metadata epoch. A
2448    /// partition without a known, dialable leader goes to the bootstrap
2449    /// connection with the epoch that its position has.
2450    #[test]
2451    fn list_offsets_partitions_group_by_leader_with_the_leader_epoch() {
2452        let position = |leader_id, leader_epoch| PartitionPosition {
2453            leader_id,
2454            leader_epoch: LeaderEpoch(leader_epoch),
2455            ..Default::default()
2456        };
2457        let positions = maplit::hashmap! {("topic-a".to_owned(), 0) => position(1, 5), ("topic-a".to_owned(), 1) => position(2, 6), ("topic-b".to_owned(), 0) => position(1, 3), ("topic-b".to_owned(), 1) => position(3, 8)};
2458        let keys = [
2459            ("topic-a".to_owned(), 0),
2460            ("topic-a".to_owned(), 1),
2461            ("topic-b".to_owned(), 0),
2462            ("topic-b".to_owned(), 1),
2463            ("topic-c".to_owned(), 0),
2464        ];
2465        let grouped = group_list_offsets(&keys, &positions, |id| id != 3);
2466        assert2::assert!(
2467            grouped
2468                == maplit::btreemap! {
2469                    BOOTSTRAP_LEADER => maplit::btreemap! {
2470                        "topic-b".to_owned() => vec![(1, LeaderEpoch(8))],
2471                        "topic-c".to_owned() => vec![(0, LeaderEpoch(-1))],
2472                    },
2473                    1 => maplit::btreemap! {
2474                        "topic-a".to_owned() => vec![(0, LeaderEpoch(5))],
2475                        "topic-b".to_owned() => vec![(0, LeaderEpoch(3))],
2476                    },
2477                    2 => maplit::btreemap! {
2478                        "topic-a".to_owned() => vec![(1, LeaderEpoch(6))],
2479                    },
2480                }
2481        );
2482    }
2483
2484    /// `OffsetFetcherUtils.handleListOffsetResponse` sets a position only for
2485    /// `NONE` with a known offset, fails for `TOPIC_AUTHORIZATION_FAILED`, and
2486    /// retries every other code.
2487    #[test]
2488    fn list_offsets_rows_map_to_kafka_consumer_actions() {
2489        for (name, error_code, offset, expected) in [
2490            ("none", 0, 40, ListOffsetsRowAction::Position(40)),
2491            ("none at zero", 0, 0, ListOffsetsRowAction::Position(0)),
2492            ("none without an offset", 0, -1, ListOffsetsRowAction::Retry),
2493            ("not leader or follower", 6, -1, ListOffsetsRowAction::Retry),
2494            (
2495                "unknown topic or partition",
2496                3,
2497                -1,
2498                ListOffsetsRowAction::Retry,
2499            ),
2500            ("leader not available", 5, -1, ListOffsetsRowAction::Retry),
2501            ("replica not available", 9, -1, ListOffsetsRowAction::Retry),
2502            ("offset not available", 78, -1, ListOffsetsRowAction::Retry),
2503            ("fenced leader epoch", 74, -1, ListOffsetsRowAction::Retry),
2504            ("unknown leader epoch", 75, -1, ListOffsetsRowAction::Retry),
2505            ("kafka storage error", 56, -1, ListOffsetsRowAction::Retry),
2506            (
2507                "unsupported for message format",
2508                43,
2509                -1,
2510                ListOffsetsRowAction::Retry,
2511            ),
2512            ("unexpected code", 2, -1, ListOffsetsRowAction::Retry),
2513            (
2514                "topic authorization failed",
2515                29,
2516                -1,
2517                ListOffsetsRowAction::Unauthorized,
2518            ),
2519        ] {
2520            check!(
2521                classify_list_offsets_row(error_code, offset) == expected,
2522                "case {name}"
2523            );
2524        }
2525    }
2526
2527    /// The answers become offsets for the requested partitions. A partition
2528    /// with an error row or without a row is retried. An answer with an
2529    /// unauthorized row gives no offsets.
2530    #[test]
2531    fn list_offsets_result_collects_offsets_retries_and_unauthorized_topics() {
2532        let row = |partition_index, error_code, offset| ListOffsetsPartitionResponse {
2533            partition_index,
2534            error_code,
2535            offset,
2536            ..Default::default()
2537        };
2538        let answer = |name: &str, rows| ListOffsetsResponse {
2539            topics: vec![ListOffsetsTopicResponse {
2540                name: name.into(),
2541                partitions: rows,
2542                ..Default::default()
2543            }],
2544            ..Default::default()
2545        };
2546        let requested = [("topic-a".to_owned(), 0), ("topic-a".to_owned(), 1)];
2547        let retry_with_first_offset = ListOffsetsResult {
2548            offsets: maplit::hashmap! {("topic-a".to_owned(), 0) => 5},
2549            retry: true,
2550            unauthorized: BTreeSet::new(),
2551        };
2552        for (name, answers, expected) in [
2553            (
2554                "every partition answered",
2555                vec![answer("topic-a", vec![row(0, 0, 5), row(1, 0, 9)])],
2556                ListOffsetsResult {
2557                    offsets: maplit::hashmap! {("topic-a".to_owned(), 0) => 5, ("topic-a".to_owned(), 1) => 9},
2558                    retry: false,
2559                    unauthorized: BTreeSet::new(),
2560                },
2561            ),
2562            (
2563                "an error row is retried",
2564                vec![answer("topic-a", vec![row(0, 0, 5), row(1, 6, -1)])],
2565                retry_with_first_offset.clone(),
2566            ),
2567            (
2568                "a missing row is retried and an unrequested row is ignored",
2569                vec![answer("topic-a", vec![row(0, 0, 5), row(2, 0, 3)])],
2570                retry_with_first_offset,
2571            ),
2572            (
2573                "an unauthorized answer gives no offsets",
2574                vec![answer("topic-a", vec![row(0, 0, 5), row(1, 29, -1)])],
2575                ListOffsetsResult {
2576                    offsets: HashMap::new(),
2577                    retry: true,
2578                    unauthorized: maplit::btreeset! {"topic-a".to_owned()},
2579                },
2580            ),
2581        ] {
2582            check!(
2583                ListOffsetsResult::collect(&requested, &answers) == expected,
2584                "case {name}"
2585            );
2586        }
2587    }
2588
2589    /// Only an unauthorized topic fails the reset.
2590    #[test]
2591    fn list_offsets_result_authorization_names_the_unauthorized_topics() {
2592        let mut result = ListOffsetsResult::default();
2593        check!(result.authorization().map_err(|e| e.to_string()) == Ok(()));
2594        result.unauthorized.insert("topic-a".to_owned());
2595        check!(
2596            result.authorization().map_err(|e| e.to_string())
2597                == Err(ConsumerError::TopicAuthorizationFailed(
2598                    maplit::btreeset! {"topic-a".to_owned()}
2599                )
2600                .to_string())
2601        );
2602    }
2603
2604    /// The recovery writes the broker's log start, and leaves a partition the
2605    /// broker did not answer for alone.
2606    ///
2607    /// Inventing a position for an unanswered partition is what the old code
2608    /// did: it wrote the -1 that an errored Fetch reports, and every following
2609    /// Fetch then asked for -1 and was out of range again.
2610    #[test]
2611    fn out_of_range_positions_uses_the_log_start_and_skips_the_unanswered() {
2612        let out_of_range = vec![
2613            (("topic-a".to_owned(), 0), 0),
2614            (("topic-a".to_owned(), 1), 3),
2615        ];
2616        let log_starts = maplit::hashmap! {("topic-a".to_owned(), 0) => 5};
2617        assert2::assert!(
2618            out_of_range_positions(&out_of_range, &log_starts)
2619                == vec![(("topic-a".to_owned(), 0), 5)]
2620        );
2621    }
2622
2623    /// Under `auto.offset.reset=none` the first out-of-range partition becomes
2624    /// the error, with the log start as the offset the caller seeks to.
2625    ///
2626    /// An answer that omits the partition gives `UNKNOWN_OFFSET`, which says
2627    /// the broker did not report one. It is not a position to fetch from.
2628    #[test]
2629    fn out_of_range_truncation_names_the_first_partition_and_its_log_start() {
2630        for (answered, expected_safe) in [(true, 5), (false, UNKNOWN_OFFSET)] {
2631            let log_starts = if answered {
2632                maplit::hashmap! {("topic-a".to_owned(), 0) => 5}
2633            } else {
2634                HashMap::new()
2635            };
2636            let out_of_range = vec![
2637                (("topic-a".to_owned(), 0), 2),
2638                (("topic-a".to_owned(), 1), 3),
2639            ];
2640            assert2::assert!(
2641                out_of_range_truncation(&out_of_range, &log_starts).map(|e| e.to_string())
2642                    == Some(
2643                        ConsumerError::LogTruncation {
2644                            topic: "topic-a".to_owned(),
2645                            partition: 0,
2646                            fetch_offset: 2,
2647                            safe_offset: expected_safe,
2648                        }
2649                        .to_string()
2650                    )
2651            );
2652        }
2653    }
2654
2655    /// Nothing out of range means nothing to report.
2656    #[test]
2657    fn out_of_range_truncation_is_none_for_an_empty_set() {
2658        assert2::assert!(out_of_range_truncation(&[], &HashMap::new()).is_none());
2659    }
2660
2661    /// `keys_at` picks exactly the partitions parked on a sentinel.
2662    #[test]
2663    fn keys_at_selects_only_the_sentinel_partitions() {
2664        let mut offsets: HashMap<(String, i32), i64> = HashMap::new();
2665        offsets.insert(("topic-a".to_owned(), 0), LATEST_SENTINEL);
2666        offsets.insert(("topic-a".to_owned(), 1), 42);
2667        assert2::assert!(keys_at(&offsets, LATEST_SENTINEL) == vec![("topic-a".to_owned(), 0)]);
2668    }
2669
2670    #[test]
2671    fn advance_target_uses_last_offset_delta_not_record_count() {
2672        // A batch spanning offsets 10..=14 (last_offset_delta = 4) but carrying
2673        // zero surviving records must still advance the fetch offset to 15.
2674        let batch = RecordBatch {
2675            base_offset: 10,
2676            last_offset_delta: 4,
2677            records: vec![],
2678            ..Default::default()
2679        };
2680        let payload = RecordsPayload::V2(vec![batch]);
2681        let batches = payload.as_v2().unwrap();
2682        assert2::assert!(super::next_offset_after(batches) == Some(15));
2683    }
2684
2685    #[test]
2686    fn advance_target_none_for_empty() {
2687        let payload = RecordsPayload::V2(vec![]);
2688        assert2::assert!(super::next_offset_after(payload.as_v2().unwrap()) == None);
2689    }
2690}
2691
2692#[cfg(test)]
2693pub mod partition_error_tests {
2694    use std::{
2695        sync::{
2696            Arc, Mutex,
2697            atomic::{AtomicU16, AtomicUsize, Ordering},
2698        },
2699        time::Duration,
2700    };
2701
2702    use assert2::check;
2703    use krabka_client_core::{Client, MockBroker};
2704    use krabka_protocol::{
2705        Encode,
2706        owned::{
2707            api_versions_request,
2708            fetch_response::{
2709                EpochEndOffset, FetchResponse, FetchableTopicResponse, LeaderIdAndEpoch,
2710                PartitionData,
2711            },
2712            metadata_request,
2713            metadata_response::{MetadataResponse, MetadataResponseBroker},
2714        },
2715        primitives::uuid::Uuid as WireUuid,
2716    };
2717    use krabka_units::secs;
2718    use tokio::{
2719        sync::watch,
2720        time::{sleep, timeout},
2721    };
2722
2723    use super::*;
2724    use crate::{Assignor, commit::AutoCommit, subscription::shared};
2725
2726    const TOPIC_ID: WireUuid = WireUuid([7; 16]);
2727
2728    fn encode(response: &impl Encode, version: i16) -> Vec<u8> {
2729        let mut buf = bytes::BytesMut::new();
2730        response.encode(&mut buf, version).unwrap();
2731        buf.to_vec()
2732    }
2733
2734    /// A mock broker that answers `ApiVersions` and a non-flexible `Metadata`,
2735    /// and counts the `Metadata` requests.
2736    async fn metadata_counting_broker(metadata_requests: Arc<AtomicUsize>) -> MockBroker {
2737        MockBroker::start(move |api_key, version, _corr_id, _body| {
2738            if api_key == api_versions_request::API_KEY {
2739                let versions = mock_test_support::api_versions_response([
2740                    (api_versions_request::API_KEY, 3),
2741                    (metadata_request::API_KEY, 8),
2742                ]);
2743                return Some(encode(&versions, 0));
2744            }
2745            if api_key == metadata_request::API_KEY {
2746                metadata_requests.fetch_add(1, Ordering::SeqCst);
2747                return Some(encode(&MetadataResponse::default(), version));
2748            }
2749            None
2750        })
2751        .await
2752        .unwrap()
2753    }
2754
2755    async fn consumer_on(broker: &MockBroker) -> Consumer {
2756        let client = Client::builder()
2757            .bootstrap(broker.addr.to_string())
2758            .build()
2759            .await
2760            .unwrap();
2761        consumer_with_client(client)
2762    }
2763
2764    /// A consumer of `group-a` on `client` that owns `orders-0` at offset 5.
2765    pub fn consumer_with_client(client: Client) -> Consumer {
2766        Consumer::test_builder(client)
2767            .commit_identity(test_support::orders_partition_commit_identity(1))
2768            .subscription(shared(vec!["orders".into()], None, true))
2769            .assigned(vec![("orders".into(), 0)])
2770            .next_offsets(maplit::hashmap! {("orders".into(), 0) => 5})
2771            .topic_ids(maplit::hashmap! {"orders".into() => TOPIC_ID})
2772            .build()
2773    }
2774
2775    /// A Fetch v13 response row: the topic by id only, one partition with
2776    /// `error_code`, no records and no leader hint.
2777    fn fetch_response(error_code: i16) -> FetchResponse {
2778        FetchResponse {
2779            responses: vec![FetchableTopicResponse {
2780                topic: String::new(),
2781                topic_id: TOPIC_ID,
2782                partitions: vec![PartitionData {
2783                    partition_index: 0,
2784                    error_code,
2785                    high_watermark: -1,
2786                    last_stable_offset: -1,
2787                    log_start_offset: -1,
2788                    diverging_epoch: EpochEndOffset {
2789                        epoch: -1,
2790                        end_offset: -1,
2791                        ..Default::default()
2792                    },
2793                    current_leader: LeaderIdAndEpoch {
2794                        leader_id: -1,
2795                        leader_epoch: -1,
2796                        ..Default::default()
2797                    },
2798                    ..Default::default()
2799                }],
2800                ..Default::default()
2801            }],
2802            ..Default::default()
2803        }
2804    }
2805
2806    /// How one poll handled a partition error row.
2807    #[derive(Debug, PartialEq)]
2808    struct PollOutcome {
2809        /// `Ok(record_count)`, `Err(Some(code))` for
2810        /// `ConsumerError::Server(code)`, or `Err(None)` for any other error.
2811        result: Result<usize, Option<i16>>,
2812        /// The `Metadata` requests that processing the response sent.
2813        metadata_requests: usize,
2814        /// Whether processing asked for a refresh before the next poll.
2815        metadata_update_requested: bool,
2816        next_offset: Option<i64>,
2817    }
2818
2819    /// Kafka's consumer raises a fatal `OffsetFetch` error from `poll()`. A
2820    /// rejoin leaves the error for `poll`, which returns it once. The next poll
2821    /// runs as usual.
2822    #[tokio::test]
2823    async fn poll_returns_a_fatal_rejoin_error_once() {
2824        let broker = metadata_counting_broker(Arc::default()).await;
2825        let consumer = consumer_on(&broker).await;
2826        let topics = maplit::btreeset! {"orders".to_owned()};
2827        coordinator::report_rejoin_error(
2828            &consumer.poll_error,
2829            ConsumerError::TopicAuthorizationFailed(topics.clone()),
2830        );
2831
2832        let first = consumer
2833            .prepare_poll(Instant::now() + Duration::from_secs(5))
2834            .await
2835            .map_err(|error| match error {
2836                ConsumerError::TopicAuthorizationFailed(topics) => Some(topics),
2837                _ => None,
2838            });
2839        let second = consumer
2840            .prepare_poll(Instant::now() + Duration::from_secs(5))
2841            .await
2842            .map_err(|_| None::<BTreeSet<String>>);
2843
2844        broker.stop();
2845        assert2::assert!((first, second) == (Err(Some(topics)), Ok(true)));
2846    }
2847
2848    /// A partition that the coordinator revokes while `poll` drains the fetch
2849    /// buffer is either still owned for the whole drain, or not drained at
2850    /// all. The drain never returns records of a partition that left the
2851    /// assignment before the drain completed.
2852    #[tokio::test]
2853    async fn drain_never_returns_records_of_a_partition_revoked_during_the_drain() {
2854        let broker = metadata_counting_broker(Arc::default()).await;
2855        let mut consumer = consumer_on(&broker).await;
2856        consumer.fetch_buffer.push(BufferedPartition {
2857            key: ("orders".into(), 0),
2858            position: 5,
2859            records: VecDeque::from([ConsumerRecord {
2860                topic: "orders".into(),
2861                partition: 0,
2862                offset: 5,
2863                leader_epoch: 0,
2864                timestamp: 0,
2865                timestamp_type: crate::TimestampType::CreateTime,
2866                key: None,
2867                value: None,
2868                headers: Vec::new(),
2869            }]),
2870            next_offset: 6,
2871            last_epoch: None,
2872        });
2873        let assigned = Arc::clone(&consumer.assigned);
2874        let next_offsets = Arc::clone(&consumer.next_offsets);
2875        // Block the drain between the ownership check and the position update.
2876        let offsets_guard = next_offsets.lock().await;
2877        let drain = tokio::spawn(async move {
2878            let records = consumer.drain_fetch_buffer().await;
2879            (consumer, records)
2880        });
2881        sleep(Duration::from_millis(50)).await;
2882        // The coordinator revokes the partition, if it can take the lock.
2883        let revoked = timeout(Duration::from_millis(50), assigned.lock())
2884            .await
2885            .is_ok_and(|mut owned| {
2886                owned.clear();
2887                true
2888            });
2889        drop(offsets_guard);
2890        let (consumer, records) = drain.await.expect("drain task");
2891        drop(consumer);
2892        broker.stop();
2893        assert2::assert!(!revoked || records.is_empty());
2894    }
2895
2896    /// While a join runs, Kafka's eager `onJoinPrepare` has revoked every
2897    /// partition, so `poll` fetches nothing until the join completes or the
2898    /// poll timeout passes. A cooperative member keeps fetching its owned
2899    /// partitions.
2900    #[tokio::test(start_paused = true)]
2901    async fn poll_waits_for_a_pending_eager_join_but_not_for_a_cooperative_one() {
2902        let broker = metadata_counting_broker(Arc::default()).await;
2903        let mut actual = Vec::new();
2904        let mut wanted = Vec::new();
2905        for (name, assignor, owns_partitions, pending, join_completes_after, expected) in [
2906            ("no join", Assignor::Range, true, false, None, (true, 0)),
2907            (
2908                "eager join completes in time",
2909                Assignor::Range,
2910                true,
2911                true,
2912                Some(100),
2913                (true, 100),
2914            ),
2915            (
2916                "eager join does not complete in time",
2917                Assignor::Range,
2918                true,
2919                true,
2920                None,
2921                (false, 500),
2922            ),
2923            (
2924                "cooperative join with owned partitions",
2925                Assignor::CooperativeSticky,
2926                true,
2927                true,
2928                None,
2929                (true, 0),
2930            ),
2931            (
2932                "cooperative join without partitions",
2933                Assignor::CooperativeSticky,
2934                false,
2935                true,
2936                Some(200),
2937                (true, 200),
2938            ),
2939        ] {
2940            let mut consumer = consumer_on(&broker).await;
2941            consumer.rebalance_protocol =
2942                assignor::rebalance_protocol_of(&[assignor]).expect("assignor protocol");
2943            if !owns_partitions {
2944                consumer.assigned.lock().await.clear();
2945            }
2946            let pending_sender = watch::Sender::new(pending);
2947            consumer.rebalance_pending = pending_sender.subscribe();
2948            let completion = join_completes_after.map(|millis| {
2949                let pending_sender = pending_sender.clone();
2950                tokio::spawn(async move {
2951                    sleep(Duration::from_millis(millis)).await;
2952                    pending_sender.send_replace(false);
2953                })
2954            });
2955            let start = Instant::now();
2956            let joined = consumer
2957                .wait_for_rebalance(start + Duration::from_millis(500))
2958                .await
2959                .expect("no listener");
2960            let elapsed_ms = u64::try_from(start.elapsed().as_millis()).expect("millis");
2961            if let Some(completion) = completion {
2962                completion.abort();
2963            }
2964            actual.push((name, (joined, elapsed_ms)));
2965            wanted.push((name, expected));
2966        }
2967        broker.stop();
2968        assert2::assert!(actual == wanted);
2969    }
2970
2971    /// One `poll` of the record cap case: the number of records, the first and
2972    /// the last offset, and the fetch position after the `poll`.
2973    #[derive(Debug, PartialEq)]
2974    struct CappedPoll {
2975        records: usize,
2976        first_offset: Option<i64>,
2977        last_offset: Option<i64>,
2978        next_offset: Option<i64>,
2979    }
2980
2981    /// Kafka's `FetchCollector.collectFetch` returns at most
2982    /// `max.poll.records` records, and keeps the rest of the completed fetch for
2983    /// the next `poll`. The position moves only past the returned records, and
2984    /// the next `poll` sends no Fetch while buffered records remain.
2985    #[tokio::test]
2986    async fn poll_returns_at_most_max_poll_records_and_keeps_the_rest_without_a_fetch() {
2987        let fetches = Arc::new(AtomicUsize::new(0));
2988        let fetches_in_mock = Arc::clone(&fetches);
2989        let broker = MockBroker::start(move |api_key, version, _corr_id, _body| {
2990            if api_key == api_versions_request::API_KEY {
2991                let versions = mock_test_support::api_versions_response_ranges([
2992                    (api_versions_request::API_KEY, 0, 3),
2993                    (metadata_request::API_KEY, 0, 8),
2994                    (fetch_request::API_KEY, 4, 11),
2995                ]);
2996                return Some(encode(&versions, 0));
2997            }
2998            if api_key == metadata_request::API_KEY {
2999                return Some(encode(&MetadataResponse::default(), version));
3000            }
3001            if api_key != fetch_request::API_KEY {
3002                return None;
3003            }
3004            fetches_in_mock.fetch_add(1, Ordering::SeqCst);
3005            let batch = RecordBatch {
3006                base_offset: 5,
3007                last_offset_delta: 1199,
3008                records: (0..1200)
3009                    .map(|offset_delta| Record {
3010                        offset_delta,
3011                        value: Some(bytes::Bytes::from_static(b"v")),
3012                        ..Default::default()
3013                    })
3014                    .collect(),
3015                ..Default::default()
3016            };
3017            let mut response = fetch_response(0);
3018            response.responses[0].topic = "orders".into();
3019            response.responses[0].partitions[0].high_watermark = 1205;
3020            response.responses[0].partitions[0].records = Some(RecordsPayload::V2(vec![batch]));
3021            Some(encode(&response, version))
3022        })
3023        .await
3024        .unwrap();
3025        let mut consumer = consumer_on(&broker).await;
3026
3027        let mut polls = Vec::new();
3028        for _ in 0..3 {
3029            let records = consumer.poll(secs(5)).await.expect("poll");
3030            polls.push(CappedPoll {
3031                records: records.len(),
3032                first_offset: records.first().map(|record| record.offset),
3033                last_offset: records.last().map(|record| record.offset),
3034                next_offset: consumer
3035                    .next_offsets
3036                    .lock()
3037                    .await
3038                    .get(&("orders".to_owned(), 0))
3039                    .copied(),
3040            });
3041        }
3042
3043        broker.stop();
3044        assert2::assert!(
3045            (polls, fetches.load(Ordering::SeqCst))
3046                == (
3047                    vec![
3048                        CappedPoll {
3049                            records: 500,
3050                            first_offset: Some(5),
3051                            last_offset: Some(504),
3052                            next_offset: Some(505),
3053                        },
3054                        CappedPoll {
3055                            records: 500,
3056                            first_offset: Some(505),
3057                            last_offset: Some(1004),
3058                            next_offset: Some(1005),
3059                        },
3060                        CappedPoll {
3061                            records: 200,
3062                            first_offset: Some(1005),
3063                            last_offset: Some(1204),
3064                            next_offset: Some(1205),
3065                        },
3066                    ],
3067                    1,
3068                )
3069        );
3070    }
3071
3072    /// What changes between two fetch-position updates.
3073    #[derive(Clone, Copy, Debug)]
3074    enum BetweenPolls {
3075        Nothing,
3076        UpdateRequested,
3077        MaxAgeElapsed,
3078        LeaderUnknown,
3079    }
3080
3081    /// Kafka's consumer refreshes metadata only when `Metadata.updateRequested`
3082    /// holds, not on every poll. A refresh on every poll pipelines a `Metadata`
3083    /// request behind a parked long-poll `Fetch` on the same connection, and
3084    /// the poll then stalls for up to `fetch.max.wait.ms` (krabka-io/krabka-broker#1394).
3085    #[tokio::test]
3086    async fn fetch_positions_refresh_metadata_only_when_it_is_due() {
3087        // (what changes between the polls, `Metadata` requests of the second)
3088        let cases = [
3089            (BetweenPolls::Nothing, 0),
3090            (BetweenPolls::UpdateRequested, 1),
3091            (BetweenPolls::MaxAgeElapsed, 1),
3092            (BetweenPolls::LeaderUnknown, 1),
3093        ];
3094        for (between, expected) in cases {
3095            let metadata_requests = Arc::new(AtomicUsize::new(0));
3096            let broker = metadata_counting_broker(Arc::clone(&metadata_requests)).await;
3097            let mut consumer = consumer_on(&broker).await;
3098            let key = ("orders".to_owned(), 0);
3099            consumer
3100                .positions
3101                .lock()
3102                .await
3103                .entry(key.clone())
3104                .or_default()
3105                .leader_id = 1;
3106            let deadline = Instant::now() + Duration::from_secs(5);
3107            check!(
3108                consumer
3109                    .update_fetch_positions(None, Some(deadline))
3110                    .await
3111                    .unwrap()
3112            );
3113            check!(
3114                metadata_requests.swap(0, Ordering::SeqCst) == 1,
3115                "{between:?}: first poll"
3116            );
3117            match between {
3118                BetweenPolls::Nothing => {}
3119                BetweenPolls::UpdateRequested => consumer.client.metadata_topics().request_update(),
3120                BetweenPolls::MaxAgeElapsed => consumer.metadata_max_age = Time::ZERO,
3121                BetweenPolls::LeaderUnknown => {
3122                    consumer
3123                        .positions
3124                        .lock()
3125                        .await
3126                        .entry(key)
3127                        .or_default()
3128                        .leader_id = UNKNOWN_LEADER_ID;
3129                }
3130            }
3131            check!(
3132                consumer
3133                    .update_fetch_positions(None, Some(deadline))
3134                    .await
3135                    .unwrap()
3136            );
3137            check!(
3138                metadata_requests.load(Ordering::SeqCst) == expected,
3139                "{between:?}"
3140            );
3141            broker.stop();
3142        }
3143    }
3144
3145    /// Kafka's `FetchCollector.handleInitializeErrors` requests a metadata
3146    /// update for 3, 6 and 100 and does not raise an error. The poll returns
3147    /// no records, keeps the fetch position, and asks for a refresh, which
3148    /// the next poll makes before it fetches. It sends no `Metadata` request
3149    /// of its own. Any code that Kafka does not handle fails the poll.
3150    #[tokio::test]
3151    async fn fetch_partition_errors_refresh_metadata_or_fail_the_poll() {
3152        for (name, error_code, expected) in [
3153            (
3154                "unknown topic or partition",
3155                3,
3156                PollOutcome {
3157                    result: Ok(0),
3158                    metadata_requests: 0,
3159                    metadata_update_requested: true,
3160                    next_offset: Some(5),
3161                },
3162            ),
3163            (
3164                "not leader or follower",
3165                6,
3166                PollOutcome {
3167                    result: Ok(0),
3168                    metadata_requests: 0,
3169                    metadata_update_requested: true,
3170                    next_offset: Some(5),
3171                },
3172            ),
3173            (
3174                "unknown topic id",
3175                100,
3176                PollOutcome {
3177                    result: Ok(0),
3178                    metadata_requests: 0,
3179                    metadata_update_requested: true,
3180                    next_offset: Some(5),
3181                },
3182            ),
3183            (
3184                "topic authorization failed",
3185                29,
3186                PollOutcome {
3187                    result: Err(Some(29)),
3188                    metadata_requests: 0,
3189                    metadata_update_requested: false,
3190                    next_offset: Some(5),
3191                },
3192            ),
3193        ] {
3194            let metadata_requests = Arc::new(AtomicUsize::new(0));
3195            let broker = metadata_counting_broker(Arc::clone(&metadata_requests)).await;
3196            let mut consumer = consumer_on(&broker).await;
3197            let topic_ids = consumer.topic_ids.lock().await.clone();
3198
3199            let result = consumer
3200                .process_fetch_responses(vec![fetch_response(error_code)], &topic_ids)
3201                .await
3202                .map(|()| consumer.fetch_buffer.len())
3203                .map_err(|error| match error {
3204                    ConsumerError::Server(code) => Some(code),
3205                    _ => None,
3206                });
3207            let outcome = PollOutcome {
3208                result,
3209                metadata_requests: metadata_requests.load(Ordering::SeqCst),
3210                metadata_update_requested: consumer.client.metadata_topics().update_requested(),
3211                next_offset: consumer
3212                    .next_offsets
3213                    .lock()
3214                    .await
3215                    .get(&("orders".to_owned(), 0))
3216                    .copied(),
3217            };
3218
3219            broker.stop();
3220            check!(outcome == expected, "case {name}");
3221        }
3222    }
3223
3224    /// How the test resets the fetch position of `orders-0` after the start of
3225    /// `poll`.
3226    #[derive(Clone, Copy, Debug)]
3227    enum PollReset {
3228        /// A Fetch row without an error and without records.
3229        None,
3230        /// A Fetch row with `OFFSET_OUT_OF_RANGE` under `auto.offset.reset=latest`.
3231        LatestOutOfRange,
3232        /// A Fetch row with `OFFSET_OUT_OF_RANGE` under
3233        /// `auto.offset.reset=earliest`. The leader answers `ListOffsets` with
3234        /// log start 2.
3235        EarliestOutOfRange,
3236        /// A Fetch row with a diverging epoch that ends at offset 3.
3237        FetchTruncation,
3238        /// The validation pass finds a truncation to offset 3.
3239        ValidationTruncation,
3240    }
3241
3242    /// Kafka's `onJoinPrepare` commits `SubscriptionState.allConsumed`, which
3243    /// reads the current fetch positions. A reset in `poll` therefore changes
3244    /// what the commit before a `JoinGroup` sends. A position that waits for
3245    /// `ListOffsets` is not valid, and `allConsumed` skips it.
3246    #[tokio::test]
3247    async fn a_reset_in_poll_moves_the_position_of_the_commit_before_join() {
3248        let orders_0 = ("orders".to_owned(), 0);
3249        let cases = [
3250            (
3251                "no reset: the position at the start of poll",
3252                PollReset::None,
3253                maplit::hashmap! {orders_0.clone() => (5, -1)},
3254            ),
3255            (
3256                "out of range under latest: no position until ListOffsets",
3257                PollReset::LatestOutOfRange,
3258                HashMap::new(),
3259            ),
3260            (
3261                "out of range under earliest: the log start",
3262                PollReset::EarliestOutOfRange,
3263                maplit::hashmap! {orders_0.clone() => (2, -1)},
3264            ),
3265            (
3266                "truncation in a Fetch response: the end offset of the diverging epoch",
3267                PollReset::FetchTruncation,
3268                maplit::hashmap! {orders_0.clone() => (3, -1)},
3269            ),
3270            (
3271                "truncation from validation: the safe offset",
3272                PollReset::ValidationTruncation,
3273                maplit::hashmap! {orders_0.clone() => (3, -1)},
3274            ),
3275        ];
3276        let mut actual = Vec::new();
3277        let mut wanted = Vec::new();
3278        for (name, reset, expected) in cases {
3279            let leader_port = Arc::new(AtomicU16::new(0));
3280            let broker = list_offsets_broker(
3281                "leader",
3282                Arc::clone(&leader_port),
3283                (0, 2),
3284                list_offsets_request::MAX_VERSION,
3285                Arc::default(),
3286                Arc::default(),
3287            )
3288            .await;
3289            leader_port.store(broker.addr.port(), Ordering::SeqCst);
3290            let mut consumer = consumer_on(&broker).await;
3291            let auto_commit = AutoCommit::new(Duration::from_secs(5));
3292            consumer.auto_commit = Some(auto_commit.clone());
3293            // The start of `poll` records the positions. The interval has not
3294            // passed, so it sends no commit.
3295            consumer.maybe_auto_commit_async().await;
3296            let topic_ids = consumer.topic_ids.lock().await.clone();
3297
3298            let result = match reset {
3299                PollReset::None => {
3300                    consumer
3301                        .process_fetch_responses(vec![fetch_response(0)], &topic_ids)
3302                        .await
3303                }
3304                PollReset::LatestOutOfRange => {
3305                    consumer.auto_offset_reset = AutoOffsetReset::Latest;
3306                    consumer
3307                        .process_fetch_responses(vec![fetch_response(1)], &topic_ids)
3308                        .await
3309                }
3310                PollReset::EarliestOutOfRange => {
3311                    consumer.auto_offset_reset = AutoOffsetReset::Earliest;
3312                    consumer
3313                        .refresh_leader_epochs()
3314                        .await
3315                        .expect("metadata for the leader route");
3316                    consumer
3317                        .process_fetch_responses(vec![fetch_response(1)], &topic_ids)
3318                        .await
3319                }
3320                PollReset::FetchTruncation => {
3321                    let mut response = fetch_response(0);
3322                    response.responses[0].partitions[0].diverging_epoch = EpochEndOffset {
3323                        epoch: 2,
3324                        end_offset: 3,
3325                        ..Default::default()
3326                    };
3327                    consumer
3328                        .process_fetch_responses(vec![response], &topic_ids)
3329                        .await
3330                }
3331                PollReset::ValidationTruncation => {
3332                    consumer
3333                        .apply_truncation(&maplit::hashmap! {orders_0.clone() => 3})
3334                        .await
3335                }
3336            };
3337            result.expect("the reset succeeds");
3338            let offsets = auto_commit
3339                .polled_offsets(&maplit::hashmap! {orders_0.clone() => 1})
3340                .await;
3341
3342            broker.stop();
3343            actual.push((name, offsets));
3344            wanted.push((name, expected));
3345        }
3346        assert2::assert!(actual == wanted);
3347    }
3348
3349    /// The broker id that the metadata names as the leader of `orders-0`.
3350    const LEADER_ID: i32 = 1;
3351    /// The leader epoch that the metadata reports for `orders-0`.
3352    const LEADER_EPOCH: i32 = 7;
3353
3354    /// The `ListOffsets` requests that the mock brokers decoded, with the
3355    /// name of the broker that received each one and the request version.
3356    type SentListOffsets = Arc<Mutex<Vec<(&'static str, i16, ListOffsetsRequest)>>>;
3357
3358    /// The response bytes for a `ListOffsets` request at `version` with one
3359    /// `orders-0` row.
3360    fn list_offsets_answer_bytes(version: i16, error_code: i16, offset: i64) -> Vec<u8> {
3361        use krabka_protocol::owned::{
3362            list_offsets_request::FLEXIBLE_MIN,
3363            list_offsets_response::{ListOffsetsPartitionResponse, ListOffsetsTopicResponse},
3364        };
3365        let answer = ListOffsetsResponse {
3366            topics: vec![ListOffsetsTopicResponse {
3367                name: "orders".into(),
3368                partitions: vec![ListOffsetsPartitionResponse {
3369                    partition_index: 0,
3370                    error_code,
3371                    offset,
3372                    ..Default::default()
3373                }],
3374                ..Default::default()
3375            }],
3376            ..Default::default()
3377        };
3378        let mut bytes = Vec::new();
3379        if version >= FLEXIBLE_MIN {
3380            // The flexible response header has an empty tagged-field section.
3381            bytes.push(0);
3382        }
3383        bytes.extend(encode(&answer, version));
3384        bytes
3385    }
3386
3387    /// A mock broker that answers `ApiVersions`, `Metadata` and `ListOffsets`.
3388    ///
3389    /// It records each decoded `ListOffsets` request under `name` and answers
3390    /// it with one `orders-0` row. The `Metadata` answer names broker
3391    /// `LEADER_ID` at `leader_port` as the leader of `orders-0`, with epoch
3392    /// `LEADER_EPOCH`.
3393    async fn list_offsets_broker(
3394        name: &'static str,
3395        leader_port: Arc<AtomicU16>,
3396        row: (i16, i64),
3397        list_offsets_max_version: i16,
3398        sent: SentListOffsets,
3399        metadata_requests: Arc<AtomicUsize>,
3400    ) -> MockBroker {
3401        use krabka_protocol::owned::list_offsets_request;
3402        MockBroker::start(move |api_key, version, _corr_id, mut body| {
3403            if api_key == api_versions_request::API_KEY {
3404                let versions = mock_test_support::api_versions_response([
3405                    (api_versions_request::API_KEY, 3),
3406                    (metadata_request::API_KEY, 8),
3407                    (list_offsets_request::API_KEY, list_offsets_max_version),
3408                ]);
3409                return Some(encode(&versions, 0));
3410            }
3411            if api_key == metadata_request::API_KEY {
3412                metadata_requests.fetch_add(1, Ordering::SeqCst);
3413                let metadata = test_support::OrdersMetadata::builder()
3414                    .topic_id(TOPIC_ID.0)
3415                    .leader_id(LEADER_ID)
3416                    .leader_epoch(LEADER_EPOCH)
3417                    .brokers(vec![MetadataResponseBroker {
3418                        node_id: LEADER_ID,
3419                        host: "127.0.0.1".into(),
3420                        port: i32::from(leader_port.load(Ordering::SeqCst)),
3421                        ..Default::default()
3422                    }])
3423                    .build();
3424                return Some(encode(&metadata, version));
3425            }
3426            if api_key == list_offsets_request::API_KEY {
3427                let request = test_support::decode_list_offsets_request(&mut body, version);
3428                sent.lock()
3429                    .expect("sent lock")
3430                    .push((name, version, request));
3431                return Some(list_offsets_answer_bytes(version, row.0, row.1));
3432            }
3433            None
3434        })
3435        .await
3436        .unwrap()
3437    }
3438
3439    /// The `ListOffsets` request that the consumer sends for `orders-0`.
3440    fn expected_list_offsets(isolation_level: i8, timestamp: i64) -> ListOffsetsRequest {
3441        ListOffsetsRequest {
3442            replica_id: -1,
3443            isolation_level,
3444            topics: vec![ListOffsetsTopic {
3445                name: "orders".into(),
3446                partitions: vec![ListOffsetsPartition {
3447                    partition_index: 0,
3448                    current_leader_epoch: LEADER_EPOCH,
3449                    timestamp,
3450                    ..Default::default()
3451                }],
3452                ..Default::default()
3453            }],
3454            ..Default::default()
3455        }
3456    }
3457
3458    /// Why a poll step failed, in a form that tests can compare.
3459    #[derive(Debug, PartialEq)]
3460    enum PollFailure {
3461        /// `ConsumerError::TopicAuthorizationFailed` with its topics.
3462        TopicAuthorizationFailed(BTreeSet<String>),
3463        /// `ClientError::IncompatibleVersion` as
3464        /// `(api_key, broker_min, broker_max, client_min, client_max)`.
3465        IncompatibleVersion(i16, i16, i16, i16, i16),
3466        /// Any other error.
3467        Other,
3468    }
3469
3470    impl From<ConsumerError> for PollFailure {
3471        fn from(error: ConsumerError) -> Self {
3472            match error {
3473                ConsumerError::TopicAuthorizationFailed(topics) => {
3474                    Self::TopicAuthorizationFailed(topics)
3475                }
3476                ConsumerError::Client(krabka_client_core::ClientError::IncompatibleVersion {
3477                    api_key,
3478                    broker_min,
3479                    broker_max,
3480                    client_min,
3481                    client_max,
3482                }) => Self::IncompatibleVersion(
3483                    api_key, broker_min, broker_max, client_min, client_max,
3484                ),
3485                _ => Self::Other,
3486            }
3487        }
3488    }
3489
3490    /// How the consumer handled one `ListOffsets` exchange.
3491    #[derive(Debug, PartialEq)]
3492    struct ListOffsetsOutcome {
3493        /// Each decoded request, with the name of the broker that received it
3494        /// and the version.
3495        requests: Vec<(&'static str, i16, ListOffsetsRequest)>,
3496        result: Result<(), PollFailure>,
3497        /// The `Metadata` requests after the setup.
3498        metadata_requests: usize,
3499        /// Whether the exchange asked for a refresh before the next poll.
3500        metadata_update_requested: bool,
3501        next_offset: Option<i64>,
3502    }
3503
3504    /// How the test reaches the `ListOffsets` path.
3505    #[derive(Clone, Copy, Debug)]
3506    enum Reset {
3507        /// `auto.offset.reset=latest`: a sentinel that `prepare_poll` resolves.
3508        Latest,
3509        /// `auto.offset.reset=by_duration:PT1H`: a sentinel that
3510        /// `prepare_poll` resolves with the timestamp of one hour ago.
3511        ByOneHour,
3512        /// `auto.offset.reset=earliest`: a Fetch row with
3513        /// `OFFSET_OUT_OF_RANGE`.
3514        EarliestOutOfRange,
3515        /// `seek_to_beginning` of every partition, then `position`.
3516        SeekToBeginning,
3517        /// `seek_to_end` of every partition, then `position`.
3518        SeekToEnd,
3519    }
3520
3521    /// One `ListOffsets` scenario on two mock brokers.
3522    #[derive(Clone, Copy, Debug)]
3523    struct Exchange {
3524        isolation: IsolationLevel,
3525        reset: Reset,
3526        /// The `(error_code, offset)` row that the leader answers.
3527        row: (i16, i64),
3528        /// The highest `ListOffsets` version that both brokers advertise.
3529        list_offsets_max_version: i16,
3530        /// `true` when a metadata refresh fills the positions map before the
3531        /// exchange, as after an earlier poll. `false` is a new assignment
3532        /// without a committed offset, whose position has no leader yet.
3533        primed: bool,
3534    }
3535
3536    /// Run `exchange` with a consumer that bootstraps from a broker that is
3537    /// not the leader of `orders-0`.
3538    ///
3539    /// The bootstrap broker answers each `ListOffsets` with
3540    /// `NOT_LEADER_OR_FOLLOWER`.
3541    async fn run_list_offsets_exchange(exchange: Exchange) -> ListOffsetsOutcome {
3542        let sent: SentListOffsets = Arc::default();
3543        let metadata_requests = Arc::new(AtomicUsize::new(0));
3544        let leader_port = Arc::new(AtomicU16::new(0));
3545        let leader = list_offsets_broker(
3546            "leader",
3547            Arc::clone(&leader_port),
3548            exchange.row,
3549            exchange.list_offsets_max_version,
3550            Arc::clone(&sent),
3551            Arc::clone(&metadata_requests),
3552        )
3553        .await;
3554        leader_port.store(leader.addr.port(), Ordering::SeqCst);
3555        let bootstrap = list_offsets_broker(
3556            "bootstrap",
3557            Arc::clone(&leader_port),
3558            (6, -1),
3559            exchange.list_offsets_max_version,
3560            Arc::clone(&sent),
3561            Arc::clone(&metadata_requests),
3562        )
3563        .await;
3564        let mut consumer = consumer_on(&bootstrap).await;
3565        consumer.isolation_level = exchange.isolation;
3566        if exchange.primed {
3567            consumer
3568                .refresh_leader_epochs()
3569                .await
3570                .expect("setup metadata");
3571            metadata_requests.store(0, Ordering::SeqCst);
3572        }
3573
3574        let result = match exchange.reset {
3575            Reset::Latest => {
3576                consumer.auto_offset_reset = AutoOffsetReset::Latest;
3577                consumer
3578                    .next_offsets
3579                    .lock()
3580                    .await
3581                    .insert(("orders".into(), 0), LATEST_SENTINEL);
3582                consumer
3583                    .prepare_poll(Instant::now() + Duration::from_secs(5))
3584                    .await
3585                    .map(|_| ())
3586            }
3587            Reset::ByOneHour => {
3588                consumer.auto_offset_reset = AutoOffsetReset::ByDuration(Duration::from_hours(1));
3589                consumer
3590                    .next_offsets
3591                    .lock()
3592                    .await
3593                    .insert(("orders".into(), 0), LATEST_SENTINEL);
3594                consumer
3595                    .prepare_poll(Instant::now() + Duration::from_secs(5))
3596                    .await
3597                    .map(|_| ())
3598            }
3599            Reset::EarliestOutOfRange => {
3600                consumer.auto_offset_reset = AutoOffsetReset::Earliest;
3601                let topic_ids = consumer.topic_ids.lock().await.clone();
3602                consumer
3603                    .process_fetch_responses(vec![fetch_response(1)], &topic_ids)
3604                    .await
3605            }
3606            Reset::SeekToBeginning | Reset::SeekToEnd => {
3607                // Kafka's `seekToBeginning` and `seekToEnd` ignore
3608                // `auto.offset.reset`.
3609                consumer.auto_offset_reset = AutoOffsetReset::ByDuration(Duration::from_hours(1));
3610                consumer.default_api_timeout = krabka_units::millis(300);
3611                let reset = if matches!(exchange.reset, Reset::SeekToBeginning) {
3612                    consumer.seek_to_beginning(&[]).await
3613                } else {
3614                    consumer.seek_to_end(&[]).await
3615                };
3616                match reset {
3617                    Ok(()) => consumer.position("orders", 0).await.map(|_| ()),
3618                    Err(error) => Err(error),
3619                }
3620            }
3621        };
3622        let requests = sent.lock().expect("sent lock").clone();
3623        let outcome = ListOffsetsOutcome {
3624            requests,
3625            result: result.map_err(PollFailure::from),
3626            metadata_requests: metadata_requests.load(Ordering::SeqCst),
3627            metadata_update_requested: consumer.client.metadata_topics().update_requested(),
3628            next_offset: consumer
3629                .next_offsets
3630                .lock()
3631                .await
3632                .get(&("orders".to_owned(), 0))
3633                .copied(),
3634        };
3635
3636        bootstrap.stop();
3637        leader.stop();
3638        outcome
3639    }
3640
3641    /// Kafka's consumer sends `ListOffsets` to the partition leader with its
3642    /// own isolation level and the leader epoch from metadata
3643    /// (`OffsetFetcher.groupListOffsetRequests` and
3644    /// `ListOffsetsRequest.Builder.forConsumer`). A row with an error never
3645    /// sets a position. `TOPIC_AUTHORIZATION_FAILED` fails the poll, and any
3646    /// other code retries after a metadata refresh
3647    /// (`OffsetFetcherUtils.handleListOffsetResponse`).
3648    #[tokio::test]
3649    async fn list_offsets_goes_to_the_leader_with_the_isolation_level_and_honours_row_errors() {
3650        let max = list_offsets_request::MAX_VERSION;
3651        let latest_request =
3652            |isolation| vec![("leader", max, expected_list_offsets(isolation, -1))];
3653        let latest = |isolation, row| Exchange {
3654            isolation,
3655            reset: Reset::Latest,
3656            row,
3657            list_offsets_max_version: max,
3658            primed: true,
3659        };
3660        let earliest = |isolation, row| Exchange {
3661            reset: Reset::EarliestOutOfRange,
3662            ..latest(isolation, row)
3663        };
3664        for (name, exchange, expected) in [
3665            (
3666                "read committed latest reads the last stable offset",
3667                latest(IsolationLevel::ReadCommitted, (0, 40)),
3668                ListOffsetsOutcome {
3669                    requests: latest_request(1),
3670                    result: Ok(()),
3671                    metadata_requests: 1,
3672                    metadata_update_requested: false,
3673                    next_offset: Some(40),
3674                },
3675            ),
3676            (
3677                "read uncommitted latest reads the high watermark",
3678                latest(IsolationLevel::ReadUncommitted, (0, 50)),
3679                ListOffsetsOutcome {
3680                    requests: latest_request(0),
3681                    result: Ok(()),
3682                    metadata_requests: 1,
3683                    metadata_update_requested: false,
3684                    next_offset: Some(50),
3685                },
3686            ),
3687            (
3688                "not leader or follower keeps the sentinel and refreshes metadata",
3689                latest(IsolationLevel::ReadCommitted, (6, -1)),
3690                ListOffsetsOutcome {
3691                    requests: latest_request(1),
3692                    result: Ok(()),
3693                    metadata_requests: 1,
3694                    metadata_update_requested: true,
3695                    next_offset: Some(LATEST_SENTINEL),
3696                },
3697            ),
3698            (
3699                "topic authorization failed fails the poll",
3700                latest(IsolationLevel::ReadCommitted, (29, -1)),
3701                ListOffsetsOutcome {
3702                    requests: latest_request(1),
3703                    result: Err(PollFailure::TopicAuthorizationFailed(
3704                        maplit::btreeset! {"orders".to_owned()},
3705                    )),
3706                    metadata_requests: 1,
3707                    metadata_update_requested: false,
3708                    next_offset: Some(LATEST_SENTINEL),
3709                },
3710            ),
3711            (
3712                "seek to beginning reads the log start",
3713                Exchange {
3714                    reset: Reset::SeekToBeginning,
3715                    ..latest(IsolationLevel::ReadUncommitted, (0, 7))
3716                },
3717                ListOffsetsOutcome {
3718                    requests: vec![("leader", max, expected_list_offsets(0, -2))],
3719                    result: Ok(()),
3720                    metadata_requests: 0,
3721                    metadata_update_requested: false,
3722                    next_offset: Some(7),
3723                },
3724            ),
3725            (
3726                "seek to end reads the last stable offset",
3727                Exchange {
3728                    reset: Reset::SeekToEnd,
3729                    ..latest(IsolationLevel::ReadCommitted, (0, 40))
3730                },
3731                ListOffsetsOutcome {
3732                    requests: latest_request(1),
3733                    result: Ok(()),
3734                    metadata_requests: 0,
3735                    metadata_update_requested: false,
3736                    next_offset: Some(40),
3737                },
3738            ),
3739            (
3740                "earliest after out of range reads the log start",
3741                earliest(IsolationLevel::ReadUncommitted, (0, 7)),
3742                ListOffsetsOutcome {
3743                    requests: vec![("leader", max, expected_list_offsets(0, -2))],
3744                    result: Ok(()),
3745                    metadata_requests: 0,
3746                    metadata_update_requested: false,
3747                    next_offset: Some(7),
3748                },
3749            ),
3750            (
3751                "earliest with a retriable row keeps the position and refreshes metadata",
3752                earliest(IsolationLevel::ReadCommitted, (74, -1)),
3753                ListOffsetsOutcome {
3754                    requests: vec![("leader", max, expected_list_offsets(1, -2))],
3755                    result: Ok(()),
3756                    metadata_requests: 0,
3757                    metadata_update_requested: true,
3758                    next_offset: Some(5),
3759                },
3760            ),
3761        ] {
3762            let outcome = run_list_offsets_exchange(exchange).await;
3763            check!(outcome == expected, "case {name}");
3764        }
3765    }
3766
3767    /// Kafka's `AutoOffsetResetStrategy.timestamp` for `by_duration` is now
3768    /// minus the duration, and the reset sends `ListOffsets` with it (KIP-1106).
3769    #[tokio::test]
3770    async fn by_duration_resets_with_the_timestamp_of_now_minus_the_duration() {
3771        let before = unix_now_ms() - 3_600_000;
3772        let mut outcome = run_list_offsets_exchange(Exchange {
3773            isolation: IsolationLevel::ReadUncommitted,
3774            reset: Reset::ByOneHour,
3775            row: (0, 40),
3776            list_offsets_max_version: list_offsets_request::MAX_VERSION,
3777            primed: true,
3778        })
3779        .await;
3780        let after = unix_now_ms() - 3_600_000;
3781        // The wall clock moves during the exchange: check the range, then
3782        // compare the whole outcome with the timestamp of the start.
3783        let mut in_range = Vec::new();
3784        for (_, _, request) in &mut outcome.requests {
3785            for topic in &mut request.topics {
3786                for partition in &mut topic.partitions {
3787                    in_range.push((before..=after).contains(&partition.timestamp));
3788                    partition.timestamp = before;
3789                }
3790            }
3791        }
3792
3793        assert2::assert!(
3794            (in_range, outcome)
3795                == (
3796                    vec![true],
3797                    ListOffsetsOutcome {
3798                        requests: vec![(
3799                            "leader",
3800                            list_offsets_request::MAX_VERSION,
3801                            expected_list_offsets(0, before),
3802                        )],
3803                        result: Ok(()),
3804                        metadata_requests: 1,
3805                        metadata_update_requested: false,
3806                        next_offset: Some(40),
3807                    }
3808                )
3809        );
3810    }
3811
3812    #[test]
3813    fn reset_timestamp_follows_the_reset_policy() {
3814        let now = 1_700_000_000_000;
3815        let actual = [
3816            AutoOffsetReset::Latest,
3817            AutoOffsetReset::None,
3818            AutoOffsetReset::ByDuration(Duration::from_hours(1)),
3819            AutoOffsetReset::ByDuration(Duration::from_millis(1500)),
3820        ]
3821        .map(|policy| reset_timestamp(policy, now));
3822        assert2::assert!(actual == [-1, -1, now - 3_600_000, now - 1500]);
3823    }
3824
3825    /// A new assignment without a committed offset has no leader in the
3826    /// positions map yet. Kafka's `OffsetFetcher.groupListOffsetRequests`
3827    /// routes with `metadata.currentLeader(tp)`, so the first `ListOffsets`
3828    /// of the first poll goes to the leader and not to the bootstrap broker.
3829    #[tokio::test]
3830    async fn first_poll_sends_list_offsets_to_the_leader_of_a_new_partition() {
3831        let outcome = run_list_offsets_exchange(Exchange {
3832            isolation: IsolationLevel::ReadCommitted,
3833            reset: Reset::Latest,
3834            row: (0, 40),
3835            list_offsets_max_version: list_offsets_request::MAX_VERSION,
3836            primed: false,
3837        })
3838        .await;
3839
3840        assert2::assert!(
3841            outcome
3842                == ListOffsetsOutcome {
3843                    requests: vec![(
3844                        "leader",
3845                        list_offsets_request::MAX_VERSION,
3846                        expected_list_offsets(1, -1),
3847                    )],
3848                    result: Ok(()),
3849                    metadata_requests: 1,
3850                    metadata_update_requested: false,
3851                    next_offset: Some(40),
3852                }
3853        );
3854    }
3855
3856    /// `isolation_level` exists from `ListOffsets` version 2. Kafka's
3857    /// `ListOffsetsRequest.Builder.forConsumer` sets the oldest allowed
3858    /// version to 2 for `READ_COMMITTED`, and to 1 (the oldest version) for
3859    /// `READ_UNCOMMITTED`. A `read_committed` consumer on a broker that
3860    /// supports only version 1 fails the poll with an unsupported version. It
3861    /// does not reset to the high watermark.
3862    #[tokio::test]
3863    async fn read_committed_list_offsets_needs_version_2() {
3864        let at = |isolation, list_offsets_max_version| Exchange {
3865            isolation,
3866            reset: Reset::Latest,
3867            row: (0, 40),
3868            list_offsets_max_version,
3869            primed: true,
3870        };
3871        // Below version 4 the request has no `current_leader_epoch`, and the
3872        // decoder gives the default.
3873        let without_epoch = |isolation_level| {
3874            let mut request = expected_list_offsets(isolation_level, -1);
3875            request.topics[0].partitions[0].current_leader_epoch =
3876                ListOffsetsPartition::default().current_leader_epoch;
3877            request
3878        };
3879        for (name, exchange, expected) in [
3880            (
3881                "read committed on a version 1 broker fails the poll",
3882                at(IsolationLevel::ReadCommitted, 1),
3883                ListOffsetsOutcome {
3884                    requests: vec![],
3885                    result: Err(PollFailure::IncompatibleVersion(
3886                        list_offsets_request::API_KEY,
3887                        0,
3888                        1,
3889                        2,
3890                        list_offsets_request::MAX_VERSION,
3891                    )),
3892                    metadata_requests: 1,
3893                    metadata_update_requested: false,
3894                    next_offset: Some(LATEST_SENTINEL),
3895                },
3896            ),
3897            (
3898                "read committed on a version 2 broker sends version 2",
3899                at(IsolationLevel::ReadCommitted, 2),
3900                ListOffsetsOutcome {
3901                    requests: vec![("leader", 2, without_epoch(1))],
3902                    result: Ok(()),
3903                    metadata_requests: 1,
3904                    metadata_update_requested: false,
3905                    next_offset: Some(40),
3906                },
3907            ),
3908            (
3909                "read uncommitted on a version 1 broker sends version 1",
3910                at(IsolationLevel::ReadUncommitted, 1),
3911                ListOffsetsOutcome {
3912                    requests: vec![("leader", 1, without_epoch(0))],
3913                    result: Ok(()),
3914                    metadata_requests: 1,
3915                    metadata_update_requested: false,
3916                    next_offset: Some(40),
3917                },
3918            ),
3919        ] {
3920            let outcome = run_list_offsets_exchange(exchange).await;
3921            check!(outcome == expected, "case {name}");
3922        }
3923    }
3924
3925    /// Kafka's consumer gives each record the timestamp and the timestamp type
3926    /// of its record batch. A `LogAppendTime` batch gives every record the
3927    /// batch `max_timestamp`.
3928    #[tokio::test]
3929    async fn fetched_records_take_the_batch_timestamp_type() {
3930        let broker = metadata_counting_broker(Arc::default()).await;
3931        let consumer = consumer_on(&broker).await;
3932        let key: (String, i32) = ("orders".into(), 0);
3933
3934        for case in &timestamp_cases::CASES {
3935            let part = PartitionData {
3936                partition_index: 0,
3937                records: Some(timestamp_cases::batch(case).into()),
3938                ..Default::default()
3939            };
3940            let offsets = maplit::hashmap! {key.clone() => timestamp_cases::BASE_OFFSET};
3941
3942            let records: Vec<ConsumerRecord> = consumer
3943                .process_partition_records(&offsets, &key, "orders", &part)
3944                .map(|partition| partition.records.into_iter().collect())
3945                .unwrap_or_default();
3946
3947            let expected: Vec<ConsumerRecord> = case
3948                .expected
3949                .iter()
3950                .zip(timestamp_cases::BASE_OFFSET..)
3951                .zip(timestamp_cases::VALUES)
3952                .map(
3953                    |(((timestamp, timestamp_type), offset), value)| ConsumerRecord {
3954                        topic: "orders".into(),
3955                        partition: 0,
3956                        offset,
3957                        leader_epoch: timestamp_cases::LEADER_EPOCH,
3958                        timestamp: *timestamp,
3959                        timestamp_type: *timestamp_type,
3960                        key: None,
3961                        value: Some(bytes::Bytes::from_static(value)),
3962                        headers: Vec::new(),
3963                    },
3964                )
3965                .collect();
3966            check!(records == expected, "case {}", case.name);
3967        }
3968    }
3969}
3970
3971/// Record batch timestamp cases that the classic and the share consumer tests
3972/// share.
3973#[cfg(test)]
3974pub mod timestamp_cases {
3975    use bytes::Bytes;
3976    use consumer::TimestampType;
3977    use krabka_protocol::records::{Attributes, Record, RecordBatch};
3978
3979    use crate::consumer;
3980
3981    pub const BASE_OFFSET: i64 = 5;
3982    pub const LEADER_EPOCH: i32 = 3;
3983    pub const VALUES: [&[u8]; 2] = [b"v0", b"v1"];
3984
3985    /// One batch with two records and the timestamps Kafka gives them.
3986    pub struct TimestampCase {
3987        pub(crate) name: &'static str,
3988        pub(crate) attributes: Attributes,
3989        pub(crate) base_timestamp: i64,
3990        pub(crate) max_timestamp: i64,
3991        pub(crate) timestamp_deltas: [i64; 2],
3992        pub(crate) expected: [(i64, TimestampType); 2],
3993    }
3994
3995    pub const CASES: [TimestampCase; 3] = [
3996        TimestampCase {
3997            name: "create time adds the delta to the base timestamp",
3998            attributes: Attributes(0),
3999            base_timestamp: 1000,
4000            max_timestamp: 1005,
4001            timestamp_deltas: [0, 5],
4002            expected: [
4003                (1000, TimestampType::CreateTime),
4004                (1005, TimestampType::CreateTime),
4005            ],
4006        },
4007        TimestampCase {
4008            name: "log append time uses the max timestamp",
4009            attributes: Attributes(Attributes::TIMESTAMP_TYPE_BIT),
4010            base_timestamp: 1000,
4011            max_timestamp: 9000,
4012            timestamp_deltas: [0, 5],
4013            expected: [
4014                (9000, TimestampType::LogAppendTime),
4015                (9000, TimestampType::LogAppendTime),
4016            ],
4017        },
4018        TimestampCase {
4019            name: "log append time in a transactional batch uses the max timestamp",
4020            attributes: Attributes(Attributes::TIMESTAMP_TYPE_BIT | Attributes::TRANSACTIONAL_BIT),
4021            base_timestamp: 1000,
4022            max_timestamp: 9000,
4023            timestamp_deltas: [0, 5],
4024            expected: [
4025                (9000, TimestampType::LogAppendTime),
4026                (9000, TimestampType::LogAppendTime),
4027            ],
4028        },
4029    ];
4030
4031    pub fn batch(case: &TimestampCase) -> RecordBatch {
4032        RecordBatch {
4033            base_offset: BASE_OFFSET,
4034            partition_leader_epoch: LEADER_EPOCH,
4035            attributes: case.attributes,
4036            last_offset_delta: 1,
4037            base_timestamp: case.base_timestamp,
4038            max_timestamp: case.max_timestamp,
4039            records: case
4040                .timestamp_deltas
4041                .iter()
4042                .zip(0..)
4043                .zip(VALUES)
4044                .map(|((timestamp_delta, offset_delta), value)| Record {
4045                    attributes: 0,
4046                    timestamp_delta: *timestamp_delta,
4047                    offset_delta,
4048                    key: None,
4049                    value: Some(Bytes::from_static(value)),
4050                    headers: Vec::new(),
4051                })
4052                .collect(),
4053            ..Default::default()
4054        }
4055    }
4056}
4057
4058#[cfg(test)]
4059mod fetch_path_tests {
4060    use std::{
4061        sync::{
4062            Arc, Mutex,
4063            atomic::{AtomicU16, Ordering},
4064        },
4065        time::Duration,
4066    };
4067
4068    use bytes::Buf as _;
4069    use krabka_client_core::{Client, MockBroker};
4070    use krabka_protocol::{
4071        Decode as _, Encode,
4072        owned::{
4073            api_versions_request, fetch_request,
4074            fetch_response::{FetchableTopicResponse, PartitionData},
4075            list_offsets_request,
4076            list_offsets_response::ListOffsetsResponse,
4077            metadata_request,
4078            metadata_response::{
4079                MetadataResponse, MetadataResponseBroker, MetadataResponsePartition,
4080                MetadataResponseTopic,
4081            },
4082        },
4083        primitives::uuid::Uuid as WireUuid,
4084    };
4085    use krabka_units::{millis, secs};
4086    use test_support::RecordingListener;
4087    use tokio::{
4088        sync::watch,
4089        time::{sleep, timeout},
4090    };
4091
4092    use super::*;
4093    use crate::rebalance_listener::AssignedCallbackGate;
4094
4095    const TOPIC_ID: WireUuid = WireUuid([7; 16]);
4096
4097    /// What a mock broker does with a Fetch request.
4098    #[derive(Clone, Copy, Debug)]
4099    enum FetchAnswer {
4100        /// No response.
4101        Silent,
4102        /// A response with this top-level error code, session id and no
4103        /// partition data.
4104        Respond { error_code: i16, session_id: i32 },
4105        /// A response with one `orders-0` row without records.
4106        Partition {
4107            error_code: i16,
4108            preferred_read_replica: i32,
4109        },
4110    }
4111
4112    /// The Fetch requests that the mock brokers received: `(broker id,
4113    /// request)`.
4114    type SentFetches = Arc<Mutex<Vec<(i32, FetchRequest)>>>;
4115
4116    fn encode(response: &impl Encode, version: i16) -> Vec<u8> {
4117        let mut buf = bytes::BytesMut::new();
4118        response.encode(&mut buf, version).expect("encode");
4119        buf.to_vec()
4120    }
4121
4122    /// Start one mock broker per entry of `answers`, with node ids from 1.
4123    /// Metadata names broker `i + 1` as the leader of `orders-i`. Each broker
4124    /// answers Fetch v4 to v11 with its entry of `answers`.
4125    async fn start_brokers(answers: &[Vec<FetchAnswer>], sent: &SentFetches) -> Vec<MockBroker> {
4126        let ports: Arc<Vec<AtomicU16>> =
4127            Arc::new(answers.iter().map(|_| AtomicU16::new(0)).collect());
4128        let mut brokers = Vec::new();
4129        for (index, answers) in answers.iter().enumerate() {
4130            let node_id = i32::try_from(index).expect("index") + 1;
4131            let ports_in_mock = Arc::clone(&ports);
4132            let sent = Arc::clone(sent);
4133            let answers = Mutex::new(answers.clone().into_iter());
4134            let last = Mutex::new(FetchAnswer::Silent);
4135            let broker = MockBroker::start(move |api_key, version, _corr_id, mut body| {
4136                if api_key == api_versions_request::API_KEY {
4137                    let versions = mock_test_support::api_versions_response_ranges([
4138                        (api_versions_request::API_KEY, 0, 3),
4139                        (metadata_request::API_KEY, 0, 8),
4140                        (fetch_request::API_KEY, 4, 11),
4141                        (list_offsets_request::API_KEY, 1, 5),
4142                    ]);
4143                    return Some(encode(&versions, 0));
4144                }
4145                if api_key == metadata_request::API_KEY {
4146                    let ports = &ports_in_mock;
4147                    let count = ports.len();
4148                    let metadata = MetadataResponse {
4149                        brokers: (0..count)
4150                            .map(|index| MetadataResponseBroker {
4151                                node_id: i32::try_from(index).expect("index") + 1,
4152                                host: "127.0.0.1".into(),
4153                                port: i32::from(ports[index].load(Ordering::SeqCst)),
4154                                ..Default::default()
4155                            })
4156                            .collect(),
4157                        topics: vec![MetadataResponseTopic {
4158                            name: Some("orders".into()),
4159                            topic_id: TOPIC_ID,
4160                            partitions: (0..count)
4161                                .map(|index| {
4162                                    let index = i32::try_from(index).expect("index");
4163                                    MetadataResponsePartition {
4164                                        partition_index: index,
4165                                        leader_id: index + 1,
4166                                        leader_epoch: 1,
4167                                        ..Default::default()
4168                                    }
4169                                })
4170                                .collect(),
4171                            ..Default::default()
4172                        }],
4173                        ..Default::default()
4174                    };
4175                    return Some(encode(&metadata, version));
4176                }
4177                if api_key == fetch_request::API_KEY {
4178                    let client_id_len = usize::try_from(body.get_i16()).expect("client id");
4179                    body.advance(client_id_len);
4180                    let request = FetchRequest::decode(&mut body, version).expect("decode fetch");
4181                    sent.lock().expect("sent lock").push((node_id, request));
4182                    let answer = answers.lock().expect("answers lock").next();
4183                    let mut last = last.lock().expect("last lock");
4184                    if let Some(answer) = answer {
4185                        *last = answer;
4186                    }
4187                    return match *last {
4188                        FetchAnswer::Silent => None,
4189                        FetchAnswer::Respond {
4190                            error_code,
4191                            session_id,
4192                        } => Some(encode(
4193                            &FetchResponse {
4194                                error_code,
4195                                session_id,
4196                                ..Default::default()
4197                            },
4198                            version,
4199                        )),
4200                        FetchAnswer::Partition {
4201                            error_code,
4202                            preferred_read_replica,
4203                        } => Some(encode(
4204                            &FetchResponse {
4205                                responses: vec![FetchableTopicResponse {
4206                                    topic: "orders".into(),
4207                                    partitions: vec![PartitionData {
4208                                        partition_index: 0,
4209                                        error_code,
4210                                        high_watermark: 5,
4211                                        last_stable_offset: 5,
4212                                        preferred_read_replica,
4213                                        ..Default::default()
4214                                    }],
4215                                    ..Default::default()
4216                                }],
4217                                ..Default::default()
4218                            },
4219                            version,
4220                        )),
4221                    };
4222                }
4223                // Every partition resolves to offset 0, whichever timestamp
4224                // the request asked for.
4225                if api_key == list_offsets_request::API_KEY {
4226                    return Some(test_support::list_offsets_mock_response(&mut body, version));
4227                }
4228                None
4229            })
4230            .await
4231            .unwrap();
4232            ports[index].store(broker.addr.port(), Ordering::SeqCst);
4233            brokers.push(broker);
4234        }
4235        brokers
4236    }
4237
4238    /// A consumer that bootstraps from the first broker and owns one partition
4239    /// of `orders` per broker, each at offset 5.
4240    async fn consumer_on(brokers: &[MockBroker]) -> Consumer {
4241        let client = Client::builder()
4242            .bootstrap(brokers[0].addr.to_string())
4243            .request_timeout(secs(30))
4244            .build()
4245            .await
4246            .expect("client");
4247        let partitions: Vec<(String, i32)> = (0..brokers.len())
4248            .map(|index| ("orders".to_owned(), i32::try_from(index).expect("index")))
4249            .collect();
4250        let consumer = poll::partition_error_tests::consumer_with_client(client);
4251        *consumer.assigned.lock().await = partitions.clone();
4252        *consumer.next_offsets.lock().await =
4253            partitions.iter().map(|key| (key.clone(), 5)).collect();
4254        consumer.commit_identity.lock().await.ownership_ids =
4255            partitions.into_iter().zip(1..).collect();
4256        consumer
4257    }
4258
4259    struct TwoPartitionFetchSetup {
4260        brokers: Vec<MockBroker>,
4261        consumer: Consumer,
4262        sent: SentFetches,
4263    }
4264
4265    async fn two_partition_fetch_setup() -> TwoPartitionFetchSetup {
4266        let answer = FetchAnswer::Respond {
4267            error_code: 0,
4268            session_id: 0,
4269        };
4270        let sent = SentFetches::default();
4271        let brokers = start_brokers(&[vec![answer], vec![answer]], &sent).await;
4272        let consumer = consumer_on(&brokers).await;
4273        consumer.refresh_leader_epochs().await.expect("metadata");
4274        TwoPartitionFetchSetup {
4275            brokers,
4276            consumer,
4277            sent,
4278        }
4279    }
4280
4281    fn stop(brokers: Vec<MockBroker>) {
4282        for broker in brokers {
4283            broker.stop();
4284        }
4285    }
4286
4287    /// Kafka's `FetchRequest.max_wait_ms` is `fetch.max.wait.ms` (default
4288    /// 500), whatever the poll timeout.
4289    #[tokio::test]
4290    async fn fetch_waits_fetch_max_wait_and_not_the_poll_timeout() {
4291        let sent = SentFetches::default();
4292        let brokers = start_brokers(
4293            &[vec![FetchAnswer::Respond {
4294                error_code: 0,
4295                session_id: 0,
4296            }]],
4297            &sent,
4298        )
4299        .await;
4300        let mut consumer = consumer_on(&brokers).await;
4301        consumer.poll(secs(30)).await.expect("poll");
4302        let max_waits: Vec<i32> = sent
4303            .lock()
4304            .expect("sent lock")
4305            .iter()
4306            .map(|(_, request)| request.max_wait_ms)
4307            .collect();
4308        drop(consumer);
4309        stop(brokers);
4310        assert2::assert!(max_waits == vec![500]);
4311    }
4312
4313    /// Kafka sends the Fetch of every leader at once, and `poll` returns by its
4314    /// timeout, also when no leader answers.
4315    #[tokio::test]
4316    async fn fetches_go_to_all_leaders_at_once_and_poll_returns_by_its_timeout() {
4317        let sent = SentFetches::default();
4318        let silent = vec![FetchAnswer::Silent];
4319        let brokers = start_brokers(&[silent.clone(), silent.clone(), silent], &sent).await;
4320        let mut consumer = consumer_on(&brokers).await;
4321        consumer.refresh_leader_epochs().await.expect("metadata");
4322        let started = Instant::now();
4323        let records = consumer.poll(millis(300)).await.expect("poll");
4324        let elapsed = started.elapsed();
4325        // Let the requests reach the brokers.
4326        sleep(Duration::from_millis(100)).await;
4327        let mut leaders: Vec<i32> = sent
4328            .lock()
4329            .expect("sent lock")
4330            .iter()
4331            .map(|(node_id, _)| *node_id)
4332            .collect();
4333        leaders.sort_unstable();
4334        drop(consumer);
4335        stop(brokers);
4336        assert2::assert!(
4337            (
4338                records.len(),
4339                elapsed < Duration::from_millis(1500),
4340                leaders
4341            ) == (0, true, vec![1, 2, 3])
4342        );
4343    }
4344
4345    /// A partition that waits for its assign callback gets no Fetch. Kafka's
4346    /// `SubscriptionState.isFetchable` is false while
4347    /// `pendingOnAssignedCallback` is set.
4348    #[tokio::test]
4349    async fn fetch_leaves_out_a_partition_that_waits_for_its_assign_callback() {
4350        let TwoPartitionFetchSetup {
4351            sent,
4352            brokers,
4353            mut consumer,
4354        } = two_partition_fetch_setup().await;
4355        let gate = AssignedCallbackGate::new(
4356            &consumer.assigned_callback_pending,
4357            &[("orders".to_owned(), 0)],
4358        );
4359        consumer.poll(millis(200)).await.expect("poll");
4360        sleep(Duration::from_millis(50)).await;
4361        let leaders = |sent: &SentFetches| {
4362            let mut leaders: Vec<i32> = sent
4363                .lock()
4364                .expect("sent lock")
4365                .drain(..)
4366                .map(|(node_id, _)| node_id)
4367                .collect();
4368            leaders.sort_unstable();
4369            leaders
4370        };
4371        let waiting = leaders(&sent);
4372        drop(gate);
4373        consumer.poll(millis(200)).await.expect("poll");
4374        sleep(Duration::from_millis(50)).await;
4375        let after = leaders(&sent);
4376        drop(consumer);
4377        stop(brokers);
4378        assert2::assert!((waiting, after) == (vec![2], vec![1, 2]));
4379    }
4380
4381    /// Kafka's `seek` sets the position of an assigned partition at once, and
4382    /// `SubscriptionState.assignedState` throws for a partition that the
4383    /// consumer does not own. The next Fetch uses the sought offset.
4384    #[tokio::test]
4385    async fn seek_moves_an_owned_position_and_rejects_an_unowned_partition() {
4386        let mut actual = Vec::new();
4387        let mut wanted = Vec::new();
4388        for (name, partition, expected_error, expected_offset) in [
4389            ("owned partition", 0, None, 2),
4390            (
4391                "unowned partition",
4392                9,
4393                Some("no current assignment for partition orders-9"),
4394                5,
4395            ),
4396        ] {
4397            let sent = SentFetches::default();
4398            let brokers = start_brokers(
4399                &[vec![FetchAnswer::Respond {
4400                    error_code: 0,
4401                    session_id: 0,
4402                }]],
4403                &sent,
4404            )
4405            .await;
4406            let mut consumer = consumer_on(&brokers).await;
4407            let error = consumer
4408                .seek("orders", partition, 2)
4409                .await
4410                .err()
4411                .map(|error| error.to_string());
4412            consumer.poll(millis(200)).await.expect("poll");
4413            let offsets: Vec<i64> = sent
4414                .lock()
4415                .expect("sent lock")
4416                .iter()
4417                .flat_map(|(_, request)| &request.topics)
4418                .flat_map(|topic| &topic.partitions)
4419                .map(|partition| partition.fetch_offset)
4420                .collect();
4421            drop(consumer);
4422            stop(brokers);
4423            actual.push((name, error, offsets));
4424            wanted.push((
4425                name,
4426                expected_error.map(str::to_owned),
4427                vec![expected_offset],
4428            ));
4429        }
4430        assert2::assert!(actual == wanted);
4431    }
4432
4433    /// Kafka's `AbstractFetch.fetchablePartitions` leaves out a paused
4434    /// partition, so its leader gets no Fetch until `resume`.
4435    #[tokio::test]
4436    async fn fetch_leaves_out_a_paused_partition_until_resume() {
4437        let TwoPartitionFetchSetup {
4438            sent,
4439            brokers,
4440            mut consumer,
4441        } = two_partition_fetch_setup().await;
4442        let fetched = |sent: &SentFetches| {
4443            let mut partitions: Vec<(i32, i32)> = sent
4444                .lock()
4445                .expect("sent lock")
4446                .drain(..)
4447                .flat_map(|(node_id, request)| {
4448                    request
4449                        .topics
4450                        .into_iter()
4451                        .flat_map(|topic| topic.partitions)
4452                        .map(move |partition| (node_id, partition.partition))
4453                })
4454                .collect();
4455            partitions.sort_unstable();
4456            partitions
4457        };
4458        consumer
4459            .pause(&[("orders".to_owned(), 0)])
4460            .await
4461            .expect("pause");
4462        consumer.poll(millis(200)).await.expect("poll");
4463        sleep(Duration::from_millis(50)).await;
4464        let paused = fetched(&sent);
4465        consumer
4466            .resume(&[("orders".to_owned(), 0)])
4467            .await
4468            .expect("resume");
4469        consumer.poll(millis(200)).await.expect("poll");
4470        sleep(Duration::from_millis(50)).await;
4471        let resumed = fetched(&sent);
4472        let unowned = consumer
4473            .position("orders", 9)
4474            .await
4475            .map_err(|error| error.to_string());
4476        drop(consumer);
4477        stop(brokers);
4478        assert2::assert!(
4479            (paused, resumed, unowned)
4480                == (
4481                    vec![(2, 1)],
4482                    vec![(1, 0), (2, 1)],
4483                    Err("no current assignment for partition orders-9".to_owned())
4484                )
4485        );
4486    }
4487
4488    /// Kafka's `pollForFetches` waits for the poll timeout when no partition
4489    /// can be fetched, for example when all are paused.
4490    #[tokio::test]
4491    async fn poll_waits_for_its_timeout_when_every_partition_is_paused() {
4492        let sent = SentFetches::default();
4493        let brokers = start_brokers(
4494            &[vec![FetchAnswer::Respond {
4495                error_code: 0,
4496                session_id: 0,
4497            }]],
4498            &sent,
4499        )
4500        .await;
4501        let mut consumer = consumer_on(&brokers).await;
4502        consumer
4503            .pause(&[("orders".to_owned(), 0)])
4504            .await
4505            .expect("pause");
4506        let started = Instant::now();
4507        let records = consumer.poll(millis(300)).await.expect("poll").len();
4508        let waited = started.elapsed() >= Duration::from_millis(250);
4509        let fetches = sent.lock().expect("sent lock").len();
4510        drop(consumer);
4511        stop(brokers);
4512        assert2::assert!((records, waited, fetches) == (0, true, 0));
4513    }
4514
4515    /// Kafka's `wakeup` makes the blocked `poll`, or the next one, throw
4516    /// `WakeupException` once.
4517    #[tokio::test]
4518    async fn wakeup_ends_the_poll_that_waits_or_the_next_one() {
4519        let sent = SentFetches::default();
4520        let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
4521        let mut consumer = consumer_on(&brokers).await;
4522        consumer.wakeup();
4523        let pending = consumer
4524            .poll(millis(100))
4525            .await
4526            .map_err(|error| error.to_string());
4527        let handle = consumer.wakeup_handle();
4528        let waker = tokio::spawn(async move {
4529            sleep(Duration::from_millis(100)).await;
4530            handle.wakeup();
4531        });
4532        let started = Instant::now();
4533        let blocked = consumer
4534            .poll(secs(10))
4535            .await
4536            .map_err(|error| error.to_string());
4537        let woken_fast = started.elapsed() < Duration::from_secs(5);
4538        waker.await.expect("waker");
4539        let next = consumer
4540            .poll(millis(100))
4541            .await
4542            .map(|records| records.len());
4543        drop(consumer);
4544        stop(brokers);
4545        let wakeup = Err("the consumer was woken up".to_owned());
4546        assert2::assert!(
4547            (
4548                pending,
4549                blocked,
4550                woken_fast,
4551                next.map_err(|error| error.to_string())
4552            ) == (wakeup.clone(), wakeup, true, Ok(0))
4553        );
4554    }
4555
4556    /// Kafka's `wakeup` ends a `poll` that waits for a metadata request of
4557    /// `updateFetchPositions`, as `ConsumerNetworkClient.poll` throws
4558    /// `WakeupException` while it waits.
4559    #[tokio::test]
4560    async fn wakeup_ends_a_poll_that_waits_for_metadata() {
4561        let silent = MockBroker::start(|api_key, _version, _corr_id, _body| {
4562            (api_key == api_versions_request::API_KEY).then(|| {
4563                encode(
4564                    &mock_test_support::api_versions_response_ranges([
4565                        (api_versions_request::API_KEY, 0, 3),
4566                        (metadata_request::API_KEY, 0, 8),
4567                    ]),
4568                    0,
4569                )
4570            })
4571        })
4572        .await
4573        .unwrap();
4574        let client = Client::builder()
4575            .bootstrap(silent.addr.to_string())
4576            .request_timeout(secs(30))
4577            .build()
4578            .await
4579            .expect("client");
4580        let mut consumer = poll::partition_error_tests::consumer_with_client(client);
4581        let handle = consumer.wakeup_handle();
4582        let waker = tokio::spawn(async move {
4583            sleep(Duration::from_millis(100)).await;
4584            handle.wakeup();
4585        });
4586        let polled = timeout(Duration::from_secs(5), consumer.poll(secs(10)))
4587            .await
4588            .map(|result| {
4589                result
4590                    .map(|records| records.len())
4591                    .map_err(|error| error.to_string())
4592            });
4593        waker.await.expect("waker");
4594        drop(consumer);
4595        silent.stop();
4596        assert2::assert!(polled == Ok(Err("the consumer was woken up".to_owned())));
4597    }
4598
4599    /// Kafka's `pollForFetches` waits at most until the coordinator needs the
4600    /// application thread: a rebalance that starts while every partition is
4601    /// paused ends the wait.
4602    #[tokio::test]
4603    async fn a_rebalance_ends_the_wait_of_a_poll_with_nothing_to_fetch() {
4604        let sent = SentFetches::default();
4605        let brokers = start_brokers(
4606            &[vec![FetchAnswer::Respond {
4607                error_code: 0,
4608                session_id: 0,
4609            }]],
4610            &sent,
4611        )
4612        .await;
4613        let mut consumer = consumer_on(&brokers).await;
4614        consumer
4615            .pause(&[("orders".to_owned(), 0)])
4616            .await
4617            .expect("pause");
4618        let (pending, pending_rx) = watch::channel(false);
4619        consumer.rebalance_pending = pending_rx;
4620        let rebalance = tokio::spawn(async move {
4621            sleep(Duration::from_millis(100)).await;
4622            pending.send_replace(true);
4623            // Keep the sender, as the coordinator task does.
4624            sleep(Duration::from_secs(10)).await;
4625        });
4626        let started = Instant::now();
4627        consumer.poll(secs(5)).await.expect("poll");
4628        let ended_early = started.elapsed() < Duration::from_secs(2);
4629        rebalance.abort();
4630        drop(consumer);
4631        stop(brokers);
4632        assert2::assert!(ended_early);
4633    }
4634
4635    /// Kafka's `ConsumerMetadata` names the current subscription, so the
4636    /// metadata of `poll` covers a topic that a pattern added in the
4637    /// coordinator task.
4638    #[tokio::test]
4639    async fn poll_names_the_topics_that_a_pattern_added() {
4640        let sent = SentFetches::default();
4641        let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
4642        let consumer = consumer_on(&brokers).await;
4643        consumer.client.metadata_topics().set(["orders".to_owned()]);
4644        consumer.subscription.send_modify(|subscription| {
4645            subscription.pattern = Some(crate::TopicPattern::new(|_| true));
4646            subscription.topics = vec!["orders".to_owned(), "payments".to_owned()];
4647        });
4648        consumer.sync_metadata_topics();
4649        let names = consumer.client.metadata_topics().names();
4650        drop(consumer);
4651        stop(brokers);
4652        assert2::assert!(names == vec!["orders".to_owned(), "payments".to_owned()]);
4653    }
4654
4655    /// Kafka's `updateFetchPositions` gives the committed offsets the timer of
4656    /// the `poll`, so an unresponsive coordinator does not hold a short poll
4657    /// for the whole `default.api.timeout.ms`.
4658    #[tokio::test]
4659    async fn a_short_poll_does_not_wait_for_the_committed_offsets() {
4660        let sent = SentFetches::default();
4661        let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
4662        let mut consumer = consumer_on(&brokers).await;
4663        consumer
4664            .next_offsets
4665            .lock()
4666            .await
4667            .insert(("orders".to_owned(), 0), COMMITTED_SENTINEL);
4668        let started = Instant::now();
4669        let result = consumer
4670            .poll(millis(200))
4671            .await
4672            .map(|records| records.len());
4673        let elapsed = started.elapsed();
4674        drop(consumer);
4675        stop(brokers);
4676        assert2::assert!((result.is_err(), elapsed < Duration::from_secs(5)) == (true, true));
4677    }
4678
4679    /// Kafka's consumer without `group.id` fetches the partitions of
4680    /// `assign` from the reset position, and throws `InvalidGroupIdException`
4681    /// for the group calls.
4682    #[tokio::test]
4683    async fn a_consumer_without_a_group_fetches_its_manual_assignment() {
4684        let sent = SentFetches::default();
4685        let brokers = start_brokers(
4686            &[vec![FetchAnswer::Respond {
4687                error_code: 0,
4688                session_id: 0,
4689            }]],
4690            &sent,
4691        )
4692        .await;
4693        let calls = Arc::new(Mutex::new(Vec::new()));
4694        let mut consumer = Consumer::builder()
4695            .bootstrap(brokers[0].addr.to_string())
4696            .auto_offset_reset(AutoOffsetReset::Earliest)
4697            .rebalance_listener(Box::new(
4698                RecordingListener::builder()
4699                    .calls(Arc::clone(&calls))
4700                    .build(),
4701            ))
4702            .build()
4703            .await
4704            .expect("build");
4705        let before_assign = consumer
4706            .poll(millis(20))
4707            .await
4708            .map(|records| records.len())
4709            .map_err(|error| error.to_string());
4710        consumer
4711            .assign(&[("orders".to_owned(), 0)])
4712            .await
4713            .expect("assign");
4714        consumer.poll(millis(300)).await.expect("poll");
4715        let fetched: Vec<(i32, String, i32, i64)> = sent
4716            .lock()
4717            .expect("sent lock")
4718            .iter()
4719            .flat_map(|(node_id, request)| {
4720                request.topics.iter().flat_map(move |topic| {
4721                    topic.partitions.iter().map(move |partition| {
4722                        (
4723                            *node_id,
4724                            topic.topic.clone(),
4725                            partition.partition,
4726                            partition.fetch_offset,
4727                        )
4728                    })
4729                })
4730            })
4731            .collect();
4732        let group_calls = (
4733            consumer
4734                .commit_sync()
4735                .await
4736                .map_err(|error| error.to_string()),
4737            consumer
4738                .subscribe(["orders"])
4739                .await
4740                .map_err(|error| error.to_string()),
4741            consumer
4742                .enforce_rebalance(None)
4743                .map_err(|error| error.to_string()),
4744        );
4745        consumer.unsubscribe().await.expect("unsubscribe");
4746        let after_unsubscribe = (
4747            consumer.assignment().await,
4748            consumer.next_offsets.lock().await.len(),
4749        );
4750        let listener_calls = calls.lock().expect("calls lock").clone();
4751        consumer.close().await.expect("close");
4752        stop(brokers);
4753        let invalid_group = Err(ConsumerError::InvalidGroupId.to_string());
4754        assert2::assert!(
4755            (
4756                before_assign,
4757                fetched.first().cloned(),
4758                group_calls,
4759                after_unsubscribe,
4760                listener_calls
4761            ) == (
4762                Err("not subscribed to any topic".to_owned()),
4763                Some((1, "orders".to_owned(), 0, 0)),
4764                (invalid_group.clone(), invalid_group.clone(), invalid_group),
4765                (Vec::new(), 0),
4766                Vec::<test_support::ListenerCall>::new()
4767            )
4768        );
4769    }
4770
4771    /// One step of a read replica case.
4772    #[derive(Clone, Copy, Debug)]
4773    enum ReplicaStep {
4774        Poll,
4775        /// Wait longer than the metadata max age of the case.
4776        WaitForExpiry,
4777    }
4778
4779    /// KIP-392: Kafka's Fetch carries `client.rack`, and a
4780    /// `preferred_read_replica` in a response sends the next Fetch of the
4781    /// partition to that replica (`AbstractFetch.selectReadReplica`) until it
4782    /// expires after `metadata.max.age.ms` or the replica answers an error
4783    /// (`FetchCollector.handleInitializeErrors`).
4784    #[tokio::test]
4785    async fn fetch_sends_the_rack_and_follows_the_preferred_read_replica() {
4786        use FetchAnswer::Partition;
4787        use ReplicaStep::{Poll, WaitForExpiry};
4788        let records = Partition {
4789            error_code: 0,
4790            preferred_read_replica: -1,
4791        };
4792        let prefer_2 = Partition {
4793            error_code: 0,
4794            preferred_read_replica: 2,
4795        };
4796        let not_leader = Partition {
4797            error_code: 6,
4798            preferred_read_replica: -1,
4799        };
4800        let mut actual = Vec::new();
4801        let mut wanted = Vec::new();
4802        for (name, rack, leader, replica, steps, expected) in [
4803            (
4804                "no rack",
4805                None,
4806                vec![records],
4807                vec![records],
4808                vec![Poll],
4809                vec![(1, String::new())],
4810            ),
4811            (
4812                "rack",
4813                Some("az-1"),
4814                vec![records],
4815                vec![records],
4816                vec![Poll],
4817                vec![(1, "az-1".to_owned())],
4818            ),
4819            (
4820                "preferred read replica",
4821                Some("az-1"),
4822                vec![prefer_2],
4823                vec![records],
4824                vec![Poll, Poll],
4825                vec![(1, "az-1".to_owned()), (2, "az-1".to_owned())],
4826            ),
4827            (
4828                "the replica answers not leader or follower",
4829                Some("az-1"),
4830                vec![prefer_2, records],
4831                vec![not_leader],
4832                vec![Poll, Poll, Poll],
4833                vec![
4834                    (1, "az-1".to_owned()),
4835                    (2, "az-1".to_owned()),
4836                    (1, "az-1".to_owned()),
4837                ],
4838            ),
4839            (
4840                "the preference expires",
4841                Some("az-1"),
4842                vec![prefer_2, records],
4843                vec![records],
4844                vec![Poll, Poll, WaitForExpiry, Poll],
4845                vec![
4846                    (1, "az-1".to_owned()),
4847                    (2, "az-1".to_owned()),
4848                    (1, "az-1".to_owned()),
4849                ],
4850            ),
4851        ] {
4852            let sent = SentFetches::default();
4853            let brokers = start_brokers(&[leader, replica], &sent).await;
4854            let mut consumer = consumer_on(&brokers).await;
4855            *consumer.assigned.lock().await = vec![("orders".to_owned(), 0)];
4856            consumer.client_rack = rack.map(str::to_owned);
4857            consumer.metadata_max_age = millis(300);
4858            for step in steps {
4859                match step {
4860                    Poll => {
4861                        let before = sent.lock().expect("sent lock").len();
4862                        consumer.poll(secs(2)).await.expect("poll");
4863                        // The fetch sessions of the two brokers are independent;
4864                        // wait until the Fetch of this poll is recorded.
4865                        let _ = timeout(Duration::from_secs(2), async {
4866                            while sent.lock().expect("sent lock").len() == before {
4867                                sleep(Duration::from_millis(5)).await;
4868                            }
4869                        })
4870                        .await;
4871                    }
4872                    WaitForExpiry => sleep(Duration::from_millis(400)).await,
4873                }
4874            }
4875            let requests: Vec<(i32, String)> = sent
4876                .lock()
4877                .expect("sent lock")
4878                .iter()
4879                .map(|(node_id, request)| (*node_id, request.rack_id.clone()))
4880                .collect();
4881            drop(consumer);
4882            stop(brokers);
4883            actual.push((name, requests));
4884            wanted.push((name, expected));
4885        }
4886        assert2::assert!(actual == wanted);
4887    }
4888
4889    /// One Fetch request of a session case: `(session_id, session_epoch,
4890    /// fetched partitions, forgotten partitions)`.
4891    type SessionFields = (i32, i32, Vec<i32>, Vec<i32>);
4892
4893    fn session_fields(request: &FetchRequest) -> SessionFields {
4894        (
4895            request.session_id,
4896            request.session_epoch,
4897            request
4898                .topics
4899                .iter()
4900                .flat_map(|topic| topic.partitions.iter().map(|p| p.partition))
4901                .collect(),
4902            request
4903                .forgotten_topics_data
4904                .iter()
4905                .flat_map(|topic| topic.partitions.iter().copied())
4906                .collect(),
4907        )
4908    }
4909
4910    /// Kafka's `FetchSessionHandler`: the first Fetch is full and creates a
4911    /// session, the next one is incremental and names no unchanged partition.
4912    /// `FETCH_SESSION_ID_NOT_FOUND` starts a new session with a full Fetch.
4913    #[tokio::test]
4914    async fn fetch_sessions_send_incremental_requests_and_recover_a_lost_session() {
4915        let mut actual = Vec::new();
4916        let mut wanted = Vec::new();
4917        for (name, answers, expected) in [
4918            (
4919                "session",
4920                vec![
4921                    FetchAnswer::Respond {
4922                        error_code: 0,
4923                        session_id: 77,
4924                    },
4925                    FetchAnswer::Respond {
4926                        error_code: 0,
4927                        session_id: 77,
4928                    },
4929                ],
4930                vec![(0, 0, vec![0], vec![]), (77, 1, vec![], vec![])],
4931            ),
4932            (
4933                "session lost",
4934                vec![
4935                    FetchAnswer::Respond {
4936                        error_code: 0,
4937                        session_id: 77,
4938                    },
4939                    FetchAnswer::Respond {
4940                        error_code: 70,
4941                        session_id: 0,
4942                    },
4943                    FetchAnswer::Respond {
4944                        error_code: 0,
4945                        session_id: 78,
4946                    },
4947                ],
4948                vec![
4949                    (0, 0, vec![0], vec![]),
4950                    (77, 1, vec![], vec![]),
4951                    (0, 0, vec![0], vec![]),
4952                ],
4953            ),
4954        ] {
4955            let sent = SentFetches::default();
4956            let rounds = answers.len();
4957            let brokers = start_brokers(&[answers], &sent).await;
4958            let mut consumer = consumer_on(&brokers).await;
4959            for _ in 0..rounds {
4960                consumer.poll(secs(5)).await.expect("poll");
4961            }
4962            let requests: Vec<SessionFields> = sent
4963                .lock()
4964                .expect("sent lock")
4965                .iter()
4966                .map(|(_, request)| session_fields(request))
4967                .collect();
4968            drop(consumer);
4969            stop(brokers);
4970            actual.push((name, requests));
4971            wanted.push((name, expected));
4972        }
4973        assert2::assert!(actual == wanted);
4974    }
4975
4976    /// A Fetch result counts as soon as the task sent it, also while the task
4977    /// has not ended yet, so `poll` does not sleep until its deadline.
4978    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4979    async fn a_fetch_result_wakes_poll_before_the_task_ends() {
4980        let sent = SentFetches::default();
4981        let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
4982        let mut consumer = consumer_on(&brokers).await;
4983        let (tx, rx) = oneshot::channel();
4984        let completed = Arc::clone(&consumer.fetches.completed);
4985        consumer.fetches.in_flight.insert(
4986            -1,
4987            InFlightFetch {
4988                handle: tokio::spawn(async move {
4989                    let _ = tx.send(Ok(FetchResponse::default()));
4990                    completed.notify_one();
4991                    // The task ends late.
4992                    sleep(Duration::from_secs(3)).await;
4993                }),
4994                result: rx,
4995                requested: HashMap::new(),
4996                from_replica: HashSet::new(),
4997            },
4998        );
4999        let started = Instant::now();
5000        consumer
5001            .wait_for_fetches(started + Duration::from_secs(2))
5002            .await;
5003        let elapsed = started.elapsed();
5004        drop(consumer);
5005        stop(brokers);
5006        assert2::assert!(elapsed < Duration::from_millis(500));
5007    }
5008
5009    /// A preferred read replica that answers `OFFSET_OUT_OF_RANGE` gives no
5010    /// offset reset, also when the preference expired while the Fetch was in
5011    /// flight: the request decides.
5012    #[tokio::test]
5013    async fn out_of_range_from_a_replica_request_gives_no_reset() {
5014        let sent = SentFetches::default();
5015        let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
5016        let mut consumer = consumer_on(&brokers).await;
5017        consumer.auto_offset_reset = AutoOffsetReset::Latest;
5018        let key = ("orders".to_owned(), 0);
5019        let (tx, rx) = oneshot::channel();
5020        consumer.fetches.in_flight.insert(
5021            2,
5022            InFlightFetch {
5023                handle: tokio::spawn(async {}),
5024                result: rx,
5025                requested: maplit::hashmap! {key.clone() => 5},
5026                from_replica: maplit::hashset! {key.clone()},
5027            },
5028        );
5029        tx.send(Ok(FetchResponse {
5030            responses: vec![FetchableTopicResponse {
5031                topic: "orders".into(),
5032                partitions: vec![PartitionData {
5033                    partition_index: 0,
5034                    error_code: 1,
5035                    ..Default::default()
5036                }],
5037                ..Default::default()
5038            }],
5039            ..Default::default()
5040        }))
5041        .expect("send response");
5042        let topic_ids = consumer.topic_ids.lock().await.clone();
5043        let responses = consumer
5044            .take_completed_fetches(&topic_ids)
5045            .await
5046            .expect("take");
5047        consumer
5048            .process_fetch_responses(responses, &topic_ids)
5049            .await
5050            .expect("process");
5051        let position = consumer.next_offsets.lock().await.get(&key).copied();
5052        drop(consumer);
5053        stop(brokers);
5054        assert2::assert!(position == Some(5));
5055    }
5056
5057    /// Kafka's `FetchCollector` discards the data of a partition whose position
5058    /// changed after the Fetch went out, for example after a seek.
5059    #[tokio::test]
5060    async fn a_fetch_response_for_a_moved_position_gives_no_records() {
5061        let sent = SentFetches::default();
5062        let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
5063        let mut consumer = consumer_on(&brokers).await;
5064        let (tx, rx) = oneshot::channel();
5065        consumer.fetches.in_flight.insert(
5066            -1,
5067            InFlightFetch {
5068                handle: tokio::spawn(async {}),
5069                result: rx,
5070                requested: maplit::hashmap! {("orders".to_owned(), 0) => 5},
5071                from_replica: HashSet::new(),
5072            },
5073        );
5074        // A seek moves the position while the Fetch is in flight.
5075        consumer
5076            .next_offsets
5077            .lock()
5078            .await
5079            .insert(("orders".to_owned(), 0), 40);
5080        let record = RecordBatch {
5081            base_offset: 5,
5082            last_offset_delta: 0,
5083            records: vec![Record {
5084                offset_delta: 0,
5085                value: Some(bytes::Bytes::from_static(b"old")),
5086                ..Default::default()
5087            }],
5088            ..Default::default()
5089        };
5090        tx.send(Ok(FetchResponse {
5091            responses: vec![FetchableTopicResponse {
5092                topic: "orders".into(),
5093                partitions: vec![PartitionData {
5094                    partition_index: 0,
5095                    high_watermark: 6,
5096                    records: Some(RecordsPayload::V2(vec![record])),
5097                    ..Default::default()
5098                }],
5099                ..Default::default()
5100            }],
5101            ..Default::default()
5102        }))
5103        .expect("send response");
5104        consumer
5105            .wait_for_fetches(Instant::now() + Duration::from_secs(5))
5106            .await;
5107        let topic_ids = consumer.topic_ids.lock().await.clone();
5108        let responses = consumer
5109            .take_completed_fetches(&topic_ids)
5110            .await
5111            .expect("take");
5112        consumer
5113            .process_fetch_responses(responses, &topic_ids)
5114            .await
5115            .expect("process");
5116        let records = consumer.drain_fetch_buffer().await;
5117        let position = consumer
5118            .next_offsets
5119            .lock()
5120            .await
5121            .get(&("orders".to_owned(), 0))
5122            .copied();
5123        drop(consumer);
5124        stop(brokers);
5125        assert2::assert!((records.len(), position) == (0, Some(40)));
5126    }
5127
5128    /// A single mock broker that serves `ApiVersions`, `Metadata`,
5129    /// `ListOffsets` (always answering `list_offsets_offset` for `orders-0`)
5130    /// and `Fetch` (an empty, error-free response). It never names a
5131    /// separately dialable leader address, so every request the consumer
5132    /// sends for `orders-0` falls back to this bootstrap connection.
5133    ///
5134    /// Captures the timestamp of every `ListOffsets` request and the
5135    /// `fetch_offset` of every `Fetch` request it decodes.
5136    async fn no_committed_offset_broker(
5137        list_offsets_offset: i64,
5138        list_offsets_timestamps: Arc<Mutex<Vec<i64>>>,
5139        fetch_offsets: Arc<Mutex<Vec<i64>>>,
5140    ) -> MockBroker {
5141        use krabka_protocol::owned::{
5142            list_offsets_request::FLEXIBLE_MIN,
5143            list_offsets_response::{ListOffsetsPartitionResponse, ListOffsetsTopicResponse},
5144        };
5145        const NODE_ID: i32 = 1;
5146        MockBroker::start(move |api_key, version, _corr_id, mut body| {
5147            if api_key == api_versions_request::API_KEY {
5148                let versions = mock_test_support::api_versions_response_ranges([
5149                    (api_versions_request::API_KEY, 0, 3),
5150                    (metadata_request::API_KEY, 0, 8),
5151                    (
5152                        list_offsets_request::API_KEY,
5153                        0,
5154                        list_offsets_request::MAX_VERSION,
5155                    ),
5156                    (fetch_request::API_KEY, 4, 11),
5157                ]);
5158                return Some(encode(&versions, 0));
5159            }
5160            if api_key == metadata_request::API_KEY {
5161                // No `brokers` entry for `NODE_ID`: the pool cannot dial it,
5162                // so `orders-0` falls back to the bootstrap connection, which
5163                // is this same mock broker.
5164                let metadata = test_support::OrdersMetadata::builder()
5165                    .topic_id([9; 16])
5166                    .leader_id(NODE_ID)
5167                    .leader_epoch(1)
5168                    .build();
5169                return Some(encode(&metadata, version));
5170            }
5171            if api_key == list_offsets_request::API_KEY {
5172                let request = test_support::decode_list_offsets_request(&mut body, version);
5173                let timestamp = request.topics[0].partitions[0].timestamp;
5174                list_offsets_timestamps
5175                    .lock()
5176                    .expect("lock")
5177                    .push(timestamp);
5178                let answer = ListOffsetsResponse {
5179                    topics: vec![ListOffsetsTopicResponse {
5180                        name: "orders".into(),
5181                        partitions: vec![ListOffsetsPartitionResponse {
5182                            partition_index: 0,
5183                            error_code: 0,
5184                            offset: list_offsets_offset,
5185                            ..Default::default()
5186                        }],
5187                        ..Default::default()
5188                    }],
5189                    ..Default::default()
5190                };
5191                let mut bytes = Vec::new();
5192                if version >= FLEXIBLE_MIN {
5193                    // The flexible response header has an empty tagged-field
5194                    // section.
5195                    bytes.push(0);
5196                }
5197                bytes.extend(encode(&answer, version));
5198                return Some(bytes);
5199            }
5200            if api_key == fetch_request::API_KEY {
5201                let client_id_len = usize::try_from(body.get_i16()).expect("client id");
5202                body.advance(client_id_len);
5203                let request = FetchRequest::decode(&mut body, version).expect("decode fetch");
5204                if let Some(partition) = request
5205                    .topics
5206                    .first()
5207                    .and_then(|topic| topic.partitions.first())
5208                {
5209                    fetch_offsets
5210                        .lock()
5211                        .expect("lock")
5212                        .push(partition.fetch_offset);
5213                }
5214                return Some(encode(&FetchResponse::default(), version));
5215            }
5216            None
5217        })
5218        .await
5219        .unwrap()
5220    }
5221
5222    /// A consumer of `group-a`, assigned `orders-0` with no committed offset:
5223    /// the priming a fresh assignment gets when `OffsetFetch` names no commit
5224    /// (`starting_offset`/`reset_starting_offset` in `consumer.rs`).
5225    async fn consumer_with_no_committed_offset(
5226        broker: &MockBroker,
5227        policy: AutoOffsetReset,
5228    ) -> Consumer {
5229        let client = Client::builder()
5230            .bootstrap(broker.addr.to_string())
5231            .request_timeout(secs(30))
5232            .build()
5233            .await
5234            .expect("client");
5235        let mut consumer = poll::partition_error_tests::consumer_with_client(client);
5236        consumer.auto_offset_reset = policy;
5237        consumer.next_offsets.lock().await.insert(
5238            ("orders".to_owned(), 0),
5239            consumer::reset_starting_offset(policy),
5240        );
5241        consumer
5242    }
5243
5244    /// Kafka's initial position for a fresh assignment with no committed
5245    /// offset, by `auto.offset.reset` policy
5246    /// (<https://github.com/krabka-io/krabka-client-rs/issues/140>):
5247    ///
5248    /// - `none` never auto-resets: `poll` fails with
5249    ///   `NoOffsetForPartition` and sends no `Fetch`.
5250    /// - `earliest` asks `ListOffsets(-2)` before its first `Fetch`, which
5251    ///   then starts from the broker's real log start (40 here), never from
5252    ///   the literal offset 0.
5253    /// - `latest` asks `ListOffsets(-1)`, unchanged by this fix.
5254    #[tokio::test]
5255    async fn initial_position_follows_the_reset_policy_with_no_committed_offset() {
5256        struct Case {
5257            name: &'static str,
5258            policy: AutoOffsetReset,
5259            expected_list_offsets: Vec<i64>,
5260            expected_fetch_offsets: Vec<i64>,
5261            expect_error: bool,
5262        }
5263        let cases = [
5264            Case {
5265                name: "none has no committed offset and no reset policy: poll fails, no Fetch",
5266                policy: AutoOffsetReset::None,
5267                expected_list_offsets: vec![],
5268                expected_fetch_offsets: vec![],
5269                expect_error: true,
5270            },
5271            Case {
5272                name: "earliest asks ListOffsets(-2) and fetches from the log start",
5273                policy: AutoOffsetReset::Earliest,
5274                expected_list_offsets: vec![EARLIEST_TIMESTAMP],
5275                expected_fetch_offsets: vec![40],
5276                expect_error: false,
5277            },
5278            Case {
5279                name: "latest asks ListOffsets(-1)",
5280                policy: AutoOffsetReset::Latest,
5281                expected_list_offsets: vec![LATEST_TIMESTAMP],
5282                expected_fetch_offsets: vec![40],
5283                expect_error: false,
5284            },
5285        ];
5286        let mut actual = Vec::new();
5287        let mut wanted = Vec::new();
5288        for case in cases {
5289            let list_offsets_timestamps: Arc<Mutex<Vec<i64>>> = Arc::default();
5290            let fetch_offsets: Arc<Mutex<Vec<i64>>> = Arc::default();
5291            let broker = no_committed_offset_broker(
5292                40,
5293                Arc::clone(&list_offsets_timestamps),
5294                Arc::clone(&fetch_offsets),
5295            )
5296            .await;
5297            let mut consumer = consumer_with_no_committed_offset(&broker, case.policy).await;
5298            let result = consumer.poll(millis(500)).await;
5299            let list_offsets_timestamps = list_offsets_timestamps.lock().expect("lock").clone();
5300            let fetch_offsets = fetch_offsets.lock().expect("lock").clone();
5301            drop(consumer);
5302            broker.stop();
5303            actual.push((
5304                case.name,
5305                list_offsets_timestamps,
5306                fetch_offsets,
5307                result.is_err(),
5308            ));
5309            wanted.push((
5310                case.name,
5311                case.expected_list_offsets,
5312                case.expected_fetch_offsets,
5313                case.expect_error,
5314            ));
5315        }
5316        assert2::assert!(actual == wanted);
5317    }
5318}