1use std::{
5 collections::{BTreeMap, BTreeSet, HashMap},
6 sync::Arc,
7};
8
9use krabka_client_core::VersionBoundedRequest;
10#[cfg(test)]
11use krabka_client_core::mock_test_support;
12use krabka_ids::LeaderEpoch;
13#[cfg(test)]
14use krabka_protocol::owned::fetch_request;
15#[cfg(test)]
16use krabka_protocol::records::RecordsPayload;
17use krabka_protocol::{
18 ProtocolRequest,
19 owned::{
20 fetch_request::{FetchPartition, FetchRequest, FetchTopic, ForgottenTopic},
21 fetch_response,
22 fetch_response::FetchResponse,
23 list_offsets_request::{self, ListOffsetsPartition, ListOffsetsRequest, ListOffsetsTopic},
24 list_offsets_response::ListOffsetsResponse,
25 },
26 primitives::uuid,
27 records::{Record, RecordBatch},
28};
29use krabka_units::{
30 ByteSize, Time,
31 convert::{ByteSizeExt as _, TimeExt as _},
32 mebibytes,
33};
34
35#[cfg(test)]
36use crate::poll;
37#[cfg(test)]
38use crate::test_support;
39use crate::{
40 assignor,
41 builder::{AutoOffsetReset, IsolationLevel},
42 consumer,
43 consumer::{Consumer, ConsumerRecord, Header, TimestampType},
44 control, coordinator,
45 error::ConsumerError,
46 fetch_buffer::BufferedPartition,
47 fetch_session::{
48 FINAL_EPOCH, FetchSession, INVALID_SESSION_ID, SessionPartition, SessionRequest,
49 },
50 offset_wire,
51 position::PartitionPosition,
52};
53
54const BOOTSTRAP_LEADER: i32 = -1;
60const UNKNOWN_FETCH_OFFSET: i64 = -1;
61const UNKNOWN_LEADER_ID: i32 = -1;
62const OFFSET_OUT_OF_RANGE: i16 = 1;
64const UNKNOWN_OFFSET: i64 = -1;
66pub(crate) const DEFAULT_FETCH_PARTITION_MAX: ByteSize = mebibytes(1);
67pub(crate) const DEFAULT_FETCH_MAX: ByteSize = mebibytes(50);
68
69type FetchSpec = (i32, i64, LeaderEpoch, LeaderEpoch);
72
73type FetchByLeader = HashMap<i32, HashMap<String, Vec<FetchSpec>>>;
75
76fn fetch_leader_id(leader_id: i32, knows_leader: bool) -> i32 {
77 if leader_id >= 0 && knows_leader {
78 leader_id
79 } else {
80 BOOTSTRAP_LEADER
81 }
82}
83
84fn should_use_bootstrap_leader(leader: i32) -> bool {
85 leader == BOOTSTRAP_LEADER
86}
87
88fn fetch_offset_or_unknown(offsets: &HashMap<(String, i32), i64>, key: &(String, i32)) -> i64 {
89 offsets.get(key).copied().unwrap_or(UNKNOWN_FETCH_OFFSET)
90}
91
92fn is_read_committed(isolation_level: IsolationLevel) -> bool {
93 isolation_level == IsolationLevel::ReadCommitted
94}
95
96pub(crate) fn is_transient_transport_error(e: &krabka_client_core::ClientError) -> bool {
97 crate::error::is_transport_error(e, true)
98}
99
100fn is_transient_poll_error(e: &ConsumerError) -> bool {
101 matches!(e, ConsumerError::Client(client) if is_transient_transport_error(client))
102}
103
104#[derive(Clone, Copy, Debug, Eq, PartialEq)]
107enum FetchPartitionAction {
108 Records,
110 ResetOffset,
112 Reroute,
114 RevalidateEpoch,
117 RefreshMetadata,
120 Fail(i16),
122}
123
124fn classify_fetch_partition_error(error_code: i16) -> FetchPartitionAction {
137 match error_code {
138 0 => FetchPartitionAction::Records,
139 1 => FetchPartitionAction::ResetOffset,
140 6 => FetchPartitionAction::Reroute,
141 74 | 75 => {
142 FetchPartitionAction::RevalidateEpoch
143 }
144 3 | 9 | 56 | 78 | 100 | 103 => FetchPartitionAction::RefreshMetadata,
150 other => FetchPartitionAction::Fail(other),
151 }
152}
153
154fn aborted_txn_started(first_offset: i64, batch_base_offset: i64) -> bool {
155 first_offset <= batch_base_offset
156}
157
158fn should_drop_aborted_batch(
159 read_committed: bool,
160 is_transactional: bool,
161 producer_is_aborted: bool,
162) -> bool {
163 read_committed && is_transactional && producer_is_aborted
164}
165
166fn record_offset(base_offset: i64, offset_delta: i32) -> i64 {
167 base_offset + i64::from(offset_delta)
168}
169
170pub(crate) fn record_timestamp(batch: &RecordBatch, record: &Record) -> (i64, TimestampType) {
176 let timestamp_type = TimestampType::from(batch.attributes.timestamp_type());
177 let timestamp = match timestamp_type {
178 TimestampType::CreateTime => batch.base_timestamp + record.timestamp_delta,
179 TimestampType::LogAppendTime => batch.max_timestamp,
180 };
181 (timestamp, timestamp_type)
182}
183
184fn session_partition(
186 topic_id: uuid::Uuid,
187 (_, fetch_offset, leader_epoch, last_fetched_epoch): FetchSpec,
188 partition_max: ByteSize,
189) -> SessionPartition {
190 SessionPartition {
191 topic_id,
192 fetch_offset,
193 current_leader_epoch: leader_epoch.get(),
196 last_fetched_epoch: last_fetched_epoch.get(),
197 partition_max_bytes: partition_max.bytes_i32(),
198 }
199}
200
201fn build_fetch_request(
207 max_wait_ms: i32,
208 isolation_level: IsolationLevel,
209 min: ByteSize,
210 max: ByteSize,
211 rack_id: &str,
212 session: SessionRequest,
213) -> FetchRequest {
214 let mut topics: Vec<FetchTopic> = Vec::new();
215 for ((topic, partition), data) in session.partitions {
216 if topics.last().is_none_or(|last| last.topic != topic) {
217 topics.push(FetchTopic {
218 topic: topic.clone(),
219 topic_id: data.topic_id,
220 ..Default::default()
221 });
222 }
223 if let Some(last) = topics.last_mut() {
224 last.partitions.push(FetchPartition {
225 partition,
226 fetch_offset: data.fetch_offset,
227 current_leader_epoch: data.current_leader_epoch,
228 last_fetched_epoch: data.last_fetched_epoch,
229 partition_max_bytes: data.partition_max_bytes,
230 ..Default::default()
231 });
232 }
233 }
234 let mut forgotten_topics_data: Vec<ForgottenTopic> = Vec::new();
235 for ((topic, partition), topic_id) in session.forgotten {
236 if forgotten_topics_data
237 .last()
238 .is_none_or(|last| last.topic != topic)
239 {
240 forgotten_topics_data.push(ForgottenTopic {
241 topic: topic.clone(),
242 topic_id,
243 ..Default::default()
244 });
245 }
246 if let Some(last) = forgotten_topics_data.last_mut() {
247 last.partitions.push(partition);
248 }
249 }
250 FetchRequest {
251 max_wait_ms,
252 min_bytes: min.bytes_i32(),
253 max_bytes: max.bytes_i32(),
254 isolation_level: isolation_level.wire(),
255 session_id: session.session_id,
256 session_epoch: session.session_epoch,
257 topics,
258 forgotten_topics_data,
259 rack_id: rack_id.to_owned(),
261 ..Default::default()
262 }
263}
264
265struct InFlightFetch {
268 handle: tokio::task::JoinHandle<()>,
269 result: tokio::sync::oneshot::Receiver<Result<FetchResponse, krabka_client_core::ClientError>>,
272 requested: HashMap<(String, i32), i64>,
273 from_replica: std::collections::HashSet<(String, i32)>,
275}
276
277struct CompletedFetch {
280 leader: i32,
281 requested: HashMap<(String, i32), i64>,
282 from_replica: std::collections::HashSet<(String, i32)>,
283 result: Option<Result<FetchResponse, krabka_client_core::ClientError>>,
284}
285
286#[derive(Default)]
293pub(crate) struct Fetches {
294 sessions: HashMap<i32, FetchSession>,
295 in_flight: HashMap<i32, InFlightFetch>,
296 ready: Vec<CompletedFetch>,
299 completed: Arc<tokio::sync::Notify>,
300 preferred_read_replicas: HashMap<(String, i32), (i32, tokio::time::Instant)>,
304}
305
306impl Fetches {
307 fn collect_ready(&mut self) {
309 let leaders: Vec<i32> = self.in_flight.keys().copied().collect();
310 for leader in leaders {
311 let Some(fetch) = self.in_flight.get_mut(&leader) else {
312 continue;
313 };
314 let result = match fetch.result.try_recv() {
315 Ok(result) => Some(result),
316 Err(tokio::sync::oneshot::error::TryRecvError::Closed) => None,
317 Err(tokio::sync::oneshot::error::TryRecvError::Empty) => continue,
318 };
319 if let Some(fetch) = self.in_flight.remove(&leader) {
320 self.ready.push(CompletedFetch {
321 leader,
322 requested: fetch.requested,
323 from_replica: fetch.from_replica,
324 result,
325 });
326 }
327 }
328 }
329}
330
331impl Drop for Fetches {
332 fn drop(&mut self) {
333 for fetch in self.in_flight.values() {
334 fetch.handle.abort();
335 }
336 }
337}
338
339pub(crate) const LATEST_TIMESTAMP: i64 = -1;
341pub(crate) const EARLIEST_TIMESTAMP: i64 = -2;
343
344pub(crate) const LATEST_SENTINEL: i64 = i64::MAX;
347
348pub(crate) const BEGINNING_SENTINEL: i64 = i64::MAX - 1;
354
355pub(crate) const END_SENTINEL: i64 = i64::MAX - 2;
358
359pub(crate) const COMMITTED_SENTINEL: i64 = i64::MAX - 3;
363
364pub(crate) const NO_OFFSET_SENTINEL: i64 = i64::MAX - 4;
373
374pub(crate) fn is_reset_sentinel(next_offset: i64) -> bool {
380 next_offset >= NO_OFFSET_SENTINEL
381}
382
383fn reset_timestamp(policy: AutoOffsetReset, now_ms: i64) -> i64 {
389 match policy {
390 AutoOffsetReset::ByDuration(duration) => {
391 now_ms.saturating_sub(i64::try_from(duration.as_millis()).unwrap_or(i64::MAX))
392 }
393 AutoOffsetReset::Earliest | AutoOffsetReset::Latest | AutoOffsetReset::None => {
394 LATEST_TIMESTAMP
395 }
396 }
397}
398
399async fn until_woken<F: std::future::Future>(
402 wakeup: Option<&control::WakeupHandle>,
403 future: F,
404) -> Result<F::Output, ConsumerError> {
405 match wakeup {
406 None => Ok(future.await),
407 Some(wakeup) => tokio::select! {
408 biased;
409 () = wakeup.woken() => Err(ConsumerError::Wakeup),
410 output = future => Ok(output),
411 },
412 }
413}
414
415fn unix_now_ms() -> i64 {
417 std::time::SystemTime::now()
418 .duration_since(std::time::UNIX_EPOCH)
419 .map_or(0, |elapsed| {
420 i64::try_from(elapsed.as_millis()).unwrap_or(i64::MAX)
421 })
422}
423
424type ListOffsetsSpecs = BTreeMap<String, Vec<(i32, LeaderEpoch)>>;
427
428fn build_offsets_request(
438 by_topic: ListOffsetsSpecs,
439 timestamp: i64,
440 isolation_level: IsolationLevel,
441) -> ListOffsetsRequest {
442 let topics: Vec<ListOffsetsTopic> = by_topic
443 .into_iter()
444 .map(|(name, partitions)| ListOffsetsTopic {
445 name,
446 partitions: partitions
447 .into_iter()
448 .map(|(partition_index, leader_epoch)| ListOffsetsPartition {
449 partition_index,
450 current_leader_epoch: leader_epoch.get(),
452 timestamp,
453 ..Default::default()
454 })
455 .collect(),
456 ..Default::default()
457 })
458 .collect();
459 ListOffsetsRequest {
460 replica_id: -1,
461 isolation_level: isolation_level.wire(),
462 topics,
463 ..Default::default()
464 }
465}
466
467pub(crate) type ReadCommittedListOffsets =
478 VersionBoundedRequest<ListOffsetsRequest, 2, { list_offsets_request::MAX_VERSION }>;
479
480fn group_list_offsets(
488 keys: &[(String, i32)],
489 positions: &HashMap<(String, i32), PartitionPosition>,
490 knows_broker: impl Fn(i32) -> bool,
491) -> BTreeMap<i32, ListOffsetsSpecs> {
492 let mut grouped: BTreeMap<i32, ListOffsetsSpecs> = BTreeMap::new();
493 for key in keys {
494 let position = positions.get(key).copied().unwrap_or_default();
495 let leader = fetch_leader_id(position.leader_id, knows_broker(position.leader_id));
496 grouped
497 .entry(leader)
498 .or_default()
499 .entry(key.0.clone())
500 .or_default()
501 .push((key.1, position.leader_epoch));
502 }
503 grouped
504}
505
506#[derive(Clone, Copy, Debug, Eq, PartialEq)]
508enum ListOffsetsRowAction {
509 Position(i64),
511 Unauthorized,
513 Retry,
516}
517
518fn classify_list_offsets_row(error_code: i16, offset: i64) -> ListOffsetsRowAction {
525 match error_code {
526 0 if offset != UNKNOWN_OFFSET => ListOffsetsRowAction::Position(offset),
527 29 => ListOffsetsRowAction::Unauthorized,
528 _ => ListOffsetsRowAction::Retry,
529 }
530}
531
532#[derive(Clone, Debug, Default, Eq, PartialEq)]
535struct ListOffsetsResult {
536 offsets: HashMap<(String, i32), i64>,
538 retry: bool,
541 unauthorized: BTreeSet<String>,
543}
544
545impl ListOffsetsResult {
546 fn collect(requested: &[(String, i32)], answers: &[ListOffsetsResponse]) -> Self {
552 let mut result = Self::default();
553 for answer in answers {
554 let mut offsets = HashMap::new();
555 let mut unauthorized = BTreeSet::new();
556 for topic in &answer.topics {
557 for row in &topic.partitions {
558 match classify_list_offsets_row(row.error_code, row.offset) {
559 ListOffsetsRowAction::Position(offset) => {
560 offsets.insert((topic.name.clone(), row.partition_index), offset);
561 }
562 ListOffsetsRowAction::Unauthorized => {
563 unauthorized.insert(topic.name.clone());
564 }
565 ListOffsetsRowAction::Retry => {}
566 }
567 }
568 }
569 if unauthorized.is_empty() {
570 result.offsets.extend(offsets);
571 } else {
572 result.unauthorized.extend(unauthorized);
573 }
574 }
575 result.offsets.retain(|key, _| requested.contains(key));
576 result.retry = requested
577 .iter()
578 .any(|key| !result.offsets.contains_key(key));
579 result
580 }
581
582 fn authorization(&self) -> Result<(), ConsumerError> {
584 if self.unauthorized.is_empty() {
585 Ok(())
586 } else {
587 Err(ConsumerError::TopicAuthorizationFailed(
588 self.unauthorized.clone(),
589 ))
590 }
591 }
592}
593
594impl Consumer {
595 #[tracing::instrument(
612 name = "consumer.poll",
613 level = "debug",
614 skip_all,
615 fields(
616 group_id = %self.group_id,
617 timeout_ms = timeout.millis_i64_trunc(),
618 assigned_partitions = tracing::field::Empty,
619 leaders = tracing::field::Empty,
620 records = tracing::field::Empty,
621 ),
622 err(level = "debug")
623 )]
624 pub async fn poll(&mut self, timeout: Time) -> Result<Vec<ConsumerRecord>, ConsumerError> {
627 let deadline = tokio::time::Instant::now()
628 + std::time::Duration::from_millis(
629 u64::try_from(timeout.millis_i64_trunc()).unwrap_or(0),
630 );
631 if self.subscription.borrow().is_none() && !self.subscription.borrow().manual_assignment {
634 return Err(ConsumerError::NotSubscribed);
635 }
636 if self.wakeup.take() {
639 return Err(ConsumerError::Wakeup);
640 }
641 if !self.prepare_poll(deadline).await? {
642 return Ok(Vec::new());
643 }
644 if !self.wait_for_rebalance(deadline).await? {
645 return Ok(Vec::new());
646 }
647 let wakeup = self.wakeup.clone();
648
649 let buffered = self.drain_fetch_buffer().await;
651 if !buffered.is_empty() {
652 tracing::Span::current().record("records", buffered.len());
653 return Ok(buffered);
654 }
655
656 let assigned = self.assigned.lock().await.clone();
658 tracing::Span::current().record("assigned_partitions", assigned.len());
659 if assigned.is_empty() {
660 tokio::select! {
661 () = tokio::time::sleep_until(deadline) => return Ok(Vec::new()),
662 () = wakeup.woken() => return Err(ConsumerError::Wakeup),
663 }
664 }
665
666 let by_leader = self.group_fetches(&assigned).await;
667 tracing::Span::current().record("leaders", by_leader.len());
668 let topic_ids = self.topic_ids.lock().await.clone();
669 self.send_fetches(by_leader, &topic_ids);
670 self.fetches.collect_ready();
671 if self.fetches.in_flight.is_empty() && self.fetches.ready.is_empty() {
672 tokio::select! {
673 () = self.wait_without_fetches(&assigned, deadline) => return Ok(Vec::new()),
674 () = wakeup.woken() => return Err(ConsumerError::Wakeup),
675 }
676 }
677 tokio::select! {
678 biased;
679 () = self.wait_for_fetches(deadline) => {}
680 () = wakeup.woken() => return Err(ConsumerError::Wakeup),
683 }
684 let responses = self.take_completed_fetches(&topic_ids).await?;
685
686 self.process_fetch_responses(responses, &topic_ids).await?;
687 let records = self.drain_fetch_buffer().await;
688 tracing::Span::current().record("records", records.len());
689 Ok(records)
690 }
691
692 pub(crate) async fn wait_for_rebalance(
710 &mut self,
711 deadline: tokio::time::Instant,
712 ) -> Result<bool, ConsumerError> {
713 let mut first_error = self.run_pending_listener_calls().await.err();
717 if !*self.rebalance_pending.borrow() {
718 return first_error.map_or(Ok(true), Err);
719 }
720 let eager = self.rebalance_protocol == assignor::RebalanceProtocol::Eager;
721 if first_error.is_none() && !eager && !self.assigned.lock().await.is_empty() {
722 return Ok(true);
723 }
724 coordinator::note_poll(&self.poll_signal);
728 let wakeup = self.wakeup.clone();
729 let joined = loop {
730 if tokio::time::Instant::now() >= deadline {
733 break false;
734 }
735 let call = tokio::select! {
736 biased;
737 call = self.listener_calls.recv() => call,
738 joined = self.rebalance_pending.wait_for(|pending| !*pending) => {
739 let _ = joined;
743 break true;
744 }
745 () = tokio::time::sleep_until(deadline) => break false,
746 () = wakeup.woken() => {
749 first_error.get_or_insert(ConsumerError::Wakeup);
750 break false;
751 }
752 };
753 match call {
754 Some(call) => {
755 if let Err(error) = self.complete_listener_call(call).await {
756 first_error.get_or_insert(error);
757 }
758 }
759 None => {
761 break tokio::select! {
762 biased;
763 joined = self.rebalance_pending.wait_for(|pending| !*pending) => {
764 let _ = joined;
765 true
766 }
767 () = tokio::time::sleep_until(deadline) => false,
768 () = wakeup.woken() => {
769 first_error.get_or_insert(ConsumerError::Wakeup);
770 false
771 }
772 };
773 }
774 }
775 };
776 if let Err(error) = self.run_pending_listener_calls().await {
779 first_error.get_or_insert(error);
780 }
781 first_error.map_or(Ok(joined), Err)
782 }
783
784 async fn drain_fetch_buffer(&mut self) -> Vec<ConsumerRecord> {
792 let ownership_ids = self.commit_identity.lock().await.ownership_ids.clone();
793 let paused = self.paused_of(&ownership_ids);
794 let assigned_guard = self.assigned.lock().await;
795 let assigned: std::collections::HashSet<(String, i32)> =
796 assigned_guard.iter().cloned().collect();
797 let mut offsets = self.next_offsets.lock().await;
798 let mut positions = self.positions.lock().await;
799 let records = self.fetch_buffer.drain(
800 self.max_poll_records,
801 &assigned,
802 &paused,
803 &mut offsets,
804 &mut positions,
805 );
806 drop(positions);
807 drop(offsets);
808 drop(assigned_guard);
809 records
810 }
811 async fn wait_without_fetches(
819 &self,
820 assigned: &[(String, i32)],
821 deadline: tokio::time::Instant,
822 ) {
823 let all_positions = {
824 let offsets = self.next_offsets.lock().await;
825 let positions = self.positions.lock().await;
826 assigned.iter().all(|key| {
827 offsets
828 .get(key)
829 .is_some_and(|offset| !is_reset_sentinel(*offset))
830 && !positions
831 .get(key)
832 .is_some_and(|position| position.awaiting_validation)
833 })
834 };
835 let until = if all_positions && self.listener_calls.is_empty() {
836 deadline
837 } else {
838 deadline.min(tokio::time::Instant::now() + self.retry_policy.initial_backoff)
839 };
840 let mut rebalance_pending = self.rebalance_pending.clone();
844 rebalance_pending.borrow_and_update();
845 let rebalance = async {
846 if rebalance_pending.changed().await.is_err() {
848 std::future::pending::<()>().await;
849 }
850 };
851 tokio::select! {
852 () = tokio::time::sleep_until(until) => {}
853 () = rebalance => {}
854 }
855 }
856
857 async fn process_fetch_responses(
860 &mut self,
861 responses: Vec<fetch_response::FetchResponse>,
862 topic_ids: &HashMap<String, uuid::Uuid>,
863 ) -> Result<(), ConsumerError> {
864 let id_to_name = offset_wire::id_to_name(topic_ids);
872
873 let still_owned: std::collections::HashSet<(String, i32)> =
880 self.assigned.lock().await.iter().cloned().collect();
881
882 let mut fetched: Vec<BufferedPartition> = Vec::new();
883 let mut refresh_after_processing = false;
884 let mut out_of_range: Vec<((String, i32), i64)> = Vec::new();
889 let mut polled_resets: Vec<((String, i32), Option<i64>)> = Vec::new();
892 let mut replica_updates: Vec<((String, i32), Option<i32>)> = Vec::new();
895 let mut offsets = self.next_offsets.lock().await;
896 for topic in responses.iter().flat_map(|resp| &resp.responses) {
897 let topic_name = if topic.topic.is_empty() {
898 id_to_name.get(&topic.topic_id).cloned().unwrap_or_default()
899 } else {
900 topic.topic.clone()
901 };
902 for part in &topic.partitions {
903 if !still_owned.contains(&(topic_name.clone(), part.partition_index)) {
906 continue;
907 }
908
909 let key = (topic_name.clone(), part.partition_index);
910
911 if part.diverging_epoch.end_offset >= 0 {
914 self.handle_truncation_in_poll(
915 &mut offsets,
916 &key,
917 part.diverging_epoch.end_offset,
918 )?;
919 polled_resets.push((key.clone(), Some(part.diverging_epoch.end_offset)));
920 continue;
921 }
922 let action = classify_fetch_partition_error(part.error_code);
924 if matches!(
927 action,
928 FetchPartitionAction::Reroute
929 | FetchPartitionAction::RevalidateEpoch
930 | FetchPartitionAction::RefreshMetadata
931 ) {
932 replica_updates.push((key.clone(), None));
933 }
934 match action {
935 FetchPartitionAction::Records => {
936 if part.preferred_read_replica >= 0 {
938 replica_updates.push((key.clone(), Some(part.preferred_read_replica)));
939 }
940 }
941 FetchPartitionAction::ResetOffset => {
942 match self.auto_offset_reset {
958 AutoOffsetReset::Latest | AutoOffsetReset::ByDuration(_) => {
959 offsets.insert(key.clone(), LATEST_SENTINEL);
960 polled_resets.push((key.clone(), None));
961 }
962 AutoOffsetReset::Earliest | AutoOffsetReset::None => {
963 let fetch_offset = fetch_offset_or_unknown(&offsets, &key);
964 out_of_range.push((key.clone(), fetch_offset));
965 }
966 }
967 continue;
968 }
969 FetchPartitionAction::Reroute => {
970 let mut positions = self.positions.lock().await;
976 if part.current_leader.leader_id >= 0 {
977 let p = positions.entry(key.clone()).or_default();
980 p.leader_id = part.current_leader.leader_id;
981 p.leader_epoch = LeaderEpoch(part.current_leader.leader_epoch);
984 } else {
985 if let Some(p) = positions.get_mut(&key) {
990 p.leader_id = UNKNOWN_LEADER_ID;
991 }
992 drop(positions);
993 refresh_after_processing = true;
994 }
995 continue;
996 }
997 FetchPartitionAction::RevalidateEpoch => {
998 let mut positions = self.positions.lock().await;
999 if let Some(p) = positions.get_mut(&key) {
1000 p.leader_epoch = LeaderEpoch(UNKNOWN_LEADER_ID);
1003 self.client.metadata_topics().request_update();
1004 if p.offset_epoch.is_known() {
1010 p.awaiting_validation = true;
1011 }
1012 }
1013 continue;
1014 }
1015 FetchPartitionAction::RefreshMetadata => {
1016 tracing::warn!(
1021 topic = %topic_name,
1022 partition = part.partition_index,
1023 error_code = part.error_code,
1024 "fetch partition names a topic the broker does not know; refreshing metadata"
1025 );
1026 refresh_after_processing = true;
1027 continue;
1028 }
1029 FetchPartitionAction::Fail(code) => {
1030 return Err(ConsumerError::Server(code));
1031 }
1032 }
1033
1034 record_readable_end(
1035 &mut *self.end_offsets.lock().await,
1036 key.clone(),
1037 self.isolation_level,
1038 part.high_watermark,
1039 part.last_stable_offset,
1040 );
1041
1042 if let Some(partition) =
1043 self.process_partition_records(&offsets, &key, &topic_name, part)
1044 {
1045 fetched.push(partition);
1046 }
1047 }
1048 }
1049 if let Some(auto_commit) = &self.auto_commit {
1050 auto_commit.reset_polled(polled_resets).await;
1051 }
1052 drop(offsets);
1055 let now = tokio::time::Instant::now();
1056 for (key, replica) in replica_updates {
1057 match replica {
1058 Some(replica) => {
1061 let expires = now + self.metadata_max_age.to_std();
1062 let entry = self
1063 .fetches
1064 .preferred_read_replicas
1065 .entry(key)
1066 .or_insert((replica, expires));
1067 if entry.0 != replica {
1068 *entry = (replica, expires);
1069 }
1070 }
1071 None => {
1072 self.fetches.preferred_read_replicas.remove(&key);
1073 }
1074 }
1075 }
1076 for partition in fetched {
1077 self.fetch_buffer.push(partition);
1078 }
1079 if !out_of_range.is_empty() && self.recover_out_of_range(&out_of_range).await? {
1080 refresh_after_processing = true;
1081 }
1082 if refresh_after_processing {
1083 self.client.metadata_topics().request_update();
1089 }
1090 Ok(())
1091 }
1092
1093 fn send_fetches(&mut self, by_leader: FetchByLeader, topic_ids: &HashMap<String, uuid::Uuid>) {
1096 let max_wait_ms = consumer::protocol_millis_i32(self.fetch_max_wait);
1097 for (leader, by_topic) in by_leader {
1098 if self.fetches.in_flight.contains_key(&leader) {
1099 continue;
1100 }
1101 let wanted: BTreeMap<(String, i32), SessionPartition> = by_topic
1102 .into_iter()
1103 .flat_map(|(topic, specs)| {
1104 let topic_id = topic_ids.get(&topic).copied().unwrap_or_default();
1105 let partition_max = self.fetch_partition_max;
1106 specs.into_iter().map(move |spec| {
1107 (
1108 (topic.clone(), spec.0),
1109 session_partition(topic_id, spec, partition_max),
1110 )
1111 })
1112 })
1113 .collect();
1114 let requested = wanted
1115 .iter()
1116 .map(|(key, partition)| (key.clone(), partition.fetch_offset))
1117 .collect();
1118 let from_replica = wanted
1119 .keys()
1120 .filter(|key| {
1121 self.fetches
1122 .preferred_read_replicas
1123 .get(*key)
1124 .is_some_and(|(replica, _)| *replica == leader)
1125 })
1126 .cloned()
1127 .collect();
1128 let session = self
1129 .fetches
1130 .sessions
1131 .entry(leader)
1132 .or_default()
1133 .build(wanted);
1134 let request = build_fetch_request(
1135 max_wait_ms,
1136 self.isolation_level,
1137 self.fetch_min,
1138 self.fetch_max,
1139 self.client_rack.as_deref().unwrap_or_default(),
1140 session,
1141 );
1142 let client = self.client.clone();
1143 let completed = Arc::clone(&self.fetches.completed);
1144 let (sender, result) = tokio::sync::oneshot::channel();
1145 let handle = tokio::spawn(async move {
1146 let response = if should_use_bootstrap_leader(leader) {
1147 client.send(request).await
1148 } else {
1149 client.broker(leader).send(request).await
1150 };
1151 let _ = sender.send(response);
1152 completed.notify_one();
1153 });
1154 self.fetches.in_flight.insert(
1155 leader,
1156 InFlightFetch {
1157 handle,
1158 result,
1159 requested,
1160 from_replica,
1161 },
1162 );
1163 }
1164 }
1165
1166 async fn wait_for_fetches(&mut self, deadline: tokio::time::Instant) {
1170 loop {
1171 self.fetches.collect_ready();
1172 if self.fetches.in_flight.is_empty() || !self.fetches.ready.is_empty() {
1173 return;
1174 }
1175 tokio::select! {
1176 () = self.fetches.completed.notified() => {}
1177 () = tokio::time::sleep_until(deadline) => return,
1178 }
1179 }
1180 }
1181
1182 async fn take_completed_fetches(
1194 &mut self,
1195 topic_ids: &HashMap<String, uuid::Uuid>,
1196 ) -> Result<Vec<FetchResponse>, ConsumerError> {
1197 self.fetches.collect_ready();
1198 if self.fetches.ready.is_empty() {
1199 return Ok(Vec::new());
1200 }
1201 let id_to_name = offset_wire::id_to_name(topic_ids);
1202 let offsets = self.next_offsets.lock().await.clone();
1203 let mut responses = Vec::new();
1204 let mut failure = None;
1205 for fetch in std::mem::take(&mut self.fetches.ready) {
1206 let leader = fetch.leader;
1207 let session = self.fetches.sessions.entry(leader).or_default();
1208 match fetch.result {
1209 Some(Ok(mut response)) => {
1210 let partitions = response
1211 .responses
1212 .iter()
1213 .map(|topic| topic.partitions.len())
1214 .sum();
1215 if !session.handle_response(
1216 response.error_code,
1217 response.session_id,
1218 partitions,
1219 response.throttle_time_ms,
1220 ) {
1221 tracing::debug!(
1222 leader,
1223 error_code = response.error_code,
1224 "fetch session answered with an error; the next fetch is a full fetch"
1225 );
1226 continue;
1227 }
1228 let mut replica_out_of_range = Vec::new();
1229 for topic in &mut response.responses {
1230 let name = if topic.topic.is_empty() {
1231 id_to_name.get(&topic.topic_id).cloned().unwrap_or_default()
1232 } else {
1233 topic.topic.clone()
1234 };
1235 topic.partitions.retain(|partition| {
1236 let key = (name.clone(), partition.partition_index);
1237 if partition.error_code == OFFSET_OUT_OF_RANGE
1243 && fetch.from_replica.contains(&key)
1244 {
1245 replica_out_of_range.push(key);
1246 return false;
1247 }
1248 fetch.requested.get(&key).is_some_and(|requested| {
1249 offsets.get(&key).copied().unwrap_or(0) == *requested
1250 })
1251 });
1252 }
1253 for key in replica_out_of_range {
1254 self.fetches.preferred_read_replicas.remove(&key);
1255 }
1256 responses.push(response);
1257 }
1258 Some(Err(error)) => {
1259 session.handle_error();
1260 for key in fetch.requested.keys() {
1263 self.fetches.preferred_read_replicas.remove(key);
1264 }
1265 if is_transient_transport_error(&error) {
1266 if should_use_bootstrap_leader(leader) {
1267 self.client.reconnect_bootstrap().await;
1268 } else {
1269 self.client.evict_broker(leader);
1270 }
1271 } else if failure.is_none() {
1272 failure = Some(error);
1273 }
1274 }
1275 None => session.handle_error(),
1277 }
1278 }
1279 match failure {
1280 Some(error) => Err(error.into()),
1281 None => Ok(responses),
1282 }
1283 }
1284
1285 pub(crate) async fn close_fetch_sessions(&mut self, timeout: std::time::Duration) {
1289 let max_wait_ms = consumer::protocol_millis_i32(self.fetch_max_wait);
1290 let closes: Vec<_> = self
1291 .fetches
1292 .sessions
1293 .iter()
1294 .filter(|(_, session)| session.session_id() != INVALID_SESSION_ID)
1295 .map(|(leader, session)| {
1296 let request = build_fetch_request(
1297 max_wait_ms,
1298 self.isolation_level,
1299 self.fetch_min,
1300 self.fetch_max,
1301 self.client_rack.as_deref().unwrap_or_default(),
1302 SessionRequest {
1303 session_id: session.session_id(),
1304 session_epoch: FINAL_EPOCH,
1305 partitions: BTreeMap::new(),
1306 forgotten: Vec::new(),
1307 },
1308 );
1309 let client = self.client.clone();
1310 let leader = *leader;
1311 async move {
1312 let _ = if should_use_bootstrap_leader(leader) {
1313 client.send(request).await
1314 } else {
1315 client.broker(leader).send(request).await
1316 };
1317 }
1318 })
1319 .collect();
1320 let _ = tokio::time::timeout(timeout, futures_util::future::join_all(closes)).await;
1321 }
1322
1323 pub(crate) async fn group_fetches(&mut self, assigned: &[(String, i32)]) -> FetchByLeader {
1324 let subscription = self.subscription.borrow().clone();
1325 let awaiting_callback = self
1326 .assigned_callback_pending
1327 .lock()
1328 .unwrap_or_else(std::sync::PoisonError::into_inner)
1329 .clone();
1330 let ownership_ids = self.commit_identity.lock().await.ownership_ids.clone();
1331 let paused = self.paused_of(&ownership_ids);
1332 let buffered = self.fetch_buffer.buffered_partitions();
1333 let mut grouped: FetchByLeader = HashMap::new();
1334 let now = tokio::time::Instant::now();
1335 let mut refresh_metadata = false;
1336 {
1337 let offsets = self.next_offsets.lock().await;
1338 let positions = self.positions.lock().await;
1339 for (t, p) in assigned {
1340 if awaiting_callback.contains(&(t.clone(), *p)) {
1341 continue;
1342 }
1343 if !subscription.manual_assignment
1350 && subscription.pattern.is_none()
1351 && subscription.regex.is_none()
1352 && !subscription.contains(t)
1353 {
1354 continue;
1355 }
1356 if positions
1359 .get(&(t.clone(), *p))
1360 .is_some_and(|x| x.awaiting_validation)
1361 {
1362 continue;
1363 }
1364 let next = offsets.get(&(t.clone(), *p)).copied().unwrap_or(0);
1365 if is_reset_sentinel(next) {
1369 continue;
1370 }
1371 if paused.contains(&(t.clone(), *p)) || buffered.contains(&(t.clone(), *p)) {
1374 continue;
1375 }
1376 let pos = positions.get(&(t.clone(), *p)).copied().unwrap_or_default();
1377 let mut leader =
1386 fetch_leader_id(pos.leader_id, self.client.knows_broker(pos.leader_id));
1387 let key = (t.clone(), *p);
1391 match self.fetches.preferred_read_replicas.get(&key).copied() {
1392 Some((_, expires)) if now > expires => {
1393 self.fetches.preferred_read_replicas.remove(&key);
1394 }
1395 Some((replica, _)) if self.client.knows_broker(replica) => leader = replica,
1396 Some(_) => {
1397 self.fetches.preferred_read_replicas.remove(&key);
1398 refresh_metadata = true;
1399 }
1400 None => {}
1401 }
1402 if self.fetches.in_flight.contains_key(&leader) {
1405 continue;
1406 }
1407 grouped
1408 .entry(leader)
1409 .or_default()
1410 .entry(t.clone())
1411 .or_default()
1412 .push((*p, next, pos.leader_epoch, pos.offset_epoch));
1413 }
1414 }
1415 if refresh_metadata {
1416 self.client.metadata_topics().request_update();
1419 }
1420
1421 grouped
1422 }
1423
1424 async fn prepare_poll(
1425 &mut self,
1426 deadline: tokio::time::Instant,
1427 ) -> Result<bool, ConsumerError> {
1428 if let Some(error) = coordinator::take_poll_error(&self.poll_error) {
1431 return Err(error);
1432 }
1433 self.maybe_auto_commit_async().await;
1436 coordinator::note_poll(&self.poll_signal);
1440 self.sync_metadata_topics();
1446 let wakeup = self.wakeup.clone();
1447 self.update_fetch_positions(Some(&wakeup), Some(deadline))
1448 .await
1449 }
1450
1451 fn sync_metadata_topics(&self) {
1456 let topics = {
1457 let subscription = self.subscription.borrow();
1458 (!subscription.is_none()).then(|| subscription.topics.clone())
1459 };
1460 if let Some(topics) = topics
1461 && self.client.metadata_topics().names() != topics
1462 {
1463 self.client.metadata_topics().set(topics);
1464 }
1465 }
1466
1467 pub(crate) async fn metadata_due(&self) -> bool {
1478 let topics = self.client.metadata_topics();
1479 if topics.update_requested()
1480 || topics
1481 .last_refreshed()
1482 .is_none_or(|at| at.elapsed() >= self.metadata_max_age.to_std())
1483 {
1484 return true;
1485 }
1486 let assigned = self.assigned.lock().await.clone();
1487 let positions = self.positions.lock().await;
1488 assigned.iter().any(|key| {
1489 positions
1490 .get(key)
1491 .is_none_or(|position| position.leader_id < 0)
1492 })
1493 }
1494
1495 pub(crate) async fn update_fetch_positions(
1507 &self,
1508 wakeup: Option<&control::WakeupHandle>,
1509 deadline: Option<tokio::time::Instant>,
1510 ) -> Result<bool, ConsumerError> {
1511 if self.metadata_due().await
1512 && let Err(error) = until_woken(wakeup, self.refresh_leader_epochs()).await?
1513 {
1514 if is_transient_poll_error(&error) {
1515 self.client.reconnect_bootstrap().await;
1516 return Ok(false);
1517 }
1518 return Err(error);
1519 }
1520 if let Err(error) = until_woken(wakeup, self.resolve_committed_sentinels(deadline)).await? {
1521 if is_transient_poll_error(&error) {
1522 self.client.reconnect_bootstrap().await;
1523 return Ok(false);
1524 }
1525 return Err(error);
1526 }
1527 self.fail_partitions_without_offset().await?;
1531 if let Err(error) = until_woken(wakeup, self.resolve_reset_sentinels()).await? {
1532 if is_transient_poll_error(&error) {
1533 self.client.reconnect_bootstrap().await;
1534 return Ok(false);
1535 }
1536 return Err(error);
1537 }
1538 let truncated = match until_woken(wakeup, self.validate_positions()).await? {
1539 Ok(truncated) => truncated,
1540 Err(error) if is_transient_poll_error(&error) => {
1541 self.client.reconnect_bootstrap().await;
1542 return Ok(false);
1543 }
1544 Err(error) => return Err(error),
1545 };
1546 if !truncated.is_empty() {
1547 self.apply_truncation(&truncated).await?;
1548 }
1549 Ok(true)
1550 }
1551
1552 fn process_partition_records(
1556 &self,
1557 offsets: &HashMap<(String, i32), i64>,
1558 key: &(String, i32),
1559 topic_name: &str,
1560 part: &fetch_response::PartitionData,
1561 ) -> Option<BufferedPartition> {
1562 let batches = part.records.as_ref()?.as_v2()?;
1565 let fetch_floor = offsets.get(key).copied().unwrap_or(0);
1573 let mut records = std::collections::VecDeque::new();
1574 let read_committed = is_read_committed(self.isolation_level);
1581 let mut aborted: std::collections::VecDeque<(i64, i64)> = if read_committed {
1584 let mut v: Vec<(i64, i64)> = part
1585 .aborted_transactions
1586 .as_deref()
1587 .unwrap_or(&[])
1588 .iter()
1589 .map(|a| (a.first_offset, a.producer_id))
1590 .collect();
1591 v.sort_unstable();
1592 v.into()
1593 } else {
1594 std::collections::VecDeque::new()
1595 };
1596 let mut aborted_pids: std::collections::HashSet<i64> = std::collections::HashSet::new();
1598 for batch in batches {
1599 if read_committed {
1602 while let Some(&(first_offset, pid)) = aborted.front() {
1603 if aborted_txn_started(first_offset, batch.base_offset) {
1604 aborted_pids.insert(pid);
1605 aborted.pop_front();
1606 } else {
1607 break;
1608 }
1609 }
1610 }
1611 if batch.attributes.is_control_batch() {
1615 if read_committed {
1616 aborted_pids.remove(&batch.producer_id);
1617 }
1618 continue;
1619 }
1620 if should_drop_aborted_batch(
1622 read_committed,
1623 batch.attributes.is_transactional(),
1624 aborted_pids.contains(&batch.producer_id),
1625 ) {
1626 continue;
1627 }
1628 for r in &batch.records {
1629 let offset = record_offset(batch.base_offset, r.offset_delta);
1630 if offset < fetch_floor {
1634 continue;
1635 }
1636 let (timestamp, timestamp_type) = record_timestamp(batch, r);
1637 records.push_back(ConsumerRecord {
1638 topic: topic_name.to_string(),
1639 partition: part.partition_index,
1640 offset,
1641 leader_epoch: batch.partition_leader_epoch,
1642 timestamp,
1643 timestamp_type,
1644 key: r.key.clone(),
1645 value: r.value.clone(),
1646 headers: r
1647 .headers
1648 .iter()
1649 .map(|h| Header {
1650 key: h.key.clone(),
1651 value: h.value.clone(),
1652 })
1653 .collect(),
1654 });
1655 }
1656 }
1657 let next_offset = next_offset_after(batches)?;
1662 Some(BufferedPartition {
1663 key: key.clone(),
1664 position: fetch_floor,
1665 records,
1666 next_offset,
1667 last_epoch: batches
1670 .iter()
1671 .map(|batch| batch.partition_leader_epoch)
1672 .max()
1673 .map(LeaderEpoch),
1674 })
1675 }
1676}
1677
1678const fn readable_end_offset(
1679 isolation: IsolationLevel,
1680 high_watermark: i64,
1681 last_stable_offset: i64,
1682) -> i64 {
1683 match isolation {
1684 IsolationLevel::ReadCommitted => last_stable_offset,
1685 IsolationLevel::ReadUncommitted => high_watermark,
1686 }
1687}
1688
1689fn record_readable_end(
1690 ends: &mut HashMap<(String, i32), i64>,
1691 key: (String, i32),
1692 isolation: IsolationLevel,
1693 high_watermark: i64,
1694 last_stable_offset: i64,
1695) {
1696 ends.insert(
1697 key,
1698 readable_end_offset(isolation, high_watermark, last_stable_offset),
1699 );
1700}
1701
1702fn next_offset_after(batches: &[krabka_protocol::records::RecordBatch]) -> Option<i64> {
1709 batches
1710 .iter()
1711 .map(|b| b.base_offset + i64::from(b.last_offset_delta) + 1)
1712 .max()
1713}
1714
1715impl Consumer {
1716 #[tracing::instrument(
1737 name = "consumer.resolve_reset_sentinels",
1738 level = "debug",
1739 skip_all,
1740 fields(group_id = %self.group_id, sentinels = tracing::field::Empty),
1741 err(level = "debug")
1742 )]
1743 async fn resolve_reset_sentinels(&self) -> Result<(), ConsumerError> {
1744 let by_sentinel = {
1745 let offsets = self.next_offsets.lock().await;
1746 [
1747 (
1748 LATEST_SENTINEL,
1749 reset_timestamp(self.auto_offset_reset, unix_now_ms()),
1750 ),
1751 (BEGINNING_SENTINEL, EARLIEST_TIMESTAMP),
1752 (END_SENTINEL, LATEST_TIMESTAMP),
1753 ]
1754 .map(|(sentinel, timestamp)| (sentinel, timestamp, keys_at(&offsets, sentinel)))
1755 };
1756 let count: usize = by_sentinel.iter().map(|(_, _, keys)| keys.len()).sum();
1757 if count == 0 {
1758 return Ok(());
1759 }
1760 tracing::Span::current().record("sentinels", count);
1761 let mut results = Vec::new();
1762 for (sentinel, timestamp, keys) in by_sentinel {
1763 if !keys.is_empty() {
1764 results.push((sentinel, self.list_offsets(&keys, timestamp).await?));
1765 }
1766 }
1767 {
1768 let mut offsets = self.next_offsets.lock().await;
1771 for (sentinel, result) in &results {
1772 for (key, offset) in &result.offsets {
1773 if let Some(next) = offsets.get_mut(key)
1774 && *next == *sentinel
1775 {
1776 *next = *offset;
1777 }
1778 }
1779 }
1780 }
1781 if results
1786 .iter()
1787 .any(|(_, result)| result.retry && result.unauthorized.is_empty())
1788 {
1789 self.client.metadata_topics().request_update();
1790 }
1791 results
1792 .into_iter()
1793 .try_for_each(|(_, result)| result.authorization())
1794 }
1795
1796 async fn fail_partitions_without_offset(&self) -> Result<(), ConsumerError> {
1810 let missing: BTreeSet<(String, i32)> = self
1811 .next_offsets
1812 .lock()
1813 .await
1814 .iter()
1815 .filter(|&(_, &offset)| offset == NO_OFFSET_SENTINEL)
1816 .map(|(key, _)| key.clone())
1817 .collect();
1818 if missing.is_empty() {
1819 Ok(())
1820 } else {
1821 Err(ConsumerError::NoOffsetForPartition(missing))
1822 }
1823 }
1824
1825 pub(crate) async fn resolve_committed_sentinels(
1834 &self,
1835 deadline: Option<tokio::time::Instant>,
1836 ) -> Result<(), ConsumerError> {
1837 let keys = keys_at(&*self.next_offsets.lock().await, COMMITTED_SENTINEL);
1838 if keys.is_empty() {
1839 return Ok(());
1840 }
1841 let deadline = deadline
1844 .unwrap_or_else(|| tokio::time::Instant::now() + self.default_api_timeout.to_std());
1845 let committed = if self.group_id.is_empty() {
1846 HashMap::new()
1847 } else {
1848 self.committed_until(&keys, deadline).await?
1849 };
1850 let reset = consumer::reset_starting_offset(self.auto_offset_reset);
1851 let mut offsets = self.next_offsets.lock().await;
1852 let mut positions = self.positions.lock().await;
1853 for key in keys {
1854 let Some(next) = offsets.get_mut(&key) else {
1855 continue;
1856 };
1857 if *next != COMMITTED_SENTINEL {
1858 continue;
1859 }
1860 match committed.get(&key).cloned().flatten() {
1861 Some(offset) => {
1862 *next = offset.offset;
1863 let position = positions.entry(key).or_default();
1864 position.offset_epoch = LeaderEpoch(offset.leader_epoch.unwrap_or(-1));
1865 position.awaiting_validation = crate::validate::should_await_validation(
1869 position.leader_epoch,
1870 position.offset_epoch,
1871 );
1872 }
1873 None => *next = reset,
1874 }
1875 }
1876 Ok(())
1877 }
1878
1879 #[cfg_attr(test, mutants::skip)] async fn list_offsets(
1893 &self,
1894 keys: &[(String, i32)],
1895 timestamp: i64,
1896 ) -> Result<ListOffsetsResult, ConsumerError> {
1897 let by_leader = group_list_offsets(keys, &*self.positions.lock().await, |id| {
1898 self.client.knows_broker(id)
1899 });
1900 let mut answers = Vec::with_capacity(by_leader.len());
1901 for (leader, by_topic) in by_leader {
1902 let request = build_offsets_request(by_topic, timestamp, self.isolation_level);
1903 let answer = if is_read_committed(self.isolation_level) {
1904 self.send_list_offsets(leader, ReadCommittedListOffsets::new(request))
1905 .await
1906 } else {
1907 self.send_list_offsets(leader, request).await
1908 };
1909 match answer {
1910 Ok(answer) => answers.push(answer),
1911 Err(e) if is_transient_transport_error(&e) => {
1912 if should_use_bootstrap_leader(leader) {
1913 self.client.reconnect_bootstrap().await;
1914 } else {
1915 self.client.evict_broker(leader);
1916 }
1917 }
1918 Err(e) => return Err(e.into()),
1919 }
1920 }
1921 Ok(ListOffsetsResult::collect(keys, &answers))
1922 }
1923
1924 async fn send_list_offsets<R>(
1927 &self,
1928 leader: i32,
1929 request: R,
1930 ) -> Result<ListOffsetsResponse, krabka_client_core::ClientError>
1931 where
1932 R: ProtocolRequest<Response = ListOffsetsResponse>,
1933 {
1934 if should_use_bootstrap_leader(leader) {
1935 self.client.send(request).await
1936 } else {
1937 self.client.broker(leader).send(request).await
1938 }
1939 }
1940
1941 #[tracing::instrument(
1960 name = "consumer.recover_out_of_range",
1961 level = "debug",
1962 skip_all,
1963 fields(group_id = %self.group_id, partitions = out_of_range.len()),
1964 err(level = "debug")
1965 )]
1966 async fn recover_out_of_range(
1967 &self,
1968 out_of_range: &[((String, i32), i64)],
1969 ) -> Result<bool, ConsumerError> {
1970 let keys: Vec<(String, i32)> = out_of_range.iter().map(|(key, _)| key.clone()).collect();
1971 let result = self.list_offsets(&keys, EARLIEST_TIMESTAMP).await?;
1972 result.authorization()?;
1973
1974 if let AutoOffsetReset::None = self.auto_offset_reset {
1975 if let Some(truncation) = out_of_range_truncation(out_of_range, &result.offsets) {
1976 return Err(truncation);
1977 }
1978 return Ok(false);
1979 }
1980
1981 let mut offsets = self.next_offsets.lock().await;
1982 let log_starts = out_of_range_positions(out_of_range, &result.offsets);
1983 for (key, log_start) in &log_starts {
1984 offsets.insert(key.clone(), *log_start);
1985 }
1986 if let Some(auto_commit) = &self.auto_commit {
1987 auto_commit
1988 .reset_polled(
1989 log_starts
1990 .into_iter()
1991 .map(|(key, log_start)| (key, Some(log_start))),
1992 )
1993 .await;
1994 }
1995 Ok(result.retry)
1996 }
1997}
1998
1999fn out_of_range_positions(
2005 out_of_range: &[((String, i32), i64)],
2006 log_starts: &HashMap<(String, i32), i64>,
2007) -> Vec<((String, i32), i64)> {
2008 out_of_range
2009 .iter()
2010 .filter_map(|(key, _)| log_starts.get(key).map(|start| (key.clone(), *start)))
2011 .collect()
2012}
2013
2014fn out_of_range_truncation(
2020 out_of_range: &[((String, i32), i64)],
2021 log_starts: &HashMap<(String, i32), i64>,
2022) -> Option<ConsumerError> {
2023 let ((topic, partition), fetch_offset) = out_of_range.first()?;
2024 Some(ConsumerError::LogTruncation {
2025 topic: topic.clone(),
2026 partition: *partition,
2027 fetch_offset: *fetch_offset,
2028 safe_offset: log_starts
2029 .get(&(topic.clone(), *partition))
2030 .copied()
2031 .unwrap_or(UNKNOWN_OFFSET),
2032 })
2033}
2034
2035fn keys_at(offsets: &HashMap<(String, i32), i64>, sentinel: i64) -> Vec<(String, i32)> {
2037 offsets
2038 .iter()
2039 .filter(|(_, value)| **value == sentinel)
2040 .map(|(key, _)| key.clone())
2041 .collect()
2042}
2043
2044impl Consumer {
2045 #[tracing::instrument(
2051 name = "consumer.apply_truncation",
2052 level = "debug",
2053 skip_all,
2054 fields(group_id = %self.group_id, truncated = truncated.len()),
2055 err(level = "debug")
2056 )]
2057 async fn apply_truncation(
2058 &self,
2059 truncated: &HashMap<(String, i32), i64>,
2060 ) -> Result<(), ConsumerError> {
2061 let mut offsets = self.next_offsets.lock().await;
2062 for (key, safe_offset) in truncated {
2063 if let AutoOffsetReset::None = self.auto_offset_reset {
2064 let fetch_offset = fetch_offset_or_unknown(&offsets, key);
2065 return Err(ConsumerError::LogTruncation {
2066 topic: key.0.clone(),
2067 partition: key.1,
2068 fetch_offset,
2069 safe_offset: *safe_offset,
2070 });
2071 }
2072 offsets.insert(key.clone(), *safe_offset);
2073 }
2074 if let Some(auto_commit) = &self.auto_commit {
2075 auto_commit
2076 .reset_polled(
2077 truncated
2078 .iter()
2079 .map(|(key, safe_offset)| (key.clone(), Some(*safe_offset))),
2080 )
2081 .await;
2082 }
2083 Ok(())
2084 }
2085
2086 fn handle_truncation_in_poll(
2089 &self,
2090 offsets: &mut HashMap<(String, i32), i64>,
2091 key: &(String, i32),
2092 safe_offset: i64,
2093 ) -> Result<(), ConsumerError> {
2094 if let AutoOffsetReset::None = self.auto_offset_reset {
2095 let fetch_offset = fetch_offset_or_unknown(offsets, key);
2096 return Err(ConsumerError::LogTruncation {
2097 topic: key.0.clone(),
2098 partition: key.1,
2099 fetch_offset,
2100 safe_offset,
2101 });
2102 }
2103 offsets.insert(key.clone(), safe_offset);
2104 Ok(())
2105 }
2106}
2107
2108#[cfg(test)]
2109mod offset_advance_tests {
2110 use std::collections::HashMap;
2111
2112 use assert2::check;
2113 use krabka_protocol::{
2114 owned::{
2115 fetch_request::ReplicaState,
2116 list_offsets_response::{ListOffsetsPartitionResponse, ListOffsetsTopicResponse},
2117 },
2118 primitives::uuid::Uuid as WireUuid,
2119 records::{RecordBatch, RecordsPayload},
2120 tagged_fields::UnknownTaggedFields,
2121 };
2122
2123 use super::*;
2124
2125 fn id(n: u8) -> WireUuid {
2126 let mut b = [0u8; 16];
2127 b[15] = n;
2128 WireUuid(b)
2129 }
2130
2131 #[test]
2132 fn fetch_partition_error_codes_map_to_kafka_consumer_actions() {
2133 for (name, error_code, expected) in [
2134 ("none", 0, FetchPartitionAction::Records),
2135 ("offset out of range", 1, FetchPartitionAction::ResetOffset),
2136 (
2137 "unknown topic or partition",
2138 3,
2139 FetchPartitionAction::RefreshMetadata,
2140 ),
2141 ("not leader or follower", 6, FetchPartitionAction::Reroute),
2142 (
2143 "fenced leader epoch",
2144 74,
2145 FetchPartitionAction::RevalidateEpoch,
2146 ),
2147 (
2148 "unknown leader epoch",
2149 75,
2150 FetchPartitionAction::RevalidateEpoch,
2151 ),
2152 (
2153 "unknown topic id",
2154 100,
2155 FetchPartitionAction::RefreshMetadata,
2156 ),
2157 (
2158 "replica not available",
2159 9,
2160 FetchPartitionAction::RefreshMetadata,
2161 ),
2162 (
2163 "kafka storage error",
2164 56,
2165 FetchPartitionAction::RefreshMetadata,
2166 ),
2167 (
2168 "offset not available",
2169 78,
2170 FetchPartitionAction::RefreshMetadata,
2171 ),
2172 (
2173 "inconsistent topic id",
2174 103,
2175 FetchPartitionAction::RefreshMetadata,
2176 ),
2177 (
2178 "topic authorization failed",
2179 29,
2180 FetchPartitionAction::Fail(29),
2181 ),
2182 ] {
2183 check!(
2184 classify_fetch_partition_error(error_code) == expected,
2185 "case {name}"
2186 );
2187 }
2188 }
2189
2190 #[test]
2191 fn fetch_leader_id_uses_known_non_negative_leader_or_bootstrap() {
2192 check!(BOOTSTRAP_LEADER == -1);
2193 check!(UNKNOWN_LEADER_ID == -1);
2194 check!(UNKNOWN_FETCH_OFFSET == -1);
2195 for (name, leader, known, expected) in [
2196 ("known non-negative", 3, true, 3),
2197 ("negative leader", -1, true, BOOTSTRAP_LEADER),
2198 ("unknown broker", 3, false, BOOTSTRAP_LEADER),
2199 ] {
2200 check!(fetch_leader_id(leader, known) == expected, "case {name}");
2201 }
2202 for (name, leader, expected) in [
2203 ("bootstrap sentinel", BOOTSTRAP_LEADER, true),
2204 ("broker leader", 3, false),
2205 ] {
2206 check!(
2207 should_use_bootstrap_leader(leader) == expected,
2208 "case {name}"
2209 );
2210 }
2211 }
2212
2213 #[test]
2214 fn readable_end_respects_consumer_isolation() {
2215 check!(readable_end_offset(IsolationLevel::ReadUncommitted, 12, 9) == 12);
2216 check!(readable_end_offset(IsolationLevel::ReadCommitted, 12, 9) == 9);
2217
2218 let mut ends = HashMap::new();
2219 record_readable_end(
2220 &mut ends,
2221 ("topic-a".into(), 0),
2222 IsolationLevel::ReadCommitted,
2223 12,
2224 9,
2225 );
2226 check!(ends.get(&("topic-a".into(), 0)) == Some(&9));
2227 }
2228
2229 #[test]
2230 fn poll_helpers_preserve_sentinel_filter_and_record_math_boundaries() {
2231 let key = ("topic-a".to_string(), 2);
2232 let mut offsets = HashMap::new();
2233
2234 assert2::assert!(fetch_offset_or_unknown(&offsets, &key) == UNKNOWN_FETCH_OFFSET);
2235 offsets.insert(key.clone(), 42);
2236 check!(fetch_offset_or_unknown(&offsets, &key) == 42);
2237
2238 for (name, isolation, expected) in [
2239 ("read committed", IsolationLevel::ReadCommitted, true),
2240 ("read uncommitted", IsolationLevel::ReadUncommitted, false),
2241 ] {
2242 check!(is_read_committed(isolation) == expected, "case {name}");
2243 }
2244 for (name, first, offset, expected) in [
2245 ("at first offset", 10, 10, true),
2246 ("after first offset", 10, 11, true),
2247 ("before first offset", 11, 10, false),
2248 ] {
2249 check!(
2250 aborted_txn_started(first, offset) == expected,
2251 "case {name}"
2252 );
2253 }
2254 for (name, committed, aborted, started, expected) in [
2255 ("drop aborted", true, true, true, true),
2256 ("uncommitted isolation", false, true, true, false),
2257 ("not aborted", true, false, true, false),
2258 ("transaction not started", true, true, false, false),
2259 ] {
2260 check!(
2261 should_drop_aborted_batch(committed, aborted, started) == expected,
2262 "case {name}"
2263 );
2264 }
2265 check!(record_offset(100, 7) == 107);
2266 }
2267
2268 #[test]
2269 fn transient_transport_error_classification_is_narrow() {
2270 use std::io;
2271
2272 use krabka_client_core::ClientError;
2273
2274 let transport_cases = vec![
2275 ("disconnected", ClientError::Disconnected, true),
2276 (
2277 "timeout",
2278 ClientError::Timeout(krabka_units::millis(10)),
2279 true,
2280 ),
2281 (
2282 "connection reset",
2283 ClientError::Io(io::Error::new(io::ErrorKind::ConnectionReset, "reset")),
2284 true,
2285 ),
2286 (
2287 "connect refused",
2288 ClientError::Connect {
2289 addr: "127.0.0.1:9092".parse().unwrap(),
2290 source: io::Error::new(io::ErrorKind::ConnectionRefused, "refused"),
2291 },
2292 true,
2293 ),
2294 ("server error", ClientError::Server { error_code: 6 }, false),
2295 (
2296 "incompatible version",
2297 ClientError::IncompatibleVersion {
2298 api_key: 1,
2299 broker_min: 0,
2300 broker_max: 1,
2301 client_min: 2,
2302 client_max: 3,
2303 },
2304 false,
2305 ),
2306 ];
2307 for (name, error, expected) in &transport_cases {
2308 check!(
2309 is_transient_transport_error(error) == *expected,
2310 "case {name}"
2311 );
2312 }
2313 let poll_cases = vec![
2314 (
2315 "client connect",
2316 ConsumerError::Client(ClientError::Connect {
2317 addr: "127.0.0.1:9092".parse().unwrap(),
2318 source: io::Error::new(io::ErrorKind::ConnectionRefused, "refused"),
2319 }),
2320 true,
2321 ),
2322 ("server", ConsumerError::Server(6), false),
2323 ];
2324 for (name, error, expected) in &poll_cases {
2325 check!(is_transient_poll_error(error) == *expected, "case {name}");
2326 }
2327 }
2328
2329 #[test]
2333 fn build_fetch_request_carries_limits_session_partitions_and_forgotten_topics() {
2334 let partition = |fetch_offset| crate::fetch_session::SessionPartition {
2335 topic_id: id(7),
2336 fetch_offset,
2337 current_leader_epoch: 5,
2338 last_fetched_epoch: 4,
2339 partition_max_bytes: 128 * 1024,
2340 };
2341 let req = build_fetch_request(
2342 500,
2343 IsolationLevel::ReadCommitted,
2344 krabka_units::bytes(7),
2345 mebibytes(2),
2346 "az-1",
2347 SessionRequest {
2348 session_id: 77,
2349 session_epoch: 3,
2350 partitions: maplit::btreemap! {("topic-a".to_owned(), 2) => partition(42), ("topic-a".to_owned(), 3) => partition(43)},
2351 forgotten: vec![
2352 (("topic-b".to_owned(), 0), id(8)),
2353 (("topic-b".to_owned(), 1), id(8)),
2354 ],
2355 },
2356 );
2357 let fetch_partition = |partition, fetch_offset| FetchPartition {
2358 partition,
2359 current_leader_epoch: 5,
2360 fetch_offset,
2361 last_fetched_epoch: 4,
2362 log_start_offset: -1,
2363 partition_max_bytes: 128 * 1024,
2364 replica_directory_id: WireUuid::default(),
2365 high_watermark: i64::MAX,
2366 unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2367 };
2368 assert2::assert!(
2369 req == FetchRequest {
2370 replica_id: -1,
2371 max_wait_ms: 500,
2372 min_bytes: 7,
2373 max_bytes: 2 * 1024 * 1024,
2374 isolation_level: 1, session_id: 77,
2376 session_epoch: 3,
2377 topics: vec![FetchTopic {
2378 topic: "topic-a".into(),
2379 topic_id: id(7),
2380 partitions: vec![fetch_partition(2, 42), fetch_partition(3, 43)],
2381 unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2382 }],
2383 forgotten_topics_data: vec![ForgottenTopic {
2384 topic: "topic-b".into(),
2385 topic_id: id(8),
2386 partitions: vec![0, 1],
2387 unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2388 }],
2389 rack_id: "az-1".into(),
2390 cluster_id: None,
2391 replica_state: ReplicaState {
2392 replica_id: -1,
2393 replica_epoch: -1,
2394 unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2395 },
2396 unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2397 }
2398 );
2399 }
2400
2401 #[test]
2409 fn offsets_request_carries_the_timestamp_isolation_and_leader_epoch() {
2410 for (timestamp, isolation, isolation_level, leader_epoch) in [
2411 (LATEST_TIMESTAMP, IsolationLevel::ReadCommitted, 1, 4),
2412 (LATEST_TIMESTAMP, IsolationLevel::ReadUncommitted, 0, -1),
2413 (EARLIEST_TIMESTAMP, IsolationLevel::ReadCommitted, 1, -1),
2414 (EARLIEST_TIMESTAMP, IsolationLevel::ReadUncommitted, 0, 9),
2415 ] {
2416 let by_topic =
2417 maplit::btreemap! {"topic-a".to_string() => vec![(3, LeaderEpoch(leader_epoch))]};
2418 let req = build_offsets_request(by_topic, timestamp, isolation);
2419
2420 check!(
2421 req == ListOffsetsRequest {
2422 replica_id: -1,
2423 isolation_level,
2424 topics: vec![ListOffsetsTopic {
2425 name: "topic-a".into(),
2426 partitions: vec![ListOffsetsPartition {
2427 partition_index: 3,
2428 current_leader_epoch: leader_epoch,
2429 timestamp,
2430 unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2431 }],
2432 unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2433 }],
2434 timeout_ms: 0,
2435 unknown_tagged_fields: UnknownTaggedFields(Vec::new()),
2436 },
2437 "timestamp {timestamp}, isolation {isolation:?}"
2438 );
2439 }
2440 }
2441
2442 #[test]
2446 fn list_offsets_partitions_group_by_leader_with_the_leader_epoch() {
2447 let position = |leader_id, leader_epoch| PartitionPosition {
2448 leader_id,
2449 leader_epoch: LeaderEpoch(leader_epoch),
2450 ..Default::default()
2451 };
2452 let positions = maplit::hashmap! {("topic-a".to_string(), 0) => position(1, 5), ("topic-a".to_string(), 1) => position(2, 6), ("topic-b".to_string(), 0) => position(1, 3), ("topic-b".to_string(), 1) => position(3, 8)};
2453 let keys = [
2454 ("topic-a".to_string(), 0),
2455 ("topic-a".to_string(), 1),
2456 ("topic-b".to_string(), 0),
2457 ("topic-b".to_string(), 1),
2458 ("topic-c".to_string(), 0),
2459 ];
2460 let grouped = group_list_offsets(&keys, &positions, |id| id != 3);
2461 assert2::assert!(
2462 grouped
2463 == maplit::btreemap! {
2464 BOOTSTRAP_LEADER => maplit::btreemap! {
2465 "topic-b".to_string() => vec![(1, LeaderEpoch(8))],
2466 "topic-c".to_string() => vec![(0, LeaderEpoch(-1))],
2467 },
2468 1 => maplit::btreemap! {
2469 "topic-a".to_string() => vec![(0, LeaderEpoch(5))],
2470 "topic-b".to_string() => vec![(0, LeaderEpoch(3))],
2471 },
2472 2 => maplit::btreemap! {
2473 "topic-a".to_string() => vec![(1, LeaderEpoch(6))],
2474 },
2475 }
2476 );
2477 }
2478
2479 #[test]
2483 fn list_offsets_rows_map_to_kafka_consumer_actions() {
2484 for (name, error_code, offset, expected) in [
2485 ("none", 0, 40, ListOffsetsRowAction::Position(40)),
2486 ("none at zero", 0, 0, ListOffsetsRowAction::Position(0)),
2487 ("none without an offset", 0, -1, ListOffsetsRowAction::Retry),
2488 ("not leader or follower", 6, -1, ListOffsetsRowAction::Retry),
2489 (
2490 "unknown topic or partition",
2491 3,
2492 -1,
2493 ListOffsetsRowAction::Retry,
2494 ),
2495 ("leader not available", 5, -1, ListOffsetsRowAction::Retry),
2496 ("replica not available", 9, -1, ListOffsetsRowAction::Retry),
2497 ("offset not available", 78, -1, ListOffsetsRowAction::Retry),
2498 ("fenced leader epoch", 74, -1, ListOffsetsRowAction::Retry),
2499 ("unknown leader epoch", 75, -1, ListOffsetsRowAction::Retry),
2500 ("kafka storage error", 56, -1, ListOffsetsRowAction::Retry),
2501 (
2502 "unsupported for message format",
2503 43,
2504 -1,
2505 ListOffsetsRowAction::Retry,
2506 ),
2507 ("unexpected code", 2, -1, ListOffsetsRowAction::Retry),
2508 (
2509 "topic authorization failed",
2510 29,
2511 -1,
2512 ListOffsetsRowAction::Unauthorized,
2513 ),
2514 ] {
2515 check!(
2516 classify_list_offsets_row(error_code, offset) == expected,
2517 "case {name}"
2518 );
2519 }
2520 }
2521
2522 #[test]
2526 fn list_offsets_result_collects_offsets_retries_and_unauthorized_topics() {
2527 let row = |partition_index, error_code, offset| ListOffsetsPartitionResponse {
2528 partition_index,
2529 error_code,
2530 offset,
2531 ..Default::default()
2532 };
2533 let answer = |name: &str, rows| ListOffsetsResponse {
2534 topics: vec![ListOffsetsTopicResponse {
2535 name: name.into(),
2536 partitions: rows,
2537 ..Default::default()
2538 }],
2539 ..Default::default()
2540 };
2541 let requested = [("topic-a".to_string(), 0), ("topic-a".to_string(), 1)];
2542 let retry_with_first_offset = ListOffsetsResult {
2543 offsets: maplit::hashmap! {("topic-a".to_string(), 0) => 5},
2544 retry: true,
2545 unauthorized: BTreeSet::new(),
2546 };
2547 for (name, answers, expected) in [
2548 (
2549 "every partition answered",
2550 vec![answer("topic-a", vec![row(0, 0, 5), row(1, 0, 9)])],
2551 ListOffsetsResult {
2552 offsets: maplit::hashmap! {("topic-a".to_string(), 0) => 5, ("topic-a".to_string(), 1) => 9},
2553 retry: false,
2554 unauthorized: BTreeSet::new(),
2555 },
2556 ),
2557 (
2558 "an error row is retried",
2559 vec![answer("topic-a", vec![row(0, 0, 5), row(1, 6, -1)])],
2560 retry_with_first_offset.clone(),
2561 ),
2562 (
2563 "a missing row is retried and an unrequested row is ignored",
2564 vec![answer("topic-a", vec![row(0, 0, 5), row(2, 0, 3)])],
2565 retry_with_first_offset.clone(),
2566 ),
2567 (
2568 "an unauthorized answer gives no offsets",
2569 vec![answer("topic-a", vec![row(0, 0, 5), row(1, 29, -1)])],
2570 ListOffsetsResult {
2571 offsets: HashMap::new(),
2572 retry: true,
2573 unauthorized: maplit::btreeset! {"topic-a".to_string()},
2574 },
2575 ),
2576 ] {
2577 check!(
2578 ListOffsetsResult::collect(&requested, &answers) == expected,
2579 "case {name}"
2580 );
2581 }
2582 }
2583
2584 #[test]
2586 fn list_offsets_result_authorization_names_the_unauthorized_topics() {
2587 let mut result = ListOffsetsResult::default();
2588 check!(result.authorization().map_err(|e| e.to_string()) == Ok(()));
2589 result.unauthorized.insert("topic-a".to_string());
2590 check!(
2591 result.authorization().map_err(|e| e.to_string())
2592 == Err(ConsumerError::TopicAuthorizationFailed(
2593 maplit::btreeset! {"topic-a".to_string()}
2594 )
2595 .to_string())
2596 );
2597 }
2598
2599 #[test]
2606 fn out_of_range_positions_uses_the_log_start_and_skips_the_unanswered() {
2607 let out_of_range = vec![
2608 (("topic-a".to_string(), 0), 0),
2609 (("topic-a".to_string(), 1), 3),
2610 ];
2611 let log_starts = maplit::hashmap! {("topic-a".to_string(), 0) => 5};
2612 assert2::assert!(
2613 out_of_range_positions(&out_of_range, &log_starts)
2614 == vec![(("topic-a".to_string(), 0), 5)]
2615 );
2616 }
2617
2618 #[test]
2624 fn out_of_range_truncation_names_the_first_partition_and_its_log_start() {
2625 for (answered, expected_safe) in [(true, 5), (false, UNKNOWN_OFFSET)] {
2626 let log_starts = if answered {
2627 maplit::hashmap! {("topic-a".to_string(), 0) => 5}
2628 } else {
2629 HashMap::new()
2630 };
2631 let out_of_range = vec![
2632 (("topic-a".to_string(), 0), 2),
2633 (("topic-a".to_string(), 1), 3),
2634 ];
2635 assert2::assert!(
2636 out_of_range_truncation(&out_of_range, &log_starts).map(|e| e.to_string())
2637 == Some(
2638 ConsumerError::LogTruncation {
2639 topic: "topic-a".to_string(),
2640 partition: 0,
2641 fetch_offset: 2,
2642 safe_offset: expected_safe,
2643 }
2644 .to_string()
2645 )
2646 );
2647 }
2648 }
2649
2650 #[test]
2652 fn out_of_range_truncation_is_none_for_an_empty_set() {
2653 assert2::assert!(out_of_range_truncation(&[], &HashMap::new()).is_none());
2654 }
2655
2656 #[test]
2658 fn keys_at_selects_only_the_sentinel_partitions() {
2659 let mut offsets: HashMap<(String, i32), i64> = HashMap::new();
2660 offsets.insert(("topic-a".to_string(), 0), LATEST_SENTINEL);
2661 offsets.insert(("topic-a".to_string(), 1), 42);
2662 assert2::assert!(keys_at(&offsets, LATEST_SENTINEL) == vec![("topic-a".to_string(), 0)]);
2663 }
2664
2665 #[test]
2666 fn advance_target_uses_last_offset_delta_not_record_count() {
2667 let batch = RecordBatch {
2670 base_offset: 10,
2671 last_offset_delta: 4,
2672 records: vec![],
2673 ..Default::default()
2674 };
2675 let payload = RecordsPayload::V2(vec![batch]);
2676 let batches = payload.as_v2().unwrap();
2677 assert2::assert!(super::next_offset_after(batches) == Some(15));
2678 }
2679
2680 #[test]
2681 fn advance_target_none_for_empty() {
2682 let payload = RecordsPayload::V2(vec![]);
2683 assert2::assert!(super::next_offset_after(payload.as_v2().unwrap()) == None);
2684 }
2685}
2686
2687#[cfg(test)]
2688pub(crate) mod partition_error_tests {
2689 use std::{
2690 sync::{
2691 Arc,
2692 atomic::{AtomicU16, AtomicUsize, Ordering},
2693 },
2694 time::Duration,
2695 };
2696
2697 use assert2::check;
2698 use krabka_client_core::{Client, MockBroker};
2699 use krabka_protocol::{
2700 Encode,
2701 owned::{
2702 api_versions_request,
2703 fetch_response::{
2704 EpochEndOffset, FetchResponse, FetchableTopicResponse, LeaderIdAndEpoch,
2705 PartitionData,
2706 },
2707 metadata_request,
2708 metadata_response::MetadataResponse,
2709 },
2710 primitives::uuid::Uuid as WireUuid,
2711 };
2712 use krabka_units::secs;
2713
2714 use super::*;
2715 use crate::Assignor;
2716
2717 const TOPIC_ID: WireUuid = WireUuid([7; 16]);
2718
2719 fn encode(response: &impl Encode, version: i16) -> Vec<u8> {
2720 let mut buf = bytes::BytesMut::new();
2721 response.encode(&mut buf, version).unwrap();
2722 buf.to_vec()
2723 }
2724
2725 async fn metadata_counting_broker(metadata_requests: Arc<AtomicUsize>) -> MockBroker {
2728 MockBroker::start(move |api_key, version, _corr_id, _body| {
2729 if api_key == api_versions_request::API_KEY {
2730 let versions = mock_test_support::api_versions_response([
2731 (api_versions_request::API_KEY, 3),
2732 (metadata_request::API_KEY, 8),
2733 ]);
2734 return Some(encode(&versions, 0));
2735 }
2736 if api_key == metadata_request::API_KEY {
2737 metadata_requests.fetch_add(1, Ordering::SeqCst);
2738 return Some(encode(&MetadataResponse::default(), version));
2739 }
2740 None
2741 })
2742 .await
2743 }
2744
2745 async fn consumer_on(broker: &MockBroker) -> Consumer {
2746 let client = Client::builder()
2747 .bootstrap(broker.addr.to_string())
2748 .build()
2749 .await
2750 .unwrap();
2751 consumer_with_client(client)
2752 }
2753
2754 pub(crate) fn consumer_with_client(client: Client) -> Consumer {
2756 Consumer::test_builder(client)
2757 .commit_identity(test_support::orders_partition_commit_identity(1))
2758 .subscription(crate::subscription::shared(
2759 vec!["orders".into()],
2760 None,
2761 true,
2762 ))
2763 .assigned(vec![("orders".into(), 0)])
2764 .next_offsets(maplit::hashmap! {("orders".into(), 0) => 5})
2765 .topic_ids(maplit::hashmap! {"orders".into() => TOPIC_ID})
2766 .build()
2767 }
2768
2769 fn fetch_response(error_code: i16) -> FetchResponse {
2772 FetchResponse {
2773 responses: vec![FetchableTopicResponse {
2774 topic: String::new(),
2775 topic_id: TOPIC_ID,
2776 partitions: vec![PartitionData {
2777 partition_index: 0,
2778 error_code,
2779 high_watermark: -1,
2780 last_stable_offset: -1,
2781 log_start_offset: -1,
2782 diverging_epoch: EpochEndOffset {
2783 epoch: -1,
2784 end_offset: -1,
2785 ..Default::default()
2786 },
2787 current_leader: LeaderIdAndEpoch {
2788 leader_id: -1,
2789 leader_epoch: -1,
2790 ..Default::default()
2791 },
2792 ..Default::default()
2793 }],
2794 ..Default::default()
2795 }],
2796 ..Default::default()
2797 }
2798 }
2799
2800 #[derive(Debug, PartialEq)]
2802 struct PollOutcome {
2803 result: Result<usize, Option<i16>>,
2806 metadata_requests: usize,
2808 metadata_update_requested: bool,
2810 next_offset: Option<i64>,
2811 }
2812
2813 #[tokio::test]
2817 async fn poll_returns_a_fatal_rejoin_error_once() {
2818 let broker = metadata_counting_broker(Arc::default()).await;
2819 let mut consumer = consumer_on(&broker).await;
2820 let topics = maplit::btreeset! {"orders".to_string()};
2821 coordinator::report_rejoin_error(
2822 &consumer.poll_error,
2823 ConsumerError::TopicAuthorizationFailed(topics.clone()),
2824 );
2825
2826 let first = consumer
2827 .prepare_poll(tokio::time::Instant::now() + Duration::from_secs(5))
2828 .await
2829 .map_err(|error| match error {
2830 ConsumerError::TopicAuthorizationFailed(topics) => Some(topics),
2831 _ => None,
2832 });
2833 let second = consumer
2834 .prepare_poll(tokio::time::Instant::now() + Duration::from_secs(5))
2835 .await
2836 .map_err(|_| None::<std::collections::BTreeSet<String>>);
2837
2838 broker.stop();
2839 assert2::assert!((first, second) == (Err(Some(topics)), Ok(true)));
2840 }
2841
2842 #[tokio::test]
2847 async fn drain_never_returns_records_of_a_partition_revoked_during_the_drain() {
2848 let broker = metadata_counting_broker(Arc::default()).await;
2849 let mut consumer = consumer_on(&broker).await;
2850 consumer
2851 .fetch_buffer
2852 .push(crate::fetch_buffer::BufferedPartition {
2853 key: ("orders".into(), 0),
2854 position: 5,
2855 records: std::collections::VecDeque::from([ConsumerRecord {
2856 topic: "orders".into(),
2857 partition: 0,
2858 offset: 5,
2859 leader_epoch: 0,
2860 timestamp: 0,
2861 timestamp_type: crate::TimestampType::CreateTime,
2862 key: None,
2863 value: None,
2864 headers: Vec::new(),
2865 }]),
2866 next_offset: 6,
2867 last_epoch: None,
2868 });
2869 let assigned = Arc::clone(&consumer.assigned);
2870 let next_offsets = Arc::clone(&consumer.next_offsets);
2871 let offsets_guard = next_offsets.lock().await;
2873 let drain = tokio::spawn(async move {
2874 let records = consumer.drain_fetch_buffer().await;
2875 (consumer, records)
2876 });
2877 tokio::time::sleep(Duration::from_millis(50)).await;
2878 let revoked = match tokio::time::timeout(Duration::from_millis(50), assigned.lock()).await {
2880 Ok(mut owned) => {
2881 owned.clear();
2882 true
2883 }
2884 Err(_) => false,
2885 };
2886 drop(offsets_guard);
2887 let (consumer, records) = drain.await.expect("drain task");
2888 drop(consumer);
2889 broker.stop();
2890 assert2::assert!(!revoked || records.is_empty());
2891 }
2892
2893 #[tokio::test(start_paused = true)]
2898 async fn poll_waits_for_a_pending_eager_join_but_not_for_a_cooperative_one() {
2899 let broker = metadata_counting_broker(Arc::default()).await;
2900 let mut actual = Vec::new();
2901 let mut wanted = Vec::new();
2902 for (name, assignor, owns_partitions, pending, join_completes_after, expected) in [
2903 ("no join", Assignor::Range, true, false, None, (true, 0)),
2904 (
2905 "eager join completes in time",
2906 Assignor::Range,
2907 true,
2908 true,
2909 Some(100),
2910 (true, 100),
2911 ),
2912 (
2913 "eager join does not complete in time",
2914 Assignor::Range,
2915 true,
2916 true,
2917 None,
2918 (false, 500),
2919 ),
2920 (
2921 "cooperative join with owned partitions",
2922 Assignor::CooperativeSticky,
2923 true,
2924 true,
2925 None,
2926 (true, 0),
2927 ),
2928 (
2929 "cooperative join without partitions",
2930 Assignor::CooperativeSticky,
2931 false,
2932 true,
2933 Some(200),
2934 (true, 200),
2935 ),
2936 ] {
2937 let mut consumer = consumer_on(&broker).await;
2938 consumer.rebalance_protocol =
2939 assignor::rebalance_protocol_of(&[assignor]).expect("assignor protocol");
2940 if !owns_partitions {
2941 consumer.assigned.lock().await.clear();
2942 }
2943 let pending_sender = tokio::sync::watch::Sender::new(pending);
2944 consumer.rebalance_pending = pending_sender.subscribe();
2945 let completion = join_completes_after.map(|millis| {
2946 let pending_sender = pending_sender.clone();
2947 tokio::spawn(async move {
2948 tokio::time::sleep(Duration::from_millis(millis)).await;
2949 pending_sender.send_replace(false);
2950 })
2951 });
2952 let start = tokio::time::Instant::now();
2953 let joined = consumer
2954 .wait_for_rebalance(start + Duration::from_millis(500))
2955 .await
2956 .expect("no listener");
2957 let elapsed_ms = u64::try_from(start.elapsed().as_millis()).expect("millis");
2958 if let Some(completion) = completion {
2959 completion.abort();
2960 }
2961 actual.push((name, (joined, elapsed_ms)));
2962 wanted.push((name, expected));
2963 }
2964 broker.stop();
2965 assert2::assert!(actual == wanted);
2966 }
2967
2968 #[derive(Debug, PartialEq)]
2971 struct CappedPoll {
2972 records: usize,
2973 first_offset: Option<i64>,
2974 last_offset: Option<i64>,
2975 next_offset: Option<i64>,
2976 }
2977
2978 #[tokio::test]
2983 async fn poll_returns_at_most_max_poll_records_and_keeps_the_rest_without_a_fetch() {
2984 let fetches = Arc::new(AtomicUsize::new(0));
2985 let fetches_in_mock = Arc::clone(&fetches);
2986 let broker = MockBroker::start(move |api_key, version, _corr_id, _body| {
2987 if api_key == api_versions_request::API_KEY {
2988 let versions = mock_test_support::api_versions_response_ranges([
2989 (api_versions_request::API_KEY, 0, 3),
2990 (metadata_request::API_KEY, 0, 8),
2991 (fetch_request::API_KEY, 4, 11),
2992 ]);
2993 return Some(encode(&versions, 0));
2994 }
2995 if api_key == metadata_request::API_KEY {
2996 return Some(encode(&MetadataResponse::default(), version));
2997 }
2998 if api_key != fetch_request::API_KEY {
2999 return None;
3000 }
3001 fetches_in_mock.fetch_add(1, Ordering::SeqCst);
3002 let batch = krabka_protocol::records::RecordBatch {
3003 base_offset: 5,
3004 last_offset_delta: 1199,
3005 records: (0..1200)
3006 .map(|offset_delta| krabka_protocol::records::Record {
3007 offset_delta,
3008 value: Some(bytes::Bytes::from_static(b"v")),
3009 ..Default::default()
3010 })
3011 .collect(),
3012 ..Default::default()
3013 };
3014 let mut response = fetch_response(0);
3015 response.responses[0].topic = "orders".into();
3016 response.responses[0].partitions[0].high_watermark = 1205;
3017 response.responses[0].partitions[0].records = Some(RecordsPayload::V2(vec![batch]));
3018 Some(encode(&response, version))
3019 })
3020 .await;
3021 let mut consumer = consumer_on(&broker).await;
3022
3023 let mut polls = Vec::new();
3024 for _ in 0..3 {
3025 let records = consumer.poll(secs(5)).await.expect("poll");
3026 polls.push(CappedPoll {
3027 records: records.len(),
3028 first_offset: records.first().map(|record| record.offset),
3029 last_offset: records.last().map(|record| record.offset),
3030 next_offset: consumer
3031 .next_offsets
3032 .lock()
3033 .await
3034 .get(&("orders".to_string(), 0))
3035 .copied(),
3036 });
3037 }
3038
3039 broker.stop();
3040 assert2::assert!(
3041 (polls, fetches.load(Ordering::SeqCst))
3042 == (
3043 vec![
3044 CappedPoll {
3045 records: 500,
3046 first_offset: Some(5),
3047 last_offset: Some(504),
3048 next_offset: Some(505),
3049 },
3050 CappedPoll {
3051 records: 500,
3052 first_offset: Some(505),
3053 last_offset: Some(1004),
3054 next_offset: Some(1005),
3055 },
3056 CappedPoll {
3057 records: 200,
3058 first_offset: Some(1005),
3059 last_offset: Some(1204),
3060 next_offset: Some(1205),
3061 },
3062 ],
3063 1,
3064 )
3065 );
3066 }
3067
3068 #[derive(Clone, Copy, Debug)]
3070 enum BetweenPolls {
3071 Nothing,
3072 UpdateRequested,
3073 MaxAgeElapsed,
3074 LeaderUnknown,
3075 }
3076
3077 #[tokio::test]
3082 async fn fetch_positions_refresh_metadata_only_when_it_is_due() {
3083 let cases = [
3085 (BetweenPolls::Nothing, 0),
3086 (BetweenPolls::UpdateRequested, 1),
3087 (BetweenPolls::MaxAgeElapsed, 1),
3088 (BetweenPolls::LeaderUnknown, 1),
3089 ];
3090 for (between, expected) in cases {
3091 let metadata_requests = Arc::new(AtomicUsize::new(0));
3092 let broker = metadata_counting_broker(Arc::clone(&metadata_requests)).await;
3093 let mut consumer = consumer_on(&broker).await;
3094 let key = ("orders".to_string(), 0);
3095 consumer
3096 .positions
3097 .lock()
3098 .await
3099 .entry(key.clone())
3100 .or_default()
3101 .leader_id = 1;
3102 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
3103 check!(
3104 consumer
3105 .update_fetch_positions(None, Some(deadline))
3106 .await
3107 .unwrap()
3108 );
3109 check!(
3110 metadata_requests.swap(0, Ordering::SeqCst) == 1,
3111 "{between:?}: first poll"
3112 );
3113 match between {
3114 BetweenPolls::Nothing => {}
3115 BetweenPolls::UpdateRequested => consumer.client.metadata_topics().request_update(),
3116 BetweenPolls::MaxAgeElapsed => consumer.metadata_max_age = Time::ZERO,
3117 BetweenPolls::LeaderUnknown => {
3118 consumer
3119 .positions
3120 .lock()
3121 .await
3122 .entry(key)
3123 .or_default()
3124 .leader_id = UNKNOWN_LEADER_ID;
3125 }
3126 }
3127 check!(
3128 consumer
3129 .update_fetch_positions(None, Some(deadline))
3130 .await
3131 .unwrap()
3132 );
3133 check!(
3134 metadata_requests.load(Ordering::SeqCst) == expected,
3135 "{between:?}"
3136 );
3137 broker.stop();
3138 }
3139 }
3140
3141 #[tokio::test]
3147 async fn fetch_partition_errors_refresh_metadata_or_fail_the_poll() {
3148 for (name, error_code, expected) in [
3149 (
3150 "unknown topic or partition",
3151 3,
3152 PollOutcome {
3153 result: Ok(0),
3154 metadata_requests: 0,
3155 metadata_update_requested: true,
3156 next_offset: Some(5),
3157 },
3158 ),
3159 (
3160 "not leader or follower",
3161 6,
3162 PollOutcome {
3163 result: Ok(0),
3164 metadata_requests: 0,
3165 metadata_update_requested: true,
3166 next_offset: Some(5),
3167 },
3168 ),
3169 (
3170 "unknown topic id",
3171 100,
3172 PollOutcome {
3173 result: Ok(0),
3174 metadata_requests: 0,
3175 metadata_update_requested: true,
3176 next_offset: Some(5),
3177 },
3178 ),
3179 (
3180 "topic authorization failed",
3181 29,
3182 PollOutcome {
3183 result: Err(Some(29)),
3184 metadata_requests: 0,
3185 metadata_update_requested: false,
3186 next_offset: Some(5),
3187 },
3188 ),
3189 ] {
3190 let metadata_requests = Arc::new(AtomicUsize::new(0));
3191 let broker = metadata_counting_broker(Arc::clone(&metadata_requests)).await;
3192 let mut consumer = consumer_on(&broker).await;
3193 let topic_ids = consumer.topic_ids.lock().await.clone();
3194
3195 let result = consumer
3196 .process_fetch_responses(vec![fetch_response(error_code)], &topic_ids)
3197 .await
3198 .map(|()| consumer.fetch_buffer.len())
3199 .map_err(|error| match error {
3200 ConsumerError::Server(code) => Some(code),
3201 _ => None,
3202 });
3203 let outcome = PollOutcome {
3204 result,
3205 metadata_requests: metadata_requests.load(Ordering::SeqCst),
3206 metadata_update_requested: consumer.client.metadata_topics().update_requested(),
3207 next_offset: consumer
3208 .next_offsets
3209 .lock()
3210 .await
3211 .get(&("orders".to_string(), 0))
3212 .copied(),
3213 };
3214
3215 broker.stop();
3216 check!(outcome == expected, "case {name}");
3217 }
3218 }
3219
3220 #[derive(Clone, Copy, Debug)]
3223 enum PollReset {
3224 None,
3226 LatestOutOfRange,
3228 EarliestOutOfRange,
3232 FetchTruncation,
3234 ValidationTruncation,
3236 }
3237
3238 #[tokio::test]
3243 async fn a_reset_in_poll_moves_the_position_of_the_commit_before_join() {
3244 let orders_0 = ("orders".to_string(), 0);
3245 let cases = [
3246 (
3247 "no reset: the position at the start of poll",
3248 PollReset::None,
3249 maplit::hashmap! {orders_0.clone() => (5, -1)},
3250 ),
3251 (
3252 "out of range under latest: no position until ListOffsets",
3253 PollReset::LatestOutOfRange,
3254 HashMap::new(),
3255 ),
3256 (
3257 "out of range under earliest: the log start",
3258 PollReset::EarliestOutOfRange,
3259 maplit::hashmap! {orders_0.clone() => (2, -1)},
3260 ),
3261 (
3262 "truncation in a Fetch response: the end offset of the diverging epoch",
3263 PollReset::FetchTruncation,
3264 maplit::hashmap! {orders_0.clone() => (3, -1)},
3265 ),
3266 (
3267 "truncation from validation: the safe offset",
3268 PollReset::ValidationTruncation,
3269 maplit::hashmap! {orders_0.clone() => (3, -1)},
3270 ),
3271 ];
3272 let mut actual = Vec::new();
3273 let mut wanted = Vec::new();
3274 for (name, reset, expected) in cases {
3275 let leader_port = Arc::new(AtomicU16::new(0));
3276 let broker = list_offsets_broker(
3277 "leader",
3278 Arc::clone(&leader_port),
3279 (0, 2),
3280 list_offsets_request::MAX_VERSION,
3281 Arc::default(),
3282 Arc::default(),
3283 )
3284 .await;
3285 leader_port.store(broker.addr.port(), Ordering::SeqCst);
3286 let mut consumer = consumer_on(&broker).await;
3287 let auto_commit = crate::commit::AutoCommit::new(std::time::Duration::from_secs(5));
3288 consumer.auto_commit = Some(auto_commit.clone());
3289 consumer.maybe_auto_commit_async().await;
3292 let topic_ids = consumer.topic_ids.lock().await.clone();
3293
3294 let result = match reset {
3295 PollReset::None => {
3296 consumer
3297 .process_fetch_responses(vec![fetch_response(0)], &topic_ids)
3298 .await
3299 }
3300 PollReset::LatestOutOfRange => {
3301 consumer.auto_offset_reset = AutoOffsetReset::Latest;
3302 consumer
3303 .process_fetch_responses(vec![fetch_response(1)], &topic_ids)
3304 .await
3305 }
3306 PollReset::EarliestOutOfRange => {
3307 consumer.auto_offset_reset = AutoOffsetReset::Earliest;
3308 consumer
3309 .refresh_leader_epochs()
3310 .await
3311 .expect("metadata for the leader route");
3312 consumer
3313 .process_fetch_responses(vec![fetch_response(1)], &topic_ids)
3314 .await
3315 }
3316 PollReset::FetchTruncation => {
3317 let mut response = fetch_response(0);
3318 response.responses[0].partitions[0].diverging_epoch = EpochEndOffset {
3319 epoch: 2,
3320 end_offset: 3,
3321 ..Default::default()
3322 };
3323 consumer
3324 .process_fetch_responses(vec![response], &topic_ids)
3325 .await
3326 }
3327 PollReset::ValidationTruncation => {
3328 consumer
3329 .apply_truncation(&maplit::hashmap! {orders_0.clone() => 3})
3330 .await
3331 }
3332 };
3333 result.expect("the reset succeeds");
3334 let offsets = auto_commit
3335 .polled_offsets(&maplit::hashmap! {orders_0.clone() => 1})
3336 .await;
3337
3338 broker.stop();
3339 actual.push((name, offsets));
3340 wanted.push((name, expected));
3341 }
3342 assert2::assert!(actual == wanted);
3343 }
3344
3345 const LEADER_ID: i32 = 1;
3347 const LEADER_EPOCH: i32 = 7;
3349
3350 type SentListOffsets = Arc<std::sync::Mutex<Vec<(&'static str, i16, ListOffsetsRequest)>>>;
3353
3354 fn list_offsets_answer_bytes(version: i16, error_code: i16, offset: i64) -> Vec<u8> {
3357 use krabka_protocol::owned::{
3358 list_offsets_request::FLEXIBLE_MIN,
3359 list_offsets_response::{ListOffsetsPartitionResponse, ListOffsetsTopicResponse},
3360 };
3361 let answer = ListOffsetsResponse {
3362 topics: vec![ListOffsetsTopicResponse {
3363 name: "orders".into(),
3364 partitions: vec![ListOffsetsPartitionResponse {
3365 partition_index: 0,
3366 error_code,
3367 offset,
3368 ..Default::default()
3369 }],
3370 ..Default::default()
3371 }],
3372 ..Default::default()
3373 };
3374 let mut bytes = Vec::new();
3375 if version >= FLEXIBLE_MIN {
3376 bytes.push(0);
3378 }
3379 bytes.extend(encode(&answer, version));
3380 bytes
3381 }
3382
3383 async fn list_offsets_broker(
3390 name: &'static str,
3391 leader_port: Arc<AtomicU16>,
3392 row: (i16, i64),
3393 list_offsets_max_version: i16,
3394 sent: SentListOffsets,
3395 metadata_requests: Arc<AtomicUsize>,
3396 ) -> MockBroker {
3397 use krabka_protocol::owned::list_offsets_request;
3398 MockBroker::start(move |api_key, version, _corr_id, mut body| {
3399 if api_key == api_versions_request::API_KEY {
3400 let versions = mock_test_support::api_versions_response([
3401 (api_versions_request::API_KEY, 3),
3402 (metadata_request::API_KEY, 8),
3403 (list_offsets_request::API_KEY, list_offsets_max_version),
3404 ]);
3405 return Some(encode(&versions, 0));
3406 }
3407 if api_key == metadata_request::API_KEY {
3408 metadata_requests.fetch_add(1, Ordering::SeqCst);
3409 let metadata = test_support::OrdersMetadata::builder()
3410 .topic_id(TOPIC_ID.0)
3411 .leader_id(LEADER_ID)
3412 .leader_epoch(LEADER_EPOCH)
3413 .brokers(vec![
3414 krabka_protocol::owned::metadata_response::MetadataResponseBroker {
3415 node_id: LEADER_ID,
3416 host: "127.0.0.1".into(),
3417 port: i32::from(leader_port.load(Ordering::SeqCst)),
3418 ..Default::default()
3419 },
3420 ])
3421 .build();
3422 return Some(encode(&metadata, version));
3423 }
3424 if api_key == list_offsets_request::API_KEY {
3425 let request = test_support::decode_list_offsets_request(&mut body, version);
3426 sent.lock()
3427 .expect("sent lock")
3428 .push((name, version, request));
3429 return Some(list_offsets_answer_bytes(version, row.0, row.1));
3430 }
3431 None
3432 })
3433 .await
3434 }
3435
3436 fn expected_list_offsets(isolation_level: i8, timestamp: i64) -> ListOffsetsRequest {
3438 ListOffsetsRequest {
3439 replica_id: -1,
3440 isolation_level,
3441 topics: vec![ListOffsetsTopic {
3442 name: "orders".into(),
3443 partitions: vec![ListOffsetsPartition {
3444 partition_index: 0,
3445 current_leader_epoch: LEADER_EPOCH,
3446 timestamp,
3447 ..Default::default()
3448 }],
3449 ..Default::default()
3450 }],
3451 ..Default::default()
3452 }
3453 }
3454
3455 #[derive(Debug, PartialEq)]
3457 enum PollFailure {
3458 TopicAuthorizationFailed(std::collections::BTreeSet<String>),
3460 IncompatibleVersion(i16, i16, i16, i16, i16),
3463 Other,
3465 }
3466
3467 impl From<ConsumerError> for PollFailure {
3468 fn from(error: ConsumerError) -> Self {
3469 match error {
3470 ConsumerError::TopicAuthorizationFailed(topics) => {
3471 Self::TopicAuthorizationFailed(topics)
3472 }
3473 ConsumerError::Client(krabka_client_core::ClientError::IncompatibleVersion {
3474 api_key,
3475 broker_min,
3476 broker_max,
3477 client_min,
3478 client_max,
3479 }) => Self::IncompatibleVersion(
3480 api_key, broker_min, broker_max, client_min, client_max,
3481 ),
3482 _ => Self::Other,
3483 }
3484 }
3485 }
3486
3487 #[derive(Debug, PartialEq)]
3489 struct ListOffsetsOutcome {
3490 requests: Vec<(&'static str, i16, ListOffsetsRequest)>,
3493 result: Result<(), PollFailure>,
3494 metadata_requests: usize,
3496 metadata_update_requested: bool,
3498 next_offset: Option<i64>,
3499 }
3500
3501 #[derive(Clone, Copy, Debug)]
3503 enum Reset {
3504 Latest,
3506 ByOneHour,
3509 EarliestOutOfRange,
3512 SeekToBeginning,
3514 SeekToEnd,
3516 }
3517
3518 #[derive(Clone, Copy, Debug)]
3520 struct Exchange {
3521 isolation: IsolationLevel,
3522 reset: Reset,
3523 row: (i16, i64),
3525 list_offsets_max_version: i16,
3527 primed: bool,
3531 }
3532
3533 async fn run_list_offsets_exchange(exchange: Exchange) -> ListOffsetsOutcome {
3539 let sent: SentListOffsets = Arc::default();
3540 let metadata_requests = Arc::new(AtomicUsize::new(0));
3541 let leader_port = Arc::new(AtomicU16::new(0));
3542 let leader = list_offsets_broker(
3543 "leader",
3544 Arc::clone(&leader_port),
3545 exchange.row,
3546 exchange.list_offsets_max_version,
3547 Arc::clone(&sent),
3548 Arc::clone(&metadata_requests),
3549 )
3550 .await;
3551 leader_port.store(leader.addr.port(), Ordering::SeqCst);
3552 let bootstrap = list_offsets_broker(
3553 "bootstrap",
3554 Arc::clone(&leader_port),
3555 (6, -1),
3556 exchange.list_offsets_max_version,
3557 Arc::clone(&sent),
3558 Arc::clone(&metadata_requests),
3559 )
3560 .await;
3561 let mut consumer = consumer_on(&bootstrap).await;
3562 consumer.isolation_level = exchange.isolation;
3563 if exchange.primed {
3564 consumer
3565 .refresh_leader_epochs()
3566 .await
3567 .expect("setup metadata");
3568 metadata_requests.store(0, Ordering::SeqCst);
3569 }
3570
3571 let result = match exchange.reset {
3572 Reset::Latest => {
3573 consumer.auto_offset_reset = AutoOffsetReset::Latest;
3574 consumer
3575 .next_offsets
3576 .lock()
3577 .await
3578 .insert(("orders".into(), 0), LATEST_SENTINEL);
3579 consumer
3580 .prepare_poll(tokio::time::Instant::now() + Duration::from_secs(5))
3581 .await
3582 .map(|_| ())
3583 }
3584 Reset::ByOneHour => {
3585 consumer.auto_offset_reset =
3586 AutoOffsetReset::ByDuration(std::time::Duration::from_hours(1));
3587 consumer
3588 .next_offsets
3589 .lock()
3590 .await
3591 .insert(("orders".into(), 0), LATEST_SENTINEL);
3592 consumer
3593 .prepare_poll(tokio::time::Instant::now() + Duration::from_secs(5))
3594 .await
3595 .map(|_| ())
3596 }
3597 Reset::EarliestOutOfRange => {
3598 consumer.auto_offset_reset = AutoOffsetReset::Earliest;
3599 let topic_ids = consumer.topic_ids.lock().await.clone();
3600 consumer
3601 .process_fetch_responses(vec![fetch_response(1)], &topic_ids)
3602 .await
3603 }
3604 Reset::SeekToBeginning | Reset::SeekToEnd => {
3605 consumer.auto_offset_reset =
3608 AutoOffsetReset::ByDuration(std::time::Duration::from_hours(1));
3609 consumer.default_api_timeout = krabka_units::millis(300);
3610 let reset = if matches!(exchange.reset, Reset::SeekToBeginning) {
3611 consumer.seek_to_beginning(&[]).await
3612 } else {
3613 consumer.seek_to_end(&[]).await
3614 };
3615 match reset {
3616 Ok(()) => consumer.position("orders", 0).await.map(|_| ()),
3617 Err(error) => Err(error),
3618 }
3619 }
3620 };
3621 let requests = sent.lock().expect("sent lock").clone();
3622 let outcome = ListOffsetsOutcome {
3623 requests,
3624 result: result.map_err(PollFailure::from),
3625 metadata_requests: metadata_requests.load(Ordering::SeqCst),
3626 metadata_update_requested: consumer.client.metadata_topics().update_requested(),
3627 next_offset: consumer
3628 .next_offsets
3629 .lock()
3630 .await
3631 .get(&("orders".to_string(), 0))
3632 .copied(),
3633 };
3634
3635 bootstrap.stop();
3636 leader.stop();
3637 outcome
3638 }
3639
3640 #[tokio::test]
3648 async fn list_offsets_goes_to_the_leader_with_the_isolation_level_and_honours_row_errors() {
3649 let max = list_offsets_request::MAX_VERSION;
3650 let latest_request =
3651 |isolation| vec![("leader", max, expected_list_offsets(isolation, -1))];
3652 let latest = |isolation, row| Exchange {
3653 isolation,
3654 reset: Reset::Latest,
3655 row,
3656 list_offsets_max_version: max,
3657 primed: true,
3658 };
3659 let earliest = |isolation, row| Exchange {
3660 reset: Reset::EarliestOutOfRange,
3661 ..latest(isolation, row)
3662 };
3663 for (name, exchange, expected) in [
3664 (
3665 "read committed latest reads the last stable offset",
3666 latest(IsolationLevel::ReadCommitted, (0, 40)),
3667 ListOffsetsOutcome {
3668 requests: latest_request(1),
3669 result: Ok(()),
3670 metadata_requests: 1,
3671 metadata_update_requested: false,
3672 next_offset: Some(40),
3673 },
3674 ),
3675 (
3676 "read uncommitted latest reads the high watermark",
3677 latest(IsolationLevel::ReadUncommitted, (0, 50)),
3678 ListOffsetsOutcome {
3679 requests: latest_request(0),
3680 result: Ok(()),
3681 metadata_requests: 1,
3682 metadata_update_requested: false,
3683 next_offset: Some(50),
3684 },
3685 ),
3686 (
3687 "not leader or follower keeps the sentinel and refreshes metadata",
3688 latest(IsolationLevel::ReadCommitted, (6, -1)),
3689 ListOffsetsOutcome {
3690 requests: latest_request(1),
3691 result: Ok(()),
3692 metadata_requests: 1,
3693 metadata_update_requested: true,
3694 next_offset: Some(LATEST_SENTINEL),
3695 },
3696 ),
3697 (
3698 "topic authorization failed fails the poll",
3699 latest(IsolationLevel::ReadCommitted, (29, -1)),
3700 ListOffsetsOutcome {
3701 requests: latest_request(1),
3702 result: Err(PollFailure::TopicAuthorizationFailed(
3703 maplit::btreeset! {"orders".to_string()},
3704 )),
3705 metadata_requests: 1,
3706 metadata_update_requested: false,
3707 next_offset: Some(LATEST_SENTINEL),
3708 },
3709 ),
3710 (
3711 "seek to beginning reads the log start",
3712 Exchange {
3713 reset: Reset::SeekToBeginning,
3714 ..latest(IsolationLevel::ReadUncommitted, (0, 7))
3715 },
3716 ListOffsetsOutcome {
3717 requests: vec![("leader", max, expected_list_offsets(0, -2))],
3718 result: Ok(()),
3719 metadata_requests: 0,
3720 metadata_update_requested: false,
3721 next_offset: Some(7),
3722 },
3723 ),
3724 (
3725 "seek to end reads the last stable offset",
3726 Exchange {
3727 reset: Reset::SeekToEnd,
3728 ..latest(IsolationLevel::ReadCommitted, (0, 40))
3729 },
3730 ListOffsetsOutcome {
3731 requests: latest_request(1),
3732 result: Ok(()),
3733 metadata_requests: 0,
3734 metadata_update_requested: false,
3735 next_offset: Some(40),
3736 },
3737 ),
3738 (
3739 "earliest after out of range reads the log start",
3740 earliest(IsolationLevel::ReadUncommitted, (0, 7)),
3741 ListOffsetsOutcome {
3742 requests: vec![("leader", max, expected_list_offsets(0, -2))],
3743 result: Ok(()),
3744 metadata_requests: 0,
3745 metadata_update_requested: false,
3746 next_offset: Some(7),
3747 },
3748 ),
3749 (
3750 "earliest with a retriable row keeps the position and refreshes metadata",
3751 earliest(IsolationLevel::ReadCommitted, (74, -1)),
3752 ListOffsetsOutcome {
3753 requests: vec![("leader", max, expected_list_offsets(1, -2))],
3754 result: Ok(()),
3755 metadata_requests: 0,
3756 metadata_update_requested: true,
3757 next_offset: Some(5),
3758 },
3759 ),
3760 ] {
3761 let outcome = run_list_offsets_exchange(exchange).await;
3762 check!(outcome == expected, "case {name}");
3763 }
3764 }
3765
3766 #[tokio::test]
3769 async fn by_duration_resets_with_the_timestamp_of_now_minus_the_duration() {
3770 let before = unix_now_ms() - 3_600_000;
3771 let mut outcome = run_list_offsets_exchange(Exchange {
3772 isolation: IsolationLevel::ReadUncommitted,
3773 reset: Reset::ByOneHour,
3774 row: (0, 40),
3775 list_offsets_max_version: list_offsets_request::MAX_VERSION,
3776 primed: true,
3777 })
3778 .await;
3779 let after = unix_now_ms() - 3_600_000;
3780 let mut in_range = Vec::new();
3783 for (_, _, request) in &mut outcome.requests {
3784 for topic in &mut request.topics {
3785 for partition in &mut topic.partitions {
3786 in_range.push((before..=after).contains(&partition.timestamp));
3787 partition.timestamp = before;
3788 }
3789 }
3790 }
3791
3792 assert2::assert!(
3793 (in_range, outcome)
3794 == (
3795 vec![true],
3796 ListOffsetsOutcome {
3797 requests: vec![(
3798 "leader",
3799 list_offsets_request::MAX_VERSION,
3800 expected_list_offsets(0, before),
3801 )],
3802 result: Ok(()),
3803 metadata_requests: 1,
3804 metadata_update_requested: false,
3805 next_offset: Some(40),
3806 }
3807 )
3808 );
3809 }
3810
3811 #[test]
3812 fn reset_timestamp_follows_the_reset_policy() {
3813 let now = 1_700_000_000_000;
3814 let actual = [
3815 AutoOffsetReset::Latest,
3816 AutoOffsetReset::None,
3817 AutoOffsetReset::ByDuration(std::time::Duration::from_hours(1)),
3818 AutoOffsetReset::ByDuration(std::time::Duration::from_millis(1500)),
3819 ]
3820 .map(|policy| reset_timestamp(policy, now));
3821 assert2::assert!(actual == [-1, -1, now - 3_600_000, now - 1500]);
3822 }
3823
3824 #[tokio::test]
3829 async fn first_poll_sends_list_offsets_to_the_leader_of_a_new_partition() {
3830 let outcome = run_list_offsets_exchange(Exchange {
3831 isolation: IsolationLevel::ReadCommitted,
3832 reset: Reset::Latest,
3833 row: (0, 40),
3834 list_offsets_max_version: list_offsets_request::MAX_VERSION,
3835 primed: false,
3836 })
3837 .await;
3838
3839 assert2::assert!(
3840 outcome
3841 == ListOffsetsOutcome {
3842 requests: vec![(
3843 "leader",
3844 list_offsets_request::MAX_VERSION,
3845 expected_list_offsets(1, -1),
3846 )],
3847 result: Ok(()),
3848 metadata_requests: 1,
3849 metadata_update_requested: false,
3850 next_offset: Some(40),
3851 }
3852 );
3853 }
3854
3855 #[tokio::test]
3862 async fn read_committed_list_offsets_needs_version_2() {
3863 let at = |isolation, list_offsets_max_version| Exchange {
3864 isolation,
3865 reset: Reset::Latest,
3866 row: (0, 40),
3867 list_offsets_max_version,
3868 primed: true,
3869 };
3870 let without_epoch = |isolation_level| {
3873 let mut request = expected_list_offsets(isolation_level, -1);
3874 request.topics[0].partitions[0].current_leader_epoch =
3875 ListOffsetsPartition::default().current_leader_epoch;
3876 request
3877 };
3878 for (name, exchange, expected) in [
3879 (
3880 "read committed on a version 1 broker fails the poll",
3881 at(IsolationLevel::ReadCommitted, 1),
3882 ListOffsetsOutcome {
3883 requests: vec![],
3884 result: Err(PollFailure::IncompatibleVersion(
3885 list_offsets_request::API_KEY,
3886 0,
3887 1,
3888 2,
3889 list_offsets_request::MAX_VERSION,
3890 )),
3891 metadata_requests: 1,
3892 metadata_update_requested: false,
3893 next_offset: Some(LATEST_SENTINEL),
3894 },
3895 ),
3896 (
3897 "read committed on a version 2 broker sends version 2",
3898 at(IsolationLevel::ReadCommitted, 2),
3899 ListOffsetsOutcome {
3900 requests: vec![("leader", 2, without_epoch(1))],
3901 result: Ok(()),
3902 metadata_requests: 1,
3903 metadata_update_requested: false,
3904 next_offset: Some(40),
3905 },
3906 ),
3907 (
3908 "read uncommitted on a version 1 broker sends version 1",
3909 at(IsolationLevel::ReadUncommitted, 1),
3910 ListOffsetsOutcome {
3911 requests: vec![("leader", 1, without_epoch(0))],
3912 result: Ok(()),
3913 metadata_requests: 1,
3914 metadata_update_requested: false,
3915 next_offset: Some(40),
3916 },
3917 ),
3918 ] {
3919 let outcome = run_list_offsets_exchange(exchange).await;
3920 check!(outcome == expected, "case {name}");
3921 }
3922 }
3923
3924 #[tokio::test]
3928 async fn fetched_records_take_the_batch_timestamp_type() {
3929 let broker = metadata_counting_broker(Arc::default()).await;
3930 let consumer = consumer_on(&broker).await;
3931 let key: (String, i32) = ("orders".into(), 0);
3932
3933 for case in ×tamp_cases::CASES {
3934 let part = PartitionData {
3935 partition_index: 0,
3936 records: Some(timestamp_cases::batch(case).into()),
3937 ..Default::default()
3938 };
3939 let offsets = maplit::hashmap! {key.clone() => timestamp_cases::BASE_OFFSET};
3940
3941 let records: Vec<ConsumerRecord> = consumer
3942 .process_partition_records(&offsets, &key, "orders", &part)
3943 .map(|partition| partition.records.into_iter().collect())
3944 .unwrap_or_default();
3945
3946 let expected: Vec<ConsumerRecord> = case
3947 .expected
3948 .iter()
3949 .zip(timestamp_cases::BASE_OFFSET..)
3950 .zip(timestamp_cases::VALUES)
3951 .map(
3952 |(((timestamp, timestamp_type), offset), value)| ConsumerRecord {
3953 topic: "orders".into(),
3954 partition: 0,
3955 offset,
3956 leader_epoch: timestamp_cases::LEADER_EPOCH,
3957 timestamp: *timestamp,
3958 timestamp_type: *timestamp_type,
3959 key: None,
3960 value: Some(bytes::Bytes::from_static(value)),
3961 headers: Vec::new(),
3962 },
3963 )
3964 .collect();
3965 check!(records == expected, "case {}", case.name);
3966 }
3967 }
3968}
3969
3970#[cfg(test)]
3973pub(crate) mod timestamp_cases {
3974 use bytes::Bytes;
3975 use consumer::TimestampType;
3976 use krabka_protocol::records::{Attributes, Record, RecordBatch};
3977
3978 use crate::consumer;
3979
3980 pub(crate) const BASE_OFFSET: i64 = 5;
3981 pub(crate) const LEADER_EPOCH: i32 = 3;
3982 pub(crate) const VALUES: [&[u8]; 2] = [b"v0", b"v1"];
3983
3984 pub(crate) struct TimestampCase {
3986 pub(crate) name: &'static str,
3987 pub(crate) attributes: Attributes,
3988 pub(crate) base_timestamp: i64,
3989 pub(crate) max_timestamp: i64,
3990 pub(crate) timestamp_deltas: [i64; 2],
3991 pub(crate) expected: [(i64, TimestampType); 2],
3992 }
3993
3994 pub(crate) const CASES: [TimestampCase; 3] = [
3995 TimestampCase {
3996 name: "create time adds the delta to the base timestamp",
3997 attributes: Attributes(0),
3998 base_timestamp: 1000,
3999 max_timestamp: 1005,
4000 timestamp_deltas: [0, 5],
4001 expected: [
4002 (1000, TimestampType::CreateTime),
4003 (1005, TimestampType::CreateTime),
4004 ],
4005 },
4006 TimestampCase {
4007 name: "log append time uses the max timestamp",
4008 attributes: Attributes(Attributes::TIMESTAMP_TYPE_BIT),
4009 base_timestamp: 1000,
4010 max_timestamp: 9000,
4011 timestamp_deltas: [0, 5],
4012 expected: [
4013 (9000, TimestampType::LogAppendTime),
4014 (9000, TimestampType::LogAppendTime),
4015 ],
4016 },
4017 TimestampCase {
4018 name: "log append time in a transactional batch uses the max timestamp",
4019 attributes: Attributes(Attributes::TIMESTAMP_TYPE_BIT | Attributes::TRANSACTIONAL_BIT),
4020 base_timestamp: 1000,
4021 max_timestamp: 9000,
4022 timestamp_deltas: [0, 5],
4023 expected: [
4024 (9000, TimestampType::LogAppendTime),
4025 (9000, TimestampType::LogAppendTime),
4026 ],
4027 },
4028 ];
4029
4030 pub(crate) fn batch(case: &TimestampCase) -> RecordBatch {
4031 RecordBatch {
4032 base_offset: BASE_OFFSET,
4033 partition_leader_epoch: LEADER_EPOCH,
4034 attributes: case.attributes,
4035 last_offset_delta: 1,
4036 base_timestamp: case.base_timestamp,
4037 max_timestamp: case.max_timestamp,
4038 records: case
4039 .timestamp_deltas
4040 .iter()
4041 .zip(0..)
4042 .zip(VALUES)
4043 .map(|((timestamp_delta, offset_delta), value)| Record {
4044 attributes: 0,
4045 timestamp_delta: *timestamp_delta,
4046 offset_delta,
4047 key: None,
4048 value: Some(Bytes::from_static(value)),
4049 headers: Vec::new(),
4050 })
4051 .collect(),
4052 ..Default::default()
4053 }
4054 }
4055}
4056
4057#[cfg(test)]
4058mod fetch_path_tests {
4059 use std::{
4060 sync::{
4061 Arc,
4062 atomic::{AtomicU16, Ordering},
4063 },
4064 time::Duration,
4065 };
4066
4067 use bytes::Buf as _;
4068 use krabka_client_core::{Client, MockBroker};
4069 use krabka_protocol::{
4070 Decode as _, Encode,
4071 owned::{
4072 api_versions_request, fetch_request,
4073 fetch_response::{FetchableTopicResponse, PartitionData},
4074 list_offsets_request,
4075 list_offsets_response::ListOffsetsResponse,
4076 metadata_request,
4077 metadata_response::{
4078 MetadataResponse, MetadataResponseBroker, MetadataResponsePartition,
4079 MetadataResponseTopic,
4080 },
4081 },
4082 primitives::uuid::Uuid as WireUuid,
4083 };
4084 use krabka_units::{millis, secs};
4085 use test_support::RecordingListener;
4086
4087 use super::*;
4088
4089 const TOPIC_ID: WireUuid = WireUuid([7; 16]);
4090
4091 #[derive(Clone, Copy, Debug)]
4093 enum FetchAnswer {
4094 Silent,
4096 Respond { error_code: i16, session_id: i32 },
4099 Partition {
4101 error_code: i16,
4102 preferred_read_replica: i32,
4103 },
4104 }
4105
4106 type SentFetches = Arc<std::sync::Mutex<Vec<(i32, FetchRequest)>>>;
4109
4110 fn encode(response: &impl Encode, version: i16) -> Vec<u8> {
4111 let mut buf = bytes::BytesMut::new();
4112 response.encode(&mut buf, version).expect("encode");
4113 buf.to_vec()
4114 }
4115
4116 async fn start_brokers(answers: &[Vec<FetchAnswer>], sent: &SentFetches) -> Vec<MockBroker> {
4120 let ports: Arc<Vec<AtomicU16>> =
4121 Arc::new(answers.iter().map(|_| AtomicU16::new(0)).collect());
4122 let mut brokers = Vec::new();
4123 for (index, answers) in answers.iter().enumerate() {
4124 let node_id = i32::try_from(index).expect("index") + 1;
4125 let ports_in_mock = Arc::clone(&ports);
4126 let sent = Arc::clone(sent);
4127 let answers = std::sync::Mutex::new(answers.clone().into_iter());
4128 let last = std::sync::Mutex::new(FetchAnswer::Silent);
4129 let broker = MockBroker::start(move |api_key, version, _corr_id, mut body| {
4130 if api_key == api_versions_request::API_KEY {
4131 let versions = mock_test_support::api_versions_response_ranges([
4132 (api_versions_request::API_KEY, 0, 3),
4133 (metadata_request::API_KEY, 0, 8),
4134 (fetch_request::API_KEY, 4, 11),
4135 (list_offsets_request::API_KEY, 1, 5),
4136 ]);
4137 return Some(encode(&versions, 0));
4138 }
4139 if api_key == metadata_request::API_KEY {
4140 let ports = &ports_in_mock;
4141 let count = ports.len();
4142 let metadata = MetadataResponse {
4143 brokers: (0..count)
4144 .map(|index| MetadataResponseBroker {
4145 node_id: i32::try_from(index).expect("index") + 1,
4146 host: "127.0.0.1".into(),
4147 port: i32::from(ports[index].load(Ordering::SeqCst)),
4148 ..Default::default()
4149 })
4150 .collect(),
4151 topics: vec![MetadataResponseTopic {
4152 name: Some("orders".into()),
4153 topic_id: TOPIC_ID,
4154 partitions: (0..count)
4155 .map(|index| {
4156 let index = i32::try_from(index).expect("index");
4157 MetadataResponsePartition {
4158 partition_index: index,
4159 leader_id: index + 1,
4160 leader_epoch: 1,
4161 ..Default::default()
4162 }
4163 })
4164 .collect(),
4165 ..Default::default()
4166 }],
4167 ..Default::default()
4168 };
4169 return Some(encode(&metadata, version));
4170 }
4171 if api_key == fetch_request::API_KEY {
4172 let client_id_len = usize::try_from(body.get_i16()).expect("client id");
4173 body.advance(client_id_len);
4174 let request = FetchRequest::decode(&mut body, version).expect("decode fetch");
4175 sent.lock().expect("sent lock").push((node_id, request));
4176 let mut last = last.lock().expect("last lock");
4177 if let Some(answer) = answers.lock().expect("answers lock").next() {
4178 *last = answer;
4179 }
4180 return match *last {
4181 FetchAnswer::Silent => None,
4182 FetchAnswer::Respond {
4183 error_code,
4184 session_id,
4185 } => Some(encode(
4186 &FetchResponse {
4187 error_code,
4188 session_id,
4189 ..Default::default()
4190 },
4191 version,
4192 )),
4193 FetchAnswer::Partition {
4194 error_code,
4195 preferred_read_replica,
4196 } => Some(encode(
4197 &FetchResponse {
4198 responses: vec![FetchableTopicResponse {
4199 topic: "orders".into(),
4200 partitions: vec![PartitionData {
4201 partition_index: 0,
4202 error_code,
4203 high_watermark: 5,
4204 last_stable_offset: 5,
4205 preferred_read_replica,
4206 ..Default::default()
4207 }],
4208 ..Default::default()
4209 }],
4210 ..Default::default()
4211 },
4212 version,
4213 )),
4214 };
4215 }
4216 if api_key == list_offsets_request::API_KEY {
4219 return Some(test_support::list_offsets_mock_response(&mut body, version));
4220 }
4221 None
4222 })
4223 .await;
4224 ports[index].store(broker.addr.port(), Ordering::SeqCst);
4225 brokers.push(broker);
4226 }
4227 brokers
4228 }
4229
4230 async fn consumer_on(brokers: &[MockBroker]) -> Consumer {
4233 let client = Client::builder()
4234 .bootstrap(brokers[0].addr.to_string())
4235 .request_timeout(secs(30))
4236 .build()
4237 .await
4238 .expect("client");
4239 let partitions: Vec<(String, i32)> = (0..brokers.len())
4240 .map(|index| ("orders".to_owned(), i32::try_from(index).expect("index")))
4241 .collect();
4242 let consumer = poll::partition_error_tests::consumer_with_client(client);
4243 *consumer.assigned.lock().await = partitions.clone();
4244 *consumer.next_offsets.lock().await =
4245 partitions.iter().map(|key| (key.clone(), 5)).collect();
4246 consumer.commit_identity.lock().await.ownership_ids =
4247 partitions.into_iter().zip(1..).collect();
4248 consumer
4249 }
4250
4251 struct TwoPartitionFetchSetup {
4252 brokers: Vec<MockBroker>,
4253 consumer: Consumer,
4254 sent: SentFetches,
4255 }
4256
4257 async fn two_partition_fetch_setup() -> TwoPartitionFetchSetup {
4258 let answer = FetchAnswer::Respond {
4259 error_code: 0,
4260 session_id: 0,
4261 };
4262 let sent = SentFetches::default();
4263 let brokers = start_brokers(&[vec![answer], vec![answer]], &sent).await;
4264 let consumer = consumer_on(&brokers).await;
4265 consumer.refresh_leader_epochs().await.expect("metadata");
4266 TwoPartitionFetchSetup {
4267 brokers,
4268 consumer,
4269 sent,
4270 }
4271 }
4272
4273 fn stop(brokers: Vec<MockBroker>) {
4274 for broker in brokers {
4275 broker.stop();
4276 }
4277 }
4278
4279 #[tokio::test]
4282 async fn fetch_waits_fetch_max_wait_and_not_the_poll_timeout() {
4283 let sent = SentFetches::default();
4284 let brokers = start_brokers(
4285 &[vec![FetchAnswer::Respond {
4286 error_code: 0,
4287 session_id: 0,
4288 }]],
4289 &sent,
4290 )
4291 .await;
4292 let mut consumer = consumer_on(&brokers).await;
4293 consumer.poll(secs(30)).await.expect("poll");
4294 let max_waits: Vec<i32> = sent
4295 .lock()
4296 .expect("sent lock")
4297 .iter()
4298 .map(|(_, request)| request.max_wait_ms)
4299 .collect();
4300 drop(consumer);
4301 stop(brokers);
4302 assert2::assert!(max_waits == vec![500]);
4303 }
4304
4305 #[tokio::test]
4308 async fn fetches_go_to_all_leaders_at_once_and_poll_returns_by_its_timeout() {
4309 let sent = SentFetches::default();
4310 let silent = vec![FetchAnswer::Silent];
4311 let brokers = start_brokers(&[silent.clone(), silent.clone(), silent], &sent).await;
4312 let mut consumer = consumer_on(&brokers).await;
4313 consumer.refresh_leader_epochs().await.expect("metadata");
4314 let started = tokio::time::Instant::now();
4315 let records = consumer.poll(millis(300)).await.expect("poll");
4316 let elapsed = started.elapsed();
4317 tokio::time::sleep(Duration::from_millis(100)).await;
4319 let mut leaders: Vec<i32> = sent
4320 .lock()
4321 .expect("sent lock")
4322 .iter()
4323 .map(|(node_id, _)| *node_id)
4324 .collect();
4325 leaders.sort_unstable();
4326 drop(consumer);
4327 stop(brokers);
4328 assert2::assert!(
4329 (
4330 records.len(),
4331 elapsed < Duration::from_millis(1500),
4332 leaders
4333 ) == (0, true, vec![1, 2, 3])
4334 );
4335 }
4336
4337 #[tokio::test]
4341 async fn fetch_leaves_out_a_partition_that_waits_for_its_assign_callback() {
4342 let TwoPartitionFetchSetup {
4343 sent,
4344 brokers,
4345 mut consumer,
4346 } = two_partition_fetch_setup().await;
4347 let gate = crate::rebalance_listener::AssignedCallbackGate::new(
4348 &consumer.assigned_callback_pending,
4349 &[("orders".to_owned(), 0)],
4350 );
4351 consumer.poll(millis(200)).await.expect("poll");
4352 tokio::time::sleep(Duration::from_millis(50)).await;
4353 let leaders = |sent: &SentFetches| {
4354 let mut leaders: Vec<i32> = sent
4355 .lock()
4356 .expect("sent lock")
4357 .drain(..)
4358 .map(|(node_id, _)| node_id)
4359 .collect();
4360 leaders.sort_unstable();
4361 leaders
4362 };
4363 let waiting = leaders(&sent);
4364 drop(gate);
4365 consumer.poll(millis(200)).await.expect("poll");
4366 tokio::time::sleep(Duration::from_millis(50)).await;
4367 let after = leaders(&sent);
4368 drop(consumer);
4369 stop(brokers);
4370 assert2::assert!((waiting, after) == (vec![2], vec![1, 2]));
4371 }
4372
4373 #[tokio::test]
4377 async fn seek_moves_an_owned_position_and_rejects_an_unowned_partition() {
4378 let mut actual = Vec::new();
4379 let mut wanted = Vec::new();
4380 for (name, partition, expected_error, expected_offset) in [
4381 ("owned partition", 0, None, 2),
4382 (
4383 "unowned partition",
4384 9,
4385 Some("no current assignment for partition orders-9"),
4386 5,
4387 ),
4388 ] {
4389 let sent = SentFetches::default();
4390 let brokers = start_brokers(
4391 &[vec![FetchAnswer::Respond {
4392 error_code: 0,
4393 session_id: 0,
4394 }]],
4395 &sent,
4396 )
4397 .await;
4398 let mut consumer = consumer_on(&brokers).await;
4399 let error = consumer
4400 .seek("orders", partition, 2)
4401 .await
4402 .err()
4403 .map(|error| error.to_string());
4404 consumer.poll(millis(200)).await.expect("poll");
4405 let offsets: Vec<i64> = sent
4406 .lock()
4407 .expect("sent lock")
4408 .iter()
4409 .flat_map(|(_, request)| &request.topics)
4410 .flat_map(|topic| &topic.partitions)
4411 .map(|partition| partition.fetch_offset)
4412 .collect();
4413 drop(consumer);
4414 stop(brokers);
4415 actual.push((name, error, offsets));
4416 wanted.push((
4417 name,
4418 expected_error.map(str::to_owned),
4419 vec![expected_offset],
4420 ));
4421 }
4422 assert2::assert!(actual == wanted);
4423 }
4424
4425 #[tokio::test]
4428 async fn fetch_leaves_out_a_paused_partition_until_resume() {
4429 let TwoPartitionFetchSetup {
4430 sent,
4431 brokers,
4432 mut consumer,
4433 } = two_partition_fetch_setup().await;
4434 let fetched = |sent: &SentFetches| {
4435 let mut partitions: Vec<(i32, i32)> = sent
4436 .lock()
4437 .expect("sent lock")
4438 .drain(..)
4439 .flat_map(|(node_id, request)| {
4440 request
4441 .topics
4442 .into_iter()
4443 .flat_map(|topic| topic.partitions)
4444 .map(move |partition| (node_id, partition.partition))
4445 })
4446 .collect();
4447 partitions.sort_unstable();
4448 partitions
4449 };
4450 consumer
4451 .pause(&[("orders".to_owned(), 0)])
4452 .await
4453 .expect("pause");
4454 consumer.poll(millis(200)).await.expect("poll");
4455 tokio::time::sleep(Duration::from_millis(50)).await;
4456 let paused = fetched(&sent);
4457 consumer
4458 .resume(&[("orders".to_owned(), 0)])
4459 .await
4460 .expect("resume");
4461 consumer.poll(millis(200)).await.expect("poll");
4462 tokio::time::sleep(Duration::from_millis(50)).await;
4463 let resumed = fetched(&sent);
4464 let unowned = consumer
4465 .position("orders", 9)
4466 .await
4467 .map_err(|error| error.to_string());
4468 drop(consumer);
4469 stop(brokers);
4470 assert2::assert!(
4471 (paused, resumed, unowned)
4472 == (
4473 vec![(2, 1)],
4474 vec![(1, 0), (2, 1)],
4475 Err("no current assignment for partition orders-9".to_owned())
4476 )
4477 );
4478 }
4479
4480 #[tokio::test]
4483 async fn poll_waits_for_its_timeout_when_every_partition_is_paused() {
4484 let sent = SentFetches::default();
4485 let brokers = start_brokers(
4486 &[vec![FetchAnswer::Respond {
4487 error_code: 0,
4488 session_id: 0,
4489 }]],
4490 &sent,
4491 )
4492 .await;
4493 let mut consumer = consumer_on(&brokers).await;
4494 consumer
4495 .pause(&[("orders".to_owned(), 0)])
4496 .await
4497 .expect("pause");
4498 let started = tokio::time::Instant::now();
4499 let records = consumer.poll(millis(300)).await.expect("poll").len();
4500 let waited = started.elapsed() >= Duration::from_millis(250);
4501 let fetches = sent.lock().expect("sent lock").len();
4502 drop(consumer);
4503 stop(brokers);
4504 assert2::assert!((records, waited, fetches) == (0, true, 0));
4505 }
4506
4507 #[tokio::test]
4510 async fn wakeup_ends_the_poll_that_waits_or_the_next_one() {
4511 let sent = SentFetches::default();
4512 let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
4513 let mut consumer = consumer_on(&brokers).await;
4514 consumer.wakeup();
4515 let pending = consumer
4516 .poll(millis(100))
4517 .await
4518 .map_err(|error| error.to_string());
4519 let handle = consumer.wakeup_handle();
4520 let waker = tokio::spawn(async move {
4521 tokio::time::sleep(Duration::from_millis(100)).await;
4522 handle.wakeup();
4523 });
4524 let started = tokio::time::Instant::now();
4525 let blocked = consumer
4526 .poll(secs(10))
4527 .await
4528 .map_err(|error| error.to_string());
4529 let woken_fast = started.elapsed() < Duration::from_secs(5);
4530 waker.await.expect("waker");
4531 let next = consumer
4532 .poll(millis(100))
4533 .await
4534 .map(|records| records.len());
4535 drop(consumer);
4536 stop(brokers);
4537 let wakeup = Err("the consumer was woken up".to_owned());
4538 assert2::assert!(
4539 (
4540 pending,
4541 blocked,
4542 woken_fast,
4543 next.map_err(|error| error.to_string())
4544 ) == (wakeup.clone(), wakeup, true, Ok(0))
4545 );
4546 }
4547
4548 #[tokio::test]
4552 async fn wakeup_ends_a_poll_that_waits_for_metadata() {
4553 let silent = MockBroker::start(|api_key, _version, _corr_id, _body| {
4554 (api_key == api_versions_request::API_KEY).then(|| {
4555 encode(
4556 &mock_test_support::api_versions_response_ranges([
4557 (api_versions_request::API_KEY, 0, 3),
4558 (metadata_request::API_KEY, 0, 8),
4559 ]),
4560 0,
4561 )
4562 })
4563 })
4564 .await;
4565 let client = Client::builder()
4566 .bootstrap(silent.addr.to_string())
4567 .request_timeout(secs(30))
4568 .build()
4569 .await
4570 .expect("client");
4571 let mut consumer = poll::partition_error_tests::consumer_with_client(client);
4572 let handle = consumer.wakeup_handle();
4573 let waker = tokio::spawn(async move {
4574 tokio::time::sleep(Duration::from_millis(100)).await;
4575 handle.wakeup();
4576 });
4577 let polled = tokio::time::timeout(Duration::from_secs(5), consumer.poll(secs(10)))
4578 .await
4579 .map(|result| {
4580 result
4581 .map(|records| records.len())
4582 .map_err(|error| error.to_string())
4583 });
4584 waker.await.expect("waker");
4585 drop(consumer);
4586 silent.stop();
4587 assert2::assert!(polled == Ok(Err("the consumer was woken up".to_owned())));
4588 }
4589
4590 #[tokio::test]
4594 async fn a_rebalance_ends_the_wait_of_a_poll_with_nothing_to_fetch() {
4595 let sent = SentFetches::default();
4596 let brokers = start_brokers(
4597 &[vec![FetchAnswer::Respond {
4598 error_code: 0,
4599 session_id: 0,
4600 }]],
4601 &sent,
4602 )
4603 .await;
4604 let mut consumer = consumer_on(&brokers).await;
4605 consumer
4606 .pause(&[("orders".to_owned(), 0)])
4607 .await
4608 .expect("pause");
4609 let (pending, pending_rx) = tokio::sync::watch::channel(false);
4610 consumer.rebalance_pending = pending_rx;
4611 let rebalance = tokio::spawn(async move {
4612 tokio::time::sleep(Duration::from_millis(100)).await;
4613 pending.send_replace(true);
4614 tokio::time::sleep(Duration::from_secs(10)).await;
4616 });
4617 let started = tokio::time::Instant::now();
4618 consumer.poll(secs(5)).await.expect("poll");
4619 let ended_early = started.elapsed() < Duration::from_secs(2);
4620 rebalance.abort();
4621 drop(consumer);
4622 stop(brokers);
4623 assert2::assert!(ended_early);
4624 }
4625
4626 #[tokio::test]
4630 async fn poll_names_the_topics_that_a_pattern_added() {
4631 let sent = SentFetches::default();
4632 let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
4633 let consumer = consumer_on(&brokers).await;
4634 consumer.client.metadata_topics().set(["orders".to_owned()]);
4635 consumer.subscription.send_modify(|subscription| {
4636 subscription.pattern = Some(crate::TopicPattern::new(|_| true));
4637 subscription.topics = vec!["orders".to_owned(), "payments".to_owned()];
4638 });
4639 consumer.sync_metadata_topics();
4640 let names = consumer.client.metadata_topics().names();
4641 drop(consumer);
4642 stop(brokers);
4643 assert2::assert!(names == vec!["orders".to_owned(), "payments".to_owned()]);
4644 }
4645
4646 #[tokio::test]
4650 async fn a_short_poll_does_not_wait_for_the_committed_offsets() {
4651 let sent = SentFetches::default();
4652 let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
4653 let mut consumer = consumer_on(&brokers).await;
4654 consumer
4655 .next_offsets
4656 .lock()
4657 .await
4658 .insert(("orders".to_owned(), 0), COMMITTED_SENTINEL);
4659 let started = tokio::time::Instant::now();
4660 let result = consumer
4661 .poll(millis(200))
4662 .await
4663 .map(|records| records.len());
4664 let elapsed = started.elapsed();
4665 drop(consumer);
4666 stop(brokers);
4667 assert2::assert!((result.is_err(), elapsed < Duration::from_secs(5)) == (true, true));
4668 }
4669
4670 #[tokio::test]
4674 async fn a_consumer_without_a_group_fetches_its_manual_assignment() {
4675 let sent = SentFetches::default();
4676 let brokers = start_brokers(
4677 &[vec![FetchAnswer::Respond {
4678 error_code: 0,
4679 session_id: 0,
4680 }]],
4681 &sent,
4682 )
4683 .await;
4684 let calls = std::sync::Arc::new(std::sync::Mutex::new(Vec::new()));
4685 let mut consumer = Consumer::builder()
4686 .bootstrap(brokers[0].addr.to_string())
4687 .auto_offset_reset(AutoOffsetReset::Earliest)
4688 .rebalance_listener(Box::new(
4689 RecordingListener::builder()
4690 .calls(std::sync::Arc::clone(&calls))
4691 .build(),
4692 ))
4693 .build()
4694 .await
4695 .expect("build");
4696 let before_assign = consumer
4697 .poll(millis(20))
4698 .await
4699 .map(|records| records.len())
4700 .map_err(|error| error.to_string());
4701 consumer
4702 .assign(&[("orders".to_owned(), 0)])
4703 .await
4704 .expect("assign");
4705 consumer.poll(millis(300)).await.expect("poll");
4706 let fetched: Vec<(i32, String, i32, i64)> = sent
4707 .lock()
4708 .expect("sent lock")
4709 .iter()
4710 .flat_map(|(node_id, request)| {
4711 request.topics.iter().flat_map(move |topic| {
4712 topic.partitions.iter().map(move |partition| {
4713 (
4714 *node_id,
4715 topic.topic.clone(),
4716 partition.partition,
4717 partition.fetch_offset,
4718 )
4719 })
4720 })
4721 })
4722 .collect();
4723 let group_calls = (
4724 consumer
4725 .commit_sync()
4726 .await
4727 .map_err(|error| error.to_string()),
4728 consumer
4729 .subscribe(["orders"])
4730 .await
4731 .map_err(|error| error.to_string()),
4732 consumer
4733 .enforce_rebalance(None)
4734 .map_err(|error| error.to_string()),
4735 );
4736 consumer.unsubscribe().await.expect("unsubscribe");
4737 let after_unsubscribe = (
4738 consumer.assignment().await,
4739 consumer.next_offsets.lock().await.len(),
4740 );
4741 let listener_calls = calls.lock().expect("calls lock").clone();
4742 consumer.close().await.expect("close");
4743 stop(brokers);
4744 let invalid_group = Err(ConsumerError::InvalidGroupId.to_string());
4745 assert2::assert!(
4746 (
4747 before_assign,
4748 fetched.first().cloned(),
4749 group_calls,
4750 after_unsubscribe,
4751 listener_calls
4752 ) == (
4753 Err("not subscribed to any topic".to_owned()),
4754 Some((1, "orders".to_owned(), 0, 0)),
4755 (invalid_group.clone(), invalid_group.clone(), invalid_group),
4756 (Vec::new(), 0),
4757 Vec::<test_support::ListenerCall>::new()
4758 )
4759 );
4760 }
4761
4762 #[derive(Clone, Copy, Debug)]
4764 enum ReplicaStep {
4765 Poll,
4766 WaitForExpiry,
4768 }
4769
4770 #[tokio::test]
4776 async fn fetch_sends_the_rack_and_follows_the_preferred_read_replica() {
4777 use FetchAnswer::Partition;
4778 use ReplicaStep::{Poll, WaitForExpiry};
4779 let records = Partition {
4780 error_code: 0,
4781 preferred_read_replica: -1,
4782 };
4783 let prefer_2 = Partition {
4784 error_code: 0,
4785 preferred_read_replica: 2,
4786 };
4787 let not_leader = Partition {
4788 error_code: 6,
4789 preferred_read_replica: -1,
4790 };
4791 let mut actual = Vec::new();
4792 let mut wanted = Vec::new();
4793 for (name, rack, leader, replica, steps, expected) in [
4794 (
4795 "no rack",
4796 None,
4797 vec![records],
4798 vec![records],
4799 vec![Poll],
4800 vec![(1, String::new())],
4801 ),
4802 (
4803 "rack",
4804 Some("az-1"),
4805 vec![records],
4806 vec![records],
4807 vec![Poll],
4808 vec![(1, "az-1".to_owned())],
4809 ),
4810 (
4811 "preferred read replica",
4812 Some("az-1"),
4813 vec![prefer_2],
4814 vec![records],
4815 vec![Poll, Poll],
4816 vec![(1, "az-1".to_owned()), (2, "az-1".to_owned())],
4817 ),
4818 (
4819 "the replica answers not leader or follower",
4820 Some("az-1"),
4821 vec![prefer_2, records],
4822 vec![not_leader],
4823 vec![Poll, Poll, Poll],
4824 vec![
4825 (1, "az-1".to_owned()),
4826 (2, "az-1".to_owned()),
4827 (1, "az-1".to_owned()),
4828 ],
4829 ),
4830 (
4831 "the preference expires",
4832 Some("az-1"),
4833 vec![prefer_2, records],
4834 vec![records],
4835 vec![Poll, Poll, WaitForExpiry, Poll],
4836 vec![
4837 (1, "az-1".to_owned()),
4838 (2, "az-1".to_owned()),
4839 (1, "az-1".to_owned()),
4840 ],
4841 ),
4842 ] {
4843 let sent = SentFetches::default();
4844 let brokers = start_brokers(&[leader, replica], &sent).await;
4845 let mut consumer = consumer_on(&brokers).await;
4846 *consumer.assigned.lock().await = vec![("orders".to_owned(), 0)];
4847 consumer.client_rack = rack.map(str::to_owned);
4848 consumer.metadata_max_age = millis(300);
4849 for step in steps {
4850 match step {
4851 Poll => {
4852 let before = sent.lock().expect("sent lock").len();
4853 consumer.poll(secs(2)).await.expect("poll");
4854 let _ = tokio::time::timeout(Duration::from_secs(2), async {
4857 while sent.lock().expect("sent lock").len() == before {
4858 tokio::time::sleep(Duration::from_millis(5)).await;
4859 }
4860 })
4861 .await;
4862 }
4863 WaitForExpiry => tokio::time::sleep(Duration::from_millis(400)).await,
4864 }
4865 }
4866 let requests: Vec<(i32, String)> = sent
4867 .lock()
4868 .expect("sent lock")
4869 .iter()
4870 .map(|(node_id, request)| (*node_id, request.rack_id.clone()))
4871 .collect();
4872 drop(consumer);
4873 stop(brokers);
4874 actual.push((name, requests));
4875 wanted.push((name, expected));
4876 }
4877 assert2::assert!(actual == wanted);
4878 }
4879
4880 type SessionFields = (i32, i32, Vec<i32>, Vec<i32>);
4883
4884 fn session_fields(request: &FetchRequest) -> SessionFields {
4885 (
4886 request.session_id,
4887 request.session_epoch,
4888 request
4889 .topics
4890 .iter()
4891 .flat_map(|topic| topic.partitions.iter().map(|p| p.partition))
4892 .collect(),
4893 request
4894 .forgotten_topics_data
4895 .iter()
4896 .flat_map(|topic| topic.partitions.iter().copied())
4897 .collect(),
4898 )
4899 }
4900
4901 #[tokio::test]
4905 async fn fetch_sessions_send_incremental_requests_and_recover_a_lost_session() {
4906 let mut actual = Vec::new();
4907 let mut wanted = Vec::new();
4908 for (name, answers, expected) in [
4909 (
4910 "session",
4911 vec![
4912 FetchAnswer::Respond {
4913 error_code: 0,
4914 session_id: 77,
4915 },
4916 FetchAnswer::Respond {
4917 error_code: 0,
4918 session_id: 77,
4919 },
4920 ],
4921 vec![(0, 0, vec![0], vec![]), (77, 1, vec![], vec![])],
4922 ),
4923 (
4924 "session lost",
4925 vec![
4926 FetchAnswer::Respond {
4927 error_code: 0,
4928 session_id: 77,
4929 },
4930 FetchAnswer::Respond {
4931 error_code: 70,
4932 session_id: 0,
4933 },
4934 FetchAnswer::Respond {
4935 error_code: 0,
4936 session_id: 78,
4937 },
4938 ],
4939 vec![
4940 (0, 0, vec![0], vec![]),
4941 (77, 1, vec![], vec![]),
4942 (0, 0, vec![0], vec![]),
4943 ],
4944 ),
4945 ] {
4946 let sent = SentFetches::default();
4947 let rounds = answers.len();
4948 let brokers = start_brokers(&[answers], &sent).await;
4949 let mut consumer = consumer_on(&brokers).await;
4950 for _ in 0..rounds {
4951 consumer.poll(secs(5)).await.expect("poll");
4952 }
4953 let requests: Vec<SessionFields> = sent
4954 .lock()
4955 .expect("sent lock")
4956 .iter()
4957 .map(|(_, request)| session_fields(request))
4958 .collect();
4959 drop(consumer);
4960 stop(brokers);
4961 actual.push((name, requests));
4962 wanted.push((name, expected));
4963 }
4964 assert2::assert!(actual == wanted);
4965 }
4966
4967 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
4970 async fn a_fetch_result_wakes_poll_before_the_task_ends() {
4971 let sent = SentFetches::default();
4972 let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
4973 let mut consumer = consumer_on(&brokers).await;
4974 let (tx, rx) = tokio::sync::oneshot::channel();
4975 let completed = Arc::clone(&consumer.fetches.completed);
4976 consumer.fetches.in_flight.insert(
4977 -1,
4978 InFlightFetch {
4979 handle: tokio::spawn(async move {
4980 let _ = tx.send(Ok(FetchResponse::default()));
4981 completed.notify_one();
4982 tokio::time::sleep(Duration::from_secs(3)).await;
4984 }),
4985 result: rx,
4986 requested: HashMap::new(),
4987 from_replica: std::collections::HashSet::new(),
4988 },
4989 );
4990 let started = tokio::time::Instant::now();
4991 consumer
4992 .wait_for_fetches(started + Duration::from_secs(2))
4993 .await;
4994 let elapsed = started.elapsed();
4995 drop(consumer);
4996 stop(brokers);
4997 assert2::assert!(elapsed < Duration::from_millis(500));
4998 }
4999
5000 #[tokio::test]
5004 async fn out_of_range_from_a_replica_request_gives_no_reset() {
5005 let sent = SentFetches::default();
5006 let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
5007 let mut consumer = consumer_on(&brokers).await;
5008 consumer.auto_offset_reset = AutoOffsetReset::Latest;
5009 let key = ("orders".to_owned(), 0);
5010 let (tx, rx) = tokio::sync::oneshot::channel();
5011 consumer.fetches.in_flight.insert(
5012 2,
5013 InFlightFetch {
5014 handle: tokio::spawn(async {}),
5015 result: rx,
5016 requested: maplit::hashmap! {key.clone() => 5},
5017 from_replica: maplit::hashset! {key.clone()},
5018 },
5019 );
5020 tx.send(Ok(FetchResponse {
5021 responses: vec![FetchableTopicResponse {
5022 topic: "orders".into(),
5023 partitions: vec![PartitionData {
5024 partition_index: 0,
5025 error_code: 1,
5026 ..Default::default()
5027 }],
5028 ..Default::default()
5029 }],
5030 ..Default::default()
5031 }))
5032 .expect("send response");
5033 let topic_ids = consumer.topic_ids.lock().await.clone();
5034 let responses = consumer
5035 .take_completed_fetches(&topic_ids)
5036 .await
5037 .expect("take");
5038 consumer
5039 .process_fetch_responses(responses, &topic_ids)
5040 .await
5041 .expect("process");
5042 let position = consumer.next_offsets.lock().await.get(&key).copied();
5043 drop(consumer);
5044 stop(brokers);
5045 assert2::assert!(position == Some(5));
5046 }
5047
5048 #[tokio::test]
5051 async fn a_fetch_response_for_a_moved_position_gives_no_records() {
5052 let sent = SentFetches::default();
5053 let brokers = start_brokers(&[vec![FetchAnswer::Silent]], &sent).await;
5054 let mut consumer = consumer_on(&brokers).await;
5055 let (tx, rx) = tokio::sync::oneshot::channel();
5056 consumer.fetches.in_flight.insert(
5057 -1,
5058 InFlightFetch {
5059 handle: tokio::spawn(async {}),
5060 result: rx,
5061 requested: maplit::hashmap! {("orders".to_owned(), 0) => 5},
5062 from_replica: std::collections::HashSet::new(),
5063 },
5064 );
5065 consumer
5067 .next_offsets
5068 .lock()
5069 .await
5070 .insert(("orders".to_owned(), 0), 40);
5071 let record = krabka_protocol::records::RecordBatch {
5072 base_offset: 5,
5073 last_offset_delta: 0,
5074 records: vec![krabka_protocol::records::Record {
5075 offset_delta: 0,
5076 value: Some(bytes::Bytes::from_static(b"old")),
5077 ..Default::default()
5078 }],
5079 ..Default::default()
5080 };
5081 tx.send(Ok(FetchResponse {
5082 responses: vec![FetchableTopicResponse {
5083 topic: "orders".into(),
5084 partitions: vec![PartitionData {
5085 partition_index: 0,
5086 high_watermark: 6,
5087 records: Some(RecordsPayload::V2(vec![record])),
5088 ..Default::default()
5089 }],
5090 ..Default::default()
5091 }],
5092 ..Default::default()
5093 }))
5094 .expect("send response");
5095 consumer
5096 .wait_for_fetches(tokio::time::Instant::now() + Duration::from_secs(5))
5097 .await;
5098 let topic_ids = consumer.topic_ids.lock().await.clone();
5099 let responses = consumer
5100 .take_completed_fetches(&topic_ids)
5101 .await
5102 .expect("take");
5103 consumer
5104 .process_fetch_responses(responses, &topic_ids)
5105 .await
5106 .expect("process");
5107 let records = consumer.drain_fetch_buffer().await;
5108 let position = consumer
5109 .next_offsets
5110 .lock()
5111 .await
5112 .get(&("orders".to_owned(), 0))
5113 .copied();
5114 drop(consumer);
5115 stop(brokers);
5116 assert2::assert!((records.len(), position) == (0, Some(40)));
5117 }
5118
5119 async fn no_committed_offset_broker(
5128 list_offsets_offset: i64,
5129 list_offsets_timestamps: Arc<std::sync::Mutex<Vec<i64>>>,
5130 fetch_offsets: Arc<std::sync::Mutex<Vec<i64>>>,
5131 ) -> MockBroker {
5132 use krabka_protocol::owned::{
5133 list_offsets_request::FLEXIBLE_MIN,
5134 list_offsets_response::{ListOffsetsPartitionResponse, ListOffsetsTopicResponse},
5135 };
5136 const NODE_ID: i32 = 1;
5137 MockBroker::start(move |api_key, version, _corr_id, mut body| {
5138 if api_key == api_versions_request::API_KEY {
5139 let versions = mock_test_support::api_versions_response_ranges([
5140 (api_versions_request::API_KEY, 0, 3),
5141 (metadata_request::API_KEY, 0, 8),
5142 (
5143 list_offsets_request::API_KEY,
5144 0,
5145 list_offsets_request::MAX_VERSION,
5146 ),
5147 (fetch_request::API_KEY, 4, 11),
5148 ]);
5149 return Some(encode(&versions, 0));
5150 }
5151 if api_key == metadata_request::API_KEY {
5152 let metadata = test_support::OrdersMetadata::builder()
5156 .topic_id([9; 16])
5157 .leader_id(NODE_ID)
5158 .leader_epoch(1)
5159 .build();
5160 return Some(encode(&metadata, version));
5161 }
5162 if api_key == list_offsets_request::API_KEY {
5163 let request = test_support::decode_list_offsets_request(&mut body, version);
5164 let timestamp = request.topics[0].partitions[0].timestamp;
5165 list_offsets_timestamps
5166 .lock()
5167 .expect("lock")
5168 .push(timestamp);
5169 let answer = ListOffsetsResponse {
5170 topics: vec![ListOffsetsTopicResponse {
5171 name: "orders".into(),
5172 partitions: vec![ListOffsetsPartitionResponse {
5173 partition_index: 0,
5174 error_code: 0,
5175 offset: list_offsets_offset,
5176 ..Default::default()
5177 }],
5178 ..Default::default()
5179 }],
5180 ..Default::default()
5181 };
5182 let mut bytes = Vec::new();
5183 if version >= FLEXIBLE_MIN {
5184 bytes.push(0);
5187 }
5188 bytes.extend(encode(&answer, version));
5189 return Some(bytes);
5190 }
5191 if api_key == fetch_request::API_KEY {
5192 let client_id_len = usize::try_from(body.get_i16()).expect("client id");
5193 body.advance(client_id_len);
5194 let request = FetchRequest::decode(&mut body, version).expect("decode fetch");
5195 if let Some(partition) = request
5196 .topics
5197 .first()
5198 .and_then(|topic| topic.partitions.first())
5199 {
5200 fetch_offsets
5201 .lock()
5202 .expect("lock")
5203 .push(partition.fetch_offset);
5204 }
5205 return Some(encode(&FetchResponse::default(), version));
5206 }
5207 None
5208 })
5209 .await
5210 }
5211
5212 async fn consumer_with_no_committed_offset(
5216 broker: &MockBroker,
5217 policy: AutoOffsetReset,
5218 ) -> Consumer {
5219 let client = Client::builder()
5220 .bootstrap(broker.addr.to_string())
5221 .request_timeout(secs(30))
5222 .build()
5223 .await
5224 .expect("client");
5225 let mut consumer = poll::partition_error_tests::consumer_with_client(client);
5226 consumer.auto_offset_reset = policy;
5227 consumer.next_offsets.lock().await.insert(
5228 ("orders".to_owned(), 0),
5229 consumer::reset_starting_offset(policy),
5230 );
5231 consumer
5232 }
5233
5234 #[tokio::test]
5245 async fn initial_position_follows_the_reset_policy_with_no_committed_offset() {
5246 struct Case {
5247 name: &'static str,
5248 policy: AutoOffsetReset,
5249 expected_list_offsets: Vec<i64>,
5250 expected_fetch_offsets: Vec<i64>,
5251 expect_error: bool,
5252 }
5253 let cases = [
5254 Case {
5255 name: "none has no committed offset and no reset policy: poll fails, no Fetch",
5256 policy: AutoOffsetReset::None,
5257 expected_list_offsets: vec![],
5258 expected_fetch_offsets: vec![],
5259 expect_error: true,
5260 },
5261 Case {
5262 name: "earliest asks ListOffsets(-2) and fetches from the log start",
5263 policy: AutoOffsetReset::Earliest,
5264 expected_list_offsets: vec![EARLIEST_TIMESTAMP],
5265 expected_fetch_offsets: vec![40],
5266 expect_error: false,
5267 },
5268 Case {
5269 name: "latest asks ListOffsets(-1)",
5270 policy: AutoOffsetReset::Latest,
5271 expected_list_offsets: vec![LATEST_TIMESTAMP],
5272 expected_fetch_offsets: vec![40],
5273 expect_error: false,
5274 },
5275 ];
5276 let mut actual = Vec::new();
5277 let mut wanted = Vec::new();
5278 for case in cases {
5279 let list_offsets_timestamps: Arc<std::sync::Mutex<Vec<i64>>> = Arc::default();
5280 let fetch_offsets: Arc<std::sync::Mutex<Vec<i64>>> = Arc::default();
5281 let broker = no_committed_offset_broker(
5282 40,
5283 Arc::clone(&list_offsets_timestamps),
5284 Arc::clone(&fetch_offsets),
5285 )
5286 .await;
5287 let mut consumer = consumer_with_no_committed_offset(&broker, case.policy).await;
5288 let result = consumer.poll(millis(500)).await;
5289 let list_offsets_timestamps = list_offsets_timestamps.lock().expect("lock").clone();
5290 let fetch_offsets = fetch_offsets.lock().expect("lock").clone();
5291 drop(consumer);
5292 broker.stop();
5293 actual.push((
5294 case.name,
5295 list_offsets_timestamps,
5296 fetch_offsets,
5297 result.is_err(),
5298 ));
5299 wanted.push((
5300 case.name,
5301 case.expected_list_offsets,
5302 case.expected_fetch_offsets,
5303 case.expect_error,
5304 ));
5305 }
5306 assert2::assert!(actual == wanted);
5307 }
5308}