1use std::{
4 collections::{BTreeSet, HashMap, HashSet},
5 future::Future,
6 sync::{Arc, atomic::Ordering},
7 time::Duration,
8};
9
10use convert::TimeExt as _;
11use krabka_client_core::ClientError;
12use krabka_protocol::owned::{
13 offset_commit_request::{OffsetCommitRequest, OffsetCommitRequestTopic},
14 offset_commit_response::OffsetCommitResponse,
15};
16use krabka_units::convert;
17use tokio::sync::Mutex;
18
19#[cfg(test)]
20use crate::test_support;
21use crate::{
22 consumer::{CommitIdentity, Consumer},
23 coordinator,
24 coordinator::{
25 COORDINATOR_LOAD_IN_PROGRESS, COORDINATOR_NOT_AVAILABLE, NOT_COORDINATOR, find_coordinator,
26 is_retriable_coordinator_code, is_retriable_transport_error, next_backoff,
27 },
28 error::ConsumerError,
29 offset_wire::{TopicNameOffsetCommit, build_commit_topics},
30 position::PartitionPosition,
31};
32
33const ASYNC_COMMIT_IDLE: u8 = 0;
34const ASYNC_COMMIT_RUNNING: u8 = 1;
35const ASYNC_COMMIT_DIRTY: u8 = 2;
36
37const UNKNOWN_TOPIC_OR_PARTITION: i16 = 3;
39const REQUEST_TIMED_OUT: i16 = 7;
41const ILLEGAL_GENERATION: i16 = 22;
43const UNKNOWN_MEMBER_ID: i16 = 25;
45const REBALANCE_IN_PROGRESS: i16 = 27;
47const TOPIC_AUTHORIZATION_FAILED: i16 = 29;
49const GROUP_AUTHORIZATION_FAILED: i16 = 30;
51const FENCED_INSTANCE_ID: i16 = 82;
53
54#[derive(Clone, Debug, Default, PartialEq, Eq)]
57pub struct OffsetAndMetadata {
58 pub offset: i64,
60 pub leader_epoch: Option<i32>,
62 pub metadata: String,
65}
66
67impl OffsetAndMetadata {
68 #[must_use]
70 pub fn new(offset: i64) -> Self {
71 Self {
72 offset,
73 leader_epoch: None,
74 metadata: String::new(),
75 }
76 }
77
78 fn of_position(offset: i64, leader_epoch: i32) -> Self {
81 Self {
82 offset,
83 leader_epoch: (leader_epoch >= 0).then_some(leader_epoch),
84 metadata: String::new(),
85 }
86 }
87}
88
89pub(crate) fn position_commits(
91 offsets: HashMap<(String, i32), (i64, i32)>,
92) -> HashMap<(String, i32), OffsetAndMetadata> {
93 offsets
94 .into_iter()
95 .map(|(partition, (offset, leader_epoch))| {
96 (
97 partition,
98 OffsetAndMetadata::of_position(offset, leader_epoch),
99 )
100 })
101 .collect()
102}
103
104fn commit_offsets(
105 raw_offsets: HashMap<(String, i32), i64>,
106 positions: &HashMap<(String, i32), PartitionPosition>,
107) -> HashMap<(String, i32), (i64, i32)> {
108 raw_offsets
109 .into_iter()
110 .map(|(k, v)| {
111 let epoch = positions.get(&k).map_or(-1, |p| p.offset_epoch.get());
114 (k, (v, epoch))
115 })
116 .collect()
117}
118
119fn validate_selected_offsets(
124 offsets: &HashMap<(String, i32), OffsetAndMetadata>,
125 assigned: &HashMap<(String, i32), u64>,
126) -> Result<(), ConsumerError> {
127 for ((topic, partition), offset) in offsets {
128 if offset.offset < 0 {
129 return Err(ConsumerError::InvalidOffset(offset.offset));
130 }
131 if !assigned.contains_key(&(topic.clone(), *partition)) {
132 return Err(ConsumerError::IllegalState(format!(
133 "cannot commit unassigned partition {topic}-{partition}"
134 )));
135 }
136 }
137 Ok(())
138}
139
140async fn snapshot_commit_offsets(
145 commit_identity: &Arc<Mutex<CommitIdentity>>,
146 offsets: &Arc<Mutex<HashMap<(String, i32), i64>>>,
147 positions: &Arc<Mutex<HashMap<(String, i32), PartitionPosition>>>,
148 auto_commit: Option<&AutoCommit>,
149) -> (HashMap<(String, i32), OffsetAndMetadata>, (i32, String)) {
150 let identity = commit_identity.lock().await.clone();
151 let mut raw_offsets = offsets.lock().await.clone();
152 raw_offsets.retain(|partition, offset| {
153 identity.ownership_ids.contains_key(partition) && has_valid_position(*offset)
154 });
155 if raw_offsets.is_empty() {
156 return (HashMap::new(), (identity.generation, identity.member_id));
157 }
158 let pos = positions.lock().await;
159 let offsets = commit_offsets(raw_offsets, &pos);
160 drop(pos);
161 if let Some(auto_commit) = auto_commit {
162 auto_commit
163 .record_sent(sent_positions(&offsets, &identity.ownership_ids))
164 .await;
165 }
166 (
167 position_commits(offsets),
168 (identity.generation, identity.member_id),
169 )
170}
171
172pub type OffsetCommitCallback =
177 Box<dyn FnOnce(&HashMap<(String, i32), OffsetAndMetadata>, Result<(), &ConsumerError>) + Send>;
178
179pub(crate) type OffsetCommitCallbacks = Arc<std::sync::Mutex<Vec<OffsetCommitCallback>>>;
181
182fn take_callbacks(callbacks: &OffsetCommitCallbacks) -> Vec<OffsetCommitCallback> {
184 std::mem::take(
185 &mut *callbacks
186 .lock()
187 .unwrap_or_else(std::sync::PoisonError::into_inner),
188 )
189}
190
191fn async_commit_result(
199 result: Result<OffsetCommitResponse, ConsumerError>,
200 group_id: &str,
201 group_instance_id: Option<&str>,
202) -> Result<(), ConsumerError> {
203 let response = match result {
204 Ok(response) => response,
205 Err(ConsumerError::Client(error))
206 if is_retriable_transport_error(&error) || matches!(error, ClientError::Timeout(_)) =>
207 {
208 return Err(ConsumerError::RetriableCommitFailed(Box::new(
209 ConsumerError::Client(error),
210 )));
211 }
212 Err(ConsumerError::CoordinatorUnavailable) => {
216 return Err(ConsumerError::RetriableCommitFailed(Box::new(
217 ConsumerError::CoordinatorUnavailable,
218 )));
219 }
220 Err(ConsumerError::Server(code)) if is_retriable_coordinator_code(code) => {
221 return Err(ConsumerError::RetriableCommitFailed(Box::new(
222 ConsumerError::Server(code),
223 )));
224 }
225 Err(error) => return Err(error),
226 };
227 let partitions = || {
228 response.topics.iter().flat_map(|topic| {
229 topic
230 .partitions
231 .iter()
232 .map(move |partition| (topic, partition.error_code))
233 })
234 };
235 let class = PartitionCommitError::of_response(&response);
236 let code = partitions()
237 .map(|(_, code)| code)
238 .find(|code| PartitionCommitError::of(*code) == class)
239 .unwrap_or_default();
240 match class {
241 PartitionCommitError::None => Ok(()),
242 PartitionCommitError::TopicAuthorization => Err(ConsumerError::TopicAuthorizationFailed(
243 partitions()
244 .filter(|(_, code)| *code == TOPIC_AUTHORIZATION_FAILED)
245 .map(|(topic, _)| topic.name.clone())
246 .collect(),
247 )),
248 PartitionCommitError::Retriable | PartitionCommitError::CoordinatorUnknown => Err(
249 ConsumerError::RetriableCommitFailed(Box::new(ConsumerError::Server(code))),
250 ),
251 PartitionCommitError::Rebalance => Err(ConsumerError::CommitFailed),
252 PartitionCommitError::FencedInstanceId => Err(ConsumerError::FencedInstanceId(
253 group_instance_id.unwrap_or_default().to_owned(),
254 )),
255 PartitionCommitError::GroupAuthorization => {
256 Err(ConsumerError::GroupAuthorizationFailed(group_id.to_owned()))
257 }
258 PartitionCommitError::Fatal => Err(ConsumerError::Server(code)),
259 }
260}
261
262fn commit_timeout(timeout: Duration, cause: Option<ConsumerError>) -> ConsumerError {
269 ConsumerError::Timeout {
270 message: format!(
271 "Timeout of {}ms expired before successfully committing offsets",
272 timeout.as_millis()
273 ),
274 cause: cause.map(Box::new),
275 }
276}
277
278#[derive(Clone, Copy)]
279struct CommitDeadline {
280 expires_at: tokio::time::Instant,
281 timeout: Duration,
282}
283
284impl CommitDeadline {
285 fn after(timeout: Duration) -> Self {
286 Self {
287 expires_at: tokio::time::Instant::now() + timeout,
288 timeout,
289 }
290 }
291
292 async fn run<F>(self, commit: F) -> Result<(), ConsumerError>
293 where
294 F: Future<Output = Result<(), ConsumerError>>,
295 {
296 tokio::time::timeout_at(self.expires_at, commit)
297 .await
298 .unwrap_or_else(|_elapsed| Err(commit_timeout(self.timeout, None)))
299 }
300}
301
302pub(crate) fn build_commit_request(
305 group_id: String,
306 generation_id_or_member_epoch: i32,
307 member_id: String,
308 group_instance_id: Option<String>,
309 topics: Vec<OffsetCommitRequestTopic>,
310) -> TopicNameOffsetCommit {
311 TopicNameOffsetCommit::new(OffsetCommitRequest {
312 group_id,
313 generation_id_or_member_epoch,
314 member_id,
315 group_instance_id,
316 topics,
317 ..Default::default()
318 })
319}
320
321#[derive(Clone, Debug, Eq, PartialEq)]
322enum CommitOutcome {
323 Acked(HashSet<(String, i32)>),
324 Retriable {
327 code: i16,
328 acknowledged: HashSet<(String, i32)>,
329 find_coordinator: bool,
332 },
333}
334
335#[derive(Clone, Copy, Debug, Eq, PartialEq)]
340enum RebalanceCommitFailure {
341 CommitFailed,
346 RebalanceInProgress,
352}
353
354#[derive(Clone, Copy, Debug)]
357struct CommitResponseContext<'a> {
358 group_id: &'a str,
359 group_instance_id: Option<&'a str>,
360 coordinator_alive: bool,
362 generation_unchanged: bool,
365}
366
367fn commit_response_outcome(
407 resp: &OffsetCommitResponse,
408 context: CommitResponseContext<'_>,
409) -> Result<CommitOutcome, ConsumerError> {
410 let mut rebalance = None;
411 let mut retriable = None;
412 let mut find_coordinator = false;
413 let mut unauthorized_topics = BTreeSet::new();
414 let mut acknowledged = HashSet::new();
415 for topic in &resp.topics {
416 for partition in &topic.partitions {
417 let code = partition.error_code;
418 match PartitionCommitError::of(code) {
419 PartitionCommitError::None => {
420 acknowledged.insert((topic.name.clone(), partition.partition_index));
421 }
422 PartitionCommitError::TopicAuthorization => {
423 unauthorized_topics.insert(topic.name.clone());
424 }
425 PartitionCommitError::Retriable => {
426 retriable.get_or_insert(code);
427 }
428 PartitionCommitError::CoordinatorUnknown => {
429 retriable.get_or_insert(code);
430 find_coordinator = true;
431 }
432 PartitionCommitError::FencedInstanceId if context.generation_unchanged => {
433 return Err(ConsumerError::FencedInstanceId(
434 context.group_instance_id.unwrap_or_default().to_owned(),
435 ));
436 }
437 PartitionCommitError::Rebalance | PartitionCommitError::FencedInstanceId => {
438 let failure = if context.coordinator_alive
444 && (code == REBALANCE_IN_PROGRESS || !context.generation_unchanged)
445 {
446 RebalanceCommitFailure::RebalanceInProgress
447 } else {
448 RebalanceCommitFailure::CommitFailed
449 };
450 rebalance.get_or_insert(failure);
451 }
452 PartitionCommitError::GroupAuthorization => {
453 return Err(ConsumerError::GroupAuthorizationFailed(
454 context.group_id.to_owned(),
455 ));
456 }
457 PartitionCommitError::Fatal => return Err(ConsumerError::Server(code)),
458 }
459 }
460 }
461 match (rebalance, retriable) {
462 (Some(RebalanceCommitFailure::CommitFailed), _) => Err(ConsumerError::CommitFailed),
463 (Some(RebalanceCommitFailure::RebalanceInProgress), _) => Err(
464 ConsumerError::RebalanceInProgress(context.group_id.to_owned()),
465 ),
466 (None, Some(code)) => Ok(CommitOutcome::Retriable {
467 code,
468 acknowledged,
469 find_coordinator,
470 }),
471 (None, None) if !unauthorized_topics.is_empty() => {
472 Err(ConsumerError::TopicAuthorizationFailed(unauthorized_topics))
473 }
474 (None, None) => Ok(CommitOutcome::Acked(acknowledged)),
475 }
476}
477
478fn retain_continuously_owned(
479 pending: &mut HashMap<(String, i32), (OffsetAndMetadata, u64)>,
480 current: &HashMap<(String, i32), u64>,
481) {
482 pending.retain(|partition, (_, ownership_id)| current.get(partition) == Some(ownership_id));
483}
484
485#[derive(Clone, Copy, Debug, Eq, PartialEq)]
487pub(crate) struct ConsumedPosition {
488 pub offset: i64,
489 pub leader_epoch: i32,
490 pub ownership_id: u64,
493}
494
495pub(crate) type ConsumedPositions = HashMap<(String, i32), ConsumedPosition>;
497
498fn has_valid_position(next_offset: i64) -> bool {
503 next_offset >= 0 && !crate::poll::is_reset_sentinel(next_offset)
504}
505
506pub(crate) fn all_consumed(
509 ownership_ids: &HashMap<(String, i32), u64>,
510 next_offsets: &HashMap<(String, i32), i64>,
511 positions: &HashMap<(String, i32), PartitionPosition>,
512) -> ConsumedPositions {
513 next_offsets
514 .iter()
515 .filter(|(_, offset)| has_valid_position(**offset))
516 .filter_map(|(partition, offset)| {
517 ownership_ids.get(partition).map(|ownership_id| {
518 let leader_epoch = positions
519 .get(partition)
520 .map_or(-1, |position| position.offset_epoch.get());
521 (
522 partition.clone(),
523 ConsumedPosition {
524 offset: *offset,
525 leader_epoch,
526 ownership_id: *ownership_id,
527 },
528 )
529 })
530 })
531 .collect()
532}
533
534fn owned_offsets(
537 consumed: &ConsumedPositions,
538 ownership_ids: &HashMap<(String, i32), u64>,
539) -> HashMap<(String, i32), (i64, i32)> {
540 consumed
541 .iter()
542 .filter(|(partition, position)| {
543 ownership_ids.get(*partition) == Some(&position.ownership_id)
544 })
545 .map(|(partition, position)| (partition.clone(), (position.offset, position.leader_epoch)))
546 .collect()
547}
548
549fn sent_positions(
552 offsets: &HashMap<(String, i32), (i64, i32)>,
553 ownership_ids: &HashMap<(String, i32), u64>,
554) -> Vec<((String, i32), ConsumedPosition)> {
555 offsets
556 .iter()
557 .filter_map(|(partition, (offset, leader_epoch))| {
558 ownership_ids.get(partition).map(|ownership_id| {
559 (
560 partition.clone(),
561 ConsumedPosition {
562 offset: *offset,
563 leader_epoch: *leader_epoch,
564 ownership_id: *ownership_id,
565 },
566 )
567 })
568 })
569 .collect()
570}
571
572fn pending_positions(
574 pending: &HashMap<(String, i32), (OffsetAndMetadata, u64)>,
575) -> Vec<((String, i32), ConsumedPosition)> {
576 pending
577 .iter()
578 .map(|(partition, (offset, ownership_id))| {
579 (
580 partition.clone(),
581 ConsumedPosition {
582 offset: offset.offset,
583 leader_epoch: offset.leader_epoch.unwrap_or(-1),
584 ownership_id: *ownership_id,
585 },
586 )
587 })
588 .collect()
589}
590
591#[derive(Clone, Copy, Debug, Eq, PartialEq)]
594enum RecordSent {
595 BeforeSend,
599 AfterAck,
604}
605
606#[derive(Clone, Copy, Debug, Eq, PartialEq)]
608pub(crate) enum AutoCommitOutcome {
609 Committed,
610 Retriable,
613 Failed,
614}
615
616#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
620enum PartitionCommitError {
621 None,
622 TopicAuthorization,
626 Retriable,
629 CoordinatorUnknown,
632 Rebalance,
635 FencedInstanceId,
638 GroupAuthorization,
641 Fatal,
645}
646
647impl PartitionCommitError {
648 fn of(code: i16) -> Self {
649 match code {
650 0 => Self::None,
651 TOPIC_AUTHORIZATION_FAILED => Self::TopicAuthorization,
652 UNKNOWN_TOPIC_OR_PARTITION | COORDINATOR_LOAD_IN_PROGRESS => Self::Retriable,
653 REQUEST_TIMED_OUT | COORDINATOR_NOT_AVAILABLE | NOT_COORDINATOR => {
654 Self::CoordinatorUnknown
655 }
656 ILLEGAL_GENERATION | UNKNOWN_MEMBER_ID | REBALANCE_IN_PROGRESS => Self::Rebalance,
657 FENCED_INSTANCE_ID => Self::FencedInstanceId,
658 GROUP_AUTHORIZATION_FAILED => Self::GroupAuthorization,
659 _ => Self::Fatal,
660 }
661 }
662}
663
664impl PartitionCommitError {
665 fn of_response(response: &OffsetCommitResponse) -> Self {
668 response
669 .topics
670 .iter()
671 .flat_map(|topic| topic.partitions.iter())
672 .map(|partition| Self::of(partition.error_code))
673 .max()
674 .unwrap_or(Self::None)
675 }
676}
677
678pub(crate) fn names_moved_coordinator(response: &OffsetCommitResponse) -> bool {
683 PartitionCommitError::of_response(response) == PartitionCommitError::CoordinatorUnknown
684}
685
686pub(crate) fn auto_commit_outcome(
698 result: &Result<OffsetCommitResponse, ConsumerError>,
699) -> AutoCommitOutcome {
700 match result {
701 Ok(response) => match PartitionCommitError::of_response(response) {
702 PartitionCommitError::None => AutoCommitOutcome::Committed,
703 PartitionCommitError::Retriable | PartitionCommitError::CoordinatorUnknown => {
704 AutoCommitOutcome::Retriable
705 }
706 PartitionCommitError::TopicAuthorization
707 | PartitionCommitError::Rebalance
708 | PartitionCommitError::FencedInstanceId
709 | PartitionCommitError::GroupAuthorization
710 | PartitionCommitError::Fatal => AutoCommitOutcome::Failed,
711 },
712 Err(ConsumerError::Client(error))
713 if is_retriable_transport_error(error) || matches!(error, ClientError::Timeout(_)) =>
714 {
715 AutoCommitOutcome::Retriable
716 }
717 Err(ConsumerError::Server(code)) if is_retriable_coordinator_code(*code) => {
718 AutoCommitOutcome::Retriable
719 }
720 Err(_) => AutoCommitOutcome::Failed,
721 }
722}
723
724#[derive(Clone, Debug)]
727pub(crate) struct AutoCommit {
728 interval: Duration,
730 next_due: Arc<Mutex<tokio::time::Instant>>,
733 polled: Arc<Mutex<ConsumedPositions>>,
747}
748
749impl AutoCommit {
750 pub(crate) fn new(interval: Duration) -> Self {
751 Self {
752 interval,
753 next_due: Arc::new(Mutex::new(tokio::time::Instant::now() + interval)),
754 polled: Arc::new(Mutex::new(HashMap::new())),
755 }
756 }
757
758 pub(crate) async fn reset_polled(
769 &self,
770 resets: impl IntoIterator<Item = ((String, i32), Option<i64>)>,
771 ) {
772 let mut polled = self.polled.lock().await;
773 for (partition, offset) in resets {
774 match offset {
775 Some(offset) => {
776 if let Some(position) = polled.get_mut(&partition) {
777 position.offset = offset;
778 position.leader_epoch = -1;
779 }
780 }
781 None => {
782 polled.remove(&partition);
783 }
784 }
785 }
786 }
787
788 pub(crate) async fn record_sent(
792 &self,
793 sent: impl IntoIterator<Item = ((String, i32), ConsumedPosition)>,
794 ) {
795 let mut polled = self.polled.lock().await;
796 for (partition, position) in sent {
797 let known = polled.entry(partition).or_insert(position);
798 if known.ownership_id != position.ownership_id || known.offset < position.offset {
799 *known = position;
800 }
801 }
802 }
803
804 pub(crate) async fn restart_interval(&self) {
807 *self.next_due.lock().await = tokio::time::Instant::now() + self.interval;
808 }
809
810 pub(crate) async fn polled_offsets(
813 &self,
814 ownership_ids: &HashMap<(String, i32), u64>,
815 ) -> HashMap<(String, i32), (i64, i32)> {
816 owned_offsets(&*self.polled.lock().await, ownership_ids)
817 }
818}
819
820#[derive(Clone)]
822struct CommitRoute {
823 client: krabka_client_core::Client,
824 group_id: String,
825 group_instance_id: Option<String>,
826 coordinator_id: Arc<std::sync::atomic::AtomicI32>,
827 retry_policy: coordinator::CoordinatorRetryPolicy,
828}
829
830impl CommitRoute {
831 async fn send(
837 &self,
838 topics: Vec<OffsetCommitRequestTopic>,
839 generation: i32,
840 member_id: &str,
841 ) -> Result<OffsetCommitResponse, ConsumerError> {
842 let request = |topics| {
843 build_commit_request(
844 self.group_id.clone(),
845 generation,
846 member_id.to_owned(),
847 self.group_instance_id.clone(),
848 topics,
849 )
850 };
851 let target = self.coordinator_id.load(Ordering::Relaxed);
852 let result = self
853 .client
854 .broker(target)
855 .send(request(topics.clone()))
856 .await;
857 let moved = match &result {
858 Ok(response) => names_moved_coordinator(response),
859 Err(error) => is_retriable_transport_error(error),
860 };
861 if !moved {
862 return result.map_err(ConsumerError::from);
863 }
864 let id = find_coordinator(&self.client, &self.group_id, self.retry_policy).await?;
865 self.coordinator_id.store(id, Ordering::Relaxed);
866 self.client
867 .broker(id)
868 .send(request(topics))
869 .await
870 .map_err(ConsumerError::from)
871 }
872}
873
874#[derive(Debug)]
875#[allow(dead_code)] struct CommitTraceContext<'a> {
877 group_id: &'a str,
878 member_id: String,
879 generation: i32,
880}
881
882fn commit_trace_context(consumer: &Consumer) -> CommitTraceContext<'_> {
883 CommitTraceContext {
884 group_id: &consumer.group_id,
885 member_id: consumer.member_id(),
886 generation: consumer.current_generation.load(Ordering::Relaxed),
887 }
888}
889
890impl Consumer {
891 fn commit_route(&self) -> CommitRoute {
892 CommitRoute {
893 client: self.client.clone(),
894 group_id: self.group_id.clone(),
895 group_instance_id: self.group_instance_id.clone(),
896 coordinator_id: Arc::clone(&self.coordinator_id),
897 retry_policy: self.retry_policy,
898 }
899 }
900
901 async fn consumed_positions(&self) -> (CommitIdentity, ConsumedPositions) {
903 let identity = self.commit_identity.lock().await.clone();
904 let next_offsets = self.next_offsets.lock().await.clone();
905 let positions = self.positions.lock().await.clone();
906 let consumed = all_consumed(&identity.ownership_ids, &next_offsets, &positions);
907 (identity, consumed)
908 }
909
910 pub(crate) async fn maybe_auto_commit_async(&self) {
919 let Some(auto_commit) = &self.auto_commit else {
920 return;
921 };
922 let (identity, consumed) = self.consumed_positions().await;
923 auto_commit.polled.lock().await.clone_from(&consumed);
924 let mut next_due = auto_commit.next_due.lock().await;
925 let now = tokio::time::Instant::now();
926 if now < *next_due {
927 return;
928 }
929 let Ok(commit_guard) = Arc::clone(&self.commit_serialization).try_lock_owned() else {
934 return;
935 };
936 *next_due = now + auto_commit.interval;
937 drop(next_due);
938 let offsets = owned_offsets(&consumed, &identity.ownership_ids);
939 if offsets.is_empty() {
940 return;
941 }
942 let route = self.commit_route();
943 let next_due = Arc::clone(&auto_commit.next_due);
944 let retry_backoff = self.retry_policy.initial_backoff;
945 tokio::spawn(async move {
946 let _commit_guard = commit_guard;
947 let result = route
948 .send(
949 build_commit_topics(position_commits(offsets)),
950 identity.generation,
951 &identity.member_id,
952 )
953 .await;
954 match auto_commit_outcome(&result) {
955 AutoCommitOutcome::Committed => {}
956 AutoCommitOutcome::Retriable => {
957 tracing::debug!(
958 group = %route.group_id,
959 ?result,
960 "asynchronous auto commit failed with a retriable error"
961 );
962 *next_due.lock().await = tokio::time::Instant::now() + retry_backoff;
963 }
964 AutoCommitOutcome::Failed => {
965 tracing::warn!(
966 group = %route.group_id,
967 ?result,
968 "asynchronous auto commit failed"
969 );
970 }
971 }
972 });
973 }
974
975 pub(crate) async fn wait_for_async_commits(&self, deadline: tokio::time::Instant) {
986 let wait = async {
987 loop {
988 if self.commit_async_state.load(Ordering::Acquire) == ASYNC_COMMIT_IDLE {
989 return;
990 }
991 drop(self.commit_serialization.lock().await);
993 tokio::time::sleep(Duration::from_millis(1)).await;
994 }
995 };
996 if tokio::time::timeout_at(deadline, wait).await.is_err() {
997 tracing::warn!(
998 group = %self.group_id,
999 "asynchronous commits were still pending when the close timeout expired"
1000 );
1001 }
1002 }
1003
1004 pub(crate) async fn auto_commit_on_close(&self, deadline: tokio::time::Instant) {
1005 if self.auto_commit.is_none() {
1006 return;
1007 }
1008 let commit = async {
1011 let _commit_guard = self.commit_serialization.lock().await;
1012 let (_, consumed) = self.consumed_positions().await;
1013 let pending = consumed
1014 .into_iter()
1015 .map(|(partition, position)| {
1016 (
1017 partition,
1018 (
1019 OffsetAndMetadata::of_position(position.offset, position.leader_epoch),
1020 position.ownership_id,
1021 ),
1022 )
1023 })
1024 .collect::<HashMap<_, _>>();
1025 if pending.is_empty() {
1026 return Ok(());
1027 }
1028 let retry_deadline = CommitDeadline::after(self.retry_policy.timeout);
1029 self.commit_pending_offsets(pending, RecordSent::BeforeSend, retry_deadline)
1030 .await
1031 };
1032 match tokio::time::timeout_at(deadline, commit).await {
1033 Ok(Ok(())) => {}
1034 Ok(Err(error)) => {
1035 tracing::warn!(group = %self.group_id, %error, "synchronous auto commit on close failed");
1036 }
1037 Err(_) => {
1038 tracing::warn!(group = %self.group_id, "synchronous auto commit on close timed out");
1039 }
1040 }
1041 }
1042}
1043
1044impl Consumer {
1045 #[cfg_attr(test, mutants::skip)] #[tracing::instrument(
1051 name = "consumer.commit_sync",
1052 level = "debug",
1053 skip_all,
1054 fields(commit_context = ?commit_trace_context(self), partitions = tracing::field::Empty),
1055 err(level = "debug")
1056 )]
1057 pub async fn commit_sync(&self) -> Result<(), ConsumerError> {
1060 self.require_group_id()?;
1061 let deadline = CommitDeadline::after(self.default_api_timeout.to_std());
1062 let commit = async {
1066 let _commit_guard = self.commit_serialization.lock().await;
1067 let pending = {
1068 let identity = self.commit_identity.lock().await;
1069 self.ensure_active_group(&identity)?;
1070 let offsets = self.next_offsets.lock().await;
1071 let positions = self.positions.lock().await;
1072 all_consumed(&identity.ownership_ids, &offsets, &positions)
1075 .into_iter()
1076 .map(|(partition, position)| {
1077 (
1078 partition,
1079 (
1080 OffsetAndMetadata::of_position(
1081 position.offset,
1082 position.leader_epoch,
1083 ),
1084 position.ownership_id,
1085 ),
1086 )
1087 })
1088 .collect::<HashMap<_, _>>()
1089 };
1090 if pending.is_empty() {
1091 return Ok(());
1092 }
1093 tracing::Span::current().record("partitions", pending.len());
1094
1095 self.commit_pending_offsets(pending, RecordSent::BeforeSend, deadline)
1096 .await
1097 };
1098 deadline.run(commit).await
1099 }
1100
1101 pub async fn commit_offsets_sync(
1128 &self,
1129 offsets: HashMap<(String, i32), OffsetAndMetadata>,
1130 ) -> Result<(), ConsumerError> {
1131 self.require_group_id()?;
1132 if offsets.is_empty() {
1133 return Ok(());
1134 }
1135 let deadline = CommitDeadline::after(self.default_api_timeout.to_std());
1136 let commit = async {
1140 let _commit_guard = self.commit_serialization.lock().await;
1141 let pending = {
1142 let identity = self.commit_identity.lock().await;
1143 self.ensure_active_group(&identity)?;
1144 validate_selected_offsets(&offsets, &identity.ownership_ids)?;
1145 offsets
1146 .into_iter()
1147 .map(|(partition, offset)| {
1148 let ownership_id = identity.ownership_ids[&partition];
1149 (partition, (offset, ownership_id))
1150 })
1151 .collect::<HashMap<_, _>>()
1152 };
1153
1154 self.commit_pending_offsets(pending, RecordSent::AfterAck, deadline)
1155 .await
1156 };
1157 deadline.run(commit).await
1158 }
1159
1160 fn ensure_active_group(&self, identity: &CommitIdentity) -> Result<(), ConsumerError> {
1173 if identity.rejoin_on_poll
1174 || self.coordinator_shutdown.is_cancelled()
1175 || self
1176 .coordinator_handle
1177 .as_ref()
1178 .is_some_and(tokio::task::JoinHandle::is_finished)
1179 {
1180 return Err(ConsumerError::CommitFailed);
1181 }
1182 Ok(())
1183 }
1184
1185 async fn commit_pending_offsets(
1204 &self,
1205 mut pending: HashMap<(String, i32), (OffsetAndMetadata, u64)>,
1206 record: RecordSent,
1207 deadline: CommitDeadline,
1208 ) -> Result<(), ConsumerError> {
1209 let mut retry_backoff = self.retry_policy.initial_backoff;
1210 loop {
1211 let identity = self.commit_identity.lock().await.clone();
1212 self.ensure_active_group(&identity)?;
1213 retain_continuously_owned(&mut pending, &identity.ownership_ids);
1214 if pending.is_empty() {
1215 return Ok(());
1216 }
1217
1218 if record == RecordSent::BeforeSend
1219 && let Some(auto_commit) = &self.auto_commit
1220 {
1221 auto_commit.record_sent(pending_positions(&pending)).await;
1222 }
1223 let topics = build_commit_topics(
1224 pending
1225 .iter()
1226 .map(|(partition, (offset, _))| (partition.clone(), offset.clone()))
1227 .collect(),
1228 );
1229 let outcome = match self
1230 .commit_topics_once(topics, (identity.generation, identity.member_id.clone()))
1231 .await
1232 {
1233 Ok(outcome) => outcome,
1234 Err(ConsumerError::Client(error))
1239 if is_retriable_transport_error(&error)
1240 || matches!(error, ClientError::Timeout(_)) =>
1241 {
1242 if tokio::time::Instant::now() >= deadline.expires_at {
1243 return Err(commit_timeout(
1244 deadline.timeout,
1245 Some(ConsumerError::Client(error)),
1246 ));
1247 }
1248 tracing::warn!(
1249 group = %self.group_id,
1250 %error,
1251 "offset commit request failed; finding the coordinator again and retrying",
1252 );
1253 self.client
1254 .evict_broker(self.coordinator_id.load(Ordering::Relaxed));
1255 self.find_coordinator_again(deadline.expires_at).await;
1256 if !self
1257 .sleep_before_retry(retry_backoff, deadline.expires_at)
1258 .await
1259 {
1260 return Err(commit_timeout(
1261 deadline.timeout,
1262 Some(ConsumerError::Client(error)),
1263 ));
1264 }
1265 retry_backoff = next_backoff(retry_backoff, self.retry_policy.max_backoff);
1266 continue;
1267 }
1268 Err(error) => return Err(error),
1269 };
1270 match outcome {
1271 CommitOutcome::Acked(acknowledged) => {
1272 if self
1273 .retain_unacknowledged(record, &mut pending, &acknowledged)
1274 .await
1275 {
1276 return Ok(());
1277 }
1278 return Err(ConsumerError::IllegalState(
1279 "offset commit response omitted requested partitions".into(),
1280 ));
1281 }
1282 CommitOutcome::Retriable {
1283 code,
1284 acknowledged,
1285 find_coordinator,
1286 } => {
1287 if self
1291 .retain_unacknowledged(record, &mut pending, &acknowledged)
1292 .await
1293 {
1294 return Ok(());
1295 }
1296 if tokio::time::Instant::now() >= deadline.expires_at {
1297 return Err(commit_timeout(
1298 deadline.timeout,
1299 Some(ConsumerError::Server(code)),
1300 ));
1301 }
1302 tracing::warn!(
1303 group = %self.group_id,
1304 error_code = code,
1305 "offset commit failed with a retriable error; retrying",
1306 );
1307 if find_coordinator {
1308 self.find_coordinator_again(deadline.expires_at).await;
1309 }
1310 if !self
1311 .sleep_before_retry(retry_backoff, deadline.expires_at)
1312 .await
1313 {
1314 return Err(commit_timeout(
1315 deadline.timeout,
1316 Some(ConsumerError::Server(code)),
1317 ));
1318 }
1319 retry_backoff = next_backoff(retry_backoff, self.retry_policy.max_backoff);
1320 }
1321 }
1322 }
1323 }
1324
1325 async fn record_acknowledged(
1329 &self,
1330 record: RecordSent,
1331 pending: &HashMap<(String, i32), (OffsetAndMetadata, u64)>,
1332 acknowledged: &HashSet<(String, i32)>,
1333 ) {
1334 if record != RecordSent::AfterAck {
1335 return;
1336 }
1337 let Some(auto_commit) = &self.auto_commit else {
1338 return;
1339 };
1340 let acked = pending
1341 .iter()
1342 .filter(|(partition, _)| acknowledged.contains(*partition))
1343 .map(|(partition, value)| (partition.clone(), value.clone()))
1344 .collect();
1345 auto_commit.record_sent(pending_positions(&acked)).await;
1346 }
1347
1348 async fn retain_unacknowledged(
1349 &self,
1350 record: RecordSent,
1351 pending: &mut HashMap<(String, i32), (OffsetAndMetadata, u64)>,
1352 acknowledged: &HashSet<(String, i32)>,
1353 ) -> bool {
1354 self.record_acknowledged(record, pending, acknowledged)
1355 .await;
1356 pending.retain(|partition, _| !acknowledged.contains(partition));
1357 pending.is_empty()
1358 }
1359
1360 async fn sleep_before_retry(&self, backoff: Duration, deadline: tokio::time::Instant) -> bool {
1365 tokio::time::sleep_until((tokio::time::Instant::now() + backoff).min(deadline)).await;
1366 tokio::time::Instant::now() < deadline
1367 }
1368
1369 async fn find_coordinator_again(&self, deadline: tokio::time::Instant) {
1375 let retry_policy = coordinator::CoordinatorRetryPolicy {
1376 timeout: deadline.saturating_duration_since(tokio::time::Instant::now()),
1377 ..self.retry_policy
1378 };
1379 match find_coordinator(&self.client, &self.group_id, retry_policy).await {
1380 Ok(id) => self.coordinator_id.store(id, Ordering::Relaxed),
1381 Err(error) => {
1382 tracing::warn!(
1383 group = %self.group_id,
1384 %error,
1385 "coordinator lookup for an offset commit failed; retrying with the last known coordinator",
1386 );
1387 }
1388 }
1389 }
1390
1391 async fn commit_topics_once(
1395 &self,
1396 topics: Vec<OffsetCommitRequestTopic>,
1397 identity: (i32, String),
1398 ) -> Result<CommitOutcome, ConsumerError> {
1399 let (generation, member_id) = identity;
1400 let target = self.coordinator_id.load(Ordering::Relaxed);
1401 let resp = self
1402 .client
1403 .broker(target)
1404 .send(build_commit_request(
1405 self.group_id.clone(),
1406 generation,
1407 member_id.clone(),
1408 self.group_instance_id.clone(),
1409 topics,
1410 ))
1411 .await?;
1412
1413 let current = self.commit_identity.lock().await.clone();
1414 let coordinator_alive = !self.subscription.borrow().is_none()
1417 && self
1418 .coordinator_handle
1419 .as_ref()
1420 .is_none_or(|h| !h.is_finished());
1421 let outcome = commit_response_outcome(
1422 &resp,
1423 CommitResponseContext {
1424 group_id: &self.group_id,
1425 group_instance_id: self.group_instance_id.as_deref(),
1426 coordinator_alive,
1427 generation_unchanged: current.generation == generation
1428 && current.member_id == member_id,
1429 },
1430 )?;
1431 Ok(outcome)
1432 }
1433
1434 #[cfg_attr(test, mutants::skip)] #[tracing::instrument(
1443 name = "consumer.commit_async",
1444 level = "debug",
1445 skip_all,
1446 fields(commit_context = ?commit_trace_context(self))
1447 )]
1448 pub fn commit_async(&self) -> Result<(), ConsumerError> {
1453 self.require_group_id()?;
1454 self.schedule_commit_async();
1455 Ok(())
1456 }
1457
1458 pub fn commit_async_with_callback<F>(&self, callback: F) -> Result<(), ConsumerError>
1471 where
1472 F: FnOnce(&HashMap<(String, i32), OffsetAndMetadata>, Result<(), &ConsumerError>)
1473 + Send
1474 + 'static,
1475 {
1476 self.require_group_id()?;
1477 self.commit_async_callbacks
1478 .lock()
1479 .unwrap_or_else(std::sync::PoisonError::into_inner)
1480 .push(Box::new(callback));
1481 self.schedule_commit_async();
1482 Ok(())
1483 }
1484
1485 fn schedule_commit_async(&self) {
1486 loop {
1487 match self.commit_async_state.load(Ordering::Acquire) {
1488 ASYNC_COMMIT_IDLE => {
1489 if self
1490 .commit_async_state
1491 .compare_exchange(
1492 ASYNC_COMMIT_IDLE,
1493 ASYNC_COMMIT_RUNNING,
1494 Ordering::AcqRel,
1495 Ordering::Acquire,
1496 )
1497 .is_ok()
1498 {
1499 break;
1500 }
1501 }
1502 ASYNC_COMMIT_RUNNING => {
1503 if self
1504 .commit_async_state
1505 .compare_exchange(
1506 ASYNC_COMMIT_RUNNING,
1507 ASYNC_COMMIT_DIRTY,
1508 Ordering::AcqRel,
1509 Ordering::Acquire,
1510 )
1511 .is_ok()
1512 {
1513 return;
1514 }
1515 }
1516 ASYNC_COMMIT_DIRTY => return,
1517 _ => unreachable!("invalid async commit state"),
1518 }
1519 }
1520
1521 let route = self.commit_route();
1522 let commit_identity = Arc::clone(&self.commit_identity);
1523 let commit_serialization = Arc::clone(&self.commit_serialization);
1524 let offsets = Arc::clone(&self.next_offsets);
1525 let positions = Arc::clone(&self.positions);
1526 let commit_async_state = Arc::clone(&self.commit_async_state);
1527 let auto_commit = self.auto_commit.clone();
1528 let callbacks = Arc::clone(&self.commit_async_callbacks);
1529 tokio::spawn(async move {
1530 loop {
1531 {
1532 let _commit_guard = commit_serialization.lock().await;
1533 commit_async_state.store(ASYNC_COMMIT_RUNNING, Ordering::Release);
1536 let waiting = take_callbacks(&callbacks);
1537 let (sent, (generation, member_id)) = snapshot_commit_offsets(
1538 &commit_identity,
1539 &offsets,
1540 &positions,
1541 auto_commit.as_ref(),
1542 )
1543 .await;
1544 let result = if sent.is_empty() {
1547 Ok(())
1548 } else {
1549 let response = route
1550 .send(build_commit_topics(sent.clone()), generation, &member_id)
1551 .await;
1552 async_commit_result(
1553 response,
1554 &route.group_id,
1555 route.group_instance_id.as_deref(),
1556 )
1557 };
1558 if let Err(error) = &result {
1561 tracing::warn!(%error, "commit_async failed");
1562 }
1563 for callback in waiting {
1564 let call = std::panic::AssertUnwindSafe(|| {
1568 callback(&sent, result.as_ref().copied());
1569 });
1570 if std::panic::catch_unwind(call).is_err() {
1571 tracing::error!("an offset commit callback panicked");
1572 }
1573 }
1574 }
1575
1576 if commit_async_state
1577 .compare_exchange(
1578 ASYNC_COMMIT_RUNNING,
1579 ASYNC_COMMIT_IDLE,
1580 Ordering::AcqRel,
1581 Ordering::Acquire,
1582 )
1583 .is_ok()
1584 {
1585 return;
1586 }
1587 if commit_async_state
1591 .compare_exchange(
1592 ASYNC_COMMIT_DIRTY,
1593 ASYNC_COMMIT_RUNNING,
1594 Ordering::AcqRel,
1595 Ordering::Acquire,
1596 )
1597 .is_err()
1598 {
1599 return;
1600 }
1601 }
1602 });
1603 }
1604}
1605
1606#[cfg(test)]
1607mod tests {
1608 use std::{
1609 sync::atomic::{AtomicI32, AtomicUsize},
1610 time::Duration,
1611 };
1612
1613 use assert2::check;
1614 use coordinator::CoordinatorRetryPolicy;
1615 use krabka_client_core::{
1616 Client, MockBroker, MockReply,
1617 mock_test_support::{group_api_versions, offset_commit_response},
1618 };
1619 use krabka_protocol::{
1620 Encode, UnknownTaggedFields,
1621 owned::{
1622 api_versions_request, find_coordinator_request, metadata_request,
1623 offset_commit_request,
1624 offset_commit_request::{OffsetCommitRequestPartition, OffsetCommitRequestTopic},
1625 offset_commit_response::{
1626 OffsetCommitResponse, OffsetCommitResponsePartition, OffsetCommitResponseTopic,
1627 },
1628 },
1629 primitives::uuid::Uuid,
1630 };
1631 use krabka_units::secs;
1632
1633 use super::*;
1634
1635 fn response(errors: &[i16]) -> OffsetCommitResponse {
1636 OffsetCommitResponse {
1637 throttle_time_ms: 0,
1638 topics: vec![OffsetCommitResponseTopic {
1639 name: "topic".into(),
1640 topic_id: Uuid::ZERO,
1641 partitions: errors
1642 .iter()
1643 .enumerate()
1644 .map(
1645 |(partition_index, error_code)| OffsetCommitResponsePartition {
1646 partition_index: i32::try_from(partition_index).unwrap(),
1647 error_code: *error_code,
1648 unknown_tagged_fields: UnknownTaggedFields::default(),
1649 },
1650 )
1651 .collect(),
1652 unknown_tagged_fields: UnknownTaggedFields::default(),
1653 }],
1654 unknown_tagged_fields: UnknownTaggedFields::default(),
1655 }
1656 }
1657
1658 async fn selected_commit(consumer: &Consumer, offset: i64) -> Result<(), ConsumerError> {
1659 consumer
1660 .commit_offsets_sync(
1661 maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata::new(offset)},
1662 )
1663 .await
1664 }
1665
1666 async fn wait_for_async_commit_idle(consumer: &Consumer, timeout: Duration) {
1667 tokio::time::timeout(timeout, async {
1668 while consumer.commit_async_state.load(Ordering::Acquire) != ASYNC_COMMIT_IDLE {
1669 tokio::task::yield_now().await;
1670 }
1671 })
1672 .await
1673 .expect("async commit worker becomes idle");
1674 }
1675
1676 fn request_string_end(body: &[u8], start: usize) -> usize {
1677 let len = i16::from_be_bytes([body[start], body[start + 1]]);
1678 start + 2 + usize::try_from(len.max(0)).unwrap()
1679 }
1680
1681 fn request_string(body: &[u8], start: usize) -> (usize, String) {
1682 let end = request_string_end(body, start);
1683 (
1684 end,
1685 String::from_utf8(body[start + 2..end].to_vec()).unwrap(),
1686 )
1687 }
1688
1689 fn request_generation(body: &[u8]) -> i32 {
1690 let group_start = request_string_end(body, 0);
1691 let generation_start = request_string_end(body, group_start);
1692 i32::from_be_bytes(
1693 body[generation_start..generation_start + 4]
1694 .try_into()
1695 .unwrap(),
1696 )
1697 }
1698
1699 fn request_member_id(body: &[u8]) -> String {
1700 let group_start = request_string_end(body, 0);
1701 let generation_start = request_string_end(body, group_start);
1702 let (_, member_id) = request_string(body, generation_start + 4);
1703 member_id
1704 }
1705
1706 fn request_offsets(body: &[u8]) -> Vec<(i32, i64)> {
1707 fn i32_at(body: &[u8], start: usize) -> i32 {
1708 i32::from_be_bytes(body[start..start + 4].try_into().unwrap())
1709 }
1710 fn i64_at(body: &[u8], start: usize) -> i64 {
1711 i64::from_be_bytes(body[start..start + 8].try_into().unwrap())
1712 }
1713
1714 let mut cursor = request_string_end(body, 0);
1715 cursor = request_string_end(body, cursor);
1716 cursor += 4;
1717 cursor = request_string_end(body, cursor);
1718 cursor += 8;
1719 let topic_count = usize::try_from(i32_at(body, cursor)).unwrap();
1720 cursor += 4;
1721 let mut offsets = Vec::new();
1722 for _ in 0..topic_count {
1723 cursor = request_string_end(body, cursor);
1724 let partition_count = usize::try_from(i32_at(body, cursor)).unwrap();
1725 cursor += 4;
1726 for _ in 0..partition_count {
1727 offsets.push((i32_at(body, cursor), i64_at(body, cursor + 4)));
1728 cursor += 12;
1729 cursor = request_string_end(body, cursor);
1730 }
1731 }
1732 offsets
1733 }
1734
1735 fn commit_identity(generation: i32, member_id: &str) -> Arc<Mutex<CommitIdentity>> {
1736 Arc::new(Mutex::new(CommitIdentity {
1737 generation,
1738 member_id: member_id.into(),
1739 ownership_ids: maplit::hashmap! {("topic".into(), 0) => 1},
1740 rejoin_on_poll: false,
1741 }))
1742 }
1743
1744 fn advance_generation_on_first_commit()
1745 -> impl Fn(&Arc<Mutex<CommitIdentity>>, &Arc<AtomicI32>, i32, &str) + Send + 'static {
1746 |identity, generation, _request_generation, _request_member_id| {
1747 identity.try_lock().unwrap().generation = 8;
1748 generation.store(8, Ordering::Relaxed);
1749 }
1750 }
1751
1752 #[derive(bon::Builder)]
1753 struct SelectedCommitConsumerStart<F> {
1754 #[builder(default = commit_identity(7, "member-a"))]
1755 commit_identity: Arc<Mutex<CommitIdentity>>,
1756 #[builder(default = Arc::new(tokio::sync::Notify::new()))]
1757 assignment_changed: Arc<tokio::sync::Notify>,
1758 #[builder(default = Arc::new(AtomicI32::new(7)))]
1759 generation: Arc<AtomicI32>,
1760 #[builder(default = Arc::new(AtomicUsize::new(0)))]
1761 requests: Arc<AtomicUsize>,
1762 #[builder(default = false)]
1763 mixed_response: bool,
1764 on_first_request: F,
1765 }
1766
1767 struct SelectedCommitConsumer {
1768 consumer: Consumer,
1769 mock: MockBroker,
1770 requests: Arc<AtomicUsize>,
1771 seen_offsets: Arc<std::sync::Mutex<Vec<(i32, i64)>>>,
1772 }
1773
1774 impl<F> SelectedCommitConsumerStart<F>
1775 where
1776 F: Fn(&Arc<Mutex<CommitIdentity>>, &Arc<AtomicI32>, i32, &str) + Send + 'static,
1777 {
1778 async fn start(self) -> SelectedCommitConsumer {
1779 let Self {
1780 commit_identity,
1781 assignment_changed,
1782 generation,
1783 requests,
1784 mixed_response,
1785 on_first_request,
1786 } = self;
1787 let identity_in_mock = Arc::clone(&commit_identity);
1788 let generation_in_mock = Arc::clone(&generation);
1789 let changed_in_mock = Arc::clone(&assignment_changed);
1790 let requests_in_mock = Arc::clone(&requests);
1791 let seen_offsets = Arc::new(std::sync::Mutex::new(Vec::new()));
1792 let seen_offsets_in_mock = Arc::clone(&seen_offsets);
1793 let mock = MockBroker::start(move |api_key, version, _corr_id, body| {
1794 if api_key == api_versions_request::API_KEY {
1795 return Some(group_api_versions(offset_commit_request::API_KEY, (2, 2)));
1796 }
1797 if api_key != offset_commit_request::API_KEY {
1798 return None;
1799 }
1800 let attempt = requests_in_mock.fetch_add(1, Ordering::SeqCst);
1801 let offsets = request_offsets(body);
1802 seen_offsets_in_mock.lock().unwrap().extend(&offsets);
1803 let partitions = offsets
1804 .iter()
1805 .map(|(partition, _)| *partition)
1806 .collect::<Vec<_>>();
1807 if attempt == 0 {
1808 on_first_request(
1809 &identity_in_mock,
1810 &generation_in_mock,
1811 request_generation(body),
1812 &request_member_id(body),
1813 );
1814 changed_in_mock.notify_waiters();
1815 let errors = if mixed_response {
1816 vec![0, 27]
1817 } else {
1818 vec![27]
1819 };
1820 Some(test_support::encode_protocol_response(
1821 &response(&errors),
1822 version,
1823 ))
1824 } else {
1825 let errors = if mixed_response && partitions != [1] {
1826 vec![42; partitions.len()]
1827 } else {
1828 vec![0; partitions.len()]
1829 };
1830 let mut response = response(&errors);
1831 for (partition, partition_index) in
1832 response.topics[0].partitions.iter_mut().zip(partitions)
1833 {
1834 partition.partition_index = partition_index;
1835 }
1836 Some(test_support::encode_protocol_response(&response, version))
1837 }
1838 })
1839 .await;
1840 let consumer = commit_consumer(
1841 CommitConsumerStart::builder()
1842 .mock(&mock)
1843 .commit_identity(commit_identity)
1844 .assignment_changed(assignment_changed)
1845 .generation(generation)
1846 .build(),
1847 )
1848 .await;
1849 SelectedCommitConsumer {
1850 consumer,
1851 mock,
1852 requests,
1853 seen_offsets,
1854 }
1855 }
1856 }
1857
1858 #[derive(Clone, Copy)]
1859 enum AutoCommitMockResponse {
1860 Success,
1861 RetryOnce,
1862 }
1863
1864 struct AutoCommitCoordinator {
1865 mock: MockBroker,
1866 requests: Arc<AtomicUsize>,
1867 }
1868
1869 async fn auto_commit_coordinator(
1870 response_mode: AutoCommitMockResponse,
1871 ) -> AutoCommitCoordinator {
1872 let requests = Arc::new(AtomicUsize::new(0));
1873 let requests_in_mock = Arc::clone(&requests);
1874 let mock = MockBroker::start(move |api_key, version, _corr_id, _body| {
1875 if api_key == api_versions_request::API_KEY {
1876 return Some(group_api_versions(offset_commit_request::API_KEY, (2, 2)));
1877 }
1878 if api_key != offset_commit_request::API_KEY {
1879 return None;
1880 }
1881 let attempt = requests_in_mock.fetch_add(1, Ordering::SeqCst);
1882 let error_code = match (response_mode, attempt) {
1883 (AutoCommitMockResponse::RetryOnce, 0) => 7,
1884 _ => 0,
1885 };
1886 Some(test_support::encode_protocol_response(
1887 &response(&[error_code]),
1888 version,
1889 ))
1890 })
1891 .await;
1892 AutoCommitCoordinator { mock, requests }
1893 }
1894
1895 #[derive(bon::Builder)]
1898 #[builder(start_fn = builder, finish_fn = build)]
1899 struct CommitConsumerStart<'a> {
1900 mock: &'a MockBroker,
1901 commit_identity: Arc<Mutex<CommitIdentity>>,
1902 #[builder(default = Arc::new(tokio::sync::Notify::new()))]
1903 assignment_changed: Arc<tokio::sync::Notify>,
1904 #[builder(default = Arc::new(AtomicI32::new(7)))]
1905 generation: Arc<AtomicI32>,
1906 }
1907
1908 async fn commit_consumer(start: CommitConsumerStart<'_>) -> Consumer {
1909 let client = Client::builder()
1910 .bootstrap(start.mock.addr.to_string())
1911 .build()
1912 .await
1913 .unwrap();
1914 let ownership = start.commit_identity.lock().await.ownership_ids.clone();
1915 let assigned = ownership.keys().cloned().collect::<Vec<_>>();
1916 let next_offsets = ownership
1917 .keys()
1918 .cloned()
1919 .map(|partition| (partition, 12))
1920 .collect::<HashMap<_, i64>>();
1921 Consumer::test_builder(client)
1922 .commit_identity(start.commit_identity)
1923 .current_generation(start.generation)
1924 .assigned(assigned)
1925 .assignment_changed(start.assignment_changed)
1926 .next_offsets(next_offsets)
1927 .build()
1928 }
1929
1930 async fn default_commit_consumer(mock: &MockBroker) -> Consumer {
1931 let start = CommitConsumerStart::builder()
1932 .mock(mock)
1933 .commit_identity(commit_identity(7, "member-a"))
1934 .build();
1935 commit_consumer(start).await
1936 }
1937
1938 #[test]
1941 fn selected_offset_validation_follows_kafka() {
1942 let assigned = maplit::hashmap! {("topic".to_string(), 2) => 1};
1943
1944 let mut actual = Vec::new();
1945 let mut wanted = Vec::new();
1946 for (name, partition, offset, expected) in [
1947 (
1948 "negative",
1949 ("topic", 2),
1950 -1,
1951 Some("invalid seek offset -1: must be non-negative"),
1952 ),
1953 (
1954 "unassigned",
1955 ("other", 2),
1956 1,
1957 Some("illegal state: cannot commit unassigned partition other-2"),
1958 ),
1959 ("at the position", ("topic", 2), 11, None),
1960 ("past the position", ("topic", 2), 100, None),
1961 ] {
1962 let offsets = maplit::hashmap! {(partition.0.to_string(), partition.1) => OffsetAndMetadata::new(offset)};
1963 let result = validate_selected_offsets(&offsets, &assigned).err();
1964 actual.push((name, result.map(|error| error.to_string())));
1965 wanted.push((name, expected.map(str::to_owned)));
1966 }
1967 assert2::assert!(actual == wanted);
1968 }
1969
1970 #[test]
1971 fn commit_offsets_use_position_epoch_or_unknown_epoch() {
1972 let mut raw = HashMap::new();
1973 raw.insert(("known".into(), 0), 11);
1974 raw.insert(("unknown".into(), 1), 22);
1975
1976 let mut positions = HashMap::new();
1977 positions.insert(
1978 ("known".into(), 0),
1979 PartitionPosition {
1980 offset_epoch: krabka_ids::LeaderEpoch(7),
1981 ..Default::default()
1982 },
1983 );
1984
1985 let offsets = commit_offsets(raw, &positions);
1986 assert2::assert!(
1987 offsets
1988 == maplit::hashmap! {("known".into(), 0) => (11, 7), ("unknown".into(), 1) => (22, -1)}
1989 );
1990 }
1991
1992 #[test]
1996 fn consumed_positions_keep_owned_valid_positions_of_the_current_ownership() {
1997 let ownership_ids = maplit::hashmap! {("orders".into(), 0) => 1, ("orders".into(), 1) => 2, ("orders".into(), 2) => 3, ("orders".into(), 3) => 4};
1998 let next_offsets = maplit::hashmap! {("orders".into(), 0) => 12, ("orders".into(), 1) => 0, ("orders".into(), 2) => i64::MAX, ("orders".into(), 3) => -1, ("unowned".into(), 0) => 5};
1999 let positions = maplit::hashmap! {("orders".into(), 0) => PartitionPosition {
2000 offset_epoch: krabka_ids::LeaderEpoch(4),
2001 ..Default::default()
2002 }};
2003
2004 let consumed = all_consumed(&ownership_ids, &next_offsets, &positions);
2005 assert2::assert!(
2006 consumed
2007 == maplit::hashmap! {("orders".into(), 0) => ConsumedPosition {
2008 offset: 12,
2009 leader_epoch: 4,
2010 ownership_id: 1,
2011 }, ("orders".into(), 1) => ConsumedPosition {
2012 offset: 0,
2013 leader_epoch: -1,
2014 ownership_id: 2,
2015 }}
2016 );
2017
2018 let reassigned = maplit::hashmap! {("orders".into(), 0) => 9, ("orders".into(), 1) => 2};
2019 assert2::assert!(
2020 owned_offsets(&consumed, &reassigned)
2021 == maplit::hashmap! {("orders".into(), 1) => (0, -1)}
2022 );
2023 }
2024
2025 #[test]
2028 fn auto_commit_outcome_follows_kafka_commit_error_classes() {
2029 let mut actual = Vec::new();
2030 let mut wanted = Vec::new();
2031 for (name, result, expected) in [
2032 (
2033 "success",
2034 Ok(response(&[0, 0])),
2035 AutoCommitOutcome::Committed,
2036 ),
2037 (
2038 "unknown topic or partition",
2039 Ok(response(&[0, 3])),
2040 AutoCommitOutcome::Retriable,
2041 ),
2042 (
2043 "request timed out",
2044 Ok(response(&[7])),
2045 AutoCommitOutcome::Retriable,
2046 ),
2047 (
2048 "coordinator loading",
2049 Ok(response(&[14])),
2050 AutoCommitOutcome::Retriable,
2051 ),
2052 (
2053 "coordinator not available",
2054 Ok(response(&[15])),
2055 AutoCommitOutcome::Retriable,
2056 ),
2057 (
2058 "not coordinator",
2059 Ok(response(&[16])),
2060 AutoCommitOutcome::Retriable,
2061 ),
2062 (
2063 "illegal generation",
2064 Ok(response(&[22])),
2065 AutoCommitOutcome::Failed,
2066 ),
2067 (
2068 "rebalance in progress",
2069 Ok(response(&[27])),
2070 AutoCommitOutcome::Failed,
2071 ),
2072 (
2073 "group authorization failed",
2074 Ok(response(&[30])),
2075 AutoCommitOutcome::Failed,
2076 ),
2077 (
2078 "disconnected",
2079 Err(ConsumerError::Client(ClientError::Disconnected)),
2080 AutoCommitOutcome::Retriable,
2081 ),
2082 (
2083 "client timeout",
2084 Err(ConsumerError::Client(ClientError::Timeout(secs(30)))),
2085 AutoCommitOutcome::Retriable,
2086 ),
2087 (
2088 "coordinator lookup not available",
2089 Err(ConsumerError::Server(15)),
2090 AutoCommitOutcome::Retriable,
2091 ),
2092 (
2093 "coordinator lookup refused",
2094 Err(ConsumerError::Server(30)),
2095 AutoCommitOutcome::Failed,
2096 ),
2097 (
2098 "topic authorization failed",
2099 Ok(response(&[0, 29])),
2100 AutoCommitOutcome::Failed,
2101 ),
2102 (
2103 "retriable partition, then a fatal partition",
2104 Ok(response(&[3, 30])),
2105 AutoCommitOutcome::Failed,
2106 ),
2107 (
2108 "fatal partition, then a retriable partition",
2109 Ok(response(&[30, 3])),
2110 AutoCommitOutcome::Failed,
2111 ),
2112 (
2113 "retriable coordinator partition, then a fatal partition",
2114 Ok(response(&[16, 22])),
2115 AutoCommitOutcome::Failed,
2116 ),
2117 (
2118 "topic authorization failure, then a retriable partition",
2119 Ok(response(&[29, 3])),
2120 AutoCommitOutcome::Retriable,
2121 ),
2122 (
2123 "retriable partition, then a topic authorization failure",
2124 Ok(response(&[3, 29])),
2125 AutoCommitOutcome::Retriable,
2126 ),
2127 ] {
2128 actual.push((name, auto_commit_outcome(&result)));
2129 wanted.push((name, expected));
2130 }
2131 assert2::assert!(actual == wanted);
2132 }
2133
2134 #[test]
2140 fn names_moved_coordinator_only_for_a_retriable_coordinator_response() {
2141 let mut actual = Vec::new();
2142 let mut wanted = Vec::new();
2143 for (name, errors, expected) in [
2144 ("success", &[0, 0][..], false),
2145 ("request timed out", &[0, 7][..], true),
2146 ("coordinator not available", &[15][..], true),
2147 ("not coordinator with a retriable code", &[3, 16][..], true),
2148 ("coordinator load in progress", &[14][..], false),
2149 ("unknown topic or partition", &[3][..], false),
2150 (
2151 "request timed out and metadata too large",
2152 &[7, 12][..],
2153 false,
2154 ),
2155 (
2156 "not coordinator and group authorization",
2157 &[30, 16][..],
2158 false,
2159 ),
2160 ("not coordinator and a rebalance", &[16, 27][..], false),
2161 (
2162 "not coordinator and topic authorization",
2163 &[29, 16][..],
2164 true,
2165 ),
2166 ] {
2167 actual.push((name, names_moved_coordinator(&response(errors))));
2168 wanted.push((name, expected));
2169 }
2170 assert2::assert!(actual == wanted);
2171 }
2172
2173 #[tokio::test(start_paused = true)]
2176 async fn busy_commit_lock_moves_the_interval_commit_to_the_next_poll() {
2177 let coordinator = auto_commit_coordinator(AutoCommitMockResponse::Success).await;
2178 let mut consumer = commit_consumer(
2179 CommitConsumerStart::builder()
2180 .mock(&coordinator.mock)
2181 .commit_identity(commit_identity(7, "member-a"))
2182 .build(),
2183 )
2184 .await;
2185 let interval = Duration::from_hours(1);
2186 consumer.auto_commit = Some(AutoCommit::new(interval));
2187 tokio::time::advance(interval).await;
2188
2189 let busy = consumer.commit_serialization.lock().await;
2190 consumer.maybe_auto_commit_async().await;
2191 let while_busy = coordinator.requests.load(Ordering::SeqCst);
2192 drop(busy);
2193 consumer.maybe_auto_commit_async().await;
2194 drop(consumer.commit_serialization.lock().await);
2195 let after = coordinator.requests.load(Ordering::SeqCst);
2196
2197 coordinator.mock.stop();
2198 assert2::assert!((while_busy, after) == (0, 1));
2199 }
2200
2201 #[tokio::test(start_paused = true)]
2205 async fn retriable_auto_commit_failure_commits_again_after_the_retry_backoff() {
2206 let coordinator = auto_commit_coordinator(AutoCommitMockResponse::RetryOnce).await;
2207 let mut consumer = commit_consumer(
2208 CommitConsumerStart::builder()
2209 .mock(&coordinator.mock)
2210 .commit_identity(commit_identity(7, "member-a"))
2211 .build(),
2212 )
2213 .await;
2214 let interval = Duration::from_hours(1);
2215 consumer.auto_commit = Some(AutoCommit::new(interval));
2216 let backoff = consumer.retry_policy.initial_backoff;
2217 let ticker = tokio::spawn(async {
2221 loop {
2222 tokio::time::sleep(Duration::from_millis(1)).await;
2223 }
2224 });
2225
2226 let mut sent = Vec::new();
2227 for advance in [interval, Duration::ZERO, backoff, Duration::ZERO] {
2228 tokio::time::advance(advance).await;
2229 consumer.maybe_auto_commit_async().await;
2230 drop(consumer.commit_serialization.lock().await);
2231 sent.push(coordinator.requests.load(Ordering::SeqCst));
2232 }
2233
2234 ticker.abort();
2235 coordinator.mock.stop();
2236 assert2::assert!(sent == vec![1, 1, 2, 2]);
2237 }
2238
2239 #[tokio::test]
2240 async fn snapshot_commit_offsets_is_empty_without_offsets() {
2241 let identity = commit_identity(7, "member-a");
2242 let offsets = Arc::new(Mutex::new(HashMap::new()));
2243 let positions = Arc::new(Mutex::new(HashMap::new()));
2244
2245 let snapshot = snapshot_commit_offsets(&identity, &offsets, &positions, None).await;
2246
2247 assert2::assert!(snapshot == (HashMap::new(), (7, "member-a".into())));
2248 }
2249
2250 #[tokio::test]
2251 async fn snapshot_commit_offsets_preserves_offsets_and_epochs() {
2252 let identity = Arc::new(Mutex::new(CommitIdentity {
2253 generation: 7,
2254 member_id: "member-a".into(),
2255 ownership_ids: maplit::hashmap! {("alpha".into(), 0) => 1, ("alpha".into(), 1) => 2},
2256 rejoin_on_poll: false,
2257 }));
2258 let offsets = Arc::new(Mutex::new(
2259 maplit::hashmap! {("alpha".to_string(), 0) => 10, ("alpha".to_string(), 1) => 20},
2260 ));
2261 let positions = Arc::new(Mutex::new(
2262 maplit::hashmap! {("alpha".to_string(), 1) => PartitionPosition {
2263 offset_epoch: krabka_ids::LeaderEpoch(7),
2264 ..Default::default()
2265 }},
2266 ));
2267 let snapshot = snapshot_commit_offsets(&identity, &offsets, &positions, None).await;
2268
2269 assert2::assert!(
2270 snapshot
2271 == (
2272 maplit::hashmap! {("alpha".to_string(), 0) => OffsetAndMetadata::new(10), ("alpha".to_string(), 1) => OffsetAndMetadata {
2273 offset: 20,
2274 leader_epoch: Some(7),
2275 metadata: String::new(),
2276 }},
2277 (7, "member-a".into())
2278 )
2279 );
2280 }
2281
2282 #[test]
2283 fn build_commit_request_preserves_group_member_generation_and_topics() {
2284 let topics = vec![
2285 krabka_protocol::owned::offset_commit_request::OffsetCommitRequestTopic {
2286 name: "topic".into(),
2287 topic_id: Uuid::ZERO,
2288 partitions: vec![OffsetCommitRequestPartition {
2289 partition_index: 3,
2290 committed_offset: 99,
2291 committed_leader_epoch: 5,
2292 committed_metadata: Some(String::new()),
2293 unknown_tagged_fields: UnknownTaggedFields::default(),
2294 }],
2295 unknown_tagged_fields: UnknownTaggedFields::default(),
2296 },
2297 ];
2298
2299 let req = build_commit_request(
2300 "group-a".into(),
2301 42,
2302 "member-a".into(),
2303 Some("instance-a".into()),
2304 topics.clone(),
2305 );
2306
2307 assert2::assert!(
2308 req == TopicNameOffsetCommit::new(OffsetCommitRequest {
2309 group_id: "group-a".into(),
2310 generation_id_or_member_epoch: 42,
2311 member_id: "member-a".into(),
2312 group_instance_id: Some("instance-a".into()),
2313 retention_time_ms: -1,
2314 topics,
2315 unknown_tagged_fields: UnknownTaggedFields::default(),
2316 })
2317 );
2318 }
2319
2320 #[test]
2329 fn commit_response_outcome_orders_errors_and_reads_the_consumer_state() {
2330 let running = CommitResponseContext {
2331 group_id: "group-a",
2332 group_instance_id: Some("instance-a"),
2333 coordinator_alive: true,
2334 generation_unchanged: true,
2335 };
2336 let stopped = CommitResponseContext {
2337 coordinator_alive: false,
2338 ..running
2339 };
2340 let rejoined = CommitResponseContext {
2341 generation_unchanged: false,
2342 ..running
2343 };
2344 let rejoined_and_stopped = CommitResponseContext {
2345 coordinator_alive: false,
2346 ..rejoined
2347 };
2348 let commit_failed = ConsumerError::CommitFailed.to_string();
2349 let rebalance_in_progress =
2350 ConsumerError::RebalanceInProgress("group-a".into()).to_string();
2351 let mut actual = Vec::new();
2352 let mut wanted = Vec::new();
2353 for (name, errors, context, expected) in [
2354 (
2355 "success",
2356 &[0, 0][..],
2357 running,
2358 Ok(CommitOutcome::Acked(
2359 maplit::hashset! {("topic".into(), 0), ("topic".into(), 1)},
2360 )),
2361 ),
2362 (
2363 "retriable code keeps the acknowledged partitions",
2364 &[0, 3][..],
2365 running,
2366 Ok(CommitOutcome::Retriable {
2367 code: 3,
2368 acknowledged: maplit::hashset! {("topic".into(), 0)},
2369 find_coordinator: false,
2370 }),
2371 ),
2372 (
2373 "a later coordinator code finds the coordinator",
2374 &[14, 7][..],
2375 running,
2376 Ok(CommitOutcome::Retriable {
2377 code: 14,
2378 acknowledged: HashSet::new(),
2379 find_coordinator: true,
2380 }),
2381 ),
2382 (
2383 "a rebalance error takes precedence over a retriable code",
2384 &[3, 27][..],
2385 running,
2386 Err(rebalance_in_progress.clone()),
2387 ),
2388 (
2389 "final code takes precedence over a rebalance error",
2390 &[27, 42][..],
2391 running,
2392 Err(ConsumerError::Server(42).to_string()),
2393 ),
2394 (
2395 "group authorization takes precedence over a rebalance error",
2396 &[22, 30][..],
2397 running,
2398 Err("not authorized to access group: group-a".into()),
2399 ),
2400 (
2401 "retriable code takes precedence over topic authorization",
2402 &[29, 16][..],
2403 running,
2404 Ok(CommitOutcome::Retriable {
2405 code: 16,
2406 acknowledged: HashSet::new(),
2407 find_coordinator: true,
2408 }),
2409 ),
2410 (
2411 "topic authorization after a success",
2412 &[0, 29][..],
2413 running,
2414 Err("not authorized to access topics: [topic]".into()),
2415 ),
2416 (
2417 "illegal generation with the request's own generation still current: a real kick",
2418 &[22][..],
2419 running,
2420 Err(commit_failed.clone()),
2421 ),
2422 (
2423 "unknown member id with the request's own generation still current: a real kick",
2424 &[25][..],
2425 running,
2426 Err(commit_failed.clone()),
2427 ),
2428 (
2429 "rebalance in progress",
2430 &[27][..],
2431 running,
2432 Err(rebalance_in_progress.clone()),
2433 ),
2434 (
2435 "illegal generation after a rejoin already moved the generation on: a stale response",
2436 &[22][..],
2437 rejoined,
2438 Err(rebalance_in_progress.clone()),
2439 ),
2440 (
2441 "unknown member id after a rejoin already moved the generation on: a stale response",
2442 &[25][..],
2443 rejoined,
2444 Err(rebalance_in_progress.clone()),
2445 ),
2446 (
2447 "rebalance in progress after a rejoin is still a rebalance error",
2448 &[27][..],
2449 rejoined,
2450 Err(rebalance_in_progress.clone()),
2451 ),
2452 (
2453 "unknown member id after the task stopped, generation still current",
2454 &[25][..],
2455 stopped,
2456 Err(commit_failed.clone()),
2457 ),
2458 (
2459 "unknown member id after the task stopped and a rejoin: no task left to rejoin on",
2460 &[25][..],
2461 CommitResponseContext {
2462 generation_unchanged: false,
2463 coordinator_alive: false,
2464 ..running
2465 },
2466 Err(commit_failed.clone()),
2467 ),
2468 (
2469 "rebalance in progress after the task stopped",
2470 &[27][..],
2471 stopped,
2472 Err(commit_failed.clone()),
2473 ),
2474 (
2475 "fenced instance id with the same generation",
2476 &[82][..],
2477 running,
2478 Err(ConsumerError::FencedInstanceId("instance-a".into()).to_string()),
2479 ),
2480 (
2481 "fenced instance id after a rejoin already moved the generation on",
2482 &[82][..],
2483 rejoined,
2484 Err(rebalance_in_progress.clone()),
2485 ),
2486 (
2487 "fenced instance id after a rejoin and a task stop",
2488 &[82][..],
2489 rejoined_and_stopped,
2490 Err(commit_failed.clone()),
2491 ),
2492 ] {
2493 actual.push((
2494 name,
2495 commit_response_outcome(&response(errors), context)
2496 .map_err(|error| error.to_string()),
2497 ));
2498 wanted.push((name, expected));
2499 }
2500 assert2::assert!(actual == wanted);
2501 }
2502
2503 #[derive(Clone, Debug, PartialEq)]
2506 struct VersionedOffsetCommit {
2507 version: i16,
2508 request: OffsetCommitRequest,
2509 }
2510
2511 type SentOffsetCommits = Vec<VersionedOffsetCommit>;
2512
2513 #[derive(Clone, Copy)]
2514 struct OffsetCommitVersionRange {
2515 min: i16,
2516 max: i16,
2517 }
2518
2519 #[derive(bon::Builder)]
2520 struct OffsetCommitCoordinatorOptions {
2521 #[builder(default = 2)]
2522 min_version: i16,
2523 #[builder(default = 7)]
2524 max_version: i16,
2525 #[builder(default)]
2526 error_code: i16,
2527 }
2528
2529 struct OffsetCommitCoordinator {
2530 mock: MockBroker,
2531 requests: Arc<std::sync::Mutex<Vec<VersionedOffsetCommit>>>,
2532 }
2533
2534 impl OffsetCommitCoordinatorOptions {
2535 async fn start(self) -> OffsetCommitCoordinator {
2536 use bytes::Buf as _;
2537 use krabka_protocol::Decode as _;
2538
2539 let Self {
2540 min_version,
2541 max_version,
2542 error_code,
2543 } = self;
2544 let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
2545 let requests_in_mock = Arc::clone(&requests);
2546 let mock = MockBroker::start(move |api_key, version, _corr_id, mut body| {
2547 if api_key == api_versions_request::API_KEY {
2548 return Some(group_api_versions(
2549 offset_commit_request::API_KEY,
2550 (min_version, max_version),
2551 ));
2552 }
2553 if api_key != offset_commit_request::API_KEY {
2554 return None;
2555 }
2556 let client_id_len = usize::try_from(body.get_i16()).expect("client id length");
2557 body.advance(client_id_len);
2558 if version >= offset_commit_request::FLEXIBLE_MIN {
2559 body.advance(1);
2560 }
2561 let request = OffsetCommitRequest::decode(&mut body, version)
2562 .expect("offset commit request decodes");
2563 requests_in_mock
2564 .lock()
2565 .expect("requests lock")
2566 .push(VersionedOffsetCommit {
2567 version,
2568 request: request.clone(),
2569 });
2570 let response = offset_commit_response(&request, |_| error_code);
2571 let mut body = Vec::new();
2572 if version >= offset_commit_request::FLEXIBLE_MIN {
2573 body.push(0);
2574 }
2575 body.extend(test_support::encode_protocol_response(&response, version));
2576 Some(body)
2577 })
2578 .await;
2579 OffsetCommitCoordinator { mock, requests }
2580 }
2581 }
2582
2583 type CommitResult = Result<(), (i16, i16, i16, i16, i16)>;
2586
2587 #[tokio::test]
2592 async fn commit_sends_offset_commit_by_topic_name_at_v9_or_lower() {
2593 let sent_request = |version| {
2594 vec![VersionedOffsetCommit {
2595 version,
2596 request: OffsetCommitRequest {
2597 group_id: "group-a".into(),
2598 generation_id_or_member_epoch: 7,
2599 member_id: "member-a".into(),
2600 group_instance_id: None,
2601 retention_time_ms: -1,
2602 topics: vec![OffsetCommitRequestTopic {
2603 name: "topic".into(),
2604 topic_id: Uuid::ZERO,
2605 partitions: vec![OffsetCommitRequestPartition {
2606 partition_index: 0,
2607 committed_offset: 12,
2608 committed_leader_epoch: -1,
2609 committed_metadata: Some(String::new()),
2610 unknown_tagged_fields: UnknownTaggedFields::default(),
2611 }],
2612 unknown_tagged_fields: UnknownTaggedFields::default(),
2613 }],
2614 unknown_tagged_fields: UnknownTaggedFields::default(),
2615 },
2616 }]
2617 };
2618 for (name, offset_commit_range, expected_requests, expected_result) in [
2619 (
2620 "coordinator stops at v7",
2621 OffsetCommitVersionRange { min: 2, max: 7 },
2622 sent_request(7),
2623 Ok(()),
2624 ),
2625 (
2626 "coordinator stops at v9",
2627 OffsetCommitVersionRange { min: 2, max: 9 },
2628 sent_request(9),
2629 Ok(()),
2630 ),
2631 (
2632 "coordinator supports v10",
2633 OffsetCommitVersionRange { min: 2, max: 10 },
2634 sent_request(9),
2635 Ok(()),
2636 ),
2637 (
2638 "coordinator supports only v10",
2639 OffsetCommitVersionRange { min: 10, max: 10 },
2640 Vec::new(),
2641 Err((offset_commit_request::API_KEY, 10, 10, 2, 9)),
2642 ),
2643 ] {
2644 let OffsetCommitCoordinator { mock, requests } =
2645 OffsetCommitCoordinatorOptions::builder()
2646 .min_version(offset_commit_range.min)
2647 .max_version(offset_commit_range.max)
2648 .build()
2649 .start()
2650 .await;
2651 let consumer = default_commit_consumer(&mock).await;
2652
2653 let result: CommitResult = consumer
2654 .commit_offsets_sync(
2655 maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata::new(12)},
2656 )
2657 .await
2658 .map_err(|error| match error {
2659 ConsumerError::Client(ClientError::IncompatibleVersion {
2660 api_key,
2661 broker_min,
2662 broker_max,
2663 client_min,
2664 client_max,
2665 }) => (api_key, broker_min, broker_max, client_min, client_max),
2666 other => panic!("case {name}: unexpected error {other:?}"),
2667 });
2668
2669 mock.stop();
2670 let requests: SentOffsetCommits = requests.lock().expect("requests lock").clone();
2671 check!(
2672 (requests, result) == (expected_requests, expected_result),
2673 "case {name}"
2674 );
2675 }
2676 }
2677
2678 #[derive(Clone, Copy, Debug)]
2680 enum CommitAnswer {
2681 Codes(&'static [(&'static str, i16)]),
2684 Close,
2686 Silent,
2688 }
2689
2690 type CommitCodeResult = (Result<(), (String, Option<String>)>, usize, usize);
2694
2695 const COMMIT_TIMED_OUT: &str =
2700 "timeout: Timeout of 1000ms expired before successfully committing offsets";
2701
2702 fn described(error: &ConsumerError) -> (String, Option<String>) {
2705 (
2706 error.to_string(),
2707 std::error::Error::source(error).map(ToString::to_string),
2708 )
2709 }
2710
2711 #[tokio::test]
2720 async fn commit_sync_handles_each_offset_commit_error_code_as_kafka_does() {
2721 let mut actual = Vec::new();
2722 let mut wanted = Vec::new();
2723 for (name, answers, rejoin, retry_policy, expected) in offset_commit_error_code_cases() {
2724 actual.push((
2725 name,
2726 scripted_commit_sync(answers, rejoin, retry_policy, Duration::ZERO).await,
2727 ));
2728 wanted.push((name, expected));
2729 }
2730 let wanted: Vec<(&str, CommitCodeResult)> = wanted;
2731 assert2::assert!(actual == wanted);
2732 }
2733
2734 const RETRY: CoordinatorRetryPolicy = CoordinatorRetryPolicy {
2736 timeout: Duration::from_secs(2),
2737 initial_backoff: Duration::from_millis(1),
2738 max_backoff: Duration::from_millis(1),
2739 };
2740
2741 const LONG_BACKOFF: CoordinatorRetryPolicy = CoordinatorRetryPolicy {
2745 timeout: Duration::from_secs(1),
2746 initial_backoff: Duration::from_secs(5),
2747 max_backoff: Duration::from_secs(5),
2748 };
2749
2750 fn offset_commit_error_code_cases() -> Vec<(
2755 &'static str,
2756 Vec<CommitAnswer>,
2757 bool,
2758 CoordinatorRetryPolicy,
2759 CommitCodeResult,
2760 )> {
2761 use CommitAnswer::{Close, Codes, Silent};
2762
2763 let fenced = described(&ConsumerError::FencedInstanceId("instance-a".into()));
2764 let commit_failed = described(&ConsumerError::CommitFailed);
2765 let rebalance_in_progress =
2766 described(&ConsumerError::RebalanceInProgress("group-a".into()));
2767 let timed_out = |cause: ConsumerError| (COMMIT_TIMED_OUT.into(), Some(cause.to_string()));
2768 let disconnected = timed_out(ConsumerError::Client(ClientError::Disconnected));
2769 let group_denied = described(&ConsumerError::GroupAuthorizationFailed("group-a".into()));
2770 let topics_denied = described(&ConsumerError::TopicAuthorizationFailed(
2771 maplit::btreeset! {"orders".into(), "payments".into()},
2772 ));
2773 vec![
2774 ("success", vec![Codes(&[])], false, RETRY, (Ok(()), 1, 0)),
2775 (
2776 "unknown topic or partition retries",
2777 vec![Codes(&[("orders", 3)]), Codes(&[])],
2778 false,
2779 RETRY,
2780 (Ok(()), 2, 0),
2781 ),
2782 (
2783 "coordinator load in progress retries",
2784 vec![Codes(&[("orders", 14)]), Codes(&[])],
2785 false,
2786 RETRY,
2787 (Ok(()), 2, 0),
2788 ),
2789 (
2790 "request timed out finds the coordinator and retries",
2791 vec![Codes(&[("orders", 7)]), Codes(&[])],
2792 false,
2793 RETRY,
2794 (Ok(()), 2, 1),
2795 ),
2796 (
2797 "coordinator not available finds the coordinator and retries",
2798 vec![Codes(&[("orders", 15)]), Codes(&[])],
2799 false,
2800 RETRY,
2801 (Ok(()), 2, 1),
2802 ),
2803 (
2804 "not coordinator finds the coordinator and retries",
2805 vec![Codes(&[("orders", 16)]), Codes(&[])],
2806 false,
2807 RETRY,
2808 (Ok(()), 2, 1),
2809 ),
2810 (
2811 "unknown topic or partition past the timeout",
2812 vec![Codes(&[("orders", 3)])],
2813 false,
2814 LONG_BACKOFF,
2815 (Err(timed_out(ConsumerError::Server(3))), 1, 0),
2816 ),
2817 (
2818 "not coordinator with a backoff past the timeout",
2819 vec![Codes(&[("orders", 16)]), Codes(&[])],
2820 false,
2821 LONG_BACKOFF,
2822 (Err(timed_out(ConsumerError::Server(16))), 1, 1),
2823 ),
2824 (
2825 "disconnect finds the coordinator and retries",
2826 vec![Close, Codes(&[])],
2827 false,
2828 RETRY,
2829 (Ok(()), 2, 1),
2830 ),
2831 (
2832 "request timeout finds the coordinator and retries",
2833 vec![Silent, Codes(&[])],
2834 false,
2835 RETRY,
2836 (Ok(()), 2, 1),
2837 ),
2838 (
2839 "disconnect past the timeout",
2840 vec![Close],
2841 false,
2842 LONG_BACKOFF,
2843 (Err(disconnected), 1, 1),
2844 ),
2845 (
2846 "group authorization failed",
2847 vec![Codes(&[("orders", 30)])],
2848 false,
2849 RETRY,
2850 (Err(group_denied), 1, 0),
2851 ),
2852 (
2853 "topic authorization failed collects every topic",
2854 vec![Codes(&[("orders", 29), ("payments", 29)])],
2855 false,
2856 RETRY,
2857 (Err(topics_denied), 1, 0),
2858 ),
2859 (
2860 "topic authorization failed with a retriable code retries",
2861 vec![Codes(&[("orders", 29), ("payments", 3)]), Codes(&[])],
2862 false,
2863 RETRY,
2864 (Ok(()), 2, 0),
2865 ),
2866 (
2867 "offset metadata too large",
2868 vec![Codes(&[("orders", 12)])],
2869 false,
2870 RETRY,
2871 (Err(described(&ConsumerError::Server(12))), 1, 0),
2872 ),
2873 (
2874 "invalid commit offset size",
2875 vec![Codes(&[("orders", 28)])],
2876 false,
2877 RETRY,
2878 (Err(described(&ConsumerError::Server(28))), 1, 0),
2879 ),
2880 (
2881 "unknown topic id is unexpected for a request by topic name",
2882 vec![Codes(&[("orders", 100)])],
2883 false,
2884 RETRY,
2885 (Err(described(&ConsumerError::Server(100))), 1, 0),
2886 ),
2887 (
2888 "not leader or follower is unexpected",
2889 vec![Codes(&[("orders", 6)])],
2890 false,
2891 RETRY,
2892 (Err(described(&ConsumerError::Server(6))), 1, 0),
2893 ),
2894 (
2895 "fenced instance id with the same generation",
2896 vec![Codes(&[("orders", 82)])],
2897 false,
2898 RETRY,
2899 (Err(fenced.clone()), 1, 0),
2900 ),
2901 (
2902 "fenced instance id is a stale artifact of an already-completed rejoin",
2903 vec![Codes(&[("orders", 82)]), Codes(&[])],
2904 true,
2905 RETRY,
2906 (Err(rebalance_in_progress.clone()), 1, 0),
2907 ),
2908 (
2909 "illegal generation with the request's own generation still current: a real kick",
2910 vec![Codes(&[("orders", 22)])],
2911 false,
2912 RETRY,
2913 (Err(commit_failed.clone()), 1, 0),
2914 ),
2915 (
2916 "illegal generation is a stale artifact of an already-completed rejoin",
2917 vec![Codes(&[("orders", 22)]), Codes(&[])],
2918 true,
2919 RETRY,
2920 (Err(rebalance_in_progress.clone()), 1, 0),
2921 ),
2922 (
2923 "unknown member id with the request's own generation still current: a real kick",
2924 vec![Codes(&[("orders", 25)])],
2925 false,
2926 RETRY,
2927 (Err(commit_failed.clone()), 1, 0),
2928 ),
2929 (
2930 "unknown member id is a stale artifact of an already-completed rejoin",
2931 vec![Codes(&[("orders", 25)]), Codes(&[])],
2932 true,
2933 RETRY,
2934 (Err(rebalance_in_progress.clone()), 1, 0),
2935 ),
2936 (
2937 "rebalance in progress fails the commit",
2938 vec![Codes(&[("orders", 27)])],
2939 false,
2940 RETRY,
2941 (Err(rebalance_in_progress.clone()), 1, 0),
2942 ),
2943 (
2944 "rebalance in progress fails the commit even after a rejoin",
2945 vec![Codes(&[("orders", 27)]), Codes(&[])],
2946 true,
2947 RETRY,
2948 (Err(rebalance_in_progress.clone()), 1, 0),
2949 ),
2950 ]
2951 }
2952
2953 #[tokio::test]
2959 async fn commit_sync_after_a_wait_for_the_commit_lock_times_out_with_the_last_retriable_error()
2960 {
2961 let retry_policy = CoordinatorRetryPolicy {
2962 timeout: Duration::from_secs(1),
2963 initial_backoff: Duration::from_secs(5),
2964 max_backoff: Duration::from_secs(5),
2965 };
2966 let result = scripted_commit_sync(
2967 vec![CommitAnswer::Codes(&[("orders", 3)])],
2968 false,
2969 retry_policy,
2970 Duration::from_millis(50),
2971 )
2972 .await;
2973 let timed_out = (
2974 COMMIT_TIMED_OUT.into(),
2975 Some(ConsumerError::Server(3).to_string()),
2976 );
2977 assert2::assert!(result == (Err(timed_out), 1, 0));
2978 }
2979
2980 async fn scripted_commit_sync(
2986 answers: Vec<CommitAnswer>,
2987 rejoin: bool,
2988 retry_policy: CoordinatorRetryPolicy,
2989 commit_lock_held: Duration,
2990 ) -> CommitCodeResult {
2991 use CommitAnswer::{Close, Codes, Silent};
2992 use krabka_protocol::owned::{
2993 find_coordinator_response::FindCoordinatorResponse, metadata_response::MetadataResponse,
2994 };
2995
2996 let identity = Arc::new(Mutex::new(CommitIdentity {
2997 generation: 7,
2998 member_id: "member-a".into(),
2999 ownership_ids: maplit::hashmap! {("orders".into(), 0) => 1, ("payments".into(), 0) => 2},
3000 rejoin_on_poll: false,
3001 }));
3002 let assignment_changed = Arc::new(tokio::sync::Notify::new());
3003 let offset_commits = Arc::new(AtomicUsize::new(0));
3004 let find_coordinators = Arc::new(AtomicUsize::new(0));
3005 let identity_in_mock = Arc::clone(&identity);
3006 let changed_in_mock = Arc::clone(&assignment_changed);
3007 let offset_commits_in_mock = Arc::clone(&offset_commits);
3008 let find_coordinators_in_mock = Arc::clone(&find_coordinators);
3009 let mock = MockBroker::start_with_replies(move |api_key, version, _corr_id, mut body| {
3010 let mut buffer = bytes::BytesMut::new();
3011 match api_key {
3012 api_versions_request::API_KEY => {
3013 MockReply::Respond(group_api_versions(offset_commit_request::API_KEY, (7, 7)))
3014 }
3015 find_coordinator_request::API_KEY => {
3016 find_coordinators_in_mock.fetch_add(1, Ordering::SeqCst);
3017 FindCoordinatorResponse::default()
3018 .encode(&mut buffer, version)
3019 .unwrap();
3020 MockReply::Respond(buffer.to_vec())
3021 }
3022 metadata_request::API_KEY => {
3023 MetadataResponse::default()
3024 .encode(&mut buffer, version)
3025 .unwrap();
3026 MockReply::Respond(buffer.to_vec())
3027 }
3028 offset_commit_request::API_KEY => {
3029 let attempt = offset_commits_in_mock.fetch_add(1, Ordering::SeqCst);
3030 if attempt == 0 && rejoin {
3031 identity_in_mock.try_lock().unwrap().generation = 8;
3032 changed_in_mock.notify_waiters();
3033 }
3034 let client_id_len = bytes::Buf::get_i16(&mut body);
3035 bytes::Buf::advance(&mut body, usize::try_from(client_id_len).unwrap());
3036 let request: OffsetCommitRequest =
3037 krabka_protocol::Decode::decode(&mut body, version).unwrap();
3038 let codes = match answers[attempt.min(answers.len() - 1)] {
3039 Codes(codes) => codes,
3040 Close => return MockReply::Close,
3041 Silent => return MockReply::Silent,
3042 };
3043 let code_of = |topic: &str| {
3044 codes
3045 .iter()
3046 .find(|(name, _)| *name == topic)
3047 .map_or(0, |(_, code)| *code)
3048 };
3049 let response = offset_commit_response(&request, code_of);
3050 MockReply::Respond(test_support::encode_protocol_response(&response, version))
3051 }
3052 _ => MockReply::Silent,
3053 }
3054 })
3055 .await;
3056 let mut consumer = commit_consumer(
3057 CommitConsumerStart::builder()
3058 .mock(&mock)
3059 .commit_identity(Arc::clone(&identity))
3060 .assignment_changed(assignment_changed)
3061 .build(),
3062 )
3063 .await;
3064 consumer.client = Client::builder()
3065 .bootstrap(mock.addr.to_string())
3066 .request_timeout(krabka_units::millis(300))
3067 .build()
3068 .await
3069 .unwrap();
3070 consumer.group_instance_id = Some("instance-a".into());
3071 consumer.retry_policy = retry_policy;
3072 consumer.default_api_timeout = convert::StdDurationExt::as_time(&retry_policy.timeout);
3073
3074 let held = consumer.commit_serialization.lock().await;
3075 let release = async move {
3076 tokio::time::sleep(commit_lock_held).await;
3077 drop(held);
3078 };
3079 let (result, ()) = tokio::join!(
3080 tokio::time::timeout(Duration::from_secs(10), consumer.commit_sync()),
3081 release,
3082 );
3083 let result = result
3084 .map_err(|_| ("commit never finished".to_string(), None))
3085 .and_then(|result| result.map_err(|error| described(&error)));
3086
3087 mock.stop();
3088 (
3089 result,
3090 offset_commits.load(Ordering::SeqCst),
3091 find_coordinators.load(Ordering::SeqCst),
3092 )
3093 }
3094
3095 #[tokio::test]
3101 async fn commit_sync_times_out_near_default_api_timeout_despite_a_silent_broker() {
3102 let mock = MockBroker::start(move |api_key, _version, _corr_id, _body| {
3103 if api_key == api_versions_request::API_KEY {
3104 return Some(group_api_versions(offset_commit_request::API_KEY, (7, 7)));
3105 }
3106 None
3109 })
3110 .await;
3111
3112 let mut consumer = commit_consumer(
3113 CommitConsumerStart::builder()
3114 .mock(&mock)
3115 .commit_identity(commit_identity(7, "member-a"))
3116 .build(),
3117 )
3118 .await;
3119 consumer.client = Client::builder()
3122 .bootstrap(mock.addr.to_string())
3123 .request_timeout(secs(30))
3124 .build()
3125 .await
3126 .unwrap();
3127 consumer.default_api_timeout = krabka_units::millis(200);
3128
3129 let started = tokio::time::Instant::now();
3130 let result = tokio::time::timeout(Duration::from_secs(5), consumer.commit_sync())
3131 .await
3132 .expect("commit_sync must return on its own well before the test's own timeout");
3133 let elapsed = started.elapsed();
3134
3135 mock.stop();
3136 check!(
3137 result.map_err(|error| described(&error))
3138 == Err((
3139 "timeout: Timeout of 200ms expired before successfully committing offsets"
3140 .into(),
3141 None
3142 ))
3143 );
3144 check!(elapsed < Duration::from_secs(2));
3145 }
3146
3147 #[tokio::test]
3153 async fn commit_fails_without_a_request_after_the_coordinator_task_stops() {
3154 #[derive(Clone, Copy)]
3155 enum Commit {
3156 All,
3157 Selected,
3158 }
3159 #[derive(Clone, Copy)]
3160 enum TaskState {
3161 Running,
3162 Fenced,
3164 Stopping,
3165 Stopped,
3166 }
3167 let commit_failed = "offset commit failed: the consumer is not part of an active group; it is likely that the consumer was kicked out of the group";
3168 for (name, commit, task_state, owned, expected) in [
3169 (
3170 "commit sync while the task runs",
3171 Commit::All,
3172 TaskState::Running,
3173 true,
3174 (Ok(()), 1),
3175 ),
3176 (
3177 "commit sync after close cancelled the task but before it ends",
3178 Commit::All,
3179 TaskState::Stopping,
3180 true,
3181 (Err(commit_failed.to_string()), 0),
3182 ),
3183 (
3184 "commit sync while the task waits for a poll after a fence",
3185 Commit::All,
3186 TaskState::Fenced,
3187 false,
3188 (Err(commit_failed.to_string()), 0),
3189 ),
3190 (
3191 "selected commit while the task waits for a poll after a fence",
3192 Commit::Selected,
3193 TaskState::Fenced,
3194 false,
3195 (Err(commit_failed.to_string()), 0),
3196 ),
3197 (
3198 "commit sync after the task stops",
3199 Commit::All,
3200 TaskState::Stopped,
3201 true,
3202 (Err(commit_failed.to_string()), 0),
3203 ),
3204 (
3205 "commit sync after a fence cleared the assignment",
3206 Commit::All,
3207 TaskState::Stopped,
3208 false,
3209 (Err(commit_failed.to_string()), 0),
3210 ),
3211 (
3212 "selected commit after the task stops",
3213 Commit::Selected,
3214 TaskState::Stopped,
3215 true,
3216 (Err(commit_failed.to_string()), 0),
3217 ),
3218 ] {
3219 let requests = Arc::new(AtomicUsize::new(0));
3220 let requests_in_mock = Arc::clone(&requests);
3221 let mock = MockBroker::start(move |api_key, version, _corr_id, _body| {
3222 if api_key == api_versions_request::API_KEY {
3223 return Some(group_api_versions(offset_commit_request::API_KEY, (2, 2)));
3224 }
3225 if api_key != offset_commit_request::API_KEY {
3226 return None;
3227 }
3228 requests_in_mock.fetch_add(1, Ordering::SeqCst);
3229 Some(test_support::encode_protocol_response(
3230 &response(&[0]),
3231 version,
3232 ))
3233 })
3234 .await;
3235 let identity = commit_identity(7, "member-a");
3236 if !owned {
3237 identity.lock().await.ownership_ids.clear();
3238 }
3239 if matches!(task_state, TaskState::Fenced) {
3240 identity.lock().await.rejoin_on_poll = true;
3241 }
3242 let mut consumer = commit_consumer(
3243 CommitConsumerStart::builder()
3244 .mock(&mock)
3245 .commit_identity(identity)
3246 .build(),
3247 )
3248 .await;
3249 let task = match task_state {
3250 TaskState::Running | TaskState::Fenced => {
3251 tokio::spawn(std::future::pending::<()>())
3252 }
3253 TaskState::Stopping => {
3254 consumer.coordinator_shutdown.cancel();
3255 tokio::spawn(std::future::pending::<()>())
3256 }
3257 TaskState::Stopped => {
3258 let task = tokio::spawn(async {});
3259 while !task.is_finished() {
3260 tokio::task::yield_now().await;
3261 }
3262 task
3263 }
3264 };
3265 consumer.coordinator_handle = Some(task);
3266
3267 let commit_future = async {
3268 match commit {
3269 Commit::All => consumer.commit_sync().await,
3270 Commit::Selected => consumer
3271 .commit_offsets_sync(
3272 maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata::new(12)},
3273 )
3274 .await,
3275 }
3276 };
3277 let result = tokio::time::timeout(Duration::from_secs(5), commit_future)
3278 .await
3279 .unwrap_or_else(|_| panic!("case {name}: commit never finished"))
3280 .map_err(|error| error.to_string());
3281
3282 if let Some(task) = consumer.coordinator_handle.take() {
3283 task.abort();
3284 }
3285 mock.stop();
3286 check!(
3287 (result, requests.load(Ordering::SeqCst)) == expected,
3288 "case {name}"
3289 );
3290 }
3291 }
3292
3293 #[tokio::test]
3297 async fn selected_commit_fails_immediately_on_rebalance_in_progress_without_retrying() {
3298 let requests = Arc::new(AtomicUsize::new(0));
3299 let SelectedCommitConsumer { consumer, mock, .. } = SelectedCommitConsumerStart::builder()
3300 .requests(Arc::clone(&requests))
3301 .on_first_request(advance_generation_on_first_commit())
3302 .build()
3303 .start()
3304 .await;
3305
3306 let result = tokio::time::timeout(
3307 std::time::Duration::from_secs(2),
3308 selected_commit(&consumer, 12),
3309 )
3310 .await
3311 .expect("commit does not wait for a rejoin");
3312
3313 mock.stop();
3314 assert2::assert!(
3315 result.map_err(|error| error.to_string())
3316 == Err(ConsumerError::RebalanceInProgress("group-a".into()).to_string())
3317 );
3318 assert2::assert!(requests.load(Ordering::SeqCst) == 1);
3319 }
3320
3321 #[tokio::test]
3325 async fn selected_commit_uses_snapshot_generation_if_ownership_changes_before_rpc() {
3326 let identity = commit_identity(7, "member-a");
3327 let generation = Arc::new(AtomicI32::new(7));
3328 let seen_generation = Arc::new(AtomicI32::new(-1));
3329 let seen_in_mock = Arc::clone(&seen_generation);
3330 let SelectedCommitConsumer { consumer, mock, .. } = SelectedCommitConsumerStart::builder()
3331 .commit_identity(Arc::clone(&identity))
3332 .generation(Arc::clone(&generation))
3333 .on_first_request(
3334 move |_ownership: &Arc<Mutex<CommitIdentity>>,
3335 _generation: &Arc<AtomicI32>,
3336 request_generation: i32,
3337 _request_member_id: &str| {
3338 seen_in_mock.store(request_generation, Ordering::Relaxed);
3339 },
3340 )
3341 .build()
3342 .start()
3343 .await;
3344 let mut changed_identity = identity.lock().await;
3345 changed_identity
3346 .ownership_ids
3347 .insert(("topic".into(), 0), 2);
3348 changed_identity.generation = 8;
3349 drop(changed_identity);
3350 generation.store(8, Ordering::Relaxed);
3351 let topics = build_commit_topics(
3352 maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata::new(12)},
3353 );
3354
3355 let error = consumer
3356 .commit_topics_once(topics, (7, "member-a".into()))
3357 .await
3358 .expect_err("rebalance in progress fails the commit");
3359
3360 mock.stop();
3361 assert2::assert!(
3362 error.to_string() == ConsumerError::RebalanceInProgress("group-a".into()).to_string()
3363 );
3364 assert2::assert!(seen_generation.load(Ordering::Relaxed) == 7);
3365 }
3366
3367 #[tokio::test]
3370 async fn selected_commit_uses_live_member_id_even_without_a_retry() {
3371 let identity = commit_identity(8, "member-new");
3372 let seen_new_member = Arc::new(std::sync::atomic::AtomicBool::new(false));
3373 let seen_in_mock = Arc::clone(&seen_new_member);
3374 let SelectedCommitConsumer { consumer, mock, .. } = SelectedCommitConsumerStart::builder()
3375 .commit_identity(identity)
3376 .on_first_request(
3377 move |_identity: &Arc<Mutex<CommitIdentity>>,
3378 generation: &Arc<AtomicI32>,
3379 _request_generation: i32,
3380 request_member_id: &str| {
3381 seen_in_mock.store(request_member_id == "member-new", Ordering::Relaxed);
3382 generation.store(8, Ordering::Release);
3383 },
3384 )
3385 .build()
3386 .start()
3387 .await;
3388 let result = consumer
3389 .commit_offsets_sync(
3390 maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata::new(12)},
3391 )
3392 .await;
3393
3394 mock.stop();
3395 assert2::assert!(result.is_err());
3396 assert2::assert!(seen_new_member.load(Ordering::Relaxed));
3397 }
3398
3399 #[tokio::test]
3403 async fn commit_sync_fails_on_any_rebalance_code_in_a_mixed_response() {
3404 let identity = Arc::new(Mutex::new(CommitIdentity {
3405 generation: 7,
3406 member_id: "member-a".into(),
3407 ownership_ids: maplit::hashmap! {("topic".into(), 0) => 1, ("topic".into(), 1) => 2},
3408 rejoin_on_poll: false,
3409 }));
3410 let requests = Arc::new(AtomicUsize::new(0));
3411 let SelectedCommitConsumer { consumer, mock, .. } = SelectedCommitConsumerStart::builder()
3412 .commit_identity(identity)
3413 .requests(Arc::clone(&requests))
3414 .mixed_response(true)
3415 .on_first_request(advance_generation_on_first_commit())
3416 .build()
3417 .start()
3418 .await;
3419
3420 let result = consumer.commit_sync().await;
3421
3422 mock.stop();
3423 assert2::assert!(
3424 result.map_err(|error| error.to_string())
3425 == Err(ConsumerError::RebalanceInProgress("group-a".into()).to_string())
3426 );
3427 assert2::assert!(requests.load(Ordering::SeqCst) == 1);
3428 }
3429
3430 #[tokio::test]
3434 async fn a_failed_commit_releases_the_lock_before_a_newer_commit_for_the_same_partition() {
3435 let identity = commit_identity(7, "member-a");
3436 let requests = Arc::new(AtomicUsize::new(0));
3437 let SelectedCommitConsumer {
3438 consumer,
3439 mock,
3440 seen_offsets,
3441 ..
3442 } = SelectedCommitConsumerStart::builder()
3443 .commit_identity(identity)
3444 .requests(requests)
3445 .on_first_request(advance_generation_on_first_commit())
3446 .build()
3447 .start()
3448 .await;
3449 let consumer = Arc::new(consumer);
3450 let blocker = Arc::clone(&consumer.commit_serialization)
3451 .lock_owned()
3452 .await;
3453 let older = {
3454 let consumer = Arc::clone(&consumer);
3455 tokio::spawn(async move {
3456 consumer
3457 .commit_offsets_sync(
3458 maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata::new(10)},
3459 )
3460 .await
3461 })
3462 };
3463 tokio::task::yield_now().await;
3464 let newer = {
3465 let consumer = Arc::clone(&consumer);
3466 tokio::spawn(async move {
3467 consumer
3468 .commit_offsets_sync(
3469 maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata::new(12)},
3470 )
3471 .await
3472 })
3473 };
3474 tokio::task::yield_now().await;
3475 drop(blocker);
3476
3477 let older_error = older
3478 .await
3479 .unwrap()
3480 .expect_err("the response for the older commit is a rebalance error");
3481 newer
3482 .await
3483 .unwrap()
3484 .expect("the newer commit sends once the older one released the lock");
3485
3486 mock.stop();
3487 assert2::assert!(
3488 older_error.to_string()
3489 == ConsumerError::RebalanceInProgress("group-a".into()).to_string()
3490 );
3491 assert2::assert!(*seen_offsets.lock().unwrap() == vec![(0, 10), (0, 12)]);
3492 }
3493
3494 #[tokio::test]
3495 async fn async_commits_queued_behind_an_rpc_are_coalesced() {
3496 let SelectedCommitConsumer {
3497 consumer,
3498 mock,
3499 requests,
3500 ..
3501 } = SelectedCommitConsumerStart::builder()
3502 .on_first_request(
3503 |_: &Arc<Mutex<CommitIdentity>>, _: &Arc<AtomicI32>, _: i32, _: &str| {},
3504 )
3505 .build()
3506 .start()
3507 .await;
3508 let blocker = consumer.commit_serialization.lock().await;
3509
3510 for _ in 0..100 {
3511 consumer.commit_async().expect("commit_async");
3512 }
3513 tokio::task::yield_now().await;
3514 assert2::assert!(consumer.commit_async_state.load(Ordering::Acquire) == ASYNC_COMMIT_DIRTY);
3515 drop(blocker);
3516
3517 wait_for_async_commit_idle(&consumer, std::time::Duration::from_secs(2)).await;
3518 mock.stop();
3519
3520 assert2::assert!(requests.load(Ordering::SeqCst) == 1);
3521 }
3522
3523 #[test]
3524 fn selected_commit_topics_exclude_unrequested_assignment_positions() {
3525 let topics = build_commit_topics(
3526 maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata::new(12)},
3527 );
3528
3529 assert2::assert!(topics.len() == 1);
3530 assert2::assert!(topics[0].partitions.len() == 1);
3531 assert2::assert!(topics[0].partitions[0].partition_index == 0);
3532 }
3533
3534 type SentCommits = Arc<std::sync::Mutex<Vec<VersionedOffsetCommit>>>;
3536
3537 async fn recording_coordinator(error_code: i16) -> OffsetCommitCoordinator {
3540 OffsetCommitCoordinatorOptions::builder()
3541 .error_code(error_code)
3542 .build()
3543 .start()
3544 .await
3545 }
3546
3547 struct CallbackCommitFixture {
3548 consumer: Consumer,
3549 mock: MockBroker,
3550 sent: SentCommits,
3551 }
3552
3553 async fn callback_commit_fixture(error_code: i16) -> CallbackCommitFixture {
3554 let OffsetCommitCoordinator {
3555 mock,
3556 requests: sent,
3557 } = recording_coordinator(error_code).await;
3558 let consumer = commit_consumer(
3559 CommitConsumerStart::builder()
3560 .mock(&mock)
3561 .commit_identity(commit_identity(7, "member-a"))
3562 .build(),
3563 )
3564 .await;
3565 CallbackCommitFixture {
3566 consumer,
3567 mock,
3568 sent,
3569 }
3570 }
3571
3572 fn sent_partitions(sent: &SentCommits) -> Vec<OffsetCommitRequestPartition> {
3573 sent.lock()
3574 .unwrap()
3575 .iter()
3576 .flat_map(|sent| sent.request.topics.iter())
3577 .flat_map(|topic| topic.partitions.iter().cloned())
3578 .collect()
3579 }
3580
3581 #[tokio::test]
3585 async fn synchronous_commits_send_kafkas_offsets_epochs_and_metadata() {
3586 let mut actual = Vec::new();
3587 let mut wanted = Vec::new();
3588 for (name, selected, expected) in [
3589 (
3590 "selected offset past the position",
3591 Some(OffsetAndMetadata {
3592 offset: 100,
3593 leader_epoch: Some(4),
3594 metadata: "note".into(),
3595 }),
3596 (100, 4, "note"),
3597 ),
3598 (
3599 "selected offset without epoch",
3600 Some(OffsetAndMetadata::new(3)),
3601 (3, -1, ""),
3602 ),
3603 ("positions", None, (12, 9, "")),
3604 ] {
3605 let OffsetCommitCoordinator {
3606 mock,
3607 requests: sent,
3608 } = recording_coordinator(0).await;
3609 let consumer = default_commit_consumer(&mock).await;
3610 consumer.positions.lock().await.insert(
3611 ("topic".into(), 0),
3612 PartitionPosition {
3613 offset_epoch: krabka_ids::LeaderEpoch(9),
3614 ..Default::default()
3615 },
3616 );
3617 let result = match selected {
3618 Some(offset) => {
3619 consumer
3620 .commit_offsets_sync(maplit::hashmap! {("topic".into(), 0) => offset})
3621 .await
3622 }
3623 None => consumer.commit_sync().await,
3624 };
3625 mock.stop();
3626 actual.push((
3627 name,
3628 result.map_err(|error| error.to_string()),
3629 sent_partitions(&sent),
3630 ));
3631 wanted.push((
3632 name,
3633 Ok(()),
3634 vec![OffsetCommitRequestPartition {
3635 partition_index: 0,
3636 committed_offset: expected.0,
3637 committed_leader_epoch: expected.1,
3638 committed_metadata: Some(expected.2.into()),
3639 ..Default::default()
3640 }],
3641 ));
3642 }
3643 assert2::assert!(actual == wanted);
3644 }
3645
3646 #[test]
3651 fn asynchronous_commit_wraps_a_failed_coordinator_lookup() {
3652 let actual = [
3653 ConsumerError::Server(15),
3654 ConsumerError::Server(14),
3655 ConsumerError::CoordinatorUnavailable,
3656 ConsumerError::Server(30),
3657 ConsumerError::CommitFailed,
3658 ]
3659 .map(|error| {
3660 async_commit_result(Err(error), "group-a", None).map_err(|error| error.to_string())
3661 });
3662 let retriable = |cause: &str| {
3663 Err(format!(
3664 "offset commit failed with a retriable exception: {cause}"
3665 ))
3666 };
3667 assert2::assert!(
3668 actual
3669 == [
3670 retriable("broker error_code 15"),
3671 retriable("broker error_code 14"),
3672 retriable("coordinator unavailable"),
3673 Err("broker error_code 30".to_owned()),
3674 Err("offset commit failed: the consumer is not part of an active group; it is likely that the consumer was kicked out of the group".to_owned()),
3675 ]
3676 );
3677 }
3678
3679 #[tokio::test]
3682 async fn asynchronous_commit_calls_the_callback_with_offsets_and_result() {
3683 let mut actual = Vec::new();
3684 let mut wanted = Vec::new();
3685 for (name, error_code, expected) in [
3686 ("success", 0, Ok(())),
3687 (
3688 "coordinator loading",
3689 14,
3690 Err("offset commit failed with a retriable exception: broker error_code 14"),
3691 ),
3692 (
3693 "group authorization",
3694 30,
3695 Err("not authorized to access group: group-a"),
3696 ),
3697 ("metadata too large", 12, Err("broker error_code 12")),
3698 ] {
3699 let CallbackCommitFixture { consumer, mock, .. } =
3700 callback_commit_fixture(error_code).await;
3701 let (tx, rx) = tokio::sync::oneshot::channel();
3702 consumer
3703 .commit_async_with_callback(move |offsets, result| {
3704 let _ = tx.send((offsets.clone(), result.map_err(ToString::to_string)));
3705 })
3706 .expect("commit_async_with_callback");
3707 let completed = tokio::time::timeout(Duration::from_secs(5), rx)
3708 .await
3709 .expect("callback runs")
3710 .expect("callback sends");
3711 mock.stop();
3712 actual.push((name, completed));
3713 wanted.push((
3714 name,
3715 (
3716 maplit::hashmap! {("topic".to_string(), 0) => OffsetAndMetadata::new(12)},
3717 expected.map_err(str::to_owned),
3718 ),
3719 ));
3720 }
3721 assert2::assert!(actual == wanted);
3722 }
3723
3724 #[tokio::test]
3727 async fn asynchronous_commit_without_positions_calls_the_callback() {
3728 let CallbackCommitFixture {
3729 consumer,
3730 mock,
3731 sent,
3732 } = callback_commit_fixture(0).await;
3733 consumer.next_offsets.lock().await.clear();
3734 let (tx, rx) = tokio::sync::oneshot::channel();
3735 consumer
3736 .commit_async_with_callback(move |offsets, result| {
3737 let _ = tx.send((offsets.clone(), result.map_err(ToString::to_string)));
3738 })
3739 .expect("commit_async_with_callback");
3740 let completed = tokio::time::timeout(Duration::from_secs(5), rx)
3741 .await
3742 .expect("callback runs")
3743 .expect("callback sends");
3744 mock.stop();
3745 assert2::assert!((completed, sent_partitions(&sent)) == ((HashMap::new(), Ok(())), vec![]));
3746 }
3747
3748 #[tokio::test]
3753 async fn explicit_offsets_reach_the_pre_join_commit_only_after_an_ack() {
3754 let mut actual = Vec::new();
3755 let mut wanted = Vec::new();
3756 for (name, error_code, expected) in [
3757 (
3758 "acknowledged",
3759 0,
3760 maplit::hashmap! {("topic".to_string(), 0) => (100, 4)},
3761 ),
3762 ("metadata too large", 12, HashMap::new()),
3763 ] {
3764 let OffsetCommitCoordinator { mock, .. } = recording_coordinator(error_code).await;
3765 let mut consumer = default_commit_consumer(&mock).await;
3766 let auto_commit = AutoCommit::new(Duration::from_secs(5));
3767 consumer.auto_commit = Some(auto_commit.clone());
3768 let _ = consumer
3769 .commit_offsets_sync(maplit::hashmap! {("topic".into(), 0) => OffsetAndMetadata {
3770 offset: 100,
3771 leader_epoch: Some(4),
3772 metadata: "note".into(),
3773 }})
3774 .await;
3775 mock.stop();
3776 let ownership = consumer.commit_identity.lock().await.ownership_ids.clone();
3777 actual.push((name, auto_commit.polled_offsets(&ownership).await));
3778 wanted.push((name, expected));
3779 }
3780 assert2::assert!(actual == wanted);
3781 }
3782
3783 #[tokio::test]
3786 async fn asynchronous_commits_continue_after_a_callback_panics() {
3787 let OffsetCommitCoordinator { mock, .. } = recording_coordinator(0).await;
3788 let consumer = default_commit_consumer(&mock).await;
3789 consumer
3790 .commit_async_with_callback(|_, _| panic!("callback panics"))
3791 .expect("commit_async_with_callback");
3792 wait_for_async_commit_idle(&consumer, Duration::from_secs(5)).await;
3793 let (tx, rx) = tokio::sync::oneshot::channel();
3794 consumer
3795 .commit_async_with_callback(move |_, result| {
3796 let _ = tx.send(result.map_err(ToString::to_string));
3797 })
3798 .expect("commit_async_with_callback");
3799 let second = tokio::time::timeout(Duration::from_secs(5), rx).await;
3800 mock.stop();
3801 assert2::assert!(let Ok(Ok(Ok(()))) = second);
3802 }
3803
3804 #[tokio::test]
3809 async fn a_rebalance_error_fails_the_commit_of_a_manual_assignment() {
3810 let OffsetCommitCoordinator { mock, .. } = recording_coordinator(ILLEGAL_GENERATION).await;
3811 let mut consumer = commit_consumer(
3812 CommitConsumerStart::builder()
3813 .mock(&mock)
3814 .commit_identity(commit_identity(-1, ""))
3815 .generation(Arc::new(AtomicI32::new(-1)))
3816 .build(),
3817 )
3818 .await;
3819 consumer.subscription = crate::subscription::shared(Vec::new(), None, true);
3820 consumer
3821 .subscription
3822 .send_modify(|subscription| subscription.manual_assignment = true);
3823 let result = tokio::time::timeout(Duration::from_secs(5), selected_commit(&consumer, 12))
3824 .await
3825 .map(|result| result.map_err(|error| error.to_string()));
3826 mock.stop();
3827 assert2::assert!(result == Ok(Err(ConsumerError::CommitFailed.to_string())));
3828 }
3829}