1use 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
66const BOOTSTRAP_LEADER: i32 = -1;
72const UNKNOWN_FETCH_OFFSET: i64 = -1;
73const UNKNOWN_LEADER_ID: i32 = -1;
74const OFFSET_OUT_OF_RANGE: i16 = 1;
76const UNKNOWN_OFFSET: i64 = -1;
78pub const DEFAULT_FETCH_PARTITION_MAX: ByteSize = mebibytes(1);
79pub const DEFAULT_FETCH_MAX: ByteSize = mebibytes(50);
80
81type FetchSpec = (i32, i64, LeaderEpoch, LeaderEpoch);
84
85type 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
119enum FetchPartitionAction {
120 Records,
122 ResetOffset,
124 Reroute,
126 RevalidateEpoch,
129 RefreshMetadata,
132 Fail(i16),
134}
135
136fn classify_fetch_partition_error(error_code: i16) -> FetchPartitionAction {
149 match error_code {
150 0 => FetchPartitionAction::Records,
151 1 => FetchPartitionAction::ResetOffset,
152 6 => FetchPartitionAction::Reroute,
153 74 | 75 => {
154 FetchPartitionAction::RevalidateEpoch
155 }
156 3 | 9 | 56 | 78 | 100 | 103 => 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
182pub 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
196fn 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 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
213fn 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 rack_id: rack_id.to_owned(),
273 ..Default::default()
274 }
275}
276
277struct InFlightFetch {
280 handle: JoinHandle<()>,
281 result: oneshot::Receiver<Result<FetchResponse, krabka_client_core::ClientError>>,
284 requested: HashMap<(String, i32), i64>,
285 from_replica: HashSet<(String, i32)>,
287}
288
289struct 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#[derive(Default)]
305pub struct Fetches {
306 sessions: HashMap<i32, FetchSession>,
307 in_flight: HashMap<i32, InFlightFetch>,
308 ready: Vec<CompletedFetch>,
311 completed: Arc<Notify>,
312 preferred_read_replicas: HashMap<(String, i32), (i32, Instant)>,
316}
317
318impl Fetches {
319 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
351pub const LATEST_TIMESTAMP: i64 = -1;
353pub const EARLIEST_TIMESTAMP: i64 = -2;
355
356pub const LATEST_SENTINEL: i64 = i64::MAX;
359
360pub const BEGINNING_SENTINEL: i64 = i64::MAX - 1;
366
367pub const END_SENTINEL: i64 = i64::MAX - 2;
370
371pub const COMMITTED_SENTINEL: i64 = i64::MAX - 3;
375
376pub const NO_OFFSET_SENTINEL: i64 = i64::MAX - 4;
385
386pub fn is_reset_sentinel(next_offset: i64) -> bool {
392 next_offset >= NO_OFFSET_SENTINEL
393}
394
395fn 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
411async 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
427fn 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
436type ListOffsetsSpecs = BTreeMap<String, Vec<(i32, LeaderEpoch)>>;
439
440fn 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 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
479pub type ReadCommittedListOffsets =
490 VersionBoundedRequest<ListOffsetsRequest, 2, { list_offsets_request::MAX_VERSION }>;
491
492fn 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#[derive(Clone, Copy, Debug, Eq, PartialEq)]
520enum ListOffsetsRowAction {
521 Position(i64),
523 Unauthorized,
525 Retry,
528}
529
530fn 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 => ListOffsetsRowAction::Unauthorized,
540 _ => ListOffsetsRowAction::Retry,
541 }
542}
543
544#[derive(Clone, Debug, Default, Eq, PartialEq)]
547struct ListOffsetsResult {
548 offsets: HashMap<(String, i32), i64>,
550 retry: bool,
553 unauthorized: BTreeSet<String>,
555}
556
557impl ListOffsetsResult {
558 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 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 #[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 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 if self.subscription.borrow().is_none() && !self.subscription.borrow().manual_assignment {
644 return Err(ConsumerError::NotSubscribed);
645 }
646 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 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 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 () = 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 pub(crate) async fn wait_for_rebalance(
720 &mut self,
721 deadline: Instant,
722 ) -> Result<bool, ConsumerError> {
723 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 coordinator::note_poll(&self.poll_signal);
738 let wakeup = self.wakeup.clone();
739 let joined = loop {
740 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 let _ = joined;
753 break true;
754 }
755 () = sleep_until(deadline) => break false,
756 () = 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 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 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 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 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 let mut rebalance_pending = self.rebalance_pending.clone();
849 rebalance_pending.borrow_and_update();
850 let rebalance = async {
851 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 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 let id_to_name = offset_wire::id_to_name(topic_ids);
877
878 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 let mut out_of_range: Vec<((String, i32), i64)> = Vec::new();
894 let mut polled_resets: Vec<((String, i32), Option<i64>)> = Vec::new();
897 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 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 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 let action = classify_fetch_partition_error(part.error_code);
929 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 if part.preferred_read_replica >= 0 {
943 replica_updates.push((key.clone(), Some(part.preferred_read_replica)));
944 }
945 }
946 FetchPartitionAction::ResetOffset => {
947 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 let mut positions = self.positions.lock().await;
981 if part.current_leader.leader_id >= 0 {
982 let p = positions.entry(key.clone()).or_default();
985 p.leader_id = part.current_leader.leader_id;
986 p.leader_epoch = LeaderEpoch(part.current_leader.leader_epoch);
989 } else {
990 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 p.leader_epoch = LeaderEpoch(UNKNOWN_LEADER_ID);
1008 self.client.metadata_topics().request_update();
1009 if p.offset_epoch.is_known() {
1015 p.awaiting_validation = true;
1016 }
1017 }
1018 drop(positions);
1019 continue;
1020 }
1021 FetchPartitionAction::RefreshMetadata => {
1022 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(offsets);
1061 let now = Instant::now();
1062 for (key, replica) in replica_updates {
1063 match replica {
1064 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 self.client.metadata_topics().request_update();
1095 }
1096 Ok(())
1097 }
1098
1099 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 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 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 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 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 None => session.handle_error(),
1283 }
1284 }
1285 if let Some(error) = failure {
1286 return Err(error.into());
1287 }
1288 Ok(responses)
1289 }
1290
1291 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 if !subscription.manual_assignment
1354 && subscription.pattern.is_none()
1355 && subscription.regex.is_none()
1356 && !subscription.contains(t)
1357 {
1358 continue;
1359 }
1360 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 if is_reset_sentinel(next) {
1373 continue;
1374 }
1375 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 let mut leader =
1390 fetch_leader_id(pos.leader_id, self.client.knows_broker(pos.leader_id));
1391 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 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 self.client.metadata_topics().request_update();
1424 }
1425
1426 grouped
1427 }
1428
1429 async fn prepare_poll(&self, deadline: Instant) -> Result<bool, ConsumerError> {
1430 if let Some(error) = coordinator::take_poll_error(&self.poll_error) {
1433 return Err(error);
1434 }
1435 self.maybe_auto_commit_async().await;
1438 coordinator::note_poll(&self.poll_signal);
1442 self.sync_metadata_topics();
1448 let wakeup = self.wakeup.clone();
1449 self.update_fetch_positions(Some(&wakeup), Some(deadline))
1450 .await
1451 }
1452
1453 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 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 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 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 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 let batches = part.records.as_ref()?.as_v2()?;
1567 let fetch_floor = offsets.get(key).copied().unwrap_or(0);
1575 let mut records = VecDeque::new();
1576 let read_committed = is_read_committed(self.isolation_level);
1583 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 let mut aborted_pids: HashSet<i64> = HashSet::new();
1600 for batch in batches {
1601 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 if batch.attributes.is_control_batch() {
1617 if read_committed {
1618 aborted_pids.remove(&batch.producer_id);
1619 }
1620 continue;
1621 }
1622 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 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 let next_offset = next_offset_after(batches)?;
1664 Some(BufferedPartition {
1665 key: key.clone(),
1666 position: fetch_floor,
1667 records,
1668 next_offset,
1669 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
1704fn 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 #[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 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 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 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 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 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 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 #[cfg_attr(test, mutants::skip)] 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 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 #[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 OptionFuture::from(reset_polled).await;
1997 drop(offsets);
1998 Ok(result.retry)
1999 }
2000}
2001
2002fn 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
2017fn 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
2038fn 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 #[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 OptionFuture::from(reset_polled).await;
2087 drop(offsets);
2088 Ok(())
2089 }
2090
2091 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 #[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, 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 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 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 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 #[derive(Debug, PartialEq)]
2808 struct PollOutcome {
2809 result: Result<usize, Option<i16>>,
2812 metadata_requests: usize,
2814 metadata_update_requested: bool,
2816 next_offset: Option<i64>,
2817 }
2818
2819 #[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 #[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 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 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 #[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 #[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 #[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 #[derive(Clone, Copy, Debug)]
3074 enum BetweenPolls {
3075 Nothing,
3076 UpdateRequested,
3077 MaxAgeElapsed,
3078 LeaderUnknown,
3079 }
3080
3081 #[tokio::test]
3086 async fn fetch_positions_refresh_metadata_only_when_it_is_due() {
3087 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 #[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 #[derive(Clone, Copy, Debug)]
3227 enum PollReset {
3228 None,
3230 LatestOutOfRange,
3232 EarliestOutOfRange,
3236 FetchTruncation,
3238 ValidationTruncation,
3240 }
3241
3242 #[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 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 const LEADER_ID: i32 = 1;
3351 const LEADER_EPOCH: i32 = 7;
3353
3354 type SentListOffsets = Arc<Mutex<Vec<(&'static str, i16, ListOffsetsRequest)>>>;
3357
3358 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 bytes.push(0);
3382 }
3383 bytes.extend(encode(&answer, version));
3384 bytes
3385 }
3386
3387 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 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 #[derive(Debug, PartialEq)]
3460 enum PollFailure {
3461 TopicAuthorizationFailed(BTreeSet<String>),
3463 IncompatibleVersion(i16, i16, i16, i16, i16),
3466 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 #[derive(Debug, PartialEq)]
3492 struct ListOffsetsOutcome {
3493 requests: Vec<(&'static str, i16, ListOffsetsRequest)>,
3496 result: Result<(), PollFailure>,
3497 metadata_requests: usize,
3499 metadata_update_requested: bool,
3501 next_offset: Option<i64>,
3502 }
3503
3504 #[derive(Clone, Copy, Debug)]
3506 enum Reset {
3507 Latest,
3509 ByOneHour,
3512 EarliestOutOfRange,
3515 SeekToBeginning,
3517 SeekToEnd,
3519 }
3520
3521 #[derive(Clone, Copy, Debug)]
3523 struct Exchange {
3524 isolation: IsolationLevel,
3525 reset: Reset,
3526 row: (i16, i64),
3528 list_offsets_max_version: i16,
3530 primed: bool,
3534 }
3535
3536 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 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 #[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 #[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 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 #[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 #[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 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 #[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 ×tamp_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#[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 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 #[derive(Clone, Copy, Debug)]
4099 enum FetchAnswer {
4100 Silent,
4102 Respond { error_code: i16, session_id: i32 },
4105 Partition {
4107 error_code: i16,
4108 preferred_read_replica: i32,
4109 },
4110 }
4111
4112 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 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 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 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 #[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 #[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 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 #[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 #[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 #[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 #[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 #[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 #[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 #[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 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 #[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 #[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 #[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 #[derive(Clone, Copy, Debug)]
4773 enum ReplicaStep {
4774 Poll,
4775 WaitForExpiry,
4777 }
4778
4779 #[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 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 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 #[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 #[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 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 #[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 #[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 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 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 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 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 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 #[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}