1use std::{
9 collections::HashMap,
10 sync::{
11 Arc,
12 atomic::{AtomicI32, AtomicU8, Ordering},
13 },
14 time::Duration,
15};
16
17use bytes::Bytes;
18use krabka_client_core::{
19 Client, ClientError, ClientFrameMax, ConnectionDispatchQueueCapacity, FetchMinBytes, security,
20 telemetry,
21};
22use krabka_protocol::{
23 owned::{
24 join_group_response::JoinGroupResponse, sync_group_request::SyncGroupRequestAssignment,
25 sync_group_response::SyncGroupResponse,
26 },
27 primitives::uuid::Uuid as WireUuid,
28 records,
29};
30use krabka_units::{
31 ByteSize, Time,
32 convert::{ByteSizeExt as _, StdDurationExt as _, TimeExt as _},
33 millis, minutes, secs,
34};
35use refined_type::{
36 Refined,
37 rule::{GreaterI32, GreaterI64, MinMaxI64, MinMaxU128},
38};
39use tokio::{
40 sync::{Mutex, Notify},
41 task::JoinHandle,
42};
43use tokio_util::sync::CancellationToken;
44
45#[cfg(test)]
46use crate::TopicPattern;
47#[cfg(test)]
48use crate::builder;
49#[cfg(test)]
50use crate::test_support;
51use crate::{
52 CloseOptions, GroupProtocol, assignor,
53 assignor::Assignor,
54 builder::{AutoOffsetReset, IsolationLevel, decode_assignment},
55 commit, consumer_group, control, coordinator,
56 coordinator::{
57 CoordinatorRefind, CoordinatorRetryPolicy, CoordinatorRetryRequest, CoordinatorState,
58 find_coordinator,
59 },
60 error::ConsumerError,
61 fetch_buffer,
62 group_metadata::ConsumerGroupMetadata,
63 offset_wire, poll, position, rebalance_listener, subscription,
64};
65
66pub(crate) fn validated_positive_i32(value: i32, name: &str) -> Result<i32, String> {
67 GreaterI32::<0>::new(value)
68 .map(Refined::into_value)
69 .map_err(|error| format!("{name}: {error}"))
70}
71
72pub(crate) fn protocol_bytes(name: &str, value: ByteSize) -> Result<i32, String> {
73 let bytes = value.bytes_f64();
74 if bytes.is_finite() && bytes.fract() == 0.0 && (1.0..=f64::from(i32::MAX)).contains(&bytes) {
75 Ok(value.bytes_i32())
76 } else {
77 Err(format!(
78 "{name} must be a positive whole-byte value that fits i32"
79 ))
80 }
81}
82
83fn validated_millisecond_duration(name: &str, value: Duration) -> Result<u64, String> {
84 let milliseconds = MinMaxU128::<1, { u64::MAX as u128 }>::new(value.as_millis())
85 .map_err(|error| format!("{name}: {error}"))?
86 .into_value();
87 let milliseconds = u64::try_from(milliseconds).map_err(|error| format!("{name}: {error}"))?;
88 if Duration::from_millis(milliseconds) == value {
89 Ok(milliseconds)
90 } else {
91 Err(format!("{name} must be a whole number of milliseconds"))
92 }
93}
94
95pub(crate) async fn client_instance_id(
96 client: &Client,
97 timeout: Time,
98) -> Result<Option<uuid::Uuid>, ConsumerError> {
99 Ok(client.client_instance_id(timeout).await?)
100}
101
102#[allow(dead_code)] pub struct Consumer {
107 pub(crate) client: Client,
108 pub(crate) group_id: String,
109 pub(crate) coordinator_id: Arc<AtomicI32>,
114 pub(crate) retry_policy: CoordinatorRetryPolicy,
115 pub(crate) member_id: tokio::sync::watch::Receiver<String>,
118 pub(crate) commit_identity: Arc<Mutex<CommitIdentity>>,
119 pub(crate) commit_serialization: Arc<Mutex<()>>,
120 pub(crate) commit_async_state: Arc<AtomicU8>,
121 pub(crate) commit_async_callbacks: commit::OffsetCommitCallbacks,
123 pub(crate) group_instance_id: Option<String>,
124 pub(crate) current_generation: Arc<AtomicI32>,
128 pub(crate) subscription: subscription::SharedSubscription,
130 pub(crate) group_protocol: crate::GroupProtocol,
132 pub(crate) assigned: Arc<Mutex<Vec<(String, i32)>>>,
134 pub(crate) assignment_changed: Arc<Notify>,
136 pub(crate) next_offsets: Arc<Mutex<HashMap<(String, i32), i64>>>,
138 pub(crate) end_offsets: Arc<Mutex<HashMap<(String, i32), i64>>>,
140 pub(crate) positions: Arc<Mutex<HashMap<(String, i32), position::PartitionPosition>>>,
142 pub(crate) topic_ids: Arc<Mutex<HashMap<String, WireUuid>>>,
145 pub(crate) session_timeout: Time,
146 pub(crate) heartbeat_interval: Time,
147 pub(crate) rebalance_protocol: assignor::RebalanceProtocol,
149 pub(crate) coordinator_shutdown: CancellationToken,
150 pub(crate) coordinator_handle: Option<JoinHandle<()>>,
151 pub(crate) isolation_level: IsolationLevel,
153 pub(crate) fetch_min: ByteSize,
154 pub(crate) fetch_max: ByteSize,
155 pub(crate) fetch_partition_max: ByteSize,
156 pub(crate) fetch_max_wait: Time,
158 pub(crate) fetches: poll::Fetches,
160 pub(crate) client_rack: Option<String>,
162 pub(crate) metadata_max_age: Time,
164 pub(crate) request_timeout: Time,
166 pub(crate) wakeup: control::WakeupHandle,
168 pub(crate) enforced_rebalances: tokio::sync::mpsc::UnboundedSender<(String, u64)>,
170 pub(crate) default_api_timeout: Time,
173 pub(crate) paused: crate::partition_state::PausedPartitions,
175 pub(crate) auto_offset_reset: AutoOffsetReset,
179 pub(crate) poll_error: coordinator::PollErrorSlot,
182 pub(crate) auto_commit: Option<commit::AutoCommit>,
185 pub(crate) poll_signal: coordinator::PollSignal,
188 pub(crate) rebalance_pending: tokio::sync::watch::Receiver<bool>,
191 pub(crate) max_poll_records: usize,
193 pub(crate) fetch_buffer: fetch_buffer::FetchBuffer,
195 pub(crate) close_operation: tokio::sync::watch::Sender<control::CloseRequest>,
198 pub(crate) rebalance_listener: Option<rebalance_listener::SharedListener>,
200 pub(crate) listener_calls:
202 tokio::sync::mpsc::UnboundedReceiver<rebalance_listener::ListenerCall>,
203 pub(crate) assigned_callback_pending: rebalance_listener::AssignedCallbackPending,
206}
207
208#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
211pub enum GroupMembershipOperation {
212 LeaveGroup,
214 RemainInGroup,
217 #[default]
221 Default,
222}
223
224#[derive(Clone)]
225pub(crate) struct CommitIdentity {
226 pub generation: i32,
227 pub member_id: String,
228 pub ownership_ids: HashMap<(String, i32), u64>,
229 pub rejoin_on_poll: bool,
233}
234
235#[derive(Clone)]
236struct StartConfig {
237 bootstrap: String,
238 client_id: String,
239 group_id: String,
240 session_timeout: Time,
241 max_poll_interval: Time,
243 max_poll_records: usize,
245 heartbeat_interval: Time,
246 subscription_metadata_refresh_interval: Time,
247 subscribe: Vec<String>,
248 subscribe_pattern: Option<crate::TopicPattern>,
251 subscribe_regex: Option<String>,
253 exclude_internal_topics: bool,
255 group_protocol: crate::GroupProtocol,
257 server_assignor: Option<String>,
259 group_instance_id: Option<String>,
260 auto_offset_reset: AutoOffsetReset,
261 isolation_level: IsolationLevel,
262 assignors: Vec<Assignor>,
264 rebalance_protocol: assignor::RebalanceProtocol,
265 fetch_min: ByteSize,
266 fetch_max: ByteSize,
267 fetch_partition_max: ByteSize,
268 fetch_max_wait: Time,
269 metadata_max_age: Time,
270 default_api_timeout: Time,
271 request_timeout: Time,
272 dns_timeout: Time,
273 send_buffer: Option<ByteSize>,
274 receive_buffer: Option<ByteSize>,
275 dispatch_queue_capacity: krabka_client_core::ConnectionDispatchQueueCapacity,
276 frame_max: krabka_client_core::ClientFrameMax,
277 metadata_recovery_strategy: krabka_client_core::MetadataRecoveryStrategy,
278 metadata_recovery_rebootstrap_trigger: Time,
279 leave_group_timeout: Time,
280 client_rack: Option<String>,
281 security: Option<security::ClientSecurity>,
282 client_dns_lookup: krabka_client_core::ClientDnsLookup,
284 reconnect_backoff: Time,
286 reconnect_backoff_max: Time,
288 connections_max_idle: Time,
290 socket_connection_setup_timeout: Time,
292 socket_connection_setup_timeout_max: Time,
294 has_rebalance_listener: bool,
296 retry_policy: ConsumerRetryPolicy,
297 auto_commit_interval: Option<Duration>,
300 allow_auto_create_topics: bool,
302 telemetry: Option<telemetry::ClientTelemetryConfig>,
305}
306
307#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
308struct RetryTime(Duration);
309
310impl RetryTime {
311 fn new(name: &str, value: Time) -> Result<Self, String> {
312 let milliseconds = GreaterI64::<0>::new(value.millis_i64())
313 .map_err(|error| format!("{name}: {error}"))?
314 .into_value();
315 if !value.secs_f64().is_finite() || Time::from_millis(milliseconds) != value {
316 return Err(format!("{name} must be a whole number of milliseconds"));
317 }
318 Ok(Self(Duration::from_millis(
319 u64::try_from(milliseconds).expect("validated retry time is positive"),
320 )))
321 }
322
323 fn time(self) -> Time {
324 Time::from_std(self.0)
325 }
326}
327
328#[derive(Clone, Copy, Debug, Eq, PartialEq)]
330pub struct ConsumerRetryPolicy {
331 startup_attempt_timeout: RetryTime,
332 startup_deadline: RetryTime,
333 startup_initial_backoff: RetryTime,
334 startup_max_backoff: RetryTime,
335 coordinator_retry_timeout: RetryTime,
336 coordinator_initial_backoff: RetryTime,
337 coordinator_max_backoff: RetryTime,
338}
339
340pub const DEFAULT_CONSUMER_RETRY_BACKOFF_MAX: Time = millis(1000);
342
343#[bon::bon]
344impl ConsumerRetryPolicy {
345 #[builder(start_fn = builder, finish_fn = build)]
352 pub fn new(
353 #[builder(default = secs(90))] startup_attempt_timeout: Time,
354 #[builder(default = minutes(5))] startup_deadline: Time,
355 #[builder(default = millis(500))] startup_initial_backoff: Time,
356 #[builder(default = DEFAULT_CONSUMER_RETRY_BACKOFF_MAX)] startup_max_backoff: Time,
357 #[builder(default = secs(30))] coordinator_retry_timeout: Time,
358 #[builder(default = millis(100))] coordinator_initial_backoff: Time,
359 #[builder(default = DEFAULT_CONSUMER_RETRY_BACKOFF_MAX)] coordinator_max_backoff: Time,
360 ) -> Result<Self, String> {
361 let policy = Self {
362 startup_attempt_timeout: RetryTime::new(
363 "consumer startup attempt timeout",
364 startup_attempt_timeout,
365 )?,
366 startup_deadline: RetryTime::new("consumer startup deadline", startup_deadline)?,
367 startup_initial_backoff: RetryTime::new(
368 "consumer startup initial backoff",
369 startup_initial_backoff,
370 )?,
371 startup_max_backoff: RetryTime::new(
372 "consumer startup maximum backoff",
373 startup_max_backoff,
374 )?,
375 coordinator_retry_timeout: RetryTime::new(
376 "consumer coordinator retry timeout",
377 coordinator_retry_timeout,
378 )?,
379 coordinator_initial_backoff: RetryTime::new(
380 "consumer coordinator initial backoff",
381 coordinator_initial_backoff,
382 )?,
383 coordinator_max_backoff: RetryTime::new(
384 "consumer coordinator maximum backoff",
385 coordinator_max_backoff,
386 )?,
387 };
388 if policy.startup_attempt_timeout > policy.startup_deadline {
389 return Err(
390 "consumer startup attempt timeout must not exceed startup deadline".to_owned(),
391 );
392 }
393 if policy.startup_initial_backoff > policy.startup_max_backoff {
394 return Err(
395 "consumer startup initial backoff must not exceed startup maximum backoff"
396 .to_owned(),
397 );
398 }
399 if policy.coordinator_initial_backoff > policy.coordinator_max_backoff {
400 return Err(
401 "consumer coordinator initial backoff must not exceed coordinator maximum backoff"
402 .to_owned(),
403 );
404 }
405 Ok(policy)
406 }
407
408 #[must_use]
410 pub fn startup_attempt_timeout(self) -> Time {
411 self.startup_attempt_timeout.time()
412 }
413
414 #[must_use]
416 pub fn startup_deadline(self) -> Time {
417 self.startup_deadline.time()
418 }
419
420 #[must_use]
422 pub fn startup_initial_backoff(self) -> Time {
423 self.startup_initial_backoff.time()
424 }
425
426 #[must_use]
428 pub fn startup_max_backoff(self) -> Time {
429 self.startup_max_backoff.time()
430 }
431
432 #[must_use]
434 pub fn coordinator_retry_timeout(self) -> Time {
435 self.coordinator_retry_timeout.time()
436 }
437
438 #[must_use]
440 pub fn coordinator_initial_backoff(self) -> Time {
441 self.coordinator_initial_backoff.time()
442 }
443
444 #[must_use]
446 pub fn coordinator_max_backoff(self) -> Time {
447 self.coordinator_max_backoff.time()
448 }
449
450 pub(crate) fn with_retry_backoff_max(retry_backoff_max: Time) -> Result<Self, String> {
463 let startup_initial_backoff = if millis(500) < retry_backoff_max {
464 millis(500)
465 } else {
466 retry_backoff_max
467 };
468 let coordinator_initial_backoff = if millis(100) < retry_backoff_max {
469 millis(100)
470 } else {
471 retry_backoff_max
472 };
473 Self::builder()
474 .startup_initial_backoff(startup_initial_backoff)
475 .startup_max_backoff(retry_backoff_max)
476 .coordinator_initial_backoff(coordinator_initial_backoff)
477 .coordinator_max_backoff(retry_backoff_max)
478 .build()
479 }
480}
481
482impl Default for ConsumerRetryPolicy {
483 fn default() -> Self {
484 Self::with_retry_backoff_max(DEFAULT_CONSUMER_RETRY_BACKOFF_MAX)
485 .expect("default consumer retry backoff max is valid")
486 }
487}
488
489pub const DEFAULT_CONSUMER_LEAVE_GROUP_TIMEOUT: Time = secs(5);
491
492#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
494pub struct ConsumerLeaveGroupTimeout(Duration);
495
496impl ConsumerLeaveGroupTimeout {
497 pub fn new(value: Duration) -> Result<Self, String> {
504 validated_millisecond_duration("consumer leave-group timeout", value)?;
505 Ok(Self(value))
506 }
507
508 #[must_use]
510 pub const fn duration(self) -> Duration {
511 self.0
512 }
513
514 #[must_use]
520 pub fn milliseconds(self) -> u64 {
521 validated_millisecond_duration("consumer leave-group timeout", self.0)
522 .expect("validated consumer leave-group timeout")
523 }
524}
525
526impl Default for ConsumerLeaveGroupTimeout {
527 fn default() -> Self {
528 Self::new(DEFAULT_CONSUMER_LEAVE_GROUP_TIMEOUT.to_std())
529 .expect("default consumer leave-group timeout is valid")
530 }
531}
532
533pub const DEFAULT_CONSUMER_SUBSCRIPTION_METADATA_REFRESH_INTERVAL: Time = secs(5);
535
536#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
538pub struct ConsumerFetchMaxBytes(i32);
539
540impl ConsumerFetchMaxBytes {
541 pub fn new(value: i32) -> Result<Self, String> {
547 validated_positive_i32(value, "consumer fetch max bytes").map(Self)
548 }
549
550 #[must_use]
552 pub const fn bytes(self) -> i32 {
553 self.0
554 }
555
556 #[must_use]
558 pub fn size(self) -> ByteSize {
559 ByteSize::from_bytes_i64(i64::from(self.0))
560 }
561}
562
563impl TryFrom<ByteSize> for ConsumerFetchMaxBytes {
564 type Error = String;
565
566 fn try_from(value: ByteSize) -> Result<Self, Self::Error> {
567 Self::new(protocol_bytes("consumer fetch max", value)?)
568 }
569}
570
571#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
573pub struct ConsumerFetchPartitionMaxBytes(i32);
574
575impl ConsumerFetchPartitionMaxBytes {
576 pub fn new(value: i32) -> Result<Self, String> {
582 validated_positive_i32(value, "consumer fetch partition max bytes").map(Self)
583 }
584
585 #[must_use]
587 pub const fn bytes(self) -> i32 {
588 self.0
589 }
590
591 #[must_use]
593 pub fn size(self) -> ByteSize {
594 ByteSize::from_bytes_i64(i64::from(self.0))
595 }
596}
597
598impl TryFrom<ByteSize> for ConsumerFetchPartitionMaxBytes {
599 type Error = String;
600
601 fn try_from(value: ByteSize) -> Result<Self, Self::Error> {
602 Self::new(protocol_bytes("consumer fetch partition max", value)?)
603 }
604}
605
606#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
608pub struct ConsumerSubscriptionMetadataRefreshInterval(Duration);
609
610impl ConsumerSubscriptionMetadataRefreshInterval {
611 pub fn new(value: Duration) -> Result<Self, String> {
618 validated_millisecond_duration("consumer subscription metadata refresh interval", value)?;
619 Ok(Self(value))
620 }
621
622 #[must_use]
624 pub const fn duration(self) -> Duration {
625 self.0
626 }
627
628 #[must_use]
634 pub fn milliseconds(self) -> u64 {
635 validated_millisecond_duration("consumer subscription metadata refresh interval", self.0)
636 .expect("validated consumer subscription metadata refresh interval")
637 }
638}
639
640impl Default for ConsumerSubscriptionMetadataRefreshInterval {
641 fn default() -> Self {
642 Self::new(DEFAULT_CONSUMER_SUBSCRIPTION_METADATA_REFRESH_INTERVAL.to_std())
643 .expect("default consumer subscription metadata refresh interval is valid")
644 }
645}
646
647impl Drop for Consumer {
648 fn drop(&mut self) {
649 self.coordinator_shutdown.cancel();
650 }
651}
652
653#[derive(Debug, Clone, PartialEq, Eq)]
656pub struct Header {
657 pub key: String,
658 pub value: Option<Bytes>,
659}
660
661#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
667pub enum TimestampType {
668 CreateTime,
670 LogAppendTime,
672}
673
674impl From<records::TimestampType> for TimestampType {
675 fn from(value: records::TimestampType) -> Self {
676 match value {
677 records::TimestampType::CreateTime => Self::CreateTime,
678 records::TimestampType::LogAppendTime => Self::LogAppendTime,
679 }
680 }
681}
682
683#[derive(Debug, Clone, PartialEq, Eq)]
689pub struct ConsumerRecord {
690 pub topic: String,
691 pub partition: i32,
692 pub offset: i64,
693 pub leader_epoch: i32,
694 pub timestamp: i64,
697 pub timestamp_type: TimestampType,
699 pub key: Option<Bytes>,
700 pub value: Option<Bytes>,
701 pub headers: Vec<Header>,
702}
703
704fn initial_join_protocols(
707 assignors: &[Assignor],
708 subscribe: &[String],
709 client_rack: Option<&str>,
710) -> Vec<(String, bytes::Bytes)> {
711 coordinator::join_protocols(assignors, subscribe, &[], -1, client_rack, None)
712}
713
714fn startup_join_fields(
717 group_id: &str,
718 group_instance_id: Option<&str>,
719 protocols: Vec<(String, bytes::Bytes)>,
720 (session_timeout_ms, rebalance_timeout_ms): (i32, i32),
721) -> coordinator::JoinRequestFields {
722 coordinator::JoinRequestFields {
723 group_id: group_id.to_owned(),
724 group_instance_id: group_instance_id.map(str::to_owned),
725 session_timeout_ms,
726 rebalance_timeout_ms,
727 protocols,
728 reason: String::new(),
729 }
730}
731
732fn first_join_member_id(resp: &JoinGroupResponse) -> Result<String, ConsumerError> {
733 let member_id = if resp.error_code == 79 || resp.error_code == 0 {
734 resp.member_id.clone()
735 } else {
736 return Err(ConsumerError::Server(resp.error_code));
737 };
738 if member_id.is_empty() {
739 return Err(ConsumerError::RebalanceFailed(
740 "broker did not assign a member_id".into(),
741 ));
742 }
743 Ok(member_id)
744}
745
746#[derive(bon::Builder, Default)]
747struct ConsumerClientOptions {
748 #[builder(default)]
749 metadata_scope: krabka_client_core::MetadataScope,
750 telemetry: Option<telemetry::ClientTelemetryConfig>,
751}
752
753async fn build_consumer_client(
754 config: &StartConfig,
755 options: ConsumerClientOptions,
756) -> Result<Client, ConsumerError> {
757 Ok(Client::builder()
758 .bootstrap(&config.bootstrap)
759 .client_id(config.client_id.clone())
760 .request_timeout(config.request_timeout)
761 .dns_timeout(config.dns_timeout)
762 .send_buffer(config.send_buffer)
763 .receive_buffer(config.receive_buffer)
764 .dispatch_queue_capacity(config.dispatch_queue_capacity.get())
765 .frame_max(config.frame_max.size())
766 .metadata_recovery_strategy(config.metadata_recovery_strategy)
767 .metadata_recovery_rebootstrap_trigger(config.metadata_recovery_rebootstrap_trigger)
768 .maybe_security(config.security.clone())
769 .client_dns_lookup(config.client_dns_lookup)
770 .reconnect_backoff(config.reconnect_backoff)
771 .reconnect_backoff_max(config.reconnect_backoff_max)
772 .connections_max_idle(config.connections_max_idle)
773 .socket_connection_setup_timeout(config.socket_connection_setup_timeout)
774 .socket_connection_setup_timeout_max(config.socket_connection_setup_timeout_max)
775 .metadata_scope(options.metadata_scope)
776 .maybe_telemetry(options.telemetry)
777 .build()
778 .await?)
779}
780
781async fn pattern_topics(config: &StartConfig) -> Result<Vec<String>, ConsumerError> {
785 let client = build_consumer_client(config, ConsumerClientOptions::default()).await?;
786 let metadata = client
787 .refresh_metadata_with(krabka_client_core::all_topics_request())
788 .await;
789 client.close();
790 let subscription = subscription::Subscription {
791 pattern: config.subscribe_pattern.clone(),
792 ..subscription::Subscription::topics(Vec::new(), config.exclude_internal_topics)
793 };
794 Ok(subscription.matching_topics(&metadata?).unwrap_or_default())
795}
796
797fn initial_commit_identity(
800 assigned: &[(String, i32)],
801 generation: i32,
802 member_id: &str,
803) -> (CommitIdentity, u64) {
804 let ownership_ids = assigned.iter().cloned().zip(1_u64..).collect();
805 let identity = CommitIdentity {
806 generation,
807 member_id: member_id.to_owned(),
808 ownership_ids,
809 rejoin_on_poll: false,
810 };
811 (identity, assigned.len() as u64 + 1)
812}
813
814fn queue_initial_assign_call(
818 calls: &rebalance_listener::ListenerCalls,
819 partitions: Vec<(String, i32)>,
820) {
821 let (done, _) = tokio::sync::oneshot::channel();
822 let _ = calls.send(rebalance_listener::ListenerCall {
823 kind: rebalance_listener::ListenerCallKind::Assigned,
824 partitions,
825 done,
826 });
827}
828
829fn subscription_metadata_scope(
833 allow_auto_create_topics: bool,
834) -> krabka_client_core::MetadataScope {
835 krabka_client_core::MetadataScope::Topics {
836 allow_auto_topic_creation: allow_auto_create_topics,
837 }
838}
839
840fn is_subscribed_topic(subscribe: &[String], name: &str) -> bool {
841 subscribe.iter().any(|s| s == name)
842}
843
844fn is_group_leader(leader: &str, member_id: &str) -> bool {
845 leader == member_id
846}
847
848async fn leave_startup_member(
849 client: &Client,
850 coordinator_id: &AtomicI32,
851 group_id: &str,
852 member_id: &str,
853 group_instance_id: Option<String>,
854 leave_group_timeout: Time,
855) {
856 if !coordinator::should_send_leave_group(
857 member_id,
858 group_instance_id.as_deref(),
859 GroupMembershipOperation::Default,
860 ) {
861 return;
862 }
863 let broker = client.broker(coordinator_id.load(Ordering::Relaxed));
864 let send = broker.send(coordinator::build_leave_group_request(
865 group_id.to_string(),
866 member_id.to_string(),
867 None,
868 ));
869 let _ = tokio::time::timeout(leave_group_timeout.to_std(), send).await;
870}
871
872fn has_assigned_partitions(assigned_partitions: &[(String, i32)]) -> bool {
873 !assigned_partitions.is_empty()
874}
875
876pub(crate) fn starting_offset(committed: i64, auto_offset_reset: AutoOffsetReset) -> i64 {
877 if committed >= 0 {
878 committed
879 } else {
880 reset_starting_offset(auto_offset_reset)
881 }
882}
883
884pub(crate) fn reset_starting_offset(auto_offset_reset: AutoOffsetReset) -> i64 {
897 match auto_offset_reset {
898 AutoOffsetReset::Earliest => poll::BEGINNING_SENTINEL,
901 AutoOffsetReset::Latest | AutoOffsetReset::ByDuration(_) => poll::LATEST_SENTINEL,
903 AutoOffsetReset::None => poll::NO_OFFSET_SENTINEL,
905 }
906}
907
908pub const DEFAULT_CONSUMER_DEFAULT_API_TIMEOUT: Time = secs(60);
910
911pub const DEFAULT_CONSUMER_MAX_POLL_INTERVAL: Time = minutes(5);
913
914static CONSUMER_CLIENT_ID_SEQUENCE: std::sync::atomic::AtomicU64 =
916 std::sync::atomic::AtomicU64::new(1);
917
918fn consumer_client_id(
923 client_id: Option<String>,
924 group_id: &str,
925 group_instance_id: Option<&str>,
926) -> String {
927 match client_id {
928 Some(client_id) if !client_id.is_empty() => client_id,
929 _ => match group_instance_id {
930 Some(group_instance_id) => format!("consumer-{group_id}-{group_instance_id}"),
931 None => format!(
932 "consumer-{group_id}-{}",
933 CONSUMER_CLIENT_ID_SEQUENCE.fetch_add(1, Ordering::Relaxed)
934 ),
935 },
936 }
937}
938
939pub const DEFAULT_CONSUMER_METADATA_MAX_AGE: Time = minutes(5);
941
942pub const DEFAULT_CONSUMER_FETCH_MAX_WAIT: Time = millis(500);
944
945fn validated_fetch_max_wait(wait: Time) -> Result<Time, String> {
950 let milliseconds = MinMaxI64::<0, { i32::MAX as i64 }>::new(wait.millis_i64())
951 .map_err(|error| format!("consumer fetch max wait: {error}"))?
952 .into_value();
953 if !wait.secs_f64().is_finite() || Time::from_millis(milliseconds) != wait {
954 return Err("consumer fetch max wait must be a whole number of milliseconds".to_owned());
955 }
956 Ok(wait)
957}
958
959pub const DEFAULT_CONSUMER_MAX_POLL_RECORDS: usize = 500;
961
962fn validated_max_poll_records(records: usize) -> Result<usize, String> {
967 if records == 0 || i32::try_from(records).is_err() {
968 return Err("consumer max poll records must be from 1 to i32::MAX".to_owned());
969 }
970 Ok(records)
971}
972
973fn validated_max_poll_interval(interval: Time) -> Result<Time, String> {
979 let milliseconds = MinMaxI64::<1, { i32::MAX as i64 }>::new(interval.millis_i64())
980 .map_err(|error| format!("consumer max poll interval: {error}"))?
981 .into_value();
982 if !interval.secs_f64().is_finite() || Time::from_millis(milliseconds) != interval {
983 return Err("consumer max poll interval must be a whole number of milliseconds".to_owned());
984 }
985 Ok(interval)
986}
987
988fn validated_auto_commit_interval(
994 enable_auto_commit: bool,
995 interval: Time,
996) -> Result<Option<Duration>, String> {
997 if !enable_auto_commit {
998 return Ok(None);
999 }
1000 let milliseconds = MinMaxI64::<0, { i32::MAX as i64 }>::new(interval.millis_i64())
1001 .map_err(|error| format!("consumer auto commit interval: {error}"))?
1002 .into_value();
1003 if !interval.secs_f64().is_finite() || Time::from_millis(milliseconds) != interval {
1004 return Err(
1005 "consumer auto commit interval must be a whole number of milliseconds".to_owned(),
1006 );
1007 }
1008 Ok(Some(interval.to_std()))
1009}
1010
1011fn primed_position(committed_epoch: i32) -> position::PartitionPosition {
1012 position::PartitionPosition {
1013 offset_epoch: krabka_ids::LeaderEpoch(committed_epoch),
1016 ..Default::default()
1017 }
1018}
1019
1020pub(crate) fn protocol_millis_i32(value: Time) -> i32 {
1027 i32::try_from(value.millis_i64_trunc()).unwrap_or(i32::MAX)
1028}
1029
1030#[bon::bon]
1031impl Consumer {
1032 #[builder(start_fn = builder, finish_fn = build)]
1060 #[tracing::instrument(
1061 name = "consumer.start",
1062 level = "info",
1063 skip_all,
1064 fields(group_id = ?group_id, client_id = ?client_id),
1065 err(level = "info")
1066 )]
1067 pub async fn start(
1070 #[builder(into)] bootstrap: String,
1071 #[builder(into)] client_id: Option<String>,
1072 #[builder(into)]
1075 group_id: Option<String>,
1076 #[builder(default = secs(45))] session_timeout: Time,
1077 #[builder(default = DEFAULT_CONSUMER_MAX_POLL_INTERVAL)] max_poll_interval: Time,
1078 #[builder(default = DEFAULT_CONSUMER_MAX_POLL_RECORDS)] max_poll_records: usize,
1079 #[builder(default = secs(3))] heartbeat_interval: Time,
1080 #[builder(default = DEFAULT_CONSUMER_SUBSCRIPTION_METADATA_REFRESH_INTERVAL)]
1081 subscription_metadata_refresh_interval: Time,
1082 #[builder(into, default)] subscribe: Vec<String>,
1083 subscribe_pattern: Option<crate::TopicPattern>,
1086 #[builder(default = true)]
1089 exclude_internal_topics: bool,
1090 #[builder(into)]
1095 subscribe_regex: Option<String>,
1096 #[builder(default)]
1100 group_protocol: crate::GroupProtocol,
1101 #[builder(into)]
1105 group_remote_assignor: Option<String>,
1106 #[builder(into)] group_instance_id: Option<String>,
1107 #[builder(default = AutoOffsetReset::Latest)] auto_offset_reset: AutoOffsetReset,
1108 #[builder(default = IsolationLevel::ReadUncommitted)] isolation_level: IsolationLevel,
1109 #[builder(into, default = Assignor::DEFAULT_LIST.to_vec())] assignors: Vec<Assignor>,
1110 #[builder(default = krabka_client_core::DEFAULT_FETCH_MIN.bytes_i32())]
1112 fetch_min: i32,
1113 #[builder(default = poll::DEFAULT_FETCH_MAX.bytes_i32())]
1115 fetch_max: i32,
1116 #[builder(default = poll::DEFAULT_FETCH_PARTITION_MAX.bytes_i32())]
1118 fetch_partition_max: i32,
1119 #[builder(default = DEFAULT_CONSUMER_FETCH_MAX_WAIT)] fetch_max_wait: Time,
1120 #[builder(default = DEFAULT_CONSUMER_METADATA_MAX_AGE)] metadata_max_age: Time,
1121 #[builder(default = DEFAULT_CONSUMER_DEFAULT_API_TIMEOUT)]
1126 default_api_timeout: Time,
1127 #[builder(default = secs(30))] request_timeout: Time,
1128 #[builder(default = krabka_client_core::DEFAULT_CLIENT_DNS_TIMEOUT)]
1130 dns_timeout: Time,
1131 #[builder(required, default = Some(krabka_client_core::DEFAULT_SEND_BUFFER))]
1133 send_buffer: Option<ByteSize>,
1134 #[builder(required, default = Some(krabka_client_core::DEFAULT_RECEIVE_BUFFER))]
1136 receive_buffer: Option<ByteSize>,
1137 #[builder(default = krabka_client_core::DEFAULT_CONNECTION_DISPATCH_QUEUE_CAPACITY)]
1138 dispatch_queue_capacity: usize,
1139 #[builder(default = krabka_client_core::DEFAULT_CLIENT_FRAME_MAX)] frame_max: ByteSize,
1140 #[builder(default)]
1141 metadata_recovery_strategy: krabka_client_core::MetadataRecoveryStrategy,
1142 #[builder(default = krabka_client_core::DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER)]
1143 metadata_recovery_rebootstrap_trigger: Time,
1144 #[builder(default = DEFAULT_CONSUMER_LEAVE_GROUP_TIMEOUT)] leave_group_timeout: Time,
1145 #[builder(into)] client_rack: Option<String>,
1146 security: Option<security::ClientSecurity>,
1147 #[builder(default)]
1149 client_dns_lookup: krabka_client_core::ClientDnsLookup,
1150 #[builder(default = krabka_client_core::DEFAULT_RECONNECT_BACKOFF)]
1152 reconnect_backoff: Time,
1153 #[builder(default = krabka_client_core::DEFAULT_RECONNECT_BACKOFF_MAX)]
1155 reconnect_backoff_max: Time,
1156 #[builder(default = krabka_client_core::DEFAULT_CONNECTIONS_MAX_IDLE)]
1158 connections_max_idle: Time,
1159 #[builder(default = krabka_client_core::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT)]
1161 socket_connection_setup_timeout: Time,
1162 #[builder(default = krabka_client_core::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX)]
1164 socket_connection_setup_timeout_max: Time,
1165 rebalance_listener: Option<Box<dyn crate::ConsumerRebalanceListener>>,
1166 #[builder(default = DEFAULT_CONSUMER_RETRY_BACKOFF_MAX)]
1169 retry_backoff_max: Time,
1170 retry_policy: Option<ConsumerRetryPolicy>,
1173 enable_auto_commit: Option<bool>,
1176 #[builder(default = secs(5))] auto_commit_interval: Time,
1177 #[builder(default = true)]
1181 allow_auto_create_topics: bool,
1182 #[builder(default = true)]
1186 enable_metrics_push: bool,
1187 ) -> Result<Self, ConsumerError> {
1188 let group_id = group_id.unwrap_or_default();
1191 if group_id.is_empty() && !(subscribe.is_empty() && subscribe_pattern.is_none()) {
1192 return Err(ConsumerError::InvalidGroupId);
1193 }
1194 if group_id.is_empty() && enable_auto_commit == Some(true) {
1196 return Err(ConsumerError::InvalidConfig(
1197 "enable.auto.commit cannot be set to true when default group id (null) is used."
1198 .to_owned(),
1199 ));
1200 }
1201 let enable_auto_commit = enable_auto_commit.unwrap_or(!group_id.is_empty());
1202 if group_protocol == GroupProtocol::Classic && group_remote_assignor.is_some() {
1206 return Err(ConsumerError::InvalidConfig(
1207 "group.remote.assignor cannot be set when group.protocol=classic".to_owned(),
1208 ));
1209 }
1210 if !subscribe.is_empty() && subscribe_pattern.is_some() {
1211 return Err(ConsumerError::InvalidConfig(
1212 "subscribe and subscribe_pattern are mutually exclusive".to_owned(),
1213 ));
1214 }
1215 if subscribe_regex.is_some() && (!subscribe.is_empty() || subscribe_pattern.is_some()) {
1216 return Err(ConsumerError::InvalidConfig(
1217 "Subscription to topics, partitions and pattern are mutually exclusive".to_owned(),
1218 ));
1219 }
1220 if let Some(regex) = &subscribe_regex {
1221 if regex.is_empty() {
1222 return Err(ConsumerError::InvalidArgument(
1223 "Topic pattern to subscribe to cannot be empty".to_owned(),
1224 ));
1225 }
1226 if group_protocol != GroupProtocol::Consumer {
1227 return Err(ConsumerError::InvalidConfig(
1228 consumer_group::RE2J_NEEDS_CONSUMER_PROTOCOL.to_owned(),
1229 ));
1230 }
1231 if group_id.is_empty() {
1232 return Err(ConsumerError::InvalidGroupId);
1233 }
1234 }
1235 if subscribe.iter().any(String::is_empty) {
1236 return Err(ConsumerError::InvalidConfig(
1237 "Topic collection to subscribe to cannot contain null or empty topic".to_owned(),
1238 ));
1239 }
1240 if group_instance_id.as_deref().is_some_and(str::is_empty) {
1241 return Err(ConsumerError::InvalidConfig(
1242 "group_instance_id must not be empty".into(),
1243 ));
1244 }
1245 krabka_client_core::ClientDnsTimeout::new(dns_timeout)
1246 .map_err(ConsumerError::InvalidConfig)?;
1247 let fetch_min = FetchMinBytes::new(fetch_min)
1248 .map_err(ConsumerError::InvalidConfig)?
1249 .size();
1250 let fetch_max = ConsumerFetchMaxBytes::new(fetch_max)
1251 .map_err(ConsumerError::InvalidConfig)?
1252 .size();
1253 if fetch_min.bytes_i32() > fetch_max.bytes_i32() {
1254 return Err(ConsumerError::InvalidConfig(
1255 "consumer fetch min must not exceed consumer fetch max".to_owned(),
1256 ));
1257 }
1258 let fetch_partition_max = ConsumerFetchPartitionMaxBytes::new(fetch_partition_max)
1259 .map_err(ConsumerError::InvalidConfig)?
1260 .size();
1261 let leave_group_timeout = ConsumerLeaveGroupTimeout::new(leave_group_timeout.to_std())
1266 .map_err(ConsumerError::InvalidConfig)?;
1267 let subscription_metadata_refresh_interval =
1268 ConsumerSubscriptionMetadataRefreshInterval::new(
1269 subscription_metadata_refresh_interval.to_std(),
1270 )
1271 .map_err(ConsumerError::InvalidConfig)?;
1272 let dispatch_queue_capacity = ConnectionDispatchQueueCapacity::new(dispatch_queue_capacity)
1273 .map_err(ConsumerError::InvalidConfig)?;
1274 let frame_max =
1275 ClientFrameMax::try_from(frame_max).map_err(ConsumerError::InvalidConfig)?;
1276 let metadata_recovery_rebootstrap_trigger =
1277 krabka_client_core::MetadataRecoveryRebootstrapTrigger::new(
1278 metadata_recovery_rebootstrap_trigger,
1279 )
1280 .map_err(ConsumerError::InvalidConfig)?
1281 .time();
1282 let auto_commit_interval =
1283 validated_auto_commit_interval(enable_auto_commit, auto_commit_interval)
1284 .map_err(ConsumerError::InvalidConfig)?;
1285 let max_poll_interval =
1286 validated_max_poll_interval(max_poll_interval).map_err(ConsumerError::InvalidConfig)?;
1287 let max_poll_records =
1288 validated_max_poll_records(max_poll_records).map_err(ConsumerError::InvalidConfig)?;
1289 let fetch_max_wait =
1290 validated_fetch_max_wait(fetch_max_wait).map_err(ConsumerError::InvalidConfig)?;
1291 if metadata_max_age.secs_f64() < 0.0 || !metadata_max_age.secs_f64().is_finite() {
1292 return Err(ConsumerError::InvalidConfig(
1293 "consumer metadata max age must not be negative".to_owned(),
1294 ));
1295 }
1296 if default_api_timeout.secs_f64() < 0.0 || !default_api_timeout.secs_f64().is_finite() {
1297 return Err(ConsumerError::InvalidConfig(
1298 "consumer default api timeout must not be negative".to_owned(),
1299 ));
1300 }
1301 let retry_policy = match retry_policy {
1304 Some(retry_policy) => retry_policy,
1305 None => ConsumerRetryPolicy::with_retry_backoff_max(retry_backoff_max)
1306 .map_err(ConsumerError::InvalidConfig)?,
1307 };
1308 let rebalance_protocol =
1309 assignor::rebalance_protocol_of(&assignors).map_err(ConsumerError::InvalidConfig)?;
1310
1311 let client_id = consumer_client_id(client_id, &group_id, group_instance_id.as_deref());
1312 let telemetry = enable_metrics_push.then(|| {
1313 telemetry::ClientTelemetryConfig::consumer(
1314 Some(group_id.as_str()),
1315 group_instance_id.as_deref(),
1316 client_rack.as_deref(),
1317 )
1318 });
1319 let config = StartConfig {
1320 bootstrap,
1321 client_id,
1322 group_id,
1323 session_timeout,
1324 max_poll_interval,
1325 max_poll_records,
1326 heartbeat_interval,
1327 subscription_metadata_refresh_interval: Time::from_std(
1328 subscription_metadata_refresh_interval.duration(),
1329 ),
1330 subscribe,
1331 subscribe_pattern,
1332 subscribe_regex,
1333 exclude_internal_topics,
1334 group_protocol,
1335 server_assignor: group_remote_assignor,
1336 group_instance_id,
1337 auto_offset_reset,
1338 isolation_level,
1339 assignors,
1340 rebalance_protocol,
1341 fetch_min,
1342 fetch_max,
1343 fetch_partition_max,
1344 fetch_max_wait,
1345 metadata_max_age,
1346 default_api_timeout,
1347 request_timeout,
1348 dns_timeout,
1349 send_buffer,
1350 receive_buffer,
1351 dispatch_queue_capacity,
1352 frame_max,
1353 metadata_recovery_strategy,
1354 metadata_recovery_rebootstrap_trigger,
1355 leave_group_timeout: Time::from_std(leave_group_timeout.duration()),
1356 client_rack,
1357 security,
1358 client_dns_lookup,
1359 reconnect_backoff,
1360 reconnect_backoff_max,
1361 connections_max_idle,
1362 socket_connection_setup_timeout,
1363 socket_connection_setup_timeout_max,
1364 has_rebalance_listener: rebalance_listener.is_some(),
1365 retry_policy,
1366 auto_commit_interval,
1367 allow_auto_create_topics,
1368 telemetry,
1369 };
1370
1371 Self::retry_start_until_success(config, rebalance_listener).await
1372 }
1373
1374 async fn retry_start_until_success(
1378 config: StartConfig,
1379 rebalance_listener: Option<Box<dyn crate::ConsumerRebalanceListener>>,
1380 ) -> Result<Self, ConsumerError> {
1381 let started = tokio::time::Instant::now();
1382 let total = config.retry_policy.startup_deadline().to_std();
1383 let mut backoff = config.retry_policy.startup_initial_backoff().to_std();
1384 loop {
1385 let remaining = total.saturating_sub(started.elapsed());
1386 if remaining.is_zero() {
1387 return Err(ConsumerError::Client(ClientError::Timeout(
1388 config.retry_policy.startup_deadline(),
1389 )));
1390 }
1391 match tokio::time::timeout(
1392 config
1393 .retry_policy
1394 .startup_attempt_timeout()
1395 .to_std()
1396 .min(remaining),
1397 Box::pin(Self::start_once(config.clone())),
1398 )
1399 .await
1400 {
1401 Ok(Ok(mut consumer)) => {
1402 consumer.rebalance_listener =
1403 rebalance_listener.map(rebalance_listener::shared);
1404 return Ok(consumer);
1405 }
1406 Ok(Err(error)) => {
1407 if started.elapsed().as_time() < config.retry_policy.startup_deadline()
1408 && is_retriable_consumer_start_error(&error)
1409 {
1410 tracing::warn!(
1411 group = %config.group_id,
1412 %error,
1413 "consumer startup failed transiently; retrying with a fresh connection"
1414 );
1415 } else {
1416 return Err(error);
1417 }
1418 }
1419 Err(_elapsed) => {
1420 if started.elapsed().as_time() >= config.retry_policy.startup_deadline() {
1421 return Err(ConsumerError::Client(ClientError::Timeout(
1422 config.retry_policy.startup_deadline(),
1423 )));
1424 }
1425 tracing::warn!(
1426 group = %config.group_id,
1427 timeout = ?config.retry_policy.startup_attempt_timeout(),
1428 "consumer startup exceeded attempt timeout \
1429 (likely a cold-boot group-join stall); \
1430 retrying with a fresh connection"
1431 );
1432 }
1433 }
1434 tokio::time::sleep(backoff.min(total.saturating_sub(started.elapsed()))).await;
1435 backoff = backoff
1436 .saturating_mul(2)
1437 .min(config.retry_policy.startup_max_backoff().to_std());
1438 }
1439 }
1440
1441 #[tracing::instrument(
1451 name = "consumer.start_once",
1452 level = "info",
1453 skip_all,
1454 fields(
1455 group_id = %config.group_id,
1456 coordinator_id = tracing::field::Empty,
1457 member_id = tracing::field::Empty,
1458 generation = tracing::field::Empty,
1459 is_leader = tracing::field::Empty,
1460 assigned_partitions = tracing::field::Empty,
1461 ),
1462 err(level = "info")
1463 )]
1464 async fn start_once(mut config: StartConfig) -> Result<Self, ConsumerError> {
1465 if config.subscribe.is_empty()
1466 && config.subscribe_pattern.is_none()
1467 && config.subscribe_regex.is_none()
1468 {
1469 return start_without_subscription(config).await;
1470 }
1471 if config.subscribe_pattern.is_some() {
1472 config.subscribe = pattern_topics(&config).await?;
1473 }
1474 if config.group_protocol == GroupProtocol::Consumer {
1475 return start_without_subscription(config).await;
1478 }
1479 let finish_config = config.clone();
1480 let telemetry = config.telemetry.clone();
1481 let StartConfig {
1482 bootstrap,
1483 client_id,
1484 group_id,
1485 session_timeout,
1486 max_poll_interval,
1487 subscribe,
1488 group_instance_id,
1489 assignors,
1490 request_timeout,
1491 dns_timeout,
1492 send_buffer,
1493 receive_buffer,
1494 dispatch_queue_capacity,
1495 frame_max,
1496 metadata_recovery_strategy,
1497 metadata_recovery_rebootstrap_trigger,
1498 client_rack,
1499 security,
1500 client_dns_lookup,
1501 reconnect_backoff,
1502 reconnect_backoff_max,
1503 connections_max_idle,
1504 socket_connection_setup_timeout,
1505 socket_connection_setup_timeout_max,
1506 allow_auto_create_topics,
1507 ..
1508 } = config;
1509 let client = Client::builder()
1510 .bootstrap(&bootstrap)
1511 .client_id(client_id.clone())
1512 .request_timeout(request_timeout)
1513 .dns_timeout(dns_timeout)
1514 .send_buffer(send_buffer)
1515 .receive_buffer(receive_buffer)
1516 .dispatch_queue_capacity(dispatch_queue_capacity.get())
1517 .frame_max(frame_max.size())
1518 .metadata_recovery_strategy(metadata_recovery_strategy)
1519 .metadata_recovery_rebootstrap_trigger(metadata_recovery_rebootstrap_trigger)
1520 .maybe_security(security.clone())
1521 .client_dns_lookup(client_dns_lookup)
1522 .reconnect_backoff(reconnect_backoff)
1523 .reconnect_backoff_max(reconnect_backoff_max)
1524 .connections_max_idle(connections_max_idle)
1525 .socket_connection_setup_timeout(socket_connection_setup_timeout)
1526 .socket_connection_setup_timeout_max(socket_connection_setup_timeout_max)
1527 .metadata_scope(subscription_metadata_scope(allow_auto_create_topics))
1528 .maybe_telemetry(telemetry)
1529 .build()
1530 .await?;
1531 client.metadata_topics().set(subscribe.iter().cloned());
1532
1533 let session_timeout_ms = protocol_millis_i32(session_timeout);
1539 let rebalance_timeout_ms = protocol_millis_i32(max_poll_interval);
1542
1543 let coordinator_retry = CoordinatorRetryPolicy::from(config.retry_policy);
1551 let coordinator_id = Arc::new(AtomicI32::new(
1552 find_coordinator(&client, &group_id, coordinator_retry).await?,
1553 ));
1554 tracing::Span::current().record("coordinator_id", coordinator_id.load(Ordering::Relaxed));
1555
1556 let protocols = initial_join_protocols(&assignors, &subscribe, client_rack.as_deref());
1560
1561 let join_fields = startup_join_fields(
1566 &group_id,
1567 group_instance_id.as_deref(),
1568 protocols.clone(),
1569 (session_timeout_ms, rebalance_timeout_ms),
1570 );
1571 let refind = CoordinatorRefind {
1572 client: &client,
1573 group_id: &group_id,
1574 coordinator_id: &coordinator_id,
1575 };
1576 let r1 = CoordinatorRetryRequest::builder()
1577 .retry(coordinator_retry)
1578 .code(|r: &JoinGroupResponse| r.error_code)
1579 .make(|| refind.send_join_group(join_fields.request(String::new())))
1580 .refind(refind)
1581 .build()
1582 .run_with_refind()
1583 .await?;
1584 let member_id = first_join_member_id(&r1)?;
1585 tracing::Span::current().record("member_id", member_id.as_str());
1586
1587 let cleanup_client = client.clone();
1588 let cleanup_coordinator_id = Arc::clone(&coordinator_id);
1589 let cleanup_group_id = group_id.clone();
1590 let cleanup_member_id = member_id.clone();
1591 let cleanup_group_instance_id = group_instance_id.clone();
1592 let cleanup_leave_group_timeout = finish_config.leave_group_timeout;
1593 let start_result = finish_startup(
1594 finish_config,
1595 client,
1596 coordinator_id,
1597 member_id,
1598 protocols,
1599 (session_timeout_ms, rebalance_timeout_ms),
1600 )
1601 .await;
1602 match start_result {
1603 Ok(consumer) => Ok(consumer),
1604 Err(error) => {
1605 leave_startup_member(
1606 &cleanup_client,
1607 &cleanup_coordinator_id,
1608 &cleanup_group_id,
1609 &cleanup_member_id,
1610 cleanup_group_instance_id,
1611 cleanup_leave_group_timeout,
1612 )
1613 .await;
1614 Err(ConsumerError::StartupAfterJoin(Box::new(error)))
1615 }
1616 }
1617 }
1618}
1619
1620async fn finish_startup(
1621 config: StartConfig,
1622 client: Client,
1623 coordinator_id: Arc<AtomicI32>,
1624 member_id: String,
1625 protocols: Vec<(String, bytes::Bytes)>,
1626 timeouts_ms: (i32, i32),
1627) -> Result<Consumer, ConsumerError> {
1628 let spawn_config = config.clone();
1629 let coordinator_retry = CoordinatorRetryPolicy::from(config.retry_policy);
1630 let StartConfig {
1631 group_id,
1632 group_instance_id,
1633 auto_offset_reset,
1634 ..
1635 } = config;
1636 let mut join_fields = startup_join_fields(
1639 &group_id,
1640 group_instance_id.as_deref(),
1641 protocols,
1642 timeouts_ms,
1643 );
1644 join_fields.reason = format!("need to re-join with the given member-id: {member_id}");
1645 let r2 = coordinator::JoinGroupRetry::builder()
1647 .client(&client)
1648 .group_id(&group_id)
1649 .coordinator_id(&coordinator_id)
1650 .retry(coordinator_retry)
1651 .fields(&join_fields)
1652 .member_id(member_id.clone())
1653 .build()
1654 .run()
1655 .await?;
1656 if r2.error_code != 0 {
1657 return Err(ConsumerError::Server(r2.error_code));
1658 }
1659
1660 let InitialAssignment {
1661 assigned_partitions,
1662 group_topics,
1663 topic_ids,
1664 topic_partitions,
1665 } = resolve_initial_assignment(&spawn_config, &client, &coordinator_id, &member_id, &r2)
1666 .await?;
1667
1668 let mut next_offsets: HashMap<(String, i32), i64> = HashMap::new();
1670 let mut positions: HashMap<(String, i32), position::PartitionPosition> = HashMap::new();
1671 if has_assigned_partitions(&assigned_partitions) {
1672 let mut by_topic: HashMap<String, Vec<i32>> = HashMap::new();
1673 for (t, p) in &assigned_partitions {
1674 by_topic.entry(t.clone()).or_default().push(*p);
1675 }
1676 let of = coordinator::send_offset_fetch(
1679 &client,
1680 &group_id,
1681 &coordinator_id,
1682 &offset_wire::build_offset_fetch(&group_id, &by_topic),
1683 coordinator_retry,
1684 )
1685 .await?;
1686 for (name, partition_index, committed, committed_epoch) in
1687 offset_wire::parse_offset_fetch(&of)
1688 {
1689 let starting = starting_offset(committed, auto_offset_reset);
1690 next_offsets.insert((name.clone(), partition_index), starting);
1691 positions.insert((name, partition_index), primed_position(committed_epoch));
1692 }
1693 }
1694
1695 spawn_consumer(
1696 spawn_config,
1697 client,
1698 coordinator_id,
1699 member_id,
1700 StartupState::builder()
1701 .generation_id(r2.generation_id)
1702 .group_topics(group_topics)
1703 .assigned_partitions(assigned_partitions)
1704 .next_offsets(next_offsets)
1705 .positions(positions)
1706 .topic_ids(topic_ids)
1707 .topic_partitions(topic_partitions)
1708 .build(),
1709 )
1710 .await
1711}
1712
1713struct InitialAssignment {
1714 assigned_partitions: Vec<(String, i32)>,
1715 group_topics: Vec<String>,
1718 topic_ids: HashMap<String, WireUuid>,
1719 topic_partitions: HashMap<String, i32>,
1720}
1721
1722async fn resolve_initial_assignment(
1723 config: &StartConfig,
1724 client: &Client,
1725 coordinator_id: &Arc<AtomicI32>,
1726 member_id: &str,
1727 r2: &JoinGroupResponse,
1728) -> Result<InitialAssignment, ConsumerError> {
1729 let group_id = &config.group_id;
1730 let subscribe = &config.subscribe;
1731 let group_instance_id = &config.group_instance_id;
1732 let protocol_name = r2.protocol_name.as_deref().unwrap_or_default();
1733 let is_leader = is_group_leader(&r2.leader, member_id);
1741 let group_topics = if is_leader {
1744 coordinator::group_subscription_topics(r2, subscribe)
1745 } else {
1746 subscribe.clone()
1747 };
1748 client.metadata_topics().set(group_topics.iter().cloned());
1749 let md = client.refresh_metadata().await;
1750 client.metadata_topics().set(subscribe.iter().cloned());
1751 let md = md?;
1752 let mut topic_ids: HashMap<String, WireUuid> = HashMap::new();
1753 let mut topic_partitions: HashMap<String, i32> = HashMap::new();
1754 for t in &md.topics {
1755 let Some(name) = &t.name else { continue };
1756 if is_subscribed_topic(&group_topics, name) {
1757 let count = i32::try_from(t.partitions.len()).unwrap_or(i32::MAX);
1758 topic_partitions.insert(name.clone(), count);
1759 topic_ids.insert(name.clone(), t.topic_id);
1760 }
1761 }
1762
1763 tracing::Span::current().record("generation", r2.generation_id);
1764 tracing::Span::current().record("is_leader", is_leader);
1765 let assignments_for_sync: Vec<SyncGroupRequestAssignment> = if is_leader {
1766 coordinator::leader_assignments(&config.assignors, r2, &md)?
1767 } else {
1768 Vec::new()
1769 };
1770
1771 let r3 = CoordinatorRetryRequest::builder()
1775 .retry(config.retry_policy.into())
1776 .code(|r: &SyncGroupResponse| r.error_code)
1777 .make(|| {
1778 let group_id = group_id.clone();
1779 let protocol_name = protocol_name.to_string();
1780 let member_id = member_id.to_string();
1781 let assignments_for_sync = assignments_for_sync.clone();
1782 let generation_id = r2.generation_id;
1783 let group_instance_id = group_instance_id.clone();
1784 let target = coordinator_id.load(Ordering::Relaxed);
1785 async move {
1786 client
1787 .broker(target)
1788 .send(
1789 coordinator::SyncGroupRequestConfig::builder()
1790 .group_id(group_id)
1791 .generation_id(generation_id)
1792 .member_id(member_id)
1793 .maybe_group_instance_id(group_instance_id)
1794 .protocol_name(protocol_name)
1795 .assignments(assignments_for_sync)
1796 .build(),
1797 )
1798 .await
1799 .map_err(ConsumerError::from)
1800 }
1801 })
1802 .refind(CoordinatorRefind {
1803 client,
1804 group_id,
1805 coordinator_id,
1806 })
1807 .build()
1808 .run_with_refind()
1809 .await?;
1810 if r3.error_code != 0 {
1811 return Err(ConsumerError::Server(r3.error_code));
1812 }
1813 let assigned_partitions = decode_assignment(&r3.assignment);
1814 tracing::Span::current().record("assigned_partitions", assigned_partitions.len());
1815
1816 Ok(InitialAssignment {
1817 assigned_partitions,
1818 group_topics,
1819 topic_ids,
1820 topic_partitions,
1821 })
1822}
1823
1824#[derive(bon::Builder, Default)]
1825struct StartupState {
1826 #[builder(default)]
1827 generation_id: i32,
1828 #[builder(default)]
1830 group_topics: Vec<String>,
1831 #[builder(default)]
1832 assigned_partitions: Vec<(String, i32)>,
1833 #[builder(default)]
1834 next_offsets: HashMap<(String, i32), i64>,
1835 #[builder(default)]
1836 positions: HashMap<(String, i32), position::PartitionPosition>,
1837 #[builder(default)]
1838 topic_ids: HashMap<String, WireUuid>,
1839 #[builder(default)]
1840 topic_partitions: HashMap<String, i32>,
1841}
1842
1843#[cfg(test)]
1844#[derive(bon::Builder)]
1845#[builder(start_fn = builder, finish_fn = build)]
1846struct TestConsumerStart {
1847 bootstrap: String,
1848 config: StartConfig,
1849 startup: StartupState,
1850}
1851
1852#[cfg(test)]
1853async fn start_test_consumer(start: TestConsumerStart) -> Consumer {
1854 let client = Client::builder()
1855 .bootstrap(start.bootstrap)
1856 .build()
1857 .await
1858 .expect("client");
1859 spawn_consumer(
1860 start.config,
1861 client,
1862 Arc::new(AtomicI32::new(0)),
1863 "member-a".into(),
1864 start.startup,
1865 )
1866 .await
1867 .expect("spawn consumer")
1868}
1869
1870async fn start_without_subscription(config: StartConfig) -> Result<Consumer, ConsumerError> {
1875 let client = build_consumer_client(
1876 &config,
1877 ConsumerClientOptions::builder()
1878 .metadata_scope(subscription_metadata_scope(config.allow_auto_create_topics))
1879 .maybe_telemetry(config.telemetry.clone())
1880 .build(),
1881 )
1882 .await?;
1883 let coordinator_id = if config.group_id.is_empty() {
1884 -1
1885 } else {
1886 find_coordinator(&client, &config.group_id, config.retry_policy.into()).await?
1887 };
1888 let member_id = if config.group_protocol == GroupProtocol::Consumer {
1891 uuid::Uuid::new_v4().to_string()
1892 } else {
1893 String::new()
1894 };
1895 client
1896 .metadata_topics()
1897 .set(config.subscribe.iter().cloned());
1898 spawn_consumer(
1899 config,
1900 client,
1901 Arc::new(AtomicI32::new(coordinator_id)),
1902 member_id,
1903 StartupState::builder().generation_id(-1).build(),
1904 )
1905 .await
1906}
1907
1908async fn coordinator_task_client(
1916 config: &StartConfig,
1917 group_topics: &[String],
1918) -> Result<Option<Client>, ConsumerError> {
1919 if config.subscribe.is_empty()
1920 && config.subscribe_pattern.is_none()
1921 && config.subscribe_regex.is_none()
1922 {
1923 return Ok(None);
1924 }
1925 let client = build_consumer_client(
1926 config,
1927 ConsumerClientOptions::builder()
1928 .metadata_scope(subscription_metadata_scope(config.allow_auto_create_topics))
1929 .build(),
1930 )
1931 .await?;
1932 let topics = if group_topics.is_empty() {
1935 &config.subscribe
1936 } else {
1937 group_topics
1938 };
1939 client.metadata_topics().set(topics.iter().cloned());
1940 Ok(Some(client))
1941}
1942
1943async fn spawn_consumer(
1944 config: StartConfig,
1945 client: Client,
1946 coordinator_id: Arc<AtomicI32>,
1947 member_id: String,
1948 startup: StartupState,
1949) -> Result<Consumer, ConsumerError> {
1950 let coordinator_client = coordinator_task_client(&config, &startup.group_topics).await?;
1951 let StartConfig {
1952 group_id,
1953 session_timeout,
1954 max_poll_interval,
1955 max_poll_records,
1956 heartbeat_interval,
1957 subscription_metadata_refresh_interval,
1958 subscribe,
1959 subscribe_pattern,
1960 subscribe_regex,
1961 exclude_internal_topics,
1962 group_protocol,
1963 server_assignor,
1964 group_instance_id,
1965 auto_offset_reset,
1966 isolation_level,
1967 assignors,
1968 rebalance_protocol,
1969 fetch_min,
1970 fetch_max,
1971 fetch_partition_max,
1972 fetch_max_wait,
1973 metadata_max_age,
1974 default_api_timeout,
1975 request_timeout,
1976 leave_group_timeout,
1977 client_rack,
1978 has_rebalance_listener,
1979 retry_policy,
1980 auto_commit_interval,
1981 ..
1982 } = config;
1983 let subscription = subscription::shared(
1984 subscribe.clone(),
1985 subscribe_pattern,
1986 exclude_internal_topics,
1987 );
1988 if subscribe_regex.is_some() {
1989 subscription.send_modify(|state| state.regex.clone_from(&subscribe_regex));
1990 }
1991 let StartupState {
1992 generation_id,
1993 group_topics: _,
1994 assigned_partitions,
1995 next_offsets,
1996 positions,
1997 topic_ids,
1998 topic_partitions,
1999 } = startup;
2000 let standalone = coordinator_client.is_none();
2002 let coordinator_client = coordinator_client.unwrap_or_else(|| client.clone());
2003
2004 let (identity, next_ownership_id) =
2005 initial_commit_identity(&assigned_partitions, generation_id, &member_id);
2006 let initial_assignment = assigned_partitions.clone();
2007 let assigned = Arc::new(Mutex::new(assigned_partitions));
2008 let assignment_changed = Arc::new(Notify::new());
2009 let commit_identity = Arc::new(Mutex::new(identity));
2010 let commit_serialization = Arc::new(Mutex::new(()));
2011 let commit_async_state = Arc::new(AtomicU8::new(0));
2012 let next_offsets = Arc::new(Mutex::new(next_offsets));
2013 let end_offsets = Arc::new(Mutex::new(HashMap::new()));
2014 let positions = Arc::new(Mutex::new(positions));
2015 let topic_ids = Arc::new(Mutex::new(topic_ids));
2016 let current_generation = Arc::new(AtomicI32::new(generation_id));
2019
2020 let poll_error = coordinator::PollErrorSlot::default();
2021 let poll_signal = coordinator::PollSignal::default();
2022 let rebalance_pending =
2025 tokio::sync::watch::Sender::new(group_protocol == GroupProtocol::Consumer);
2026 let rebalance_pending_receiver = rebalance_pending.subscribe();
2027 let close_operation = tokio::sync::watch::Sender::new(control::CloseRequest::default());
2028 let (listener_sender, listener_calls) = tokio::sync::mpsc::unbounded_channel();
2029 let (enforced_rebalances, enforced_rebalance_reasons) = tokio::sync::mpsc::unbounded_channel();
2030 let assigned_callback_pending = rebalance_listener::AssignedCallbackPending::default();
2031 if has_rebalance_listener && !standalone {
2033 queue_initial_assign_call(&listener_sender, initial_assignment);
2034 }
2035 let auto_commit = auto_commit_interval.map(commit::AutoCommit::new);
2036
2037 let shutdown = CancellationToken::new();
2038 let published_member_id = tokio::sync::watch::Sender::new(member_id.clone());
2039 let live_member_id = published_member_id.subscribe();
2040 let state = CoordinatorState {
2041 client: coordinator_client,
2042 group_id: group_id.clone(),
2043 coordinator_id: Arc::clone(&coordinator_id),
2044 member_id,
2045 published_member_id,
2046 commit_identity: Arc::clone(&commit_identity),
2047 group_instance_id: group_instance_id.clone(),
2048 generation_id,
2049 current_generation: Arc::clone(¤t_generation),
2050 assignors,
2051 rebalance_protocol,
2052 subscription: Arc::clone(&subscription),
2053 subscription_changes: subscription.subscribe(),
2054 joined_topics: subscribe.clone(),
2055 seen_unsubscribes: 0,
2056 group_protocol,
2057 server_assignor,
2058 member_epoch: 0,
2059 sent_heartbeat_fields: consumer_group::SentFields::default(),
2060 assigned: Arc::clone(&assigned),
2061 assignment_changed: Arc::clone(&assignment_changed),
2062 next_ownership_id,
2063 next_offsets: Arc::clone(&next_offsets),
2064 end_offsets: Arc::clone(&end_offsets),
2065 positions: Arc::clone(&positions),
2066 topic_ids: Arc::clone(&topic_ids),
2067 session_timeout,
2068 max_poll_interval,
2069 heartbeat_interval,
2070 subscription_metadata_refresh_interval,
2071 leave_group_timeout,
2072 auto_offset_reset,
2073 client_rack: client_rack.clone(),
2074 initial_subscribed_counts: topic_partitions,
2080 retry_policy: retry_policy.into(),
2081 poll_error: Arc::clone(&poll_error),
2082 auto_commit: auto_commit.clone(),
2083 commit_serialization: Arc::clone(&commit_serialization),
2084 join_prepared: false,
2085 polls: poll_signal.subscribe(),
2086 rebalance_pending,
2087 close_operation: close_operation.subscribe(),
2088 rejoin_reason: String::new(),
2089 listener_calls: has_rebalance_listener.then_some(listener_sender),
2090 enforced_rebalances: enforced_rebalance_reasons,
2091 lost_partitions: Vec::new(),
2092 assigned_callback_pending: Arc::clone(&assigned_callback_pending),
2093 poll_timer: coordinator::PollTimer::new(max_poll_interval),
2094 poll_timeout_in_callback: false,
2095 };
2096 let coord_handle = (!standalone).then(|| match group_protocol {
2100 GroupProtocol::Classic => tokio::spawn(coordinator::run(state, shutdown.clone())),
2101 GroupProtocol::Consumer => tokio::spawn(consumer_group::run(state, shutdown.clone())),
2102 });
2103
2104 Ok(Consumer {
2105 client,
2106 group_id,
2107 coordinator_id,
2108 retry_policy: retry_policy.into(),
2109 member_id: live_member_id,
2110 commit_identity,
2111 commit_serialization,
2112 commit_async_state,
2113 commit_async_callbacks: Arc::default(),
2114 group_instance_id: group_instance_id.clone(),
2115 current_generation,
2116 subscription,
2117 group_protocol,
2118 assigned,
2119 assignment_changed,
2120 next_offsets,
2121 end_offsets,
2122 positions,
2123 topic_ids,
2124 session_timeout,
2125 heartbeat_interval,
2126 rebalance_protocol,
2127 coordinator_shutdown: shutdown,
2128 coordinator_handle: coord_handle,
2129 isolation_level,
2130 fetch_min,
2131 fetch_max,
2132 fetch_partition_max,
2133 fetch_max_wait,
2134 fetches: poll::Fetches::default(),
2135 client_rack,
2136 metadata_max_age,
2137 request_timeout,
2138 wakeup: control::WakeupHandle::default(),
2139 enforced_rebalances,
2140 default_api_timeout,
2141 paused: std::sync::Mutex::default(),
2142 auto_offset_reset,
2143 poll_error,
2144 auto_commit,
2145 poll_signal,
2146 rebalance_pending: rebalance_pending_receiver,
2147 max_poll_records,
2148 fetch_buffer: fetch_buffer::FetchBuffer::default(),
2149 close_operation,
2150 rebalance_listener: None,
2151 listener_calls,
2152 assigned_callback_pending,
2153 })
2154}
2155
2156fn is_retriable_consumer_start_error(error: &ConsumerError) -> bool {
2162 matches!(
2172 error,
2173 ConsumerError::Server(14 | 15 | 16 | 22 | 25 | 27 | 79)
2177 | ConsumerError::Client(ClientError::Disconnected)
2178 ) || matches!(
2179 error,
2180 ConsumerError::StartupAfterJoin(inner)
2181 if is_retriable_consumer_start_error(inner)
2182 || matches!(
2183 inner.as_ref(),
2184 ConsumerError::Client(ClientError::Timeout(_))
2185 )
2186 )
2187}
2188
2189impl Consumer {
2190 #[must_use]
2192 pub fn group_id(&self) -> &str {
2193 &self.group_id
2194 }
2195
2196 #[must_use]
2199 pub fn member_id(&self) -> String {
2200 self.member_id.borrow().clone()
2201 }
2202
2203 #[must_use]
2206 pub fn generation_id(&self) -> i32 {
2207 self.current_generation.load(Ordering::Relaxed)
2208 }
2209
2210 #[must_use]
2216 pub fn group_metadata(&self) -> ConsumerGroupMetadata {
2217 ConsumerGroupMetadata {
2218 group_id: self.group_id.clone(),
2219 generation_id: self.current_generation.load(Ordering::Relaxed),
2220 member_id: self.member_id(),
2221 group_instance_id: self.group_instance_id.clone(),
2222 }
2223 }
2224
2225 pub async fn assignment(&self) -> Vec<(String, i32)> {
2227 self.assigned.lock().await.clone()
2228 }
2229
2230 pub async fn at_log_end(&self) -> bool {
2233 let assigned = self.assigned.lock().await;
2234 if assigned.is_empty() {
2235 return false;
2236 }
2237 let positions = self.next_offsets.lock().await;
2238 let ends = self.end_offsets.lock().await;
2239 assigned.iter().all(|partition| {
2240 positions
2241 .get(partition)
2242 .filter(|position| !poll::is_reset_sentinel(**position))
2243 .zip(ends.get(partition))
2244 .is_some_and(|(position, end)| position >= end)
2245 })
2246 }
2247
2248 pub async fn close(self) -> Result<(), ConsumerError> {
2256 self.close_with(CloseOptions::builder().build()).await
2257 }
2258
2259 pub async fn client_instance_id(
2276 &self,
2277 timeout: Time,
2278 ) -> Result<Option<uuid::Uuid>, ConsumerError> {
2279 client_instance_id(&self.client, timeout).await
2280 }
2281
2282 #[tracing::instrument(
2304 name = "consumer.close",
2305 level = "info",
2306 skip_all,
2307 fields(group_id = %self.group_id, member_id = %self.member_id(), ?options),
2308 err(level = "info")
2309 )]
2310 pub async fn close_with(mut self, options: crate::CloseOptions) -> Result<(), ConsumerError> {
2313 let timeout = control::close_timeout(options)?.to_std();
2314 let started = tokio::time::Instant::now();
2315 let request_timeout = self.request_timeout.to_std();
2316 let coordinator_deadline = started + timeout.min(request_timeout);
2319 self.auto_commit_on_close(coordinator_deadline).await;
2322 self.wait_for_async_commits(coordinator_deadline).await;
2325 let listener_result = self.leave_prepare().await;
2329 self.close_operation.send_replace(control::CloseRequest {
2330 operation: options.group_membership_operation,
2331 });
2332 self.coordinator_shutdown.cancel();
2333 if let Some(h) = self.coordinator_handle.take() {
2334 let _ = h.await;
2335 }
2336 let left = timeout
2339 .saturating_sub(started.elapsed())
2340 .min(request_timeout);
2341 self.close_fetch_sessions(left).await;
2342 let left = timeout.saturating_sub(started.elapsed());
2344 let _ = tokio::time::timeout(left, self.client.close_telemetry()).await;
2345 listener_result
2346 }
2347
2348 pub(crate) async fn leave_prepare(&mut self) -> Result<(), ConsumerError> {
2351 if self.rebalance_listener.is_none() {
2352 return Ok(());
2353 }
2354 let pending = self.run_pending_listener_calls().await;
2355 let owned = self.assigned.lock().await.clone();
2356 if owned.is_empty() {
2357 return pending;
2358 }
2359 let kind = if self.member_id.borrow().is_empty() || *self.rebalance_pending.borrow() {
2360 rebalance_listener::ListenerCallKind::Lost
2361 } else {
2362 rebalance_listener::ListenerCallKind::Revoked
2363 };
2364 let result = self.run_listener(kind, &owned).await;
2365 pending.and(result)
2366 }
2367}
2368
2369#[cfg(test)]
2370impl Consumer {
2371 pub(crate) fn test_builder(
2372 client: krabka_client_core::Client,
2373 ) -> test_support::ConsumerTestBuilder {
2374 test_support::ConsumerTest::builder(client)
2375 }
2376}
2377
2378#[cfg(test)]
2379mod protocol_millis_tests {
2380 use assert2::check;
2381 use krabka_units::{Time, convert::TimeExt as _, millis, minutes, secs};
2382
2383 use super::protocol_millis_i32;
2384
2385 #[test]
2391 fn protocol_millis_truncates_rather_than_rounding() {
2392 for (_name, value, expected) in [
2393 ("whole default session timeout", secs(45), 45_000),
2394 ("whole default rebalance timeout", minutes(1), 60_000),
2395 ("whole millisecond", millis(37), 37),
2396 (
2397 "fractional millisecond truncates down",
2398 Time::from_secs_f64(0.0016),
2399 1,
2400 ),
2401 (
2402 "fractional millisecond that would round up",
2403 Time::from_secs_f64(1.9999),
2404 1_999,
2405 ),
2406 ("below one millisecond", Time::from_secs_f64(0.0009), 0),
2407 ] {
2408 check!(protocol_millis_i32(value) == expected);
2409 }
2410 }
2411
2412 #[test]
2416 fn protocol_millis_saturates_at_i32_max() {
2417 check!(protocol_millis_i32(Time::from_secs(10_000_000_000)) == i32::MAX);
2418 }
2419}
2420
2421#[cfg(test)]
2422mod consumer_retry_policy_tests {
2423 use assert2::{assert, check};
2424 use krabka_units::{Time, convert::TimeExt as _, millis, minutes, secs};
2425
2426 use super::ConsumerRetryPolicy;
2427
2428 #[test]
2429 fn consumer_retry_policy_defaults_and_overrides() {
2430 let defaults = ConsumerRetryPolicy::default();
2431 check!(defaults.startup_attempt_timeout() == secs(90));
2432 check!(defaults.startup_deadline() == minutes(5));
2433 check!(defaults.startup_initial_backoff() == millis(500));
2434 check!(defaults.startup_max_backoff() == millis(1000));
2435 check!(defaults.coordinator_retry_timeout() == secs(30));
2436 check!(defaults.coordinator_initial_backoff() == millis(100));
2437 check!(defaults.coordinator_max_backoff() == secs(1));
2438
2439 let configured = ConsumerRetryPolicy::builder()
2440 .startup_attempt_timeout(secs(11))
2441 .startup_deadline(secs(12))
2442 .startup_initial_backoff(millis(13))
2443 .startup_max_backoff(millis(14))
2444 .coordinator_retry_timeout(secs(15))
2445 .coordinator_initial_backoff(millis(16))
2446 .coordinator_max_backoff(millis(17))
2447 .build()
2448 .expect("valid retry policy");
2449 check!(configured.startup_attempt_timeout() == secs(11));
2450 check!(configured.startup_deadline() == secs(12));
2451 check!(configured.startup_initial_backoff() == millis(13));
2452 check!(configured.startup_max_backoff() == millis(14));
2453 check!(configured.coordinator_retry_timeout() == secs(15));
2454 check!(configured.coordinator_initial_backoff() == millis(16));
2455 check!(configured.coordinator_max_backoff() == millis(17));
2456 }
2457
2458 #[test]
2463 fn with_retry_backoff_max_caps_startup_and_coordinator_backoff() {
2464 for (name, retry_backoff_max) in [
2465 ("Kafka's default", millis(1000)),
2466 ("custom, above both fixed initial backoffs", secs(3)),
2467 ("below both fixed initial backoffs", millis(100)),
2468 ("between the two fixed initial backoffs", millis(250)),
2469 ] {
2470 let policy = ConsumerRetryPolicy::with_retry_backoff_max(retry_backoff_max)
2471 .unwrap_or_else(|error| panic!("{name}: {error}"));
2472 check!(policy.startup_max_backoff() == retry_backoff_max);
2473 check!(policy.coordinator_max_backoff() == retry_backoff_max);
2474 check!(
2475 policy.startup_initial_backoff() <= retry_backoff_max,
2476 "{name}: startup initial backoff must not exceed retry_backoff_max"
2477 );
2478 check!(
2479 policy.coordinator_initial_backoff() <= retry_backoff_max,
2480 "{name}: coordinator initial backoff must not exceed retry_backoff_max"
2481 );
2482 }
2483 }
2484
2485 #[test]
2486 fn with_retry_backoff_max_rejects_invalid_values() {
2487 for retry_backoff_max in [
2488 Time::ZERO,
2489 Time::from_secs_f64(0.0005),
2490 Time::from_secs_f64(f64::INFINITY),
2491 ] {
2492 assert!(ConsumerRetryPolicy::with_retry_backoff_max(retry_backoff_max).is_err());
2493 }
2494 }
2495
2496 #[test]
2497 fn consumer_retry_policy_rejects_invalid_values() {
2498 for invalid in [
2499 Time::ZERO,
2500 Time::from_secs_f64(0.0005),
2501 Time::from_secs_f64(f64::INFINITY),
2502 ] {
2503 assert!(
2504 ConsumerRetryPolicy::builder()
2505 .startup_attempt_timeout(invalid)
2506 .build()
2507 .is_err()
2508 );
2509 assert!(
2510 ConsumerRetryPolicy::builder()
2511 .startup_deadline(invalid)
2512 .build()
2513 .is_err()
2514 );
2515 assert!(
2516 ConsumerRetryPolicy::builder()
2517 .startup_initial_backoff(invalid)
2518 .build()
2519 .is_err()
2520 );
2521 assert!(
2522 ConsumerRetryPolicy::builder()
2523 .startup_max_backoff(invalid)
2524 .build()
2525 .is_err()
2526 );
2527 assert!(
2528 ConsumerRetryPolicy::builder()
2529 .coordinator_retry_timeout(invalid)
2530 .build()
2531 .is_err()
2532 );
2533 assert!(
2534 ConsumerRetryPolicy::builder()
2535 .coordinator_initial_backoff(invalid)
2536 .build()
2537 .is_err()
2538 );
2539 assert!(
2540 ConsumerRetryPolicy::builder()
2541 .coordinator_max_backoff(invalid)
2542 .build()
2543 .is_err()
2544 );
2545 }
2546
2547 let invalid_orderings = [
2548 ConsumerRetryPolicy::builder()
2549 .startup_attempt_timeout(secs(13))
2550 .startup_deadline(secs(12))
2551 .build(),
2552 ConsumerRetryPolicy::builder()
2553 .startup_initial_backoff(millis(3))
2554 .startup_max_backoff(millis(2))
2555 .build(),
2556 ConsumerRetryPolicy::builder()
2557 .coordinator_initial_backoff(millis(3))
2558 .coordinator_max_backoff(millis(2))
2559 .build(),
2560 ];
2561 for policy in invalid_orderings {
2562 assert!(policy.is_err());
2563 }
2564 }
2565}
2566
2567#[cfg(test)]
2568mod security_arg_tests {
2569 use assert2::check;
2570 use krabka_client_core::{
2571 ClientError, MockBroker,
2572 security::{ClientSecurity, SaslCredentials},
2573 };
2574 use krabka_protocol::{
2575 Encode,
2576 owned::{
2577 api_versions_request, leave_group_request,
2578 leave_group_response::{self, LeaveGroupResponse},
2579 },
2580 };
2581 use krabka_security::ListenerProtocol;
2582
2583 use super::*;
2584
2585 async fn invalid_consumer_config<F, E>(build: F) -> E
2586 where
2587 F: std::future::Future<Output = Result<Consumer, E>>,
2588 E: std::fmt::Debug,
2589 {
2590 build.await.err().expect("invalid configuration")
2591 }
2592
2593 fn leave_group_response_at(version: i16) -> Vec<u8> {
2594 let mut buf = bytes::BytesMut::new();
2595 if version >= leave_group_response::FLEXIBLE_MIN {
2596 buf.extend_from_slice(&[0x00]);
2597 }
2598 LeaveGroupResponse::default()
2599 .encode(&mut buf, version)
2600 .unwrap();
2601 buf.to_vec()
2602 }
2603
2604 #[test]
2605 fn leave_group_timeout_uses_default_and_valid_override() {
2606 let default = ConsumerLeaveGroupTimeout::default();
2607 assert2::assert!(default.duration() == Duration::from_secs(5));
2608 assert2::assert!(default.milliseconds() == 5_000);
2609
2610 let timeout = ConsumerLeaveGroupTimeout::new(Duration::from_millis(37))
2611 .expect("positive whole milliseconds");
2612 assert2::assert!(timeout.duration() == Duration::from_millis(37));
2613 assert2::assert!(timeout.milliseconds() == 37);
2614 }
2615
2616 #[test]
2617 fn leave_group_timeout_validates_millisecond_boundaries() {
2618 assert2::assert!(ConsumerLeaveGroupTimeout::new(Duration::ZERO).is_err());
2619 assert2::assert!(
2620 ConsumerLeaveGroupTimeout::new(Duration::from_millis(1) + Duration::from_nanos(1))
2621 .is_err()
2622 );
2623 assert2::assert!(ConsumerLeaveGroupTimeout::new(Duration::from_millis(u64::MAX)).is_ok());
2624 assert2::assert!(ConsumerLeaveGroupTimeout::new(Duration::from_secs(u64::MAX)).is_err());
2625 }
2626
2627 #[test]
2628 fn subscription_metadata_refresh_interval_uses_default_and_valid_override() {
2629 let default = ConsumerSubscriptionMetadataRefreshInterval::default();
2630 assert2::assert!(default.duration() == Duration::from_secs(5));
2631 assert2::assert!(default.milliseconds() == 5_000);
2632
2633 let interval = ConsumerSubscriptionMetadataRefreshInterval::new(Duration::from_millis(37))
2634 .expect("positive whole milliseconds");
2635 assert2::assert!(interval.duration() == Duration::from_millis(37));
2636 assert2::assert!(interval.milliseconds() == 37);
2637 }
2638
2639 #[test]
2640 fn subscription_metadata_refresh_interval_validates_millisecond_boundaries() {
2641 assert2::assert!(ConsumerSubscriptionMetadataRefreshInterval::new(Duration::ZERO).is_err());
2642 assert2::assert!(
2643 ConsumerSubscriptionMetadataRefreshInterval::new(
2644 Duration::from_millis(1) + Duration::from_nanos(1)
2645 )
2646 .is_err()
2647 );
2648 assert2::assert!(
2649 ConsumerSubscriptionMetadataRefreshInterval::new(Duration::from_millis(u64::MAX))
2650 .is_ok()
2651 );
2652 assert2::assert!(
2653 ConsumerSubscriptionMetadataRefreshInterval::new(Duration::from_secs(u64::MAX))
2654 .is_err()
2655 );
2656 }
2657
2658 #[tokio::test]
2659 async fn invalid_subscription_metadata_refresh_interval_fails_before_broker_lookup() {
2660 let error = invalid_consumer_config(
2661 Consumer::builder()
2662 .bootstrap("invalid.invalid:9092")
2663 .group_id("metadata-refresh-validation")
2664 .subscribe(["topic".to_owned()])
2665 .subscription_metadata_refresh_interval(krabka_units::secs(0))
2666 .build(),
2667 )
2668 .await;
2669
2670 assert2::assert!(
2671 error
2672 .to_string()
2673 .contains("consumer subscription metadata refresh interval")
2674 );
2675 }
2676
2677 #[tokio::test]
2680 async fn negative_sub_millisecond_times_fail_before_broker_lookup() {
2681 let negative = Time::from_secs_f64(-0.0001);
2682 let mut actual = Vec::new();
2683 for (name, default_api_timeout, metadata_max_age) in [
2684 (
2685 "default api timeout",
2686 negative,
2687 DEFAULT_CONSUMER_METADATA_MAX_AGE,
2688 ),
2689 (
2690 "metadata max age",
2691 DEFAULT_CONSUMER_DEFAULT_API_TIMEOUT,
2692 negative,
2693 ),
2694 ] {
2695 let build = Consumer::builder()
2696 .bootstrap("invalid.invalid:9092")
2697 .group_id("negative-times")
2698 .subscribe(["topic".to_owned()])
2699 .default_api_timeout(default_api_timeout)
2700 .metadata_max_age(metadata_max_age)
2701 .build();
2702 let error = tokio::time::timeout(Duration::from_secs(5), build)
2703 .await
2704 .map(|result| result.err().map(|error| error.to_string()));
2705 actual.push((name, error.ok().flatten()));
2706 }
2707 assert2::assert!(
2708 actual
2709 == vec![
2710 (
2711 "default api timeout",
2712 Some("invalid configuration: consumer default api timeout must not be negative".to_owned())
2713 ),
2714 (
2715 "metadata max age",
2716 Some("invalid configuration: consumer metadata max age must not be negative".to_owned())
2717 ),
2718 ]
2719 );
2720 }
2721
2722 #[tokio::test]
2726 async fn low_retry_backoff_max_clamps_initial_backoffs_instead_of_failing() {
2727 let mock = MockBroker::start(|_api_key, _version, _corr_id, _body| None).await;
2728
2729 let build = Consumer::builder()
2730 .bootstrap(mock.addr.to_string())
2731 .client_id("low-retry-backoff-max-consumer")
2732 .group_id("low-retry-backoff-max-group")
2733 .subscribe(vec!["orders".to_string()])
2734 .request_timeout(krabka_units::millis(50))
2735 .retry_backoff_max(millis(100))
2736 .build();
2737 let res = tokio::time::timeout(Duration::from_millis(500), build).await;
2738
2739 mock.stop();
2740
2741 let Err(err) = res.expect("a valid low retry_backoff_max must not hang the build") else {
2742 panic!("silent broker must time out during build")
2743 };
2744 assert2::assert!(
2745 matches!(err, ConsumerError::Client(ClientError::Timeout(_))),
2746 "build must fail on the connection timeout, not on config validation: {err}"
2747 );
2748 }
2749
2750 #[tokio::test]
2753 async fn invalid_retry_backoff_max_fails_before_broker_lookup() {
2754 let build = Consumer::builder()
2755 .bootstrap("invalid.invalid:9092")
2756 .group_id("retry-backoff-max")
2757 .subscribe(["topic".to_owned()])
2758 .retry_backoff_max(Time::ZERO)
2759 .build();
2760 let error = tokio::time::timeout(Duration::from_secs(5), build)
2761 .await
2762 .expect("build must fail before any broker lookup")
2763 .err()
2764 .expect("a zero retry backoff max must be rejected");
2765 assert2::assert!(error.to_string().contains("backoff"));
2766 }
2767
2768 #[tokio::test]
2772 async fn the_subscription_of_a_build_is_checked_before_broker_lookup() {
2773 let mut actual = Vec::new();
2774 let mut wanted = Vec::new();
2775 for (name, topics, pattern, expected) in [
2776 (
2777 "a subscription without a group id",
2778 vec!["orders".to_owned()],
2779 None,
2780 "To use the group management or offset commit APIs, you must provide a valid group.id in the consumer configuration.",
2781 ),
2782 (
2783 "both",
2784 vec!["orders".to_owned()],
2785 Some(TopicPattern::new(|_| true)),
2786 "invalid configuration: subscribe and subscribe_pattern are mutually exclusive",
2787 ),
2788 (
2789 "empty topic name",
2790 vec![String::new()],
2791 None,
2792 "invalid configuration: Topic collection to subscribe to cannot contain null or empty topic",
2793 ),
2794 ] {
2795 let error = Consumer::builder()
2796 .bootstrap("invalid.invalid:9092")
2797 .maybe_group_id(
2798 (name != "a subscription without a group id")
2799 .then_some("subscription-validation"),
2800 )
2801 .subscribe(topics)
2802 .maybe_subscribe_pattern(pattern)
2803 .build()
2804 .await
2805 .err()
2806 .map(|error| error.to_string());
2807 actual.push((name, error));
2808 wanted.push((name, Some(expected.to_owned())));
2809 }
2810 assert2::assert!(actual == wanted);
2811 }
2812
2813 #[tokio::test]
2818 async fn the_consumer_group_protocol_options_are_checked_before_broker_lookup() {
2819 let mut actual = Vec::new();
2820 let mut wanted = Vec::new();
2821 for (name, protocol, assignor, regex, topics, expected) in [
2822 (
2823 "a remote assignor with the classic protocol",
2824 GroupProtocol::Classic,
2825 Some("uniform"),
2826 None,
2827 Vec::new(),
2828 "invalid configuration: group.remote.assignor cannot be set when group.protocol=classic",
2829 ),
2830 (
2831 "a regular expression with the classic protocol",
2832 GroupProtocol::Classic,
2833 None,
2834 Some("orders-.*"),
2835 Vec::new(),
2836 "invalid configuration: Subscribe to RE2/J pattern is not supported when using the CLASSIC protocol defined in config group.protocol",
2837 ),
2838 (
2839 "an empty regular expression",
2840 GroupProtocol::Consumer,
2841 None,
2842 Some(""),
2843 Vec::new(),
2844 "invalid argument: Topic pattern to subscribe to cannot be empty",
2845 ),
2846 (
2847 "a regular expression and topics",
2848 GroupProtocol::Consumer,
2849 None,
2850 Some("orders-.*"),
2851 vec!["orders".to_owned()],
2852 "invalid configuration: Subscription to topics, partitions and pattern are mutually exclusive",
2853 ),
2854 ] {
2855 let error = Consumer::builder()
2856 .bootstrap("invalid.invalid:9092")
2857 .group_id("protocol-validation")
2858 .subscribe(topics)
2859 .group_protocol(protocol)
2860 .maybe_group_remote_assignor(assignor)
2861 .maybe_subscribe_regex(regex)
2862 .build()
2863 .await
2864 .err()
2865 .map(|error| error.to_string());
2866 actual.push((name, error));
2867 wanted.push((name, Some(expected.to_owned())));
2868 }
2869 assert2::assert!(actual == wanted);
2870 }
2871
2872 #[tokio::test]
2873 async fn invalid_leave_group_timeout_fails_before_broker_lookup() {
2874 let error = invalid_consumer_config(
2875 Consumer::builder()
2876 .bootstrap("invalid.invalid:9092")
2877 .group_id("leave-validation")
2878 .subscribe(["topic".to_owned()])
2879 .leave_group_timeout(krabka_units::secs(0))
2880 .build(),
2881 )
2882 .await;
2883
2884 assert2::assert!(error.to_string().contains("consumer leave-group timeout"));
2885 }
2886
2887 #[tokio::test]
2888 async fn invalid_client_resource_policy_fails_before_broker_lookup() {
2889 let error = Consumer::builder()
2890 .bootstrap("invalid.invalid:9092")
2891 .group_id("client-policy-validation")
2892 .subscribe(["topic".to_owned()])
2893 .dispatch_queue_capacity(0)
2894 .build()
2895 .await
2896 .err()
2897 .expect("invalid configuration");
2898
2899 assert2::assert!(error.to_string().contains("client dispatch queue capacity"));
2900 }
2901
2902 #[tokio::test]
2903 async fn startup_member_cleanup_leaves_the_group_for_a_dynamic_member_only() {
2904 let mut actual = Vec::new();
2905 let mut wanted = Vec::new();
2906 for (name, group_instance_id, expected) in [
2907 ("dynamic member", None, true),
2908 ("static member", Some("instance-a".to_owned()), false),
2909 ] {
2910 let saw_leave = Arc::new(std::sync::atomic::AtomicBool::new(false));
2911 let saw_leave_in_mock = Arc::clone(&saw_leave);
2912 let mock = MockBroker::start(move |api_key, version, _corr_id, _body| {
2913 if api_key == api_versions_request::API_KEY {
2914 Some(test_support::leave_group_api_versions())
2915 } else if api_key == leave_group_request::API_KEY {
2916 saw_leave_in_mock.store(true, Ordering::SeqCst);
2917 Some(leave_group_response_at(version))
2918 } else {
2919 None
2920 }
2921 })
2922 .await;
2923
2924 let client = Client::builder()
2925 .bootstrap(mock.addr.to_string())
2926 .request_timeout(krabka_units::millis(100))
2927 .build()
2928 .await
2929 .unwrap();
2930 let coordinator_id = AtomicI32::new(0);
2931
2932 leave_startup_member(
2933 &client,
2934 &coordinator_id,
2935 "group-a",
2936 "member-a",
2937 group_instance_id,
2938 krabka_units::millis(37),
2939 )
2940 .await;
2941
2942 mock.stop();
2943 actual.push((name, saw_leave.load(Ordering::SeqCst)));
2944 wanted.push((name, expected));
2945 }
2946 assert2::assert!(actual == wanted);
2947 }
2948
2949 #[tokio::test]
2950 async fn startup_member_cleanup_bounds_stalled_leave_with_configured_timeout() {
2951 let fixture = test_support::StalledLeaveGroup::start().await;
2952
2953 tokio::time::timeout(
2954 Duration::from_secs(1),
2955 leave_startup_member(
2956 &fixture.client,
2957 &fixture.coordinator_id,
2958 "group-a",
2959 "member-a",
2960 None,
2961 krabka_units::millis(37),
2962 ),
2963 )
2964 .await
2965 .expect("configured leave deadline bounds cleanup");
2966
2967 fixture.broker.stop();
2968 assert2::assert!(fixture.saw_leave.load(Ordering::SeqCst));
2969 }
2970
2971 #[test]
2976 fn initial_join_protocols_follow_the_assignor_list() {
2977 use builder::DecodedSubscription;
2978
2979 let protocols = initial_join_protocols(
2980 &[
2981 Assignor::Range,
2982 Assignor::Sticky,
2983 Assignor::CooperativeSticky,
2984 ],
2985 &["topic".into()],
2986 Some("rack-a"),
2987 );
2988 let decoded: Vec<(String, DecodedSubscription)> = protocols
2989 .iter()
2990 .map(|(name, metadata)| (name.clone(), builder::decode_subscription(metadata)))
2991 .collect();
2992 let subscription = |user_data: Option<Bytes>| DecodedSubscription {
2993 topics: vec!["topic".into()],
2994 owned: Vec::new(),
2995 generation_id: -1,
2996 rack_id: Some("rack-a".into()),
2997 user_data,
2998 };
2999 check!(
3000 decoded
3001 == vec![
3002 ("range".to_owned(), subscription(None)),
3003 ("sticky".to_owned(), subscription(None)),
3004 (
3005 "cooperative-sticky".to_owned(),
3006 subscription(Some(Bytes::from_static(&[0xff, 0xff, 0xff, 0xff]))),
3007 ),
3008 ]
3009 );
3010 }
3011
3012 #[test]
3013 fn startup_join_fields_preserve_member_timeouts_protocols_and_empty_reason() {
3014 let metadata = bytes::Bytes::from_static(b"metadata");
3015 let fields = startup_join_fields(
3016 "group-a",
3017 Some("instance-a"),
3018 vec![
3019 ("range".to_owned(), metadata.clone()),
3020 ("cooperative-sticky".to_owned(), metadata.clone()),
3021 ],
3022 (45_000, 60_000),
3023 );
3024
3025 assert2::assert!(fields.group_id == "group-a");
3026 assert2::assert!(fields.group_instance_id.as_deref() == Some("instance-a"));
3027 assert2::assert!(fields.session_timeout_ms == 45_000);
3028 assert2::assert!(fields.rebalance_timeout_ms == 60_000);
3029 assert2::assert!(
3030 fields.protocols
3031 == vec![
3032 ("range".to_owned(), metadata.clone()),
3033 ("cooperative-sticky".to_owned(), metadata),
3034 ]
3035 );
3036 assert2::assert!(fields.reason.is_empty());
3037 }
3038
3039 #[test]
3040 fn first_join_member_id_accepts_required_or_success_and_rejects_errors() {
3041 enum Expected<'a> {
3042 Member(&'a str),
3043 Server(i16),
3044 RebalanceFailed,
3045 }
3046 for (_name, error_code, member_id, expected) in [
3047 (
3048 "member id required",
3049 79,
3050 "member-a",
3051 Expected::Member("member-a"),
3052 ),
3053 ("success", 0, "member-b", Expected::Member("member-b")),
3054 ("server error", 42, "member-c", Expected::Server(42)),
3055 ("empty member", 0, "", Expected::RebalanceFailed),
3056 ] {
3057 let response = JoinGroupResponse {
3058 error_code,
3059 member_id: member_id.into(),
3060 ..Default::default()
3061 };
3062 let actual = first_join_member_id(&response);
3063 let matches = match expected {
3064 Expected::Member(expected) => matches!(actual, Ok(actual) if actual == expected),
3065 Expected::Server(expected) => {
3066 matches!(actual, Err(ConsumerError::Server(code)) if code == expected)
3067 }
3068 Expected::RebalanceFailed => {
3069 matches!(actual, Err(ConsumerError::RebalanceFailed(_)))
3070 }
3071 };
3072 assert2::assert!(matches);
3073 }
3074 }
3075
3076 #[test]
3077 fn is_subscribed_topic_matches_exact_topic_names() {
3078 let subscribe = vec!["orders".to_string(), "payments".to_string()];
3079
3080 for (name, expected) in [
3081 ("orders", true),
3082 ("payments", true),
3083 ("shipments", false),
3084 ("order", false),
3085 ] {
3086 assert2::assert!(is_subscribed_topic(&subscribe, name) == expected);
3087 }
3088 }
3089
3090 #[test]
3091 fn is_group_leader_matches_exact_member_id() {
3092 for (_name, leader, member_id, expected) in [
3093 ("exact leader", "member-a", "member-a", true),
3094 ("different member", "member-a", "member-b", false),
3095 ("empty member", "member-a", "", false),
3096 ] {
3097 assert2::assert!(is_group_leader(leader, member_id) == expected);
3098 }
3099 }
3100
3101 #[test]
3102 fn build_sync_assignment_preserves_member_and_assignment_payload() {
3103 let assignment =
3104 coordinator::build_sync_group_assignment("member-a".into(), &[("orders".into(), 3)]);
3105 let decoded = decode_assignment(&assignment.assignment);
3106
3107 assert2::assert!(
3108 (assignment.member_id.as_str(), decoded) == ("member-a", vec![("orders".into(), 3)])
3109 );
3110 }
3111
3112 #[test]
3113 fn offset_prime_helpers_preserve_assignment_presence_offsets_and_epochs() {
3114 for (_name, partitions, expected) in [
3115 ("empty assignment", vec![], false),
3116 ("assigned partition", vec![("orders".to_string(), 0)], true),
3117 ] {
3118 assert2::assert!(has_assigned_partitions(&partitions) == expected);
3119 }
3120
3121 test_support::assert_starting_offsets(starting_offset);
3122
3123 let position = primed_position(9);
3124 assert2::assert!(
3125 position
3126 == position::PartitionPosition {
3127 offset_epoch: krabka_ids::LeaderEpoch(9),
3128 leader_id: -1,
3129 leader_epoch: krabka_ids::LeaderEpoch(-1),
3130 awaiting_validation: false,
3131 }
3132 );
3133 }
3134
3135 #[tokio::test]
3136 async fn consumer_builder_uses_request_timeout_for_initial_connect_handshake() {
3137 let mock = MockBroker::start(|_api_key, _version, _corr_id, _body| None).await;
3138
3139 let build = Consumer::builder()
3140 .bootstrap(mock.addr.to_string())
3141 .client_id("timeout-consumer")
3142 .group_id("timeout-group")
3143 .subscribe(vec!["orders".to_string()])
3144 .request_timeout(krabka_units::millis(100))
3145 .build();
3146 let res = tokio::time::timeout(Duration::from_millis(500), build).await;
3147
3148 mock.stop();
3149
3150 let Err(err) = res.expect("consumer build should not retain the 30s connect timeout")
3151 else {
3152 panic!("silent broker must time out during build")
3153 };
3154 assert2::assert!(matches!(
3155 err,
3156 ConsumerError::Client(ClientError::Timeout(d))
3157 if d == krabka_units::millis(100)
3158 ));
3159 }
3160
3161 #[tokio::test]
3165 async fn a_refused_start_records_no_error() {
3166 use krabka_client_core::mock_test_support::recorded_levels;
3167
3168 let (invalid, levels) = recorded_levels(Box::pin(
3169 Consumer::builder()
3170 .bootstrap("127.0.0.1:1")
3171 .group_id("log-level-group")
3172 .subscribe(vec!["orders".to_string()])
3173 .fetch_min(0)
3174 .build(),
3175 ))
3176 .await;
3177 assert2::assert!(matches!(invalid, Err(ConsumerError::InvalidConfig(_))));
3178 check!(levels == vec![tracing::Level::INFO]);
3179
3180 let mock = MockBroker::start(|_api_key, _version, _corr_id, _body| None).await;
3181 let (silent, levels) = recorded_levels(Box::pin(
3182 Consumer::builder()
3183 .bootstrap(mock.addr.to_string())
3184 .group_id("log-level-group")
3185 .subscribe(vec!["orders".to_string()])
3186 .request_timeout(millis(100))
3187 .build(),
3188 ))
3189 .await;
3190 mock.stop();
3191 assert2::assert!(silent.is_err());
3192 check!(!levels.is_empty());
3193 check!(!levels.contains(&tracing::Level::ERROR), "{levels:?}");
3194 }
3195
3196 #[tokio::test]
3197 async fn consumer_builder_rejects_non_positive_fetch_budgets_before_network_io() {
3198 let min = Consumer::builder()
3199 .bootstrap("127.0.0.1:1")
3200 .group_id("timeout-group")
3201 .subscribe(vec!["orders".to_string()])
3202 .fetch_min(0)
3203 .build()
3204 .await;
3205 assert2::assert!(matches!(min, Err(ConsumerError::InvalidConfig(_))));
3206
3207 let max = Consumer::builder()
3208 .bootstrap("127.0.0.1:1")
3209 .group_id("timeout-group")
3210 .subscribe(vec!["orders".to_string()])
3211 .fetch_max(0)
3212 .build()
3213 .await;
3214 assert2::assert!(matches!(max, Err(ConsumerError::InvalidConfig(_))));
3215
3216 let partition_max = Consumer::builder()
3217 .bootstrap("127.0.0.1:1")
3218 .group_id("timeout-group")
3219 .subscribe(vec!["orders".to_string()])
3220 .fetch_partition_max(0)
3221 .build()
3222 .await;
3223 assert2::assert!(matches!(
3224 partition_max,
3225 Err(ConsumerError::InvalidConfig(_))
3226 ));
3227
3228 let inverted = Consumer::builder()
3229 .bootstrap("127.0.0.1:1")
3230 .group_id("timeout-group")
3231 .subscribe(vec!["orders".to_string()])
3232 .fetch_min(2)
3233 .fetch_max(1)
3234 .build()
3235 .await;
3236 assert2::assert!(matches!(
3237 inverted,
3238 Err(ConsumerError::InvalidConfig(message))
3239 if message == "consumer fetch min must not exceed consumer fetch max"
3240 ));
3241 }
3242
3243 #[test]
3244 fn consumer_fetch_byte_settings_validate_and_convert_to_uom() {
3245 for value in [1, i32::MAX] {
3246 let size = ByteSize::from_bytes_i64(i64::from(value));
3247 let max = ConsumerFetchMaxBytes::try_from(size).expect("positive fetch maximum");
3248 let partition_max = ConsumerFetchPartitionMaxBytes::try_from(size)
3249 .expect("positive partition fetch maximum");
3250
3251 check!(max.bytes() == value);
3252 check!(partition_max.bytes() == value);
3253 check!(max.size().bytes_i32() == value);
3254 check!(partition_max.size().bytes_i32() == value);
3255 }
3256
3257 for invalid in [
3258 ByteSize::ZERO,
3259 ByteSize::from_bytes_f64(-1.0),
3260 ByteSize::from_bytes_f64(1.5),
3261 ByteSize::from_bytes_f64(f64::from(i32::MAX) + 1.0),
3262 ] {
3263 check!(ConsumerFetchMaxBytes::try_from(invalid).is_err());
3264 check!(ConsumerFetchPartitionMaxBytes::try_from(invalid).is_err());
3265 }
3266 }
3267
3268 async fn test_consumer() -> Consumer {
3269 let client = Client::builder()
3270 .bootstrap("127.0.0.1:1")
3271 .client_id("test-client")
3272 .build()
3273 .await
3274 .unwrap();
3275 Consumer::test_builder(client)
3276 .coordinator_id(Arc::new(AtomicI32::new(3)))
3277 .generation_id(7)
3278 .commit_identity(test_support::orders_partition_commit_identity(7))
3279 .group_instance_id("instance-a")
3280 .subscription(subscription::shared(
3281 vec!["orders".into(), "payments".into()],
3282 None,
3283 true,
3284 ))
3285 .assigned(vec![("orders".into(), 0)])
3286 .build()
3287 }
3288
3289 #[tokio::test]
3290 async fn accessors_return_consumer_identity_subscription_and_assignment() {
3291 let consumer = test_consumer().await;
3292
3293 check!(
3294 (
3295 consumer.group_id(),
3296 consumer.member_id(),
3297 consumer.generation_id(),
3298 consumer.subscription(),
3299 consumer.assignment().await,
3300 ) == (
3301 "group-a",
3302 "member-a".to_owned(),
3303 7,
3304 vec!["orders".to_string(), "payments".to_string()],
3305 vec![("orders".into(), 0)],
3306 )
3307 );
3308
3309 let metadata = consumer.group_metadata();
3310 assert2::assert!(
3311 metadata
3312 == ConsumerGroupMetadata {
3313 group_id: "group-a".into(),
3314 generation_id: 7,
3315 member_id: "member-a".into(),
3316 group_instance_id: Some("instance-a".into()),
3317 }
3318 );
3319 }
3320
3321 #[tokio::test]
3322 async fn log_end_requires_every_assigned_position_to_reach_a_reported_end() {
3323 let consumer = test_consumer().await;
3324 assert2::assert!(!consumer.at_log_end().await);
3325
3326 consumer
3327 .next_offsets
3328 .lock()
3329 .await
3330 .insert(("orders".into(), 0), 12);
3331 consumer
3332 .end_offsets
3333 .lock()
3334 .await
3335 .insert(("orders".into(), 0), 12);
3336 assert2::assert!(consumer.at_log_end().await);
3337
3338 consumer.seek("orders", 0, 0).await.unwrap();
3339 let end = consumer
3342 .end_offsets
3343 .lock()
3344 .await
3345 .get(&("orders".to_owned(), 0))
3346 .copied();
3347 assert2::assert!((consumer.at_log_end().await, end) == (false, Some(12)));
3348 }
3349
3350 #[tokio::test]
3358 async fn generation_tracks_coordinator_rejoins_via_shared_atomic() {
3359 let consumer = test_consumer().await;
3360 assert2::assert!(
3361 (
3362 consumer.generation_id(),
3363 consumer.group_metadata().generation_id
3364 ) == (7, 7)
3365 );
3366
3367 consumer.current_generation.store(11, Ordering::Relaxed);
3369
3370 assert2::assert!(
3371 (
3372 consumer.generation_id(),
3373 consumer.group_metadata().generation_id
3374 ) == (11, 11)
3375 );
3376 }
3377
3378 #[tokio::test]
3379 async fn close_cancels_shutdown_token_even_without_spawned_handle() {
3380 let consumer = test_consumer().await;
3381 let shutdown = consumer.coordinator_shutdown.clone();
3382
3383 consumer.close().await.unwrap();
3384
3385 assert2::assert!(shutdown.is_cancelled());
3386 }
3387
3388 #[test]
3391 fn retriable_error_classification() {
3392 use std::net::{IpAddr, Ipv4Addr, SocketAddr};
3393
3394 use krabka_client_core::ClientError;
3395
3396 let transient_codes: &[i16] = &[14, 15, 16, 22, 25, 27, 79];
3398 for &code in transient_codes {
3399 assert2::assert!(is_retriable_consumer_start_error(&ConsumerError::Server(
3400 code
3401 )));
3402 }
3403
3404 assert2::assert!(!is_retriable_consumer_start_error(&ConsumerError::Server(
3406 42
3407 )));
3408
3409 for (_name, error, expected) in [
3410 (
3412 "disconnected",
3413 ConsumerError::Client(ClientError::Disconnected),
3414 true,
3415 ),
3416 (
3421 "timeout",
3422 ConsumerError::Client(ClientError::Timeout(krabka_units::secs(1))),
3423 false,
3424 ),
3425 (
3426 "startup after join",
3427 ConsumerError::StartupAfterJoin(Box::new(ConsumerError::Client(
3428 ClientError::Timeout(krabka_units::secs(1)),
3429 ))),
3430 true,
3431 ),
3432 (
3433 "connect refused",
3434 ConsumerError::Client(ClientError::Connect {
3435 addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9092),
3436 source: std::io::Error::new(std::io::ErrorKind::ConnectionRefused, "refused"),
3437 }),
3438 false,
3439 ),
3440 ("not subscribed", ConsumerError::NotSubscribed, false),
3442 (
3443 "invalid config",
3444 ConsumerError::InvalidConfig("group_id required".into()),
3445 false,
3446 ),
3447 (
3448 "incompatible version",
3449 ConsumerError::Client(ClientError::IncompatibleVersion {
3450 api_key: 0,
3451 broker_min: 0,
3452 broker_max: 5,
3453 client_min: 7,
3454 client_max: 10,
3455 }),
3456 false,
3457 ),
3458 ] {
3459 assert2::assert!(is_retriable_consumer_start_error(&error) == expected);
3460 }
3461 }
3462
3463 #[tokio::test]
3464 async fn consumer_builder_accepts_security() {
3465 let security = ClientSecurity {
3466 protocol: ListenerProtocol::SaslPlaintext,
3467 tls: None,
3468 sasl: Some(SaslCredentials::Plain {
3469 username: "u".into(),
3470 password: "p".into(),
3471 }),
3472 sasl_host: None,
3473 };
3474 let res = Consumer::builder()
3478 .bootstrap("127.0.0.1:1")
3479 .group_id("g")
3480 .subscribe(vec!["t".to_string()])
3481 .security(security)
3482 .build()
3483 .await;
3484 assert2::assert!(res.is_err());
3485 }
3486}
3487
3488#[cfg(test)]
3489mod auto_commit_tests {
3490 use std::{collections::VecDeque, sync::atomic::AtomicI16};
3491
3492 use bytes::Buf as _;
3493 use krabka_client_core::{MockBroker, mock_test_support::offset_commit_response};
3494 use krabka_protocol::{
3495 Decode as _, Encode, UnknownTaggedFields,
3496 owned::{
3497 api_versions_request, fetch_request,
3498 fetch_response::FetchResponse,
3499 find_coordinator_request,
3500 find_coordinator_response::FindCoordinatorResponse,
3501 heartbeat_request,
3502 heartbeat_response::HeartbeatResponse,
3503 join_group_request::{self, JoinGroupRequest},
3504 leave_group_request::{self, LeaveGroupRequest},
3505 leave_group_response::LeaveGroupResponse,
3506 list_offsets_request, metadata_request,
3507 metadata_response::MetadataResponse,
3508 offset_commit_request::{
3509 self, OffsetCommitRequest, OffsetCommitRequestPartition, OffsetCommitRequestTopic,
3510 },
3511 offset_fetch_request,
3512 offset_fetch_response::OffsetFetchResponse,
3513 sync_group_request,
3514 },
3515 };
3516
3517 use super::*;
3518
3519 const GROUP: &str = "group-a";
3520 pub(super) const MEMBER: &str = "member-a";
3521 pub(super) const TOPIC: &str = "orders";
3522 const REBALANCE_IN_PROGRESS: i16 = 27;
3524 const AUTO_COMMIT_INTERVAL: Duration = Duration::from_secs(5);
3525 const SESSION_TIMEOUT: Duration = Duration::from_secs(45);
3526 const REBALANCE_WAIT: Duration = Duration::from_secs(30);
3529
3530 #[derive(Clone, Debug, PartialEq)]
3533 pub(super) enum GroupRequest {
3534 OffsetCommit(OffsetCommitRequest),
3535 JoinGroup,
3536 SyncGroup,
3537 LeaveGroup,
3538 CommitSyncReturned,
3540 SyncGroupLate,
3543 CloseTimedOut,
3545 SessionExpired,
3548 }
3549
3550 #[derive(Clone, Copy, Debug)]
3552 pub(super) enum CommitReply {
3553 Drop,
3555 Error(i16),
3557 }
3558
3559 const API_VERSIONS: [(i16, i16, i16); 11] = [
3563 (api_versions_request::API_KEY, 0, 3),
3564 (metadata_request::API_KEY, 0, 8),
3565 (find_coordinator_request::API_KEY, 0, 2),
3566 (join_group_request::API_KEY, 0, 8),
3567 (sync_group_request::API_KEY, 0, 3),
3568 (heartbeat_request::API_KEY, 0, 3),
3569 (leave_group_request::API_KEY, 0, 5),
3570 (offset_commit_request::API_KEY, 2, 7),
3571 (offset_fetch_request::API_KEY, 1, 5),
3572 (fetch_request::API_KEY, 4, 11),
3573 (list_offsets_request::API_KEY, 1, 5),
3574 ];
3575
3576 fn encoded_api_versions() -> Vec<u8> {
3577 encode(
3578 &krabka_client_core::mock_test_support::api_versions_response_ranges(
3579 API_VERSIONS.iter().copied(),
3580 ),
3581 0,
3582 )
3583 }
3584
3585 pub(super) struct MockCoordinator {
3589 protocol: &'static str,
3590 requests: std::sync::Mutex<Vec<GroupRequest>>,
3591 pub(super) heartbeat_error: AtomicI16,
3594 assignments: std::sync::Mutex<VecDeque<Vec<(String, i32)>>>,
3597 generation: AtomicI32,
3598 pub(super) commit_replies: std::sync::Mutex<VecDeque<CommitReply>>,
3601 rebalance_on_commit: std::sync::atomic::AtomicBool,
3604 pub(super) last_heartbeat: std::sync::Mutex<tokio::time::Instant>,
3606 drop_heartbeats: std::sync::atomic::AtomicBool,
3609 pub(super) joins: std::sync::Mutex<Vec<JoinGroupRequest>>,
3611 pub(super) leaves: std::sync::Mutex<Vec<LeaveGroupRequest>>,
3613 pub(super) drop_leaves: std::sync::atomic::AtomicBool,
3615 }
3616
3617 pub(super) async fn mock_coordinator_broker(coordinator: &Arc<MockCoordinator>) -> MockBroker {
3618 let in_mock = Arc::clone(coordinator);
3619 MockBroker::start(move |api_key, version, _corr_id, body| {
3620 in_mock.respond(api_key, version, body)
3621 })
3622 .await
3623 }
3624
3625 pub(super) fn single_partition_startup(topic_partition_count: i32) -> StartupState {
3626 StartupState::builder()
3627 .generation_id(1)
3628 .assigned_partitions(vec![partition(0)])
3629 .next_offsets(maplit::hashmap! {partition(0) => 12})
3630 .topic_partitions(maplit::hashmap! {TOPIC.to_owned() => topic_partition_count})
3631 .build()
3632 }
3633
3634 pub(super) async fn start_consumer_with_startup(
3635 config: StartConfig,
3636 startup: StartupState,
3637 ) -> Consumer {
3638 start_test_consumer(
3639 TestConsumerStart::builder()
3640 .bootstrap(config.bootstrap.clone())
3641 .config(config)
3642 .startup(startup)
3643 .build(),
3644 )
3645 .await
3646 }
3647
3648 pub(super) async fn start_single_partition_consumer(config: StartConfig) -> Consumer {
3649 start_consumer_with_startup(config, single_partition_startup(1)).await
3650 }
3651
3652 pub(super) fn two_partition_startup() -> StartupState {
3653 StartupState::builder()
3654 .generation_id(1)
3655 .assigned_partitions(vec![partition(0), partition(1)])
3656 .next_offsets(maplit::hashmap! {partition(0) => 12, partition(1) => 7})
3657 .topic_partitions(maplit::hashmap! {TOPIC.to_owned() => 2})
3658 .build()
3659 }
3660
3661 fn encode(response: &impl Encode, version: i16) -> Vec<u8> {
3662 let mut buf = bytes::BytesMut::new();
3663 response.encode(&mut buf, version).expect("encode response");
3664 buf.to_vec()
3665 }
3666
3667 impl MockCoordinator {
3668 pub(super) fn range_partition(index: i32) -> Arc<Self> {
3669 Self::new(Assignor::Range, vec![vec![partition(index)]])
3670 }
3671
3672 pub(super) fn new(assignor: Assignor, assignments: Vec<Vec<(String, i32)>>) -> Arc<Self> {
3675 Arc::new(Self {
3676 protocol: assignor.protocol_name(),
3677 requests: std::sync::Mutex::new(Vec::new()),
3678 heartbeat_error: AtomicI16::new(0),
3679 assignments: std::sync::Mutex::new(assignments.into()),
3680 generation: AtomicI32::new(1),
3681 commit_replies: std::sync::Mutex::new(VecDeque::new()),
3682 rebalance_on_commit: std::sync::atomic::AtomicBool::new(false),
3683 last_heartbeat: std::sync::Mutex::new(tokio::time::Instant::now()),
3684 drop_heartbeats: std::sync::atomic::AtomicBool::new(false),
3685 joins: std::sync::Mutex::new(Vec::new()),
3686 leaves: std::sync::Mutex::new(Vec::new()),
3687 drop_leaves: std::sync::atomic::AtomicBool::new(false),
3688 })
3689 }
3690
3691 pub(super) fn set_heartbeat_error(&self, error_code: i16) {
3692 self.heartbeat_error.store(error_code, Ordering::SeqCst);
3693 }
3694
3695 fn record(&self, request: GroupRequest) {
3696 self.requests.lock().expect("requests lock").push(request);
3697 }
3698
3699 pub(super) fn requests(&self) -> Vec<GroupRequest> {
3700 self.requests.lock().expect("requests lock").clone()
3701 }
3702
3703 pub(super) fn sync_groups(&self) -> usize {
3704 self.requests()
3705 .iter()
3706 .filter(|request| **request == GroupRequest::SyncGroup)
3707 .count()
3708 }
3709
3710 pub(super) fn respond(
3711 &self,
3712 api_key: i16,
3713 version: i16,
3714 mut body: &[u8],
3715 ) -> Option<Vec<u8>> {
3716 match api_key {
3717 api_versions_request::API_KEY => Some(encoded_api_versions()),
3718 metadata_request::API_KEY => Some(encode(&MetadataResponse::default(), version)),
3719 heartbeat_request::API_KEY => {
3720 *self.last_heartbeat.lock().expect("last heartbeat lock") =
3721 tokio::time::Instant::now();
3722 let error_code = self.heartbeat_error.swap(0, Ordering::SeqCst);
3723 if error_code == 0 && self.drop_heartbeats.load(Ordering::SeqCst) {
3724 return None;
3725 }
3726 Some(encode(
3727 &HeartbeatResponse {
3728 error_code,
3729 ..Default::default()
3730 },
3731 version,
3732 ))
3733 }
3734 join_group_request::API_KEY => {
3735 let mut last_heartbeat =
3736 self.last_heartbeat.lock().expect("last heartbeat lock");
3737 if last_heartbeat.elapsed() > SESSION_TIMEOUT {
3738 self.record(GroupRequest::SessionExpired);
3739 }
3740 *last_heartbeat = tokio::time::Instant::now();
3741 drop(last_heartbeat);
3742 self.record(GroupRequest::JoinGroup);
3743 let client_id_len = body.get_i16();
3744 body.advance(usize::try_from(client_id_len.max(0)).expect("client id length"));
3745 let flexible = version >= join_group_request::FLEXIBLE_MIN;
3746 if flexible {
3747 body.advance(1);
3749 }
3750 self.joins
3751 .lock()
3752 .expect("joins lock")
3753 .push(JoinGroupRequest::decode(&mut body, version).expect("decode join"));
3754 let mut response = if flexible { vec![0] } else { Vec::new() };
3755 response.extend(encode(
3756 &JoinGroupResponse {
3757 generation_id: self.generation.fetch_add(1, Ordering::SeqCst) + 1,
3758 protocol_name: Some(self.protocol.to_owned()),
3759 protocol_type: Some("consumer".into()),
3760 leader: "leader".into(),
3761 member_id: MEMBER.into(),
3762 ..Default::default()
3763 },
3764 version,
3765 ));
3766 Some(response)
3767 }
3768 sync_group_request::API_KEY => {
3769 self.record(GroupRequest::SyncGroup);
3770 let mut assignments = self.assignments.lock().expect("assignments lock");
3771 let assignment = if assignments.len() > 1 {
3772 assignments.pop_front().expect("assignment")
3773 } else {
3774 assignments.front().cloned().expect("assignment")
3775 };
3776 drop(assignments);
3777 Some(encode(
3778 &SyncGroupResponse {
3779 assignment: builder::encode_assignment(&assignment),
3780 ..Default::default()
3781 },
3782 version,
3783 ))
3784 }
3785 fetch_request::API_KEY => Some(encode(&FetchResponse::default(), version)),
3787 list_offsets_request::API_KEY => {
3790 Some(test_support::list_offsets_mock_response(&mut body, version))
3791 }
3792 offset_fetch_request::API_KEY => {
3794 Some(encode(&OffsetFetchResponse::default(), version))
3795 }
3796 leave_group_request::API_KEY => {
3797 self.record(GroupRequest::LeaveGroup);
3798 let client_id_len = body.get_i16();
3799 body.advance(usize::try_from(client_id_len.max(0)).expect("client id length"));
3800 let flexible = version >= leave_group_request::FLEXIBLE_MIN;
3801 if flexible {
3802 body.advance(1);
3804 }
3805 self.leaves
3806 .lock()
3807 .expect("leaves lock")
3808 .push(LeaveGroupRequest::decode(&mut body, version).expect("decode leave"));
3809 if self.drop_leaves.load(Ordering::SeqCst) {
3810 return None;
3811 }
3812 let mut response = if flexible { vec![0] } else { Vec::new() };
3813 response.extend(encode(&LeaveGroupResponse::default(), version));
3814 Some(response)
3815 }
3816 find_coordinator_request::API_KEY => Some(encode(
3817 &FindCoordinatorResponse {
3818 node_id: 0,
3819 ..Default::default()
3820 },
3821 version,
3822 )),
3823 offset_commit_request::API_KEY => self.respond_offset_commit(version, body),
3824 _ => None,
3825 }
3826 }
3827
3828 fn respond_offset_commit(&self, version: i16, mut body: &[u8]) -> Option<Vec<u8>> {
3829 let client_id_len = body.get_i16();
3830 body.advance(usize::try_from(client_id_len.max(0)).expect("client id length"));
3831 let mut request =
3832 OffsetCommitRequest::decode(&mut body, version).expect("decode commit");
3833 request.topics.sort_by(|a, b| a.name.cmp(&b.name));
3834 for topic in &mut request.topics {
3835 topic.partitions.sort_by_key(|p| p.partition_index);
3836 }
3837 if self.rebalance_on_commit.swap(false, Ordering::SeqCst) {
3838 self.heartbeat_error
3839 .store(REBALANCE_IN_PROGRESS, Ordering::SeqCst);
3840 }
3841 let error_code = match self
3842 .commit_replies
3843 .lock()
3844 .expect("commit replies lock")
3845 .pop_front()
3846 {
3847 None => 0,
3848 Some(CommitReply::Error(code)) => code,
3849 Some(CommitReply::Drop) => {
3850 self.record(GroupRequest::OffsetCommit(request));
3851 return None;
3852 }
3853 };
3854 let response = offset_commit_response(&request, |_| error_code);
3855 self.record(GroupRequest::OffsetCommit(request));
3856 Some(encode(&response, version))
3857 }
3858 }
3859
3860 pub(super) fn partition(index: i32) -> (String, i32) {
3861 (TOPIC.to_owned(), index)
3862 }
3863
3864 fn commit(generation: i32, offsets: &[(i32, i64, i32)]) -> GroupRequest {
3867 GroupRequest::OffsetCommit(OffsetCommitRequest {
3868 group_id: GROUP.into(),
3869 generation_id_or_member_epoch: generation,
3870 member_id: MEMBER.into(),
3871 group_instance_id: None,
3872 retention_time_ms: -1,
3873 topics: vec![OffsetCommitRequestTopic {
3874 name: TOPIC.into(),
3875 partitions: offsets
3876 .iter()
3877 .map(
3878 |(partition_index, committed_offset, committed_leader_epoch)| {
3879 OffsetCommitRequestPartition {
3880 partition_index: *partition_index,
3881 committed_offset: *committed_offset,
3882 committed_leader_epoch: *committed_leader_epoch,
3883 committed_metadata: Some(String::new()),
3884 unknown_tagged_fields: UnknownTaggedFields::default(),
3885 }
3886 },
3887 )
3888 .collect(),
3889 ..Default::default()
3890 }],
3891 unknown_tagged_fields: UnknownTaggedFields::default(),
3892 })
3893 }
3894
3895 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3899 async fn metadata_requests_name_the_subscribed_topics() {
3900 use krabka_protocol::owned::{
3901 find_coordinator_response::FindCoordinatorResponse, metadata_request::MetadataRequest,
3902 };
3903
3904 let named = |topics: &[&str], allow| {
3905 krabka_client_core::topics_request(topics.iter().map(|&t| t.to_owned()), allow)
3906 };
3907 for (name, subscribe, allow_auto_create_topics, expected) in [
3908 (
3909 "subscribe a and b",
3910 vec!["a", "b"],
3911 true,
3912 named(&["a", "b"], true),
3913 ),
3914 ("auto creation off", vec!["a"], false, named(&["a"], false)),
3915 ] {
3916 let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
3917 let handler_requests = Arc::clone(&requests);
3918 let mock = MockBroker::start(move |api_key, version, _corr_id, mut body| {
3919 match api_key {
3920 api_versions_request::API_KEY => Some(encoded_api_versions()),
3921 find_coordinator_request::API_KEY => Some(encode(
3922 &FindCoordinatorResponse {
3923 node_id: 1,
3924 host: "127.0.0.1".into(),
3925 port: 0,
3926 ..Default::default()
3927 },
3928 version,
3929 )),
3930 metadata_request::API_KEY => {
3931 let client_id_len = body.get_i16();
3932 body.advance(usize::try_from(client_id_len.max(0)).expect("length"));
3933 let request =
3934 MetadataRequest::decode(&mut body, version).expect("decode Metadata");
3935 handler_requests
3936 .lock()
3937 .expect("requests lock")
3938 .push(request);
3939 Some(encode(&MetadataResponse::default(), version))
3940 }
3941 _ => None,
3943 }
3944 })
3945 .await;
3946 let mut config =
3947 start_config(mock.addr.to_string(), Assignor::Range, false, minutes(1));
3948 config.subscribe = subscribe.into_iter().map(str::to_owned).collect();
3949 config.allow_auto_create_topics = allow_auto_create_topics;
3950 let start = tokio::spawn(Consumer::start_once(config));
3951 let first = test_support::FirstRecorded::builder()
3952 .requests(Arc::clone(&requests))
3953 .timeout(Duration::from_secs(10))
3954 .poll_interval(Duration::from_millis(5))
3955 .build()
3956 .wait()
3957 .await
3958 .expect("the consumer sends Metadata");
3959 start.abort();
3960 mock.stop();
3961 assert2::check!(first == expected, "{name}");
3962 }
3963 }
3964
3965 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
3970 async fn the_coordinator_task_of_a_leader_build_watches_the_group_topics() {
3971 use krabka_protocol::owned::{
3972 find_coordinator_response::FindCoordinatorResponse,
3973 heartbeat_response::HeartbeatResponse,
3974 };
3975
3976 let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
3977 let handler_requests = Arc::clone(&requests);
3978 let mock = MockBroker::start(move |api_key, version, _corr_id, mut body| match api_key {
3979 api_versions_request::API_KEY => Some(encoded_api_versions()),
3980 find_coordinator_request::API_KEY => {
3981 Some(encode(&FindCoordinatorResponse::default(), version))
3982 }
3983 heartbeat_request::API_KEY => Some(encode(&HeartbeatResponse::default(), version)),
3984 metadata_request::API_KEY => {
3985 let names = test_support::metadata_topic_names(&mut body, version);
3986 handler_requests.lock().expect("requests lock").push(names);
3987 Some(encode(&MetadataResponse::default(), version))
3988 }
3989 _ => None,
3990 })
3991 .await;
3992 let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, minutes(1));
3993 config.heartbeat_interval = millis(50);
3994 config.subscription_metadata_refresh_interval = millis(50);
3995 let consumer = start_test_consumer(
3996 TestConsumerStart::builder()
3997 .bootstrap(mock.addr.to_string())
3998 .config(config)
3999 .startup(
4000 StartupState::builder()
4001 .generation_id(1)
4002 .group_topics(vec![TOPIC.to_owned(), "payments".to_owned()])
4003 .assigned_partitions(vec![partition(0)])
4004 .next_offsets(maplit::hashmap! {partition(0) => 12})
4005 .topic_partitions(maplit::hashmap! {TOPIC.to_owned() => 1})
4006 .build(),
4007 )
4008 .build(),
4009 )
4010 .await;
4011 let first = test_support::FirstRecorded::builder()
4012 .requests(Arc::clone(&requests))
4013 .timeout(Duration::from_secs(5))
4014 .poll_interval(Duration::from_millis(10))
4015 .build()
4016 .wait()
4017 .await;
4018 drop(consumer);
4019 mock.stop();
4020 assert2::assert!(first == Ok(vec![TOPIC.to_owned(), "payments".to_owned()]));
4021 }
4022
4023 pub(super) fn start_config(
4024 bootstrap: String,
4025 assignor: Assignor,
4026 auto_commit: bool,
4027 rebalance_timeout: Time,
4028 ) -> StartConfig {
4029 StartConfig {
4030 bootstrap,
4031 client_id: "auto-commit-test".into(),
4032 group_id: GROUP.into(),
4033 session_timeout: Time::from_std(SESSION_TIMEOUT),
4034 max_poll_interval: rebalance_timeout,
4035 max_poll_records: DEFAULT_CONSUMER_MAX_POLL_RECORDS,
4036 heartbeat_interval: secs(3),
4037 subscription_metadata_refresh_interval: minutes(60),
4038 subscribe: vec![TOPIC.into()],
4039 subscribe_pattern: None,
4040 subscribe_regex: None,
4041 exclude_internal_topics: true,
4042 group_protocol: GroupProtocol::Classic,
4043 server_assignor: None,
4044 group_instance_id: None,
4045 auto_offset_reset: AutoOffsetReset::Earliest,
4046 isolation_level: IsolationLevel::ReadUncommitted,
4047 assignors: vec![assignor],
4048 rebalance_protocol: assignor::rebalance_protocol_of(&[assignor])
4049 .expect("assignor protocol"),
4050 fetch_min: krabka_client_core::DEFAULT_FETCH_MIN,
4051 fetch_max: poll::DEFAULT_FETCH_MAX,
4052 fetch_partition_max: poll::DEFAULT_FETCH_PARTITION_MAX,
4053 fetch_max_wait: DEFAULT_CONSUMER_FETCH_MAX_WAIT,
4054 metadata_max_age: DEFAULT_CONSUMER_METADATA_MAX_AGE,
4055 default_api_timeout: DEFAULT_CONSUMER_DEFAULT_API_TIMEOUT,
4056 request_timeout: secs(30),
4057 dns_timeout: krabka_client_core::DEFAULT_CLIENT_DNS_TIMEOUT,
4058 send_buffer: Some(krabka_client_core::DEFAULT_SEND_BUFFER),
4059 receive_buffer: Some(krabka_client_core::DEFAULT_RECEIVE_BUFFER),
4060 dispatch_queue_capacity: ConnectionDispatchQueueCapacity::new(
4061 krabka_client_core::DEFAULT_CONNECTION_DISPATCH_QUEUE_CAPACITY,
4062 )
4063 .expect("dispatch queue capacity"),
4064 frame_max: ClientFrameMax::try_from(krabka_client_core::DEFAULT_CLIENT_FRAME_MAX)
4065 .expect("frame max"),
4066 metadata_recovery_strategy: krabka_client_core::MetadataRecoveryStrategy::default(),
4067 metadata_recovery_rebootstrap_trigger:
4068 krabka_client_core::DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER,
4069 leave_group_timeout: secs(5),
4070 client_rack: None,
4071 security: None,
4072 client_dns_lookup: krabka_client_core::ClientDnsLookup::default(),
4073 reconnect_backoff: krabka_client_core::DEFAULT_RECONNECT_BACKOFF,
4074 reconnect_backoff_max: krabka_client_core::DEFAULT_RECONNECT_BACKOFF_MAX,
4075 connections_max_idle: krabka_client_core::DEFAULT_CONNECTIONS_MAX_IDLE,
4076 socket_connection_setup_timeout:
4077 krabka_client_core::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT,
4078 socket_connection_setup_timeout_max:
4079 krabka_client_core::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX,
4080 has_rebalance_listener: false,
4081 retry_policy: ConsumerRetryPolicy::default(),
4082 auto_commit_interval: auto_commit.then_some(AUTO_COMMIT_INTERVAL),
4083 allow_auto_create_topics: true,
4084 telemetry: None,
4085 }
4086 }
4087
4088 #[derive(Clone, Copy, Debug)]
4090 enum Step {
4091 Poll,
4093 AdvanceInterval,
4095 ReceiveRecords,
4098 Rebalance { syncs: usize, within: Duration },
4102 CommitSync,
4104 CommitSyncDuringRebalance {
4108 reply: CommitReply,
4109 syncs: usize,
4110 within: Duration,
4111 },
4112 ReplyToNextCommit(CommitReply),
4114 HoldCommitLock,
4116 DropHeartbeats,
4119 Close,
4121 }
4122
4123 pub(super) async fn wait_for_sync_groups(
4128 coordinator: &MockCoordinator,
4129 polls: &coordinator::PollSignal,
4130 syncs: usize,
4131 within: Duration,
4132 ) {
4133 let waited = tokio::time::timeout(within, async {
4134 while coordinator.sync_groups() < syncs {
4135 coordinator::note_poll(polls);
4136 tokio::time::sleep(Duration::from_millis(100)).await;
4137 }
4138 })
4139 .await;
4140 if waited.is_err() {
4141 coordinator.record(GroupRequest::SyncGroupLate);
4142 }
4143 tokio::time::sleep(Duration::from_secs(1)).await;
4145 }
4146
4147 async fn run(
4151 auto_commit: bool,
4152 assignor: Assignor,
4153 assignments: Vec<Vec<(String, i32)>>,
4154 steps: &[Step],
4155 ) -> Vec<GroupRequest> {
4156 run_with_rebalance_timeout(auto_commit, assignor, assignments, minutes(1), steps).await
4157 }
4158
4159 async fn run_with_rebalance_timeout(
4161 auto_commit: bool,
4162 assignor: Assignor,
4163 assignments: Vec<Vec<(String, i32)>>,
4164 rebalance_timeout: Time,
4165 steps: &[Step],
4166 ) -> Vec<GroupRequest> {
4167 let coordinator = MockCoordinator::new(assignor, assignments);
4168 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
4169 let config = start_config(
4170 mock.addr.to_string(),
4171 assignor,
4172 auto_commit,
4173 rebalance_timeout,
4174 );
4175 let client = Client::builder()
4176 .bootstrap(mock.addr.to_string())
4177 .build()
4178 .await
4179 .expect("client");
4180 let mut consumer = Some(
4181 spawn_consumer(
4182 config,
4183 client,
4184 Arc::new(AtomicI32::new(0)),
4185 MEMBER.into(),
4186 StartupState::builder()
4187 .generation_id(1)
4188 .assigned_partitions(vec![partition(0), partition(1)])
4189 .next_offsets(maplit::hashmap! {partition(0) => 12, partition(1) => 7})
4190 .positions(maplit::hashmap! {partition(0) => primed_position(3)})
4191 .topic_partitions(maplit::hashmap! {TOPIC.to_owned() => 2})
4192 .build(),
4193 )
4194 .await
4195 .expect("spawn consumer"),
4196 );
4197 let mut held_commit_lock = None;
4198
4199 for step in steps {
4200 match step {
4201 Step::Poll => {
4202 let consumer = consumer.as_ref().expect("open consumer");
4203 consumer.maybe_auto_commit_async().await;
4204 drop(consumer.commit_serialization.lock().await);
4205 }
4206 Step::AdvanceInterval => tokio::time::advance(AUTO_COMMIT_INTERVAL).await,
4207 Step::ReceiveRecords => {
4208 let consumer = consumer.as_ref().expect("open consumer");
4209 consumer.next_offsets.lock().await.insert(partition(0), 20);
4210 }
4211 Step::Rebalance { syncs, within } => {
4212 let consumer = consumer.as_ref().expect("open consumer");
4213 coordinator
4214 .heartbeat_error
4215 .store(REBALANCE_IN_PROGRESS, Ordering::SeqCst);
4216 wait_for_sync_groups(&coordinator, &consumer.poll_signal, *syncs, *within)
4217 .await;
4218 }
4219 Step::CommitSync => {
4220 let consumer = consumer.as_ref().expect("open consumer");
4221 consumer.commit_sync().await.expect("commit_sync");
4222 coordinator.record(GroupRequest::CommitSyncReturned);
4223 }
4224 Step::CommitSyncDuringRebalance {
4225 reply,
4226 syncs,
4227 within,
4228 } => {
4229 let consumer = consumer.as_ref().expect("open consumer");
4230 coordinator
4231 .commit_replies
4232 .lock()
4233 .expect("commit replies lock")
4234 .push_back(*reply);
4235 coordinator
4236 .rebalance_on_commit
4237 .store(true, Ordering::SeqCst);
4238 let commit = async {
4239 let _ = consumer.commit_sync().await;
4242 coordinator.record(GroupRequest::CommitSyncReturned);
4243 };
4244 tokio::join!(
4245 commit,
4246 wait_for_sync_groups(&coordinator, &consumer.poll_signal, *syncs, *within)
4247 );
4248 }
4249 Step::ReplyToNextCommit(reply) => {
4250 coordinator
4251 .commit_replies
4252 .lock()
4253 .expect("commit replies lock")
4254 .push_back(*reply);
4255 }
4256 Step::DropHeartbeats => {
4257 coordinator.drop_heartbeats.store(true, Ordering::SeqCst);
4258 }
4259 Step::HoldCommitLock => {
4260 let consumer = consumer.as_ref().expect("open consumer");
4261 held_commit_lock = Some(
4262 Arc::clone(&consumer.commit_serialization)
4263 .lock_owned()
4264 .await,
4265 );
4266 }
4267 Step::Close => {
4268 let close = consumer.take().expect("open consumer").close();
4269 match tokio::time::timeout(Duration::from_mins(2), close).await {
4270 Ok(result) => result.expect("close"),
4271 Err(_) => coordinator.record(GroupRequest::CloseTimedOut),
4272 }
4273 }
4274 }
4275 }
4276
4277 let requests = coordinator.requests();
4278 drop(held_commit_lock);
4279 drop(consumer);
4280 mock.stop();
4281 requests
4282 }
4283
4284 #[derive(bon::Builder)]
4286 struct AutoCommitCase {
4287 name: &'static str,
4288 #[builder(default = true)]
4289 auto_commit: bool,
4290 #[builder(default = Assignor::Range)]
4291 assignor: Assignor,
4292 #[builder(default)]
4293 assignments: Vec<Vec<(String, i32)>>,
4294 steps: Vec<Step>,
4295 expected: Vec<GroupRequest>,
4296 }
4297
4298 fn rebalance_steps(syncs: usize) -> Vec<Step> {
4299 use Step::{Poll, Rebalance, ReceiveRecords};
4300
4301 vec![
4302 Poll,
4303 ReceiveRecords,
4304 Rebalance {
4305 syncs,
4306 within: REBALANCE_WAIT,
4307 },
4308 ]
4309 }
4310
4311 #[derive(bon::Builder)]
4312 struct AutoCommitTimeoutCase {
4313 name: &'static str,
4314 #[builder(default = minutes(1))]
4315 rebalance_timeout: Time,
4316 steps: Vec<Step>,
4317 expected: Vec<GroupRequest>,
4318 }
4319
4320 trait NamedGroupRequestCase {
4321 fn name(&self) -> &'static str;
4322 fn expected(&self) -> &[GroupRequest];
4323 }
4324
4325 impl NamedGroupRequestCase for AutoCommitCase {
4326 fn name(&self) -> &'static str {
4327 self.name
4328 }
4329
4330 fn expected(&self) -> &[GroupRequest] {
4331 &self.expected
4332 }
4333 }
4334
4335 impl NamedGroupRequestCase for AutoCommitTimeoutCase {
4336 fn name(&self) -> &'static str {
4337 self.name
4338 }
4339
4340 fn expected(&self) -> &[GroupRequest] {
4341 &self.expected
4342 }
4343 }
4344
4345 #[derive(Debug)]
4346 struct GroupRequestObservation {
4347 name: &'static str,
4348 actual: Vec<GroupRequest>,
4349 expected: Vec<GroupRequest>,
4350 }
4351
4352 async fn assert_group_request_cases<C, F, Fut>(cases: Vec<C>, mut run: F)
4353 where
4354 C: NamedGroupRequestCase,
4355 F: FnMut(C) -> Fut,
4356 Fut: std::future::Future<Output = Vec<GroupRequest>>,
4357 {
4358 let mut observations = Vec::with_capacity(cases.len());
4359 for case in cases {
4360 let name = case.name();
4361 let expected = case.expected().to_vec();
4362 observations.push(GroupRequestObservation {
4363 name,
4364 actual: run(case).await,
4365 expected,
4366 });
4367 }
4368 let failed_cases: Vec<_> = observations
4369 .iter()
4370 .filter(|observation| observation.actual != observation.expected)
4371 .map(|observation| observation.name)
4372 .collect();
4373 assert2::assert!(failed_cases.is_empty());
4374 }
4375
4376 fn rebalance_step(within: Duration) -> Step {
4377 Step::Rebalance { syncs: 1, within }
4378 }
4379
4380 fn commit_sync_then_rebalance_requests(received: &[(i32, i64, i32)]) -> Vec<GroupRequest> {
4381 use GroupRequest::{CommitSyncReturned, JoinGroup, SyncGroup};
4382
4383 vec![
4384 commit(1, received),
4385 CommitSyncReturned,
4386 commit(1, received),
4387 JoinGroup,
4388 SyncGroup,
4389 ]
4390 }
4391
4392 #[tokio::test(start_paused = true)]
4397 async fn auto_commit_sends_offset_commit_where_kafka_does() {
4398 use GroupRequest::{JoinGroup, LeaveGroup, SyncGroup};
4399 use Step::{AdvanceInterval, Close, Poll, ReceiveRecords};
4400
4401 let both = [(0, 12, 3), (1, 7, -1)];
4402 let cases = vec![
4403 AutoCommitCase::builder()
4404 .name("on: a poll after the interval commits the positions")
4405 .steps(vec![Poll, AdvanceInterval, Poll])
4406 .expected(vec![commit(1, &both)])
4407 .build(),
4408 AutoCommitCase::builder()
4409 .name("on: a poll before the interval commits nothing")
4410 .steps(vec![Poll, Poll])
4411 .expected(vec![])
4412 .build(),
4413 AutoCommitCase::builder()
4414 .name("off: a poll after the interval commits nothing")
4415 .auto_commit(false)
4416 .steps(vec![Poll, AdvanceInterval, Poll])
4417 .expected(vec![])
4418 .build(),
4419 AutoCommitCase::builder()
4420 .name("on, eager: a rebalance commits the positions of the last poll before JoinGroup")
4421 .assignments(vec![vec![partition(0), partition(1)]])
4422 .steps(rebalance_steps(1))
4423 .expected(vec![commit(1, &both), JoinGroup, SyncGroup])
4424 .build(),
4425 AutoCommitCase::builder()
4426 .name("off, eager: a rebalance commits nothing")
4427 .auto_commit(false)
4428 .assignments(vec![vec![partition(0), partition(1)]])
4429 .steps(rebalance_steps(1))
4430 .expected(vec![JoinGroup, SyncGroup])
4431 .build(),
4432 AutoCommitCase::builder()
4433 .name("on, cooperative: a revoke of partition 1 commits before each JoinGroup")
4434 .assignor(Assignor::CooperativeSticky)
4435 .assignments(vec![vec![partition(0)]])
4436 .steps(rebalance_steps(2))
4437 .expected(vec![
4438 commit(1, &both),
4439 JoinGroup,
4440 SyncGroup,
4441 commit(2, &[(0, 12, 3)]),
4442 JoinGroup,
4443 SyncGroup,
4444 ])
4445 .build(),
4446 AutoCommitCase::builder()
4447 .name("off, cooperative: a revoke of partition 1 commits nothing")
4448 .auto_commit(false)
4449 .assignor(Assignor::CooperativeSticky)
4450 .assignments(vec![vec![partition(0)]])
4451 .steps(rebalance_steps(2))
4452 .expected(vec![JoinGroup, SyncGroup, JoinGroup, SyncGroup])
4453 .build(),
4454 AutoCommitCase::builder()
4455 .name("on: close commits the current positions, then leaves the group")
4456 .steps(vec![Poll, ReceiveRecords, Close])
4457 .expected(vec![commit(1, &[(0, 20, 3), (1, 7, -1)]), LeaveGroup])
4458 .build(),
4459 AutoCommitCase::builder()
4460 .name("off: close only leaves the group")
4461 .auto_commit(false)
4462 .assignor(Assignor::CooperativeSticky)
4463 .steps(vec![Poll, ReceiveRecords, Close])
4464 .expected(vec![LeaveGroup])
4465 .build(),
4466 ];
4467 assert_group_request_cases(cases, |case| async move {
4468 run(
4469 case.auto_commit,
4470 case.assignor,
4471 case.assignments,
4472 &case.steps,
4473 )
4474 .await
4475 })
4476 .await;
4477 }
4478
4479 #[tokio::test(start_paused = true)]
4485 async fn auto_commit_keeps_the_commit_order_within_the_group_timeouts() {
4486 use CommitReply::{Drop, Error};
4487 use GroupRequest::{CommitSyncReturned, JoinGroup, LeaveGroup, SyncGroup};
4488 use Step::{
4489 Close, CommitSync, CommitSyncDuringRebalance, DropHeartbeats, HoldCommitLock, Poll,
4490 ReceiveRecords, ReplyToNextCommit,
4491 };
4492
4493 let polled = [(0, 12, 3), (1, 7, -1)];
4494 let received = [(0, 20, 3), (1, 7, -1)];
4495 let both_partitions = vec![vec![partition(0), partition(1)]];
4496 let cases = vec![
4497 AutoCommitTimeoutCase::builder()
4498 .name("commit_sync, then a rebalance: the commit before JoinGroup does not go back")
4499 .steps(vec![
4500 Poll,
4501 ReceiveRecords,
4502 CommitSync,
4503 rebalance_step(REBALANCE_WAIT),
4504 ])
4505 .expected(commit_sync_then_rebalance_requests(&received))
4506 .build(),
4507 AutoCommitTimeoutCase::builder()
4508 .name(
4509 "a rebalance while commit_sync waits for its response: the commit before \
4510 JoinGroup comes after it",
4511 )
4512 .steps(vec![
4513 Poll,
4514 ReceiveRecords,
4515 CommitSyncDuringRebalance {
4516 reply: Drop,
4517 syncs: 1,
4518 within: Duration::from_mins(1),
4519 },
4520 ])
4521 .expected(vec![
4522 commit(1, &received),
4527 commit(1, &received),
4528 CommitSyncReturned,
4529 commit(1, &received),
4530 JoinGroup,
4531 SyncGroup,
4532 ])
4533 .build(),
4534 AutoCommitTimeoutCase::builder()
4535 .name(
4536 "a rebalance while commit_sync waits for its response, with a \
4537 REBALANCE_IN_PROGRESS reply: commit_sync fails immediately instead of \
4538 waiting for the rejoin, so the commit before JoinGroup comes after it",
4539 )
4540 .steps(vec![
4541 Poll,
4542 ReceiveRecords,
4543 CommitSyncDuringRebalance {
4544 reply: Error(REBALANCE_IN_PROGRESS),
4545 syncs: 1,
4546 within: REBALANCE_WAIT,
4547 },
4548 ])
4549 .expected(commit_sync_then_rebalance_requests(&received))
4550 .build(),
4551 AutoCommitTimeoutCase::builder()
4552 .name(
4553 "another commit holds the commit lock: JoinGroup at the rebalance timeout, \
4554 without a commit",
4555 )
4556 .rebalance_timeout(secs(10))
4557 .steps(vec![Poll, HoldCommitLock, rebalance_step(REBALANCE_WAIT)])
4558 .expected(vec![JoinGroup, SyncGroup])
4559 .build(),
4560 AutoCommitTimeoutCase::builder()
4561 .name(
4562 "another commit holds the commit lock past the session timeout: heartbeats keep \
4563 the member in the group",
4564 )
4565 .steps(vec![
4566 Poll,
4567 HoldCommitLock,
4568 rebalance_step(Duration::from_secs(90)),
4569 ])
4570 .expected(vec![JoinGroup, SyncGroup])
4571 .build(),
4572 AutoCommitTimeoutCase::builder()
4573 .name(
4574 "the coordinator does not answer a heartbeat during the commit before JoinGroup: \
4575 JoinGroup at the rebalance timeout",
4576 )
4577 .rebalance_timeout(secs(10))
4578 .steps(vec![
4579 Poll,
4580 HoldCommitLock,
4581 DropHeartbeats,
4582 rebalance_step(Duration::from_secs(20)),
4583 ])
4584 .expected(vec![JoinGroup, SyncGroup])
4585 .build(),
4586 AutoCommitTimeoutCase::builder()
4587 .name(
4588 "the coordinator does not answer the commit before JoinGroup: JoinGroup at the \
4589 rebalance timeout",
4590 )
4591 .rebalance_timeout(secs(10))
4592 .steps(vec![Poll, ReplyToNextCommit(Drop), rebalance_step(REBALANCE_WAIT)])
4593 .expected(vec![commit(1, &polled), JoinGroup, SyncGroup])
4594 .build(),
4595 AutoCommitTimeoutCase::builder()
4596 .name(
4597 "another commit holds the commit lock: close leaves the group after the close \
4598 timeout",
4599 )
4600 .steps(vec![Poll, ReceiveRecords, HoldCommitLock, Close])
4601 .expected(vec![LeaveGroup])
4602 .build(),
4603 ];
4604 assert_group_request_cases(cases, |case| {
4605 let assignments = both_partitions.clone();
4606 async move {
4607 run_with_rebalance_timeout(
4608 true,
4609 Assignor::Range,
4610 assignments,
4611 case.rebalance_timeout,
4612 &case.steps,
4613 )
4614 .await
4615 }
4616 })
4617 .await;
4618 }
4619
4620 #[test]
4621 fn auto_commit_interval_follows_kafka_config_bounds() {
4622 let mut actual = Vec::new();
4623 let mut wanted = Vec::new();
4624 for (name, enable, interval, expected) in [
4625 ("off", false, Time::from_millis(-1), Ok(None)),
4626 ("default", true, secs(5), Ok(Some(Duration::from_secs(5)))),
4627 ("zero", true, Time::from_millis(0), Ok(Some(Duration::ZERO))),
4628 (
4629 "i32::MAX milliseconds",
4630 true,
4631 Time::from_millis(i64::from(i32::MAX)),
4632 Ok(Some(Duration::from_millis(u64::from(
4633 i32::MAX.unsigned_abs(),
4634 )))),
4635 ),
4636 ("negative", true, Time::from_millis(-1), Err(())),
4637 (
4638 "above i32::MAX milliseconds",
4639 true,
4640 Time::from_millis(i64::from(i32::MAX) + 1),
4641 Err(()),
4642 ),
4643 (
4644 "half a millisecond",
4645 true,
4646 Time::from_secs_f64(0.0005),
4647 Err(()),
4648 ),
4649 (
4650 "i32::MAX milliseconds and a fraction",
4651 true,
4652 Time::from_secs_f64((f64::from(i32::MAX) + 0.25) / 1e3),
4653 Err(()),
4654 ),
4655 ("not a number", true, Time::from_secs_f64(f64::NAN), Err(())),
4656 (
4657 "infinite",
4658 true,
4659 Time::from_secs_f64(f64::INFINITY),
4660 Err(()),
4661 ),
4662 ] {
4663 let result = validated_auto_commit_interval(enable, interval).map_err(|_| ());
4664 actual.push((name, result));
4665 wanted.push((name, expected));
4666 }
4667 assert2::assert!(actual == wanted);
4668 }
4669}
4670
4671#[cfg(test)]
4672mod group_rejoin_tests {
4673 use super::{
4674 auto_commit_tests::{
4675 GroupRequest, MockCoordinator, partition, start_config, wait_for_sync_groups,
4676 },
4677 *,
4678 };
4679
4680 const FENCED_INSTANCE_ID: i16 = 82;
4682
4683 #[derive(Debug, PartialEq)]
4685 struct FenceObservation {
4686 first_poll: Result<usize, String>,
4687 commit_after_fence: Result<(), String>,
4688 requests_before_next_poll: Vec<GroupRequest>,
4689 next_poll: Result<usize, String>,
4690 requests_after_next_poll: Vec<GroupRequest>,
4691 assignment: Vec<(String, i32)>,
4692 generation: i32,
4693 }
4694
4695 #[tokio::test(start_paused = true)]
4701 async fn poll_after_a_fenced_instance_id_joins_the_group_again() {
4702 let coordinator = MockCoordinator::range_partition(1);
4703 coordinator.set_heartbeat_error(FENCED_INSTANCE_ID);
4704 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
4705 let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, minutes(1));
4706 config.group_instance_id = Some("instance-a".into());
4707 let mut consumer = start_test_consumer(
4708 TestConsumerStart::builder()
4709 .bootstrap(mock.addr.to_string())
4710 .config(config)
4711 .startup(super::auto_commit_tests::two_partition_startup())
4712 .build(),
4713 )
4714 .await;
4715 tokio::time::timeout(Duration::from_secs(30), async {
4717 while !coordinator::poll_error_pending(&consumer.poll_error) {
4718 tokio::time::sleep(Duration::from_millis(100)).await;
4719 }
4720 })
4721 .await
4722 .expect("the coordinator fences the member");
4723
4724 let first_poll = consumer
4725 .poll(millis(100))
4726 .await
4727 .map(|records| records.len())
4728 .map_err(|error| error.to_string());
4729 tokio::time::sleep(Duration::from_secs(20)).await;
4731 let commit_after_fence = consumer
4732 .commit_sync()
4733 .await
4734 .map_err(|error| error.to_string());
4735 let requests_before_next_poll = coordinator.requests();
4736 let next_poll = consumer
4737 .poll(millis(100))
4738 .await
4739 .map(|records| records.len())
4740 .map_err(|error| error.to_string());
4741 wait_for_sync_groups(
4742 &coordinator,
4743 &consumer.poll_signal,
4744 1,
4745 Duration::from_secs(30),
4746 )
4747 .await;
4748 let observation = FenceObservation {
4749 first_poll,
4750 commit_after_fence,
4751 requests_before_next_poll,
4752 next_poll,
4753 requests_after_next_poll: coordinator
4754 .requests()
4755 .into_iter()
4756 .filter(|request| *request != GroupRequest::SessionExpired)
4757 .collect(),
4758 assignment: consumer.assignment().await,
4759 generation: consumer.generation_id(),
4760 };
4761 consumer.close().await.expect("close");
4762 mock.stop();
4763
4764 assert2::assert!(
4765 observation
4766 == FenceObservation {
4767 first_poll: Err(
4768 ConsumerError::FencedInstanceId("instance-a".into()).to_string()
4769 ),
4770 commit_after_fence: Err(ConsumerError::CommitFailed.to_string()),
4771 requests_before_next_poll: Vec::new(),
4772 next_poll: Ok(0),
4773 requests_after_next_poll: vec![
4774 GroupRequest::JoinGroup,
4775 GroupRequest::SyncGroup
4776 ],
4777 assignment: vec![partition(1)],
4778 generation: 2,
4779 }
4780 );
4781 }
4782}
4783
4784#[cfg(test)]
4785mod poll_interval_tests {
4786 use krabka_protocol::owned::leave_group_request::LeaveGroupRequest;
4787
4788 use super::{
4789 auto_commit_tests::{MEMBER, MockCoordinator, TOPIC, start_config},
4790 *,
4791 };
4792
4793 #[derive(Clone, Copy, Debug)]
4799 enum Polls {
4800 EachTenthOfASecond,
4802 Late,
4804 }
4805
4806 #[derive(Debug, PartialEq)]
4808 struct PollTimeoutObservation {
4809 leaves: Vec<LeaveGroupRequest>,
4810 heartbeats_after_timeout: bool,
4811 commit: Result<(), String>,
4812 member_id_before_next_poll: String,
4814 joins_after_next_poll: Vec<(String, Option<String>)>,
4817 member_id_after_next_poll: String,
4819 }
4820
4821 async fn run_poll_timeout_case(
4822 group_instance_id: Option<&str>,
4823 polls: Polls,
4824 ) -> PollTimeoutObservation {
4825 let coordinator = MockCoordinator::range_partition(0);
4826 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
4827 let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, millis(500));
4828 config.group_instance_id = group_instance_id.map(str::to_owned);
4829 config.heartbeat_interval = millis(100);
4830 let mut consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
4831 let spawned_at = *coordinator
4833 .last_heartbeat
4834 .lock()
4835 .expect("last heartbeat lock");
4836 while *coordinator
4837 .last_heartbeat
4838 .lock()
4839 .expect("last heartbeat lock")
4840 == spawned_at
4841 {
4842 tokio::time::sleep(Duration::from_millis(10)).await;
4843 }
4844 coordinator::note_poll(&consumer.poll_signal);
4845
4846 let poller = matches!(polls, Polls::EachTenthOfASecond).then(|| {
4847 let poll_signal = consumer.poll_signal.clone();
4848 tokio::spawn(async move {
4849 loop {
4850 tokio::time::sleep(Duration::from_millis(100)).await;
4851 coordinator::note_poll(&poll_signal);
4852 }
4853 })
4854 });
4855 tokio::time::sleep(Duration::from_millis(700)).await;
4856 let heartbeat_at_timeout = *coordinator
4857 .last_heartbeat
4858 .lock()
4859 .expect("last heartbeat lock");
4860 tokio::time::sleep(Duration::from_millis(400)).await;
4862 let heartbeats_after_timeout = *coordinator
4863 .last_heartbeat
4864 .lock()
4865 .expect("last heartbeat lock")
4866 != heartbeat_at_timeout;
4867 let commit = consumer
4868 .commit_sync()
4869 .await
4870 .map_err(|error| error.to_string());
4871 let leaves = coordinator.leaves.lock().expect("leaves lock").clone();
4872 let member_id_before_next_poll = consumer.member_id();
4873 let joins_before = coordinator.joins.lock().expect("joins lock").len();
4874 consumer.poll(millis(0)).await.expect("next poll");
4875 let _ = tokio::time::timeout(Duration::from_secs(5), async {
4877 while consumer.member_id().is_empty() {
4878 coordinator::note_poll(&consumer.poll_signal);
4879 tokio::time::sleep(Duration::from_millis(10)).await;
4880 }
4881 })
4882 .await;
4883 let joins_after_next_poll = coordinator.joins.lock().expect("joins lock")[joins_before..]
4884 .iter()
4885 .map(|join| (join.member_id.clone(), join.group_instance_id.clone()))
4886 .collect();
4887 let member_id_after_next_poll = consumer.member_id();
4888 if let Some(poller) = poller {
4889 poller.abort();
4890 }
4891 coordinator.set_heartbeat_error(0);
4892 drop(consumer);
4893 mock.stop();
4894 PollTimeoutObservation {
4895 leaves,
4896 heartbeats_after_timeout,
4897 commit,
4898 member_id_before_next_poll,
4899 joins_after_next_poll,
4900 member_id_after_next_poll,
4901 }
4902 }
4903
4904 #[tokio::test]
4912 async fn a_member_leaves_the_group_when_poll_does_not_come_within_max_poll_interval() {
4913 let commit_failed = Err(ConsumerError::CommitFailed.to_string());
4914 let leave =
4915 test_support::dynamic_member_leave_request("consumer poll timeout has expired.");
4916 let mut actual = Vec::new();
4917 let mut wanted = Vec::new();
4918 for (name, group_instance_id, polls, expected) in [
4919 (
4920 "dynamic member polls in time",
4921 None,
4922 Polls::EachTenthOfASecond,
4923 PollTimeoutObservation {
4924 leaves: Vec::new(),
4925 heartbeats_after_timeout: true,
4926 commit: Ok(()),
4927 member_id_before_next_poll: MEMBER.into(),
4928 joins_after_next_poll: Vec::new(),
4929 member_id_after_next_poll: MEMBER.into(),
4930 },
4931 ),
4932 (
4933 "dynamic member polls late",
4934 None,
4935 Polls::Late,
4936 PollTimeoutObservation {
4937 leaves: vec![leave],
4938 heartbeats_after_timeout: false,
4939 commit: commit_failed.clone(),
4940 member_id_before_next_poll: String::new(),
4941 joins_after_next_poll: vec![(String::new(), None)],
4942 member_id_after_next_poll: MEMBER.into(),
4943 },
4944 ),
4945 (
4946 "static member polls late",
4947 Some("instance-a"),
4948 Polls::Late,
4949 PollTimeoutObservation {
4950 leaves: Vec::new(),
4951 heartbeats_after_timeout: false,
4952 commit: commit_failed,
4953 member_id_before_next_poll: String::new(),
4954 joins_after_next_poll: vec![(String::new(), Some("instance-a".into()))],
4955 member_id_after_next_poll: MEMBER.into(),
4956 },
4957 ),
4958 ] {
4959 actual.push((name, run_poll_timeout_case(group_instance_id, polls).await));
4960 wanted.push((name, expected));
4961 }
4962 assert2::assert!(actual == wanted);
4963 }
4964
4965 #[tokio::test]
4969 async fn a_poll_during_the_poll_timeout_leave_fetches_nothing_and_starts_the_join() {
4970 let coordinator = MockCoordinator::range_partition(0);
4971 coordinator
4972 .drop_leaves
4973 .store(true, std::sync::atomic::Ordering::SeqCst);
4974 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
4975 let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, millis(300));
4976 config.heartbeat_interval = millis(50);
4977 config.leave_group_timeout = secs(1);
4978 let mut consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
4979 tokio::time::timeout(Duration::from_secs(10), async {
4980 while coordinator.leaves.lock().expect("leaves lock").is_empty() {
4981 tokio::time::sleep(Duration::from_millis(10)).await;
4982 }
4983 })
4984 .await
4985 .expect("the member sends LeaveGroup");
4986
4987 let records = consumer.poll(millis(100)).await.expect("poll");
4988 let joined = tokio::time::timeout(Duration::from_secs(5), async {
4989 while coordinator.joins.lock().expect("joins lock").is_empty() {
4990 tokio::time::sleep(Duration::from_millis(10)).await;
4991 }
4992 })
4993 .await
4994 .is_ok();
4995 let observation = (
4996 records.len(),
4997 joined,
4998 coordinator.leaves.lock().expect("leaves lock").len(),
4999 );
5000 drop(consumer);
5001 mock.stop();
5002 assert2::assert!(observation == (0, true, 1));
5003 }
5004
5005 #[tokio::test]
5008 async fn join_group_sends_max_poll_interval_as_the_rebalance_timeout() {
5009 let mut actual = Vec::new();
5010 let mut wanted = Vec::new();
5011 for (name, max_poll_interval, expected) in [
5012 ("default", None, 300_000),
5013 ("configured", Some(secs(10)), 10_000),
5014 ] {
5015 let coordinator = MockCoordinator::new(Assignor::Range, vec![Vec::new()]);
5016 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5017 let consumer = Consumer::builder()
5018 .bootstrap(mock.addr.to_string())
5019 .group_id("group-a")
5020 .subscribe(vec![TOPIC.to_owned()])
5021 .maybe_max_poll_interval(max_poll_interval)
5022 .build()
5023 .await
5024 .expect("build");
5025 let rebalance_timeouts: Vec<i32> = coordinator
5026 .joins
5027 .lock()
5028 .expect("joins lock")
5029 .iter()
5030 .map(|join| join.rebalance_timeout_ms)
5031 .collect();
5032 drop(consumer);
5033 mock.stop();
5034 actual.push((name, rebalance_timeouts));
5035 wanted.push((name, vec![expected, expected]));
5036 }
5037 assert2::assert!(actual == wanted);
5038 }
5039
5040 #[test]
5044 fn client_id_is_generated_as_kafka_does() {
5045 let first = consumer_client_id(None, "g", None);
5046 let sequence: u64 = first
5047 .strip_prefix("consumer-g-")
5048 .and_then(|n| n.parse().ok())
5049 .expect("generated sequence");
5050 let actual = [
5051 consumer_client_id(None, "g", None),
5052 consumer_client_id(Some(String::new()), "g", None),
5053 consumer_client_id(None, "g", Some("i-1")),
5054 consumer_client_id(Some(String::new()), "g", Some("i-1")),
5055 consumer_client_id(Some("app".into()), "g", Some("i-1")),
5056 ];
5057 let next = |step: u64| format!("consumer-g-{}", sequence + step);
5058 assert2::assert!(
5059 actual
5060 == [
5061 next(1),
5062 next(2),
5063 "consumer-g-i-1".to_owned(),
5064 "consumer-g-i-1".to_owned(),
5065 "app".to_owned(),
5066 ]
5067 );
5068 }
5069
5070 #[test]
5071 fn max_poll_settings_follow_kafka_config_bounds() {
5072 let interval = |value: Time| validated_max_poll_interval(value).map_err(|_| ());
5073 let records = |value: usize| validated_max_poll_records(value).map_err(|_| ());
5074 let actual = (
5075 [
5076 interval(Time::from_millis(0)),
5077 interval(Time::from_millis(1)),
5078 interval(DEFAULT_CONSUMER_MAX_POLL_INTERVAL),
5079 interval(Time::from_millis(i64::from(i32::MAX))),
5080 interval(Time::from_millis(i64::from(i32::MAX) + 1)),
5081 interval(Time::from_secs_f64(0.0015)),
5082 ],
5083 [
5084 records(0),
5085 records(1),
5086 records(DEFAULT_CONSUMER_MAX_POLL_RECORDS),
5087 records(usize::try_from(i32::MAX).expect("fits")),
5088 records(usize::try_from(i32::MAX).expect("fits") + 1),
5089 ],
5090 );
5091 let expected = (
5092 [
5093 Err(()),
5094 Ok(Time::from_millis(1)),
5095 Ok(minutes(5)),
5096 Ok(Time::from_millis(i64::from(i32::MAX))),
5097 Err(()),
5098 Err(()),
5099 ],
5100 [
5101 Err(()),
5102 Ok(1),
5103 Ok(500),
5104 Ok(usize::try_from(i32::MAX).expect("fits")),
5105 Err(()),
5106 ],
5107 );
5108 assert2::assert!(actual == expected);
5109 }
5110}
5111
5112#[cfg(test)]
5113mod group_membership_tests {
5114 use std::sync::{Arc, atomic::Ordering};
5115
5116 use krabka_client_core::MockBroker;
5117
5118 use super::{
5119 auto_commit_tests::{
5120 GroupRequest, MEMBER, MockCoordinator, TOPIC, partition, start_config,
5121 },
5122 *,
5123 };
5124
5125 async fn started_consumer(mock: &MockBroker, group_instance_id: Option<&str>) -> Consumer {
5126 started_consumer_with(mock, group_instance_id, &[Assignor::Range]).await
5127 }
5128
5129 async fn started_consumer_with(
5130 mock: &MockBroker,
5131 group_instance_id: Option<&str>,
5132 assignors: &[Assignor],
5133 ) -> Consumer {
5134 let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, minutes(1));
5135 config.assignors = assignors.to_vec();
5136 config.rebalance_protocol =
5137 assignor::rebalance_protocol_of(assignors).expect("assignor protocol");
5138 config.group_instance_id = group_instance_id.map(str::to_owned);
5139 config.heartbeat_interval = millis(50);
5140 let client = Client::builder()
5141 .bootstrap(mock.addr.to_string())
5142 .build()
5143 .await
5144 .expect("client");
5145 spawn_consumer(
5146 config,
5147 client,
5148 Arc::new(AtomicI32::new(0)),
5149 MEMBER.into(),
5150 super::auto_commit_tests::single_partition_startup(1),
5151 )
5152 .await
5153 .expect("spawn consumer")
5154 }
5155
5156 struct StartedRangeConsumer {
5157 coordinator: Arc<MockCoordinator>,
5158 mock: MockBroker,
5159 consumer: Consumer,
5160 }
5161
5162 impl StartedRangeConsumer {
5163 async fn start(group_instance_id: Option<&str>) -> Self {
5164 let coordinator = MockCoordinator::range_partition(0);
5165 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5166 let consumer = started_consumer(&mock, group_instance_id).await;
5167 Self {
5168 coordinator,
5169 mock,
5170 consumer,
5171 }
5172 }
5173 }
5174
5175 #[tokio::test]
5180 async fn close_leaves_the_group_as_the_membership_operation_says() {
5181 let leave = test_support::dynamic_member_leave_request("the consumer is being closed");
5182 let mut actual = Vec::new();
5183 let mut wanted = Vec::new();
5184 for (name, group_instance_id, operation, expected) in [
5185 (
5186 "dynamic member, default",
5187 None,
5188 GroupMembershipOperation::Default,
5189 vec![leave.clone()],
5190 ),
5191 (
5192 "static member, default",
5193 Some("i-1"),
5194 GroupMembershipOperation::Default,
5195 Vec::new(),
5196 ),
5197 (
5198 "static member, leave group",
5199 Some("i-1"),
5200 GroupMembershipOperation::LeaveGroup,
5201 vec![leave.clone()],
5202 ),
5203 (
5204 "dynamic member, remain in group",
5205 None,
5206 GroupMembershipOperation::RemainInGroup,
5207 Vec::new(),
5208 ),
5209 ] {
5210 let fixture = StartedRangeConsumer::start(group_instance_id).await;
5211 fixture
5212 .consumer
5213 .close_with(CloseOptions::group_membership_operation(operation))
5214 .await
5215 .expect("close");
5216 fixture.mock.stop();
5217 actual.push((
5218 name,
5219 fixture
5220 .coordinator
5221 .leaves
5222 .lock()
5223 .expect("leaves lock")
5224 .clone(),
5225 ));
5226 wanted.push((name, expected));
5227 }
5228 assert2::assert!(actual == wanted);
5229 }
5230
5231 #[tokio::test]
5236 async fn a_join_after_the_member_was_reset_during_preparation_owns_nothing() {
5237 let coordinator =
5238 MockCoordinator::new(Assignor::CooperativeSticky, vec![vec![partition(0)]]);
5239 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5240 let mut config = start_config(
5241 mock.addr.to_string(),
5242 Assignor::CooperativeSticky,
5243 true,
5244 minutes(1),
5245 );
5246 config.heartbeat_interval = millis(50);
5247 let mut consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
5248 consumer.poll(millis(20)).await.expect("first poll");
5249 coordinator
5252 .commit_replies
5253 .lock()
5254 .expect("commit replies lock")
5255 .push_back(super::auto_commit_tests::CommitReply::Drop);
5256 coordinator.set_heartbeat_error(27);
5257 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
5258 let mut reset = false;
5259 while coordinator.joins.lock().expect("joins lock").is_empty()
5260 && tokio::time::Instant::now() < deadline
5261 {
5262 let _ = consumer.poll(millis(20)).await;
5263 let committing = coordinator
5264 .requests()
5265 .iter()
5266 .any(|request| matches!(request, GroupRequest::OffsetCommit(_)));
5267 if committing && !reset {
5268 coordinator.set_heartbeat_error(25);
5269 reset = true;
5270 }
5271 }
5272 let subscriptions: Vec<(Vec<(String, i32)>, i32)> = coordinator
5273 .joins
5274 .lock()
5275 .expect("joins lock")
5276 .iter()
5277 .flat_map(|join| join.protocols.iter())
5278 .map(|protocol| {
5279 let subscription = builder::decode_subscription(&protocol.metadata);
5280 (subscription.owned, subscription.generation_id)
5281 })
5282 .collect();
5283 drop(consumer);
5284 mock.stop();
5285 assert2::assert!(subscriptions.first() == Some(&(Vec::new(), -1)));
5286 }
5287
5288 #[tokio::test]
5292 async fn close_timeout_bounds_the_auto_commit() {
5293 let mut actual = Vec::new();
5294 let mut wanted = Vec::new();
5295 for (name, timeout, expected) in [("300 ms", millis(300), (Ok(()), true, 1))] {
5296 let coordinator = MockCoordinator::range_partition(0);
5297 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5298 let mut config = start_config(mock.addr.to_string(), Assignor::Range, true, minutes(1));
5299 config.heartbeat_interval = millis(50);
5300 let consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
5301 let held = Arc::clone(&consumer.commit_serialization)
5304 .lock_owned()
5305 .await;
5306 let started = tokio::time::Instant::now();
5307 let closed = tokio::time::timeout(
5308 Duration::from_secs(10),
5309 consumer.close_with(CloseOptions::timeout(timeout)),
5310 )
5311 .await
5312 .expect("close returns")
5313 .map_err(|error| error.to_string());
5314 let fast = started.elapsed() < Duration::from_secs(3);
5315 drop(held);
5316 mock.stop();
5317 let leaves = coordinator.leaves.lock().expect("leaves lock").len();
5318 actual.push((name, (closed, fast, leaves)));
5319 wanted.push((name, expected));
5320 }
5321 assert2::assert!(actual == wanted);
5322 }
5323
5324 #[tokio::test]
5327 async fn close_waits_for_a_pending_asynchronous_commit() {
5328 let coordinator = MockCoordinator::range_partition(0);
5329 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5330 let consumer = started_consumer(&mock, None).await;
5331 let held = Arc::clone(&consumer.commit_serialization)
5334 .lock_owned()
5335 .await;
5336 let release = tokio::spawn(async move {
5337 tokio::time::sleep(Duration::from_millis(200)).await;
5338 drop(held);
5339 });
5340 let called = Arc::new(std::sync::atomic::AtomicBool::new(false));
5341 let called_in_callback = Arc::clone(&called);
5342 consumer
5343 .commit_async_with_callback(move |_, _| {
5344 called_in_callback.store(true, Ordering::SeqCst);
5345 })
5346 .expect("commit_async_with_callback");
5347 consumer.close().await.expect("close");
5348 let called_before_close_returned = called.load(Ordering::SeqCst);
5349 release.await.expect("release");
5350 mock.stop();
5351 assert2::assert!(called_before_close_returned);
5352 }
5353
5354 #[derive(Clone, Copy, Debug)]
5356 enum Change {
5357 Subscribe(&'static str),
5358 Unsubscribe,
5359 UnsubscribeThenSubscribe,
5360 PatternWithoutMatch,
5361 }
5362
5363 type SubscriptionOutcome = (
5367 Vec<(Vec<String>, Option<String>)>,
5368 Vec<Option<String>>,
5369 Vec<String>,
5370 );
5371
5372 #[tokio::test]
5377 async fn subscription_changes_follow_kafka() {
5378 let changed = |from: &str, to: &str| {
5379 Some(format!(
5380 "the subscription has changed from [{from}] to [{to}] since the last group join"
5381 ))
5382 };
5383 let mut actual = Vec::new();
5384 let mut wanted = Vec::new();
5385 for (name, change, expected) in [
5386 (
5387 "subscribe to another topic",
5388 Change::Subscribe("payments"),
5389 (
5390 vec![(vec!["payments".to_owned()], changed("orders", "payments"))],
5391 vec![],
5392 vec![],
5393 ),
5394 ),
5395 (
5396 "subscribe to the same topic",
5397 Change::Subscribe("orders"),
5398 (vec![], vec![], vec![]),
5399 ),
5400 (
5401 "unsubscribe",
5402 Change::Unsubscribe,
5403 (
5404 vec![],
5405 vec![Some("the consumer unsubscribed from all topics".to_owned())],
5406 vec!["not subscribed to any topic".to_owned()],
5407 ),
5408 ),
5409 (
5410 "unsubscribe, then subscribe again",
5411 Change::UnsubscribeThenSubscribe,
5412 (
5413 vec![(
5414 vec!["orders".to_owned()],
5415 Some("consumer pro-actively leaving the group".to_owned()),
5416 )],
5417 vec![Some("the consumer unsubscribed from all topics".to_owned())],
5418 vec!["not subscribed to any topic".to_owned()],
5419 ),
5420 ),
5421 (
5422 "a pattern that matches no topic",
5423 Change::PatternWithoutMatch,
5424 (vec![(vec![], changed("orders", ""))], vec![], vec![]),
5425 ),
5426 ] {
5427 let StartedRangeConsumer {
5428 coordinator,
5429 mock,
5430 mut consumer,
5431 } = StartedRangeConsumer::start(None).await;
5432 let mut errors = Vec::new();
5433 match change {
5434 Change::Subscribe(topic) => consumer.subscribe([topic]).await.expect("subscribe"),
5435 Change::Unsubscribe => consumer.unsubscribe().await.expect("unsubscribe"),
5436 Change::UnsubscribeThenSubscribe => {
5437 consumer.unsubscribe().await.expect("unsubscribe");
5438 if let Err(error) = consumer.poll(millis(20)).await {
5439 errors.push(error.to_string());
5440 }
5441 consumer.subscribe(["orders"]).await.expect("subscribe");
5442 }
5443 Change::PatternWithoutMatch => {
5444 consumer
5445 .subscribe_pattern(TopicPattern::new(|_| false))
5446 .expect("subscribe_pattern");
5447 }
5448 }
5449 let deadline = tokio::time::Instant::now() + Duration::from_millis(500);
5450 while tokio::time::Instant::now() < deadline {
5451 if let Err(error) = consumer.poll(millis(20)).await
5452 && !errors.contains(&error.to_string())
5453 {
5454 errors.push(error.to_string());
5455 }
5456 if !coordinator.joins.lock().expect("joins lock").is_empty() {
5457 break;
5458 }
5459 tokio::time::sleep(Duration::from_millis(10)).await;
5460 }
5461 let joins = coordinator
5462 .joins
5463 .lock()
5464 .expect("joins lock")
5465 .iter()
5466 .map(|join| {
5467 let subscription = builder::decode_subscription(&join.protocols[0].metadata);
5468 (subscription.topics, join.reason.clone())
5469 })
5470 .collect();
5471 let leaves = coordinator
5472 .leaves
5473 .lock()
5474 .expect("leaves lock")
5475 .iter()
5476 .flat_map(|leave| leave.members.iter().map(|member| member.reason.clone()))
5477 .collect();
5478 drop(consumer);
5479 mock.stop();
5480 let outcome: SubscriptionOutcome = (joins, leaves, errors);
5481 actual.push((name, outcome));
5482 wanted.push((name, expected));
5483 }
5484 assert2::assert!(actual == wanted);
5485 }
5486
5487 #[tokio::test]
5490 async fn unsubscribe_gives_up_the_assignment_before_it_returns() {
5491 let coordinator = MockCoordinator::range_partition(0);
5492 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5493 let mut consumer = started_consumer(&mock, None).await;
5494 let before = consumer.assignment().await;
5495 consumer.coordinator_shutdown.cancel();
5498 if let Some(handle) = consumer.coordinator_handle.take() {
5499 handle.await.expect("coordinator task");
5500 }
5501 consumer.unsubscribe().await.expect("unsubscribe");
5502 let after = consumer.assignment().await;
5503 let owned = consumer.commit_identity.lock().await.ownership_ids.len();
5504 drop(consumer);
5505 mock.stop();
5506 assert2::assert!((before, after, owned) == (vec![partition(0)], vec![], 0));
5507 }
5508
5509 #[tokio::test]
5513 async fn a_manual_assignment_after_unsubscribe_survives_and_blocks_enforce_rebalance() {
5514 let coordinator = MockCoordinator::range_partition(0);
5515 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5516 let mut consumer = started_consumer(&mock, None).await;
5517 consumer.unsubscribe().await.expect("unsubscribe");
5518 consumer
5519 .assign(&[(TOPIC.to_owned(), 1)])
5520 .await
5521 .expect("assign");
5522 tokio::time::sleep(Duration::from_millis(300)).await;
5524 let assignment = consumer.assignment().await;
5525 let enforced = consumer
5526 .enforce_rebalance(None)
5527 .map_err(|error| error.to_string());
5528 let rejected = consumer
5529 .assign(&[(TOPIC.to_owned(), -1)])
5530 .await
5531 .map_err(|error| error.to_string());
5532 drop(consumer);
5533 mock.stop();
5534 assert2::assert!(
5535 (assignment, enforced, rejected)
5536 == (
5537 vec![(TOPIC.to_owned(), 1)],
5538 Err("illegal state: Tried to force a rebalance but the consumer has a manual assignment.".to_owned()),
5539 Err("invalid argument: Partition index of orders--1 is negative".to_owned())
5540 )
5541 );
5542 }
5543
5544 #[tokio::test]
5547 async fn enforce_rebalance_joins_on_the_next_poll_with_the_reason() {
5548 let mut actual = Vec::new();
5549 let mut wanted = Vec::new();
5550 for (name, reason, expected) in [
5551 ("no reason", None, "rebalance enforced by user"),
5552 ("empty reason", Some(""), "rebalance enforced by user"),
5553 ("reason", Some("scale out"), "scale out"),
5554 ] {
5555 let StartedRangeConsumer {
5556 coordinator,
5557 mock,
5558 mut consumer,
5559 } = StartedRangeConsumer::start(None).await;
5560 consumer
5561 .enforce_rebalance(reason)
5562 .expect("enforce_rebalance");
5563 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
5564 while coordinator.joins.lock().expect("joins lock").is_empty()
5565 && tokio::time::Instant::now() < deadline
5566 {
5567 consumer.poll(millis(20)).await.expect("poll");
5568 }
5569 let reasons: Vec<Option<String>> = coordinator
5570 .joins
5571 .lock()
5572 .expect("joins lock")
5573 .iter()
5574 .map(|join| join.reason.clone())
5575 .collect();
5576 drop(consumer);
5577 mock.stop();
5578 actual.push((name, reasons));
5579 wanted.push((name, vec![Some(expected.to_owned())]));
5580 }
5581 assert2::assert!(actual == wanted);
5582 }
5583
5584 #[tokio::test]
5588 async fn a_poll_that_waits_for_an_eager_join_starts_a_join_requested_after_it_began() {
5589 let coordinator = MockCoordinator::range_partition(0);
5590 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5591 let mut consumer = started_consumer(&mock, None).await;
5592 coordinator::note_poll(&consumer.poll_signal);
5594 coordinator.set_heartbeat_error(27);
5595 tokio::time::timeout(Duration::from_secs(5), async {
5596 while !*consumer.rebalance_pending.borrow() {
5597 tokio::time::sleep(Duration::from_millis(5)).await;
5598 }
5599 })
5600 .await
5601 .expect("the heartbeat requests a rebalance");
5602
5603 let started = tokio::time::Instant::now();
5604 let joined = consumer
5605 .wait_for_rebalance(started + Duration::from_secs(5))
5606 .await
5607 .expect("no listener");
5608 let within_timeout = started.elapsed() < Duration::from_secs(4);
5609 let joins = coordinator.joins.lock().expect("joins lock").len();
5610 drop(consumer);
5611 mock.stop();
5612 assert2::assert!((joined, within_timeout, joins) == (true, true, 1));
5613 }
5614
5615 #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
5619 async fn continuous_polls_do_not_stop_the_heartbeats() {
5620 let coordinator = MockCoordinator::range_partition(0);
5621 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5622 let consumer = started_consumer(&mock, None).await;
5623 let first_heartbeat = *coordinator
5624 .last_heartbeat
5625 .lock()
5626 .expect("last heartbeat lock");
5627 let poll_signal = consumer.poll_signal.clone();
5629 tokio::task::spawn_blocking(move || {
5630 let end = std::time::Instant::now() + Duration::from_millis(500);
5631 while std::time::Instant::now() < end {
5632 coordinator::note_poll(&poll_signal);
5633 }
5634 })
5635 .await
5636 .expect("poller");
5637 let heartbeat_during_polls = *coordinator
5638 .last_heartbeat
5639 .lock()
5640 .expect("last heartbeat lock")
5641 > first_heartbeat + Duration::from_millis(200);
5642 drop(consumer);
5643 mock.stop();
5644 assert2::assert!(heartbeat_during_polls);
5645 }
5646
5647 #[tokio::test]
5651 async fn the_poll_timeout_leave_goes_out_before_the_join_of_the_next_poll() {
5652 let coordinator = MockCoordinator::range_partition(0);
5653 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5654 let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, millis(200));
5655 config.heartbeat_interval = millis(50);
5656 let consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
5657 tokio::time::timeout(Duration::from_secs(5), async {
5659 while !*consumer.rebalance_pending.borrow() {
5660 tokio::task::yield_now().await;
5661 }
5662 })
5663 .await
5664 .expect("the poll timeout expires");
5665 coordinator::note_poll(&consumer.poll_signal);
5666 tokio::time::timeout(Duration::from_secs(5), async {
5667 while coordinator.sync_groups() == 0 {
5668 tokio::time::sleep(Duration::from_millis(10)).await;
5669 }
5670 })
5671 .await
5672 .expect("the member joins again");
5673 let order: Vec<GroupRequest> = coordinator
5674 .requests()
5675 .into_iter()
5676 .filter(|request| matches!(request, GroupRequest::LeaveGroup | GroupRequest::JoinGroup))
5677 .take(2)
5678 .collect();
5679 drop(consumer);
5680 mock.stop();
5681 assert2::assert!(order == vec![GroupRequest::LeaveGroup, GroupRequest::JoinGroup]);
5682 }
5683
5684 #[tokio::test]
5689 async fn rejoin_sends_a_subscription_per_assignor_with_the_owned_partitions_of_the_protocol() {
5690 use builder::{DecodedSubscription, decode_subscription};
5691 let subscription =
5692 |owned: Vec<(String, i32)>, user_data: Option<Bytes>| DecodedSubscription {
5693 topics: vec![TOPIC.to_owned()],
5694 owned,
5695 generation_id: 1,
5696 rack_id: None,
5697 user_data,
5698 };
5699 let generation_one = Some(Bytes::from_static(&[0, 0, 0, 1]));
5700 let mut actual = Vec::new();
5701 let mut wanted = Vec::new();
5702 for (name, assignors, expected) in [
5703 (
5704 "default list, eager",
5705 Assignor::DEFAULT_LIST.to_vec(),
5706 vec![
5707 ("range".to_owned(), subscription(Vec::new(), None)),
5708 (
5709 "cooperative-sticky".to_owned(),
5710 subscription(Vec::new(), generation_one.clone()),
5711 ),
5712 ],
5713 ),
5714 (
5715 "cooperative sticky, cooperative",
5716 vec![Assignor::CooperativeSticky],
5717 vec![(
5718 "cooperative-sticky".to_owned(),
5719 subscription(vec![partition(0)], generation_one),
5720 )],
5721 ),
5722 (
5723 "sticky, eager",
5724 vec![Assignor::Sticky],
5725 vec![(
5726 "sticky".to_owned(),
5727 subscription(
5728 Vec::new(),
5729 assignor::subscription_user_data(
5730 Assignor::Sticky,
5731 Some(&[partition(0)]),
5732 1,
5733 ),
5734 ),
5735 )],
5736 ),
5737 ] {
5738 let coordinator = MockCoordinator::range_partition(0);
5739 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5740 let consumer = started_consumer_with(&mock, None, &assignors).await;
5741 coordinator.set_heartbeat_error(27);
5742 let protocols = tokio::time::timeout(Duration::from_secs(10), async {
5743 loop {
5744 coordinator::note_poll(&consumer.poll_signal);
5745 if let Some(join) = coordinator.joins.lock().expect("joins lock").first() {
5746 break join
5747 .protocols
5748 .iter()
5749 .map(|protocol| {
5750 (
5751 protocol.name.clone(),
5752 decode_subscription(&protocol.metadata),
5753 )
5754 })
5755 .collect::<Vec<_>>();
5756 }
5757 tokio::time::sleep(Duration::from_millis(20)).await;
5758 }
5759 })
5760 .await
5761 .expect("the consumer joins again");
5762 drop(consumer);
5763 mock.stop();
5764 actual.push((name, protocols));
5765 wanted.push((name, expected));
5766 }
5767 assert2::assert!(actual == wanted);
5768 }
5769
5770 #[tokio::test]
5775 async fn join_group_carries_the_rejoin_reason() {
5776 let mut actual = Vec::new();
5777 let mut wanted = Vec::new();
5778 for (name, heartbeat_error, expected) in [
5779 (
5780 "rebalance in progress",
5781 27,
5782 ("member-a", "group is already rebalancing"),
5783 ),
5784 (
5785 "illegal generation",
5786 22,
5787 (
5788 "member-a",
5789 "encountered ILLEGAL_GENERATION from HEARTBEAT response",
5790 ),
5791 ),
5792 (
5793 "unknown member id",
5794 25,
5795 ("", "encountered UNKNOWN_MEMBER_ID from HEARTBEAT response"),
5796 ),
5797 ] {
5798 let fixture = StartedRangeConsumer::start(None).await;
5799 fixture
5800 .coordinator
5801 .heartbeat_error
5802 .store(heartbeat_error, Ordering::SeqCst);
5803 let join = tokio::time::timeout(Duration::from_secs(10), async {
5804 loop {
5805 coordinator::note_poll(&fixture.consumer.poll_signal);
5806 if let Some(join) = fixture
5807 .coordinator
5808 .joins
5809 .lock()
5810 .expect("joins lock")
5811 .first()
5812 {
5813 break (join.member_id.clone(), join.reason.clone());
5814 }
5815 tokio::time::sleep(Duration::from_millis(20)).await;
5816 }
5817 })
5818 .await
5819 .expect("the consumer joins again");
5820 drop(fixture.consumer);
5821 fixture.mock.stop();
5822 actual.push((name, join));
5823 wanted.push((name, (expected.0.to_owned(), Some(expected.1.to_owned()))));
5824 }
5825 assert2::assert!(actual == wanted);
5826 }
5827}
5828
5829#[cfg(test)]
5830mod rebalance_listener_tests {
5831 use std::sync::atomic::Ordering;
5832
5833 use super::{
5834 auto_commit_tests::{GroupRequest, MockCoordinator, TOPIC, partition, start_config},
5835 *,
5836 };
5837 use crate::{
5838 ConsumerRebalanceListener, RebalanceListenerError,
5839 test_support::{ListenerStep as Step, RecordingListener},
5840 };
5841
5842 #[derive(Clone, Debug, PartialEq)]
5844 enum Call {
5845 Revoked(Vec<(String, i32)>),
5846 Assigned(Vec<(String, i32)>),
5847 Lost(Vec<(String, i32)>),
5848 CommitInRevoke(Result<(), String>),
5849 AwaitingCallback(Vec<(String, i32)>),
5852 PollError(String),
5854 Position(Option<i64>),
5856 }
5857
5858 struct Recorder {
5859 calls: Arc<std::sync::Mutex<Vec<Call>>>,
5860 commit_on_revoke: bool,
5861 fail_revoke: bool,
5862 seek_on_assign: Option<i64>,
5863 }
5864
5865 impl Recorder {
5866 fn record(&self, call: Call) {
5867 self.calls.lock().expect("calls lock").push(call);
5868 }
5869 }
5870
5871 #[async_trait::async_trait]
5872 impl ConsumerRebalanceListener for Recorder {
5873 async fn on_partitions_revoked(
5874 &mut self,
5875 consumer: &Consumer,
5876 partitions: &[(String, i32)],
5877 ) -> Result<(), RebalanceListenerError> {
5878 self.record(Call::Revoked(partitions.to_vec()));
5879 if self.commit_on_revoke {
5880 let result = consumer
5881 .commit_sync()
5882 .await
5883 .map_err(|error| error.to_string());
5884 self.record(Call::CommitInRevoke(result));
5885 }
5886 if self.fail_revoke {
5887 return Err("revoke failed".into());
5888 }
5889 Ok(())
5890 }
5891
5892 async fn on_partitions_assigned(
5893 &mut self,
5894 consumer_for_assignment: &Consumer,
5895 assigned_partitions: &[(String, i32)],
5896 ) -> Result<(), RebalanceListenerError> {
5897 self.record(Call::Assigned(assigned_partitions.to_vec()));
5898 let mut awaiting: Vec<_> = consumer_for_assignment
5899 .assigned_callback_pending
5900 .lock()
5901 .expect("pending lock")
5902 .iter()
5903 .cloned()
5904 .collect();
5905 awaiting.sort();
5906 if !assigned_partitions.is_empty() {
5907 self.record(Call::AwaitingCallback(awaiting));
5908 }
5909 if let (Some(offset), Some(partition)) =
5910 (self.seek_on_assign, assigned_partitions.first())
5911 {
5912 consumer_for_assignment
5913 .seek(partition.0.clone(), partition.1, offset)
5914 .await
5915 .expect("seek");
5916 }
5917 Ok(())
5918 }
5919
5920 async fn on_partitions_lost(
5921 &mut self,
5922 _consumer: &Consumer,
5923 partitions: &[(String, i32)],
5924 ) -> Result<(), RebalanceListenerError> {
5925 self.record(Call::Lost(partitions.to_vec()));
5926 Ok(())
5927 }
5928 }
5929
5930 #[derive(Clone, Copy, Debug)]
5932 enum Event {
5933 Heartbeat(i16),
5935 Close,
5937 Subscribe(&'static str),
5939 }
5940
5941 struct Case {
5942 assignor: Assignor,
5943 owned: Vec<(String, i32)>,
5944 assignments: Vec<Vec<(String, i32)>>,
5946 event: Event,
5947 commit_on_revoke: bool,
5948 fail_revoke: bool,
5949 seek_on_assign: Option<i64>,
5950 calls: usize,
5952 }
5953
5954 #[derive(Debug, PartialEq)]
5957 struct CaseResult {
5958 calls: Vec<Call>,
5959 commit_generations: Vec<i32>,
5960 }
5961
5962 #[derive(bon::Builder)]
5963 struct ListenerScenario {
5964 name: &'static str,
5965 case: Case,
5966 expected_calls: Vec<Call>,
5967 expected_commit_generations: Vec<i32>,
5968 }
5969
5970 async fn run_case(case: Case) -> CaseResult {
5971 let coordinator = MockCoordinator::new(case.assignor, case.assignments);
5972 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
5973 let mut config = start_config(mock.addr.to_string(), case.assignor, false, minutes(1));
5974 config.heartbeat_interval = millis(50);
5975 config.has_rebalance_listener = true;
5976 let mut consumer = super::auto_commit_tests::start_consumer_with_startup(
5977 config,
5978 StartupState::builder()
5979 .generation_id(1)
5980 .assigned_partitions(case.owned.clone())
5981 .next_offsets(case.owned.iter().map(|key| (key.clone(), 12)).collect())
5982 .topic_partitions(maplit::hashmap! {TOPIC.to_owned() => 2})
5983 .build(),
5984 )
5985 .await;
5986 let calls = Arc::new(std::sync::Mutex::new(Vec::new()));
5987 consumer.rebalance_listener = Some(rebalance_listener::shared(Box::new(Recorder {
5988 calls: Arc::clone(&calls),
5989 commit_on_revoke: case.commit_on_revoke,
5990 fail_revoke: case.fail_revoke,
5991 seek_on_assign: case.seek_on_assign,
5992 })));
5993
5994 consumer.poll(millis(20)).await.expect("first poll");
5996 if case.seek_on_assign.is_some() {
5997 let position = consumer
5998 .next_offsets
5999 .lock()
6000 .await
6001 .get(&partition(0))
6002 .copied();
6003 calls
6004 .lock()
6005 .expect("calls lock")
6006 .push(Call::Position(position));
6007 }
6008 match case.event {
6009 event @ (Event::Heartbeat(_) | Event::Subscribe(_)) => {
6010 match event {
6011 Event::Heartbeat(error_code) => coordinator
6012 .heartbeat_error
6013 .store(error_code, Ordering::SeqCst),
6014 Event::Subscribe(topic) => {
6015 consumer.subscribe([topic]).await.expect("subscribe");
6016 }
6017 Event::Close => {}
6018 }
6019 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
6020 while calls.lock().expect("calls lock").len() < case.calls
6021 && tokio::time::Instant::now() < deadline
6022 {
6023 poll_recording(&mut consumer, &calls).await;
6024 }
6025 for _ in 0..10 {
6027 poll_recording(&mut consumer, &calls).await;
6028 }
6029 drop(consumer);
6030 }
6031 Event::Close => consumer.close().await.expect("close"),
6032 }
6033 let generations = coordinator
6034 .requests()
6035 .into_iter()
6036 .filter_map(|request| match request {
6037 GroupRequest::OffsetCommit(commit) => Some(commit.generation_id_or_member_epoch),
6038 _ => None,
6039 })
6040 .collect();
6041 mock.stop();
6042 let calls = calls.lock().expect("calls lock").clone();
6043 CaseResult {
6044 calls,
6045 commit_generations: generations,
6046 }
6047 }
6048
6049 async fn assert_listener_scenarios(scenarios: impl IntoIterator<Item = ListenerScenario>) {
6050 for scenario in scenarios {
6051 let expected = CaseResult {
6052 calls: scenario.expected_calls,
6053 commit_generations: scenario.expected_commit_generations,
6054 };
6055 assert2::assert!(
6056 run_case(scenario.case).await == expected,
6057 "{}",
6058 scenario.name
6059 );
6060 }
6061 }
6062
6063 async fn poll_recording(consumer: &mut Consumer, calls: &std::sync::Mutex<Vec<Call>>) {
6065 if let Err(error) = consumer.poll(millis(20)).await {
6066 calls
6067 .lock()
6068 .expect("calls lock")
6069 .push(Call::PollError(error.to_string()));
6070 }
6071 }
6072
6073 #[tokio::test]
6077 async fn a_cancelled_poll_keeps_the_listener_and_runs_the_callback_again() {
6078 let coordinator = MockCoordinator::new(
6079 Assignor::CooperativeSticky,
6080 vec![vec![partition(0), partition(1)]],
6081 );
6082 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
6083 let mut config = start_config(
6084 mock.addr.to_string(),
6085 Assignor::CooperativeSticky,
6086 false,
6087 minutes(1),
6088 );
6089 config.heartbeat_interval = millis(50);
6090 config.has_rebalance_listener = true;
6091 let mut consumer = super::auto_commit_tests::start_consumer_with_startup(
6092 config,
6093 super::auto_commit_tests::single_partition_startup(2),
6094 )
6095 .await;
6096 let steps = Arc::new(std::sync::Mutex::new(Vec::new()));
6097 let listener = RecordingListener::builder()
6098 .steps(Arc::clone(&steps))
6099 .slow_partition(1)
6100 .slow_for(Duration::from_millis(500))
6101 .build();
6102 consumer.rebalance_listener = Some(rebalance_listener::shared(Box::new(listener)));
6103 consumer.poll(millis(20)).await.expect("first poll");
6104 coordinator.set_heartbeat_error(27);
6105 let ended = |steps: &std::sync::Mutex<Vec<Step>>| {
6106 steps
6107 .lock()
6108 .expect("steps lock")
6109 .contains(&Step::Ended("assigned", vec![partition(1)]))
6110 };
6111 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
6112 while !ended(&steps) && tokio::time::Instant::now() < deadline {
6113 let _ =
6115 tokio::time::timeout(Duration::from_millis(100), consumer.poll(millis(20))).await;
6116 }
6117 let awaiting = consumer
6118 .assigned_callback_pending
6119 .lock()
6120 .expect("pending lock")
6121 .len();
6122 consumer.close().await.expect("close");
6123 mock.stop();
6124 let steps = steps.lock().expect("steps lock").clone();
6125 assert2::assert!(
6126 (steps, awaiting)
6127 == (
6128 vec![
6129 Step::Started("assigned", vec![partition(0)]),
6130 Step::Ended("assigned", vec![partition(0)]),
6131 Step::Started("assigned", vec![partition(1)]),
6132 Step::Started("assigned", vec![partition(1)]),
6133 Step::Ended("assigned", vec![partition(1)]),
6134 Step::Started("revoked", vec![partition(0), partition(1)]),
6135 Step::Ended("revoked", vec![partition(0), partition(1)]),
6136 ],
6137 0
6138 )
6139 );
6140 }
6141
6142 #[tokio::test]
6146 async fn poll_timeout_during_a_callback_leaves_the_group() {
6147 let coordinator = MockCoordinator::new(
6148 Assignor::CooperativeSticky,
6149 vec![vec![partition(0), partition(1)]],
6150 );
6151 let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
6152 let mut config = start_config(
6153 mock.addr.to_string(),
6154 Assignor::CooperativeSticky,
6155 false,
6156 millis(400),
6157 );
6158 config.heartbeat_interval = millis(50);
6159 config.has_rebalance_listener = true;
6160 let mut consumer = start_test_consumer(
6161 TestConsumerStart::builder()
6162 .bootstrap(mock.addr.to_string())
6163 .config(config)
6164 .startup(super::auto_commit_tests::single_partition_startup(2))
6165 .build(),
6166 )
6167 .await;
6168 let listener = RecordingListener::builder()
6169 .slow_partition(1)
6170 .slow_for(Duration::from_millis(1200))
6171 .build();
6172 consumer.rebalance_listener = Some(rebalance_listener::shared(Box::new(listener)));
6173 consumer.poll(millis(20)).await.expect("first poll");
6174 coordinator.set_heartbeat_error(27);
6175 let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
6177 while coordinator.leaves.lock().expect("leaves lock").is_empty()
6178 && tokio::time::Instant::now() < deadline
6179 {
6180 let _ = consumer.poll(millis(20)).await;
6181 }
6182 let reasons: Vec<Option<String>> = coordinator
6183 .leaves
6184 .lock()
6185 .expect("leaves lock")
6186 .iter()
6187 .flat_map(|leave| leave.members.iter().map(|member| member.reason.clone()))
6188 .collect();
6189 drop(consumer);
6190 mock.stop();
6191 assert2::assert!(reasons == vec![Some("consumer poll timeout has expired.".to_owned())]);
6192 }
6193
6194 fn case(
6196 assignor: Assignor,
6197 owned: Vec<(String, i32)>,
6198 assignment: Vec<(String, i32)>,
6199 event: Event,
6200 calls: usize,
6201 ) -> Case {
6202 Case {
6203 assignor,
6204 owned,
6205 assignments: vec![assignment],
6206 event,
6207 commit_on_revoke: false,
6208 fail_revoke: false,
6209 seek_on_assign: None,
6210 calls,
6211 }
6212 }
6213
6214 fn awaiting(partitions: Vec<(String, i32)>) -> Call {
6215 Call::AwaitingCallback(partitions)
6216 }
6217
6218 fn empty_assignment_call() -> Call {
6219 Call::Assigned(Vec::new())
6220 }
6221
6222 fn range_revoke_calls() -> Vec<Call> {
6223 vec![
6224 Call::Assigned(vec![partition(0)]),
6225 awaiting(vec![]),
6226 Call::Revoked(vec![partition(0)]),
6227 ]
6228 }
6229
6230 fn range_revoke_commit_calls() -> Vec<Call> {
6231 let mut calls = range_revoke_calls();
6232 calls.extend([
6233 Call::CommitInRevoke(Ok(())),
6234 Call::Assigned(vec![partition(0)]),
6235 awaiting(vec![partition(0)]),
6236 ]);
6237 calls
6238 }
6239
6240 #[tokio::test]
6241 async fn listener_calls_follow_kafkas_order() {
6242 let p = partition;
6243 Box::pin(assert_listener_scenarios([
6244 ListenerScenario::builder()
6245 .name("eager rebalance")
6246 .case(case(
6247 Assignor::Range,
6248 vec![p(0), p(1)],
6249 vec![p(0)],
6250 Event::Heartbeat(27),
6251 5,
6252 ))
6253 .expected_calls(vec![
6254 Call::Assigned(vec![p(0), p(1)]),
6255 awaiting(vec![]),
6256 Call::Revoked(vec![p(0), p(1)]),
6257 Call::Assigned(vec![p(0)]),
6258 awaiting(vec![p(0)]),
6259 ])
6260 .expected_commit_generations(vec![])
6261 .build(),
6262 ListenerScenario::builder()
6263 .name("cooperative rebalance that revokes")
6264 .case(case(
6265 Assignor::CooperativeSticky,
6266 vec![p(0), p(1)],
6267 vec![p(0)],
6268 Event::Heartbeat(27),
6269 5,
6270 ))
6271 .expected_calls(vec![
6272 Call::Assigned(vec![p(0), p(1)]),
6273 awaiting(vec![]),
6274 Call::Revoked(vec![p(1)]),
6275 empty_assignment_call(),
6276 empty_assignment_call(),
6277 ])
6278 .expected_commit_generations(vec![])
6279 .build(),
6280 ListenerScenario::builder()
6281 .name("cooperative rebalance that adds")
6282 .case(case(
6283 Assignor::CooperativeSticky,
6284 vec![p(0)],
6285 vec![p(0), p(1)],
6286 Event::Heartbeat(27),
6287 4,
6288 ))
6289 .expected_calls(vec![
6290 Call::Assigned(vec![p(0)]),
6291 awaiting(vec![]),
6292 Call::Assigned(vec![p(1)]),
6293 awaiting(vec![p(1)]),
6294 ])
6295 .expected_commit_generations(vec![])
6296 .build(),
6297 ListenerScenario::builder()
6298 .name("unknown member id")
6299 .case(case(
6300 Assignor::Range,
6301 vec![p(0)],
6302 vec![p(0)],
6303 Event::Heartbeat(25),
6304 5,
6305 ))
6306 .expected_calls(vec![
6307 Call::Assigned(vec![p(0)]),
6308 awaiting(vec![]),
6309 Call::Lost(vec![p(0)]),
6310 Call::Assigned(vec![p(0)]),
6311 awaiting(vec![p(0)]),
6312 ])
6313 .expected_commit_generations(vec![])
6314 .build(),
6315 ListenerScenario::builder()
6316 .name("close")
6317 .case(case(
6318 Assignor::Range,
6319 vec![p(0)],
6320 vec![p(0)],
6321 Event::Close,
6322 3,
6323 ))
6324 .expected_calls(range_revoke_calls())
6325 .expected_commit_generations(vec![])
6326 .build(),
6327 ListenerScenario::builder()
6328 .name("a commit in the eager revoke callback")
6329 .case(Case {
6330 commit_on_revoke: true,
6331 ..case(
6332 Assignor::Range,
6333 vec![p(0)],
6334 vec![p(0)],
6335 Event::Heartbeat(27),
6336 6,
6337 )
6338 })
6339 .expected_calls(range_revoke_commit_calls())
6340 .expected_commit_generations(vec![1])
6341 .build(),
6342 ]))
6343 .await;
6344 }
6345
6346 #[tokio::test]
6349 async fn listener_calls_follow_kafka_around_commits_errors_and_seeks() {
6350 let p = partition;
6351 assert_listener_scenarios([
6352 ListenerScenario::builder()
6353 .name("a commit in the cooperative revoke callback carries the new generation")
6354 .case(Case {
6355 commit_on_revoke: true,
6356 ..case(
6357 Assignor::CooperativeSticky,
6358 vec![p(0), p(1)],
6359 vec![p(0)],
6360 Event::Heartbeat(27),
6361 6,
6362 )
6363 })
6364 .expected_calls(vec![
6365 Call::Assigned(vec![p(0), p(1)]),
6366 awaiting(vec![]),
6367 Call::Revoked(vec![p(1)]),
6368 Call::CommitInRevoke(Ok(())),
6369 empty_assignment_call(),
6370 empty_assignment_call(),
6371 ])
6372 .expected_commit_generations(vec![2])
6373 .build(),
6374 ListenerScenario::builder()
6375 .name("a failed cooperative revoke callback fails the poll after the rebalance")
6376 .case(Case {
6377 fail_revoke: true,
6378 ..case(
6379 Assignor::CooperativeSticky,
6380 vec![p(0), p(1)],
6381 vec![p(0)],
6382 Event::Heartbeat(27),
6383 6,
6384 )
6385 })
6386 .expected_calls(vec![
6387 Call::Assigned(vec![p(0), p(1)]),
6388 awaiting(vec![]),
6389 Call::Revoked(vec![p(1)]),
6390 empty_assignment_call(),
6391 empty_assignment_call(),
6392 Call::PollError("rebalance listener failed: revoke failed".into()),
6393 ])
6394 .expected_commit_generations(vec![])
6395 .build(),
6396 ListenerScenario::builder()
6397 .name("a cooperative join after a subscription change revokes the partitions of the old topic")
6398 .case(case(
6399 Assignor::CooperativeSticky,
6400 vec![p(0), p(1)],
6401 vec![],
6402 Event::Subscribe("payments"),
6403 4,
6404 ))
6405 .expected_calls(vec![
6406 Call::Assigned(vec![p(0), p(1)]),
6407 awaiting(vec![]),
6408 Call::Revoked(vec![p(0), p(1)]),
6409 empty_assignment_call(),
6410 ])
6411 .expected_commit_generations(vec![])
6412 .build(),
6413 ListenerScenario::builder()
6414 .name("a seek in the assign callback moves the position of the same poll")
6415 .case(Case {
6416 seek_on_assign: Some(3),
6417 ..case(Assignor::Range, vec![p(0)], vec![p(0)], Event::Close, 4)
6418 })
6419 .expected_calls(vec![
6420 Call::Assigned(vec![p(0)]),
6421 awaiting(vec![]),
6422 Call::Position(Some(3)),
6423 Call::Revoked(vec![p(0)]),
6424 ])
6425 .expected_commit_generations(vec![])
6426 .build(),
6427 ])
6428 .await;
6429 }
6430}
6431
6432#[cfg(test)]
6433mod transport_configuration_tests {
6434 use assert2::assert;
6435 use bytes::BytesMut;
6436 use krabka_client_core::MockBroker;
6437 use krabka_protocol::{
6438 Encode,
6439 owned::{api_versions_request, api_versions_response::ApiVersionsResponse},
6440 };
6441 use krabka_units::{kibibytes, millis};
6442
6443 use super::{ClientError, Consumer};
6444
6445 #[tokio::test]
6446 async fn consumer_retains_dns_and_socket_overrides_in_its_client() {
6447 let mock = MockBroker::start(|api, _, _, _| {
6448 if api != api_versions_request::API_KEY {
6449 return None;
6450 }
6451 let mut body = BytesMut::new();
6452 ApiVersionsResponse::default()
6453 .encode(&mut body, 0)
6454 .expect("encode ApiVersions");
6455 Some(body.to_vec())
6456 })
6457 .await;
6458 let consumer = Consumer::builder()
6459 .bootstrap(mock.addr.to_string())
6460 .request_timeout(millis(10))
6461 .dns_timeout(millis(95))
6462 .send_buffer(None)
6463 .receive_buffer(Some(kibibytes(19)))
6464 .enable_metrics_push(false)
6465 .build()
6466 .await
6467 .expect("build consumer");
6468 let options = consumer.client.connection_options();
6469 assert!(
6470 (
6471 options.dns_timeout.time(),
6472 options.send_buffer,
6473 options.receive_buffer
6474 ) == (millis(95), None, Some(kibibytes(19)))
6475 );
6476 assert!(consumer.client.refresh_metadata().await.is_err());
6477 consumer.close().await.expect("close consumer");
6478 mock.stop();
6479 }
6480
6481 #[tokio::test(start_paused = true)]
6482 async fn consumer_total_startup_deadline_bounds_attempts_and_retry_sleeps() {
6483 use krabka_units::{convert::TimeExt as _, secs};
6484
6485 use super::ConsumerRetryPolicy;
6486 for backoff in [millis(100), secs(1), secs(2)] {
6487 let mock = MockBroker::start(|_, _, _, _| None).await;
6488 let policy = ConsumerRetryPolicy::builder()
6489 .startup_attempt_timeout(secs(2))
6490 .startup_deadline(secs(3))
6491 .startup_initial_backoff(backoff)
6492 .startup_max_backoff(backoff)
6493 .coordinator_retry_timeout(secs(1))
6494 .coordinator_initial_backoff(millis(100))
6495 .coordinator_max_backoff(secs(1))
6496 .build()
6497 .expect("policy");
6498 let started = tokio::time::Instant::now();
6499 let result = Consumer::builder()
6500 .bootstrap(mock.addr.to_string())
6501 .group_id("startup-guard")
6502 .subscribe(["startup-topic".to_owned()])
6503 .retry_policy(policy)
6504 .enable_metrics_push(false)
6505 .build()
6506 .await;
6507 assert!(matches!(result, Err(super::ConsumerError::Client(
6508 ClientError::Timeout(deadline)
6509 )) if deadline == secs(3)));
6510 assert!(started.elapsed() == secs(3).to_std());
6511 mock.stop();
6512 }
6513 }
6514
6515 #[tokio::test]
6516 async fn consumer_rejects_invalid_dns_before_retrying_bootstrap() {
6517 let result = Consumer::builder()
6518 .bootstrap("127.0.0.1:1")
6519 .dns_timeout(millis(0))
6520 .build()
6521 .await;
6522 assert!(matches!(
6523 result,
6524 Err(super::ConsumerError::InvalidConfig(_))
6525 ));
6526 }
6527}