use std::{
collections::HashMap,
sync::{
Arc,
atomic::{AtomicI32, AtomicU8, Ordering},
},
time::Duration,
};
use bytes::Bytes;
use krabka_client_core::{
Client, ClientError, ClientFrameMax, ConnectionDispatchQueueCapacity, FetchMinBytes, security,
telemetry,
};
use krabka_protocol::{
owned::{
join_group_response::JoinGroupResponse, sync_group_request::SyncGroupRequestAssignment,
sync_group_response::SyncGroupResponse,
},
primitives::uuid::Uuid as WireUuid,
records,
};
use krabka_units::{
ByteSize, Time,
convert::{ByteSizeExt as _, StdDurationExt as _, TimeExt as _},
millis, minutes, secs,
};
use refined_type::{
Refined,
rule::{GreaterI32, GreaterI64, MinMaxI64, MinMaxU128},
};
use tokio::{
sync::{Mutex, Notify},
task::JoinHandle,
};
use tokio_util::sync::CancellationToken;
#[cfg(test)]
use crate::TopicPattern;
#[cfg(test)]
use crate::builder;
#[cfg(test)]
use crate::test_support;
use crate::{
CloseOptions, GroupProtocol, assignor,
assignor::Assignor,
builder::{AutoOffsetReset, IsolationLevel, decode_assignment},
commit, consumer_group, control, coordinator,
coordinator::{
CoordinatorRefind, CoordinatorRetryPolicy, CoordinatorRetryRequest, CoordinatorState,
find_coordinator,
},
error::ConsumerError,
fetch_buffer,
group_metadata::ConsumerGroupMetadata,
offset_wire, poll, position, rebalance_listener, subscription,
};
pub(crate) fn validated_positive_i32(value: i32, name: &str) -> Result<i32, String> {
GreaterI32::<0>::new(value)
.map(Refined::into_value)
.map_err(|error| format!("{name}: {error}"))
}
pub(crate) fn protocol_bytes(name: &str, value: ByteSize) -> Result<i32, String> {
let bytes = value.bytes_f64();
if bytes.is_finite() && bytes.fract() == 0.0 && (1.0..=f64::from(i32::MAX)).contains(&bytes) {
Ok(value.bytes_i32())
} else {
Err(format!(
"{name} must be a positive whole-byte value that fits i32"
))
}
}
fn validated_millisecond_duration(name: &str, value: Duration) -> Result<u64, String> {
let milliseconds = MinMaxU128::<1, { u64::MAX as u128 }>::new(value.as_millis())
.map_err(|error| format!("{name}: {error}"))?
.into_value();
let milliseconds = u64::try_from(milliseconds).map_err(|error| format!("{name}: {error}"))?;
if Duration::from_millis(milliseconds) == value {
Ok(milliseconds)
} else {
Err(format!("{name} must be a whole number of milliseconds"))
}
}
pub(crate) async fn client_instance_id(
client: &Client,
timeout: Time,
) -> Result<Option<uuid::Uuid>, ConsumerError> {
Ok(client.client_instance_id(timeout).await?)
}
#[allow(dead_code)] pub struct Consumer {
pub(crate) client: Client,
pub(crate) group_id: String,
pub(crate) coordinator_id: Arc<AtomicI32>,
pub(crate) retry_policy: CoordinatorRetryPolicy,
pub(crate) member_id: tokio::sync::watch::Receiver<String>,
pub(crate) commit_identity: Arc<Mutex<CommitIdentity>>,
pub(crate) commit_serialization: Arc<Mutex<()>>,
pub(crate) commit_async_state: Arc<AtomicU8>,
pub(crate) commit_async_callbacks: commit::OffsetCommitCallbacks,
pub(crate) group_instance_id: Option<String>,
pub(crate) current_generation: Arc<AtomicI32>,
pub(crate) subscription: subscription::SharedSubscription,
pub(crate) group_protocol: crate::GroupProtocol,
pub(crate) assigned: Arc<Mutex<Vec<(String, i32)>>>,
pub(crate) assignment_changed: Arc<Notify>,
pub(crate) next_offsets: Arc<Mutex<HashMap<(String, i32), i64>>>,
pub(crate) end_offsets: Arc<Mutex<HashMap<(String, i32), i64>>>,
pub(crate) positions: Arc<Mutex<HashMap<(String, i32), position::PartitionPosition>>>,
pub(crate) topic_ids: Arc<Mutex<HashMap<String, WireUuid>>>,
pub(crate) session_timeout: Time,
pub(crate) heartbeat_interval: Time,
pub(crate) rebalance_protocol: assignor::RebalanceProtocol,
pub(crate) coordinator_shutdown: CancellationToken,
pub(crate) coordinator_handle: Option<JoinHandle<()>>,
pub(crate) isolation_level: IsolationLevel,
pub(crate) fetch_min: ByteSize,
pub(crate) fetch_max: ByteSize,
pub(crate) fetch_partition_max: ByteSize,
pub(crate) fetch_max_wait: Time,
pub(crate) fetches: poll::Fetches,
pub(crate) client_rack: Option<String>,
pub(crate) metadata_max_age: Time,
pub(crate) request_timeout: Time,
pub(crate) wakeup: control::WakeupHandle,
pub(crate) enforced_rebalances: tokio::sync::mpsc::UnboundedSender<(String, u64)>,
pub(crate) default_api_timeout: Time,
pub(crate) paused: crate::partition_state::PausedPartitions,
pub(crate) auto_offset_reset: AutoOffsetReset,
pub(crate) poll_error: coordinator::PollErrorSlot,
pub(crate) auto_commit: Option<commit::AutoCommit>,
pub(crate) poll_signal: coordinator::PollSignal,
pub(crate) rebalance_pending: tokio::sync::watch::Receiver<bool>,
pub(crate) max_poll_records: usize,
pub(crate) fetch_buffer: fetch_buffer::FetchBuffer,
pub(crate) close_operation: tokio::sync::watch::Sender<control::CloseRequest>,
pub(crate) rebalance_listener: Option<rebalance_listener::SharedListener>,
pub(crate) listener_calls:
tokio::sync::mpsc::UnboundedReceiver<rebalance_listener::ListenerCall>,
pub(crate) assigned_callback_pending: rebalance_listener::AssignedCallbackPending,
}
#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
pub enum GroupMembershipOperation {
LeaveGroup,
RemainInGroup,
#[default]
Default,
}
#[derive(Clone)]
pub(crate) struct CommitIdentity {
pub generation: i32,
pub member_id: String,
pub ownership_ids: HashMap<(String, i32), u64>,
pub rejoin_on_poll: bool,
}
#[derive(Clone)]
struct StartConfig {
bootstrap: String,
client_id: String,
group_id: String,
session_timeout: Time,
max_poll_interval: Time,
max_poll_records: usize,
heartbeat_interval: Time,
subscription_metadata_refresh_interval: Time,
subscribe: Vec<String>,
subscribe_pattern: Option<crate::TopicPattern>,
subscribe_regex: Option<String>,
exclude_internal_topics: bool,
group_protocol: crate::GroupProtocol,
server_assignor: Option<String>,
group_instance_id: Option<String>,
auto_offset_reset: AutoOffsetReset,
isolation_level: IsolationLevel,
assignors: Vec<Assignor>,
rebalance_protocol: assignor::RebalanceProtocol,
fetch_min: ByteSize,
fetch_max: ByteSize,
fetch_partition_max: ByteSize,
fetch_max_wait: Time,
metadata_max_age: Time,
default_api_timeout: Time,
request_timeout: Time,
dns_timeout: Time,
send_buffer: Option<ByteSize>,
receive_buffer: Option<ByteSize>,
dispatch_queue_capacity: krabka_client_core::ConnectionDispatchQueueCapacity,
frame_max: krabka_client_core::ClientFrameMax,
metadata_recovery_strategy: krabka_client_core::MetadataRecoveryStrategy,
metadata_recovery_rebootstrap_trigger: Time,
leave_group_timeout: Time,
client_rack: Option<String>,
security: Option<security::ClientSecurity>,
client_dns_lookup: krabka_client_core::ClientDnsLookup,
reconnect_backoff: Time,
reconnect_backoff_max: Time,
connections_max_idle: Time,
socket_connection_setup_timeout: Time,
socket_connection_setup_timeout_max: Time,
has_rebalance_listener: bool,
retry_policy: ConsumerRetryPolicy,
auto_commit_interval: Option<Duration>,
allow_auto_create_topics: bool,
telemetry: Option<telemetry::ClientTelemetryConfig>,
}
#[derive(Clone, Copy, Debug, Eq, Ord, PartialEq, PartialOrd)]
struct RetryTime(Duration);
impl RetryTime {
fn new(name: &str, value: Time) -> Result<Self, String> {
let milliseconds = GreaterI64::<0>::new(value.millis_i64())
.map_err(|error| format!("{name}: {error}"))?
.into_value();
if !value.secs_f64().is_finite() || Time::from_millis(milliseconds) != value {
return Err(format!("{name} must be a whole number of milliseconds"));
}
Ok(Self(Duration::from_millis(
u64::try_from(milliseconds).expect("validated retry time is positive"),
)))
}
fn time(self) -> Time {
Time::from_std(self.0)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct ConsumerRetryPolicy {
startup_attempt_timeout: RetryTime,
startup_deadline: RetryTime,
startup_initial_backoff: RetryTime,
startup_max_backoff: RetryTime,
coordinator_retry_timeout: RetryTime,
coordinator_initial_backoff: RetryTime,
coordinator_max_backoff: RetryTime,
}
pub const DEFAULT_CONSUMER_RETRY_BACKOFF_MAX: Time = millis(1000);
#[bon::bon]
impl ConsumerRetryPolicy {
#[builder(start_fn = builder, finish_fn = build)]
pub fn new(
#[builder(default = secs(90))] startup_attempt_timeout: Time,
#[builder(default = minutes(5))] startup_deadline: Time,
#[builder(default = millis(500))] startup_initial_backoff: Time,
#[builder(default = DEFAULT_CONSUMER_RETRY_BACKOFF_MAX)] startup_max_backoff: Time,
#[builder(default = secs(30))] coordinator_retry_timeout: Time,
#[builder(default = millis(100))] coordinator_initial_backoff: Time,
#[builder(default = DEFAULT_CONSUMER_RETRY_BACKOFF_MAX)] coordinator_max_backoff: Time,
) -> Result<Self, String> {
let policy = Self {
startup_attempt_timeout: RetryTime::new(
"consumer startup attempt timeout",
startup_attempt_timeout,
)?,
startup_deadline: RetryTime::new("consumer startup deadline", startup_deadline)?,
startup_initial_backoff: RetryTime::new(
"consumer startup initial backoff",
startup_initial_backoff,
)?,
startup_max_backoff: RetryTime::new(
"consumer startup maximum backoff",
startup_max_backoff,
)?,
coordinator_retry_timeout: RetryTime::new(
"consumer coordinator retry timeout",
coordinator_retry_timeout,
)?,
coordinator_initial_backoff: RetryTime::new(
"consumer coordinator initial backoff",
coordinator_initial_backoff,
)?,
coordinator_max_backoff: RetryTime::new(
"consumer coordinator maximum backoff",
coordinator_max_backoff,
)?,
};
if policy.startup_attempt_timeout > policy.startup_deadline {
return Err(
"consumer startup attempt timeout must not exceed startup deadline".to_owned(),
);
}
if policy.startup_initial_backoff > policy.startup_max_backoff {
return Err(
"consumer startup initial backoff must not exceed startup maximum backoff"
.to_owned(),
);
}
if policy.coordinator_initial_backoff > policy.coordinator_max_backoff {
return Err(
"consumer coordinator initial backoff must not exceed coordinator maximum backoff"
.to_owned(),
);
}
Ok(policy)
}
#[must_use]
pub fn startup_attempt_timeout(self) -> Time {
self.startup_attempt_timeout.time()
}
#[must_use]
pub fn startup_deadline(self) -> Time {
self.startup_deadline.time()
}
#[must_use]
pub fn startup_initial_backoff(self) -> Time {
self.startup_initial_backoff.time()
}
#[must_use]
pub fn startup_max_backoff(self) -> Time {
self.startup_max_backoff.time()
}
#[must_use]
pub fn coordinator_retry_timeout(self) -> Time {
self.coordinator_retry_timeout.time()
}
#[must_use]
pub fn coordinator_initial_backoff(self) -> Time {
self.coordinator_initial_backoff.time()
}
#[must_use]
pub fn coordinator_max_backoff(self) -> Time {
self.coordinator_max_backoff.time()
}
pub(crate) fn with_retry_backoff_max(retry_backoff_max: Time) -> Result<Self, String> {
let startup_initial_backoff = if millis(500) < retry_backoff_max {
millis(500)
} else {
retry_backoff_max
};
let coordinator_initial_backoff = if millis(100) < retry_backoff_max {
millis(100)
} else {
retry_backoff_max
};
Self::builder()
.startup_initial_backoff(startup_initial_backoff)
.startup_max_backoff(retry_backoff_max)
.coordinator_initial_backoff(coordinator_initial_backoff)
.coordinator_max_backoff(retry_backoff_max)
.build()
}
}
impl Default for ConsumerRetryPolicy {
fn default() -> Self {
Self::with_retry_backoff_max(DEFAULT_CONSUMER_RETRY_BACKOFF_MAX)
.expect("default consumer retry backoff max is valid")
}
}
pub const DEFAULT_CONSUMER_LEAVE_GROUP_TIMEOUT: Time = secs(5);
#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
pub struct ConsumerLeaveGroupTimeout(Duration);
impl ConsumerLeaveGroupTimeout {
pub fn new(value: Duration) -> Result<Self, String> {
validated_millisecond_duration("consumer leave-group timeout", value)?;
Ok(Self(value))
}
#[must_use]
pub const fn duration(self) -> Duration {
self.0
}
#[must_use]
pub fn milliseconds(self) -> u64 {
validated_millisecond_duration("consumer leave-group timeout", self.0)
.expect("validated consumer leave-group timeout")
}
}
impl Default for ConsumerLeaveGroupTimeout {
fn default() -> Self {
Self::new(DEFAULT_CONSUMER_LEAVE_GROUP_TIMEOUT.to_std())
.expect("default consumer leave-group timeout is valid")
}
}
pub const DEFAULT_CONSUMER_SUBSCRIPTION_METADATA_REFRESH_INTERVAL: Time = secs(5);
#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
pub struct ConsumerFetchMaxBytes(i32);
impl ConsumerFetchMaxBytes {
pub fn new(value: i32) -> Result<Self, String> {
validated_positive_i32(value, "consumer fetch max bytes").map(Self)
}
#[must_use]
pub const fn bytes(self) -> i32 {
self.0
}
#[must_use]
pub fn size(self) -> ByteSize {
ByteSize::from_bytes_i64(i64::from(self.0))
}
}
impl TryFrom<ByteSize> for ConsumerFetchMaxBytes {
type Error = String;
fn try_from(value: ByteSize) -> Result<Self, Self::Error> {
Self::new(protocol_bytes("consumer fetch max", value)?)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
pub struct ConsumerFetchPartitionMaxBytes(i32);
impl ConsumerFetchPartitionMaxBytes {
pub fn new(value: i32) -> Result<Self, String> {
validated_positive_i32(value, "consumer fetch partition max bytes").map(Self)
}
#[must_use]
pub const fn bytes(self) -> i32 {
self.0
}
#[must_use]
pub fn size(self) -> ByteSize {
ByteSize::from_bytes_i64(i64::from(self.0))
}
}
impl TryFrom<ByteSize> for ConsumerFetchPartitionMaxBytes {
type Error = String;
fn try_from(value: ByteSize) -> Result<Self, Self::Error> {
Self::new(protocol_bytes("consumer fetch partition max", value)?)
}
}
#[derive(Clone, Copy, Debug, Eq, PartialEq, derive_more::Into)]
pub struct ConsumerSubscriptionMetadataRefreshInterval(Duration);
impl ConsumerSubscriptionMetadataRefreshInterval {
pub fn new(value: Duration) -> Result<Self, String> {
validated_millisecond_duration("consumer subscription metadata refresh interval", value)?;
Ok(Self(value))
}
#[must_use]
pub const fn duration(self) -> Duration {
self.0
}
#[must_use]
pub fn milliseconds(self) -> u64 {
validated_millisecond_duration("consumer subscription metadata refresh interval", self.0)
.expect("validated consumer subscription metadata refresh interval")
}
}
impl Default for ConsumerSubscriptionMetadataRefreshInterval {
fn default() -> Self {
Self::new(DEFAULT_CONSUMER_SUBSCRIPTION_METADATA_REFRESH_INTERVAL.to_std())
.expect("default consumer subscription metadata refresh interval is valid")
}
}
impl Drop for Consumer {
fn drop(&mut self) {
self.coordinator_shutdown.cancel();
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct Header {
pub key: String,
pub value: Option<Bytes>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum TimestampType {
CreateTime,
LogAppendTime,
}
impl From<records::TimestampType> for TimestampType {
fn from(value: records::TimestampType) -> Self {
match value {
records::TimestampType::CreateTime => Self::CreateTime,
records::TimestampType::LogAppendTime => Self::LogAppendTime,
}
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ConsumerRecord {
pub topic: String,
pub partition: i32,
pub offset: i64,
pub leader_epoch: i32,
pub timestamp: i64,
pub timestamp_type: TimestampType,
pub key: Option<Bytes>,
pub value: Option<Bytes>,
pub headers: Vec<Header>,
}
fn initial_join_protocols(
assignors: &[Assignor],
subscribe: &[String],
client_rack: Option<&str>,
) -> Vec<(String, bytes::Bytes)> {
coordinator::join_protocols(assignors, subscribe, &[], -1, client_rack, None)
}
fn startup_join_fields(
group_id: &str,
group_instance_id: Option<&str>,
protocols: Vec<(String, bytes::Bytes)>,
(session_timeout_ms, rebalance_timeout_ms): (i32, i32),
) -> coordinator::JoinRequestFields {
coordinator::JoinRequestFields {
group_id: group_id.to_owned(),
group_instance_id: group_instance_id.map(str::to_owned),
session_timeout_ms,
rebalance_timeout_ms,
protocols,
reason: String::new(),
}
}
fn first_join_member_id(resp: &JoinGroupResponse) -> Result<String, ConsumerError> {
let member_id = if resp.error_code == 79 || resp.error_code == 0 {
resp.member_id.clone()
} else {
return Err(ConsumerError::Server(resp.error_code));
};
if member_id.is_empty() {
return Err(ConsumerError::RebalanceFailed(
"broker did not assign a member_id".into(),
));
}
Ok(member_id)
}
#[derive(bon::Builder, Default)]
struct ConsumerClientOptions {
#[builder(default)]
metadata_scope: krabka_client_core::MetadataScope,
telemetry: Option<telemetry::ClientTelemetryConfig>,
}
async fn build_consumer_client(
config: &StartConfig,
options: ConsumerClientOptions,
) -> Result<Client, ConsumerError> {
Ok(Client::builder()
.bootstrap(&config.bootstrap)
.client_id(config.client_id.clone())
.request_timeout(config.request_timeout)
.dns_timeout(config.dns_timeout)
.send_buffer(config.send_buffer)
.receive_buffer(config.receive_buffer)
.dispatch_queue_capacity(config.dispatch_queue_capacity.get())
.frame_max(config.frame_max.size())
.metadata_recovery_strategy(config.metadata_recovery_strategy)
.metadata_recovery_rebootstrap_trigger(config.metadata_recovery_rebootstrap_trigger)
.maybe_security(config.security.clone())
.client_dns_lookup(config.client_dns_lookup)
.reconnect_backoff(config.reconnect_backoff)
.reconnect_backoff_max(config.reconnect_backoff_max)
.connections_max_idle(config.connections_max_idle)
.socket_connection_setup_timeout(config.socket_connection_setup_timeout)
.socket_connection_setup_timeout_max(config.socket_connection_setup_timeout_max)
.metadata_scope(options.metadata_scope)
.maybe_telemetry(options.telemetry)
.build()
.await?)
}
async fn pattern_topics(config: &StartConfig) -> Result<Vec<String>, ConsumerError> {
let client = build_consumer_client(config, ConsumerClientOptions::default()).await?;
let metadata = client
.refresh_metadata_with(krabka_client_core::all_topics_request())
.await;
client.close();
let subscription = subscription::Subscription {
pattern: config.subscribe_pattern.clone(),
..subscription::Subscription::topics(Vec::new(), config.exclude_internal_topics)
};
Ok(subscription.matching_topics(&metadata?).unwrap_or_default())
}
fn initial_commit_identity(
assigned: &[(String, i32)],
generation: i32,
member_id: &str,
) -> (CommitIdentity, u64) {
let ownership_ids = assigned.iter().cloned().zip(1_u64..).collect();
let identity = CommitIdentity {
generation,
member_id: member_id.to_owned(),
ownership_ids,
rejoin_on_poll: false,
};
(identity, assigned.len() as u64 + 1)
}
fn queue_initial_assign_call(
calls: &rebalance_listener::ListenerCalls,
partitions: Vec<(String, i32)>,
) {
let (done, _) = tokio::sync::oneshot::channel();
let _ = calls.send(rebalance_listener::ListenerCall {
kind: rebalance_listener::ListenerCallKind::Assigned,
partitions,
done,
});
}
fn subscription_metadata_scope(
allow_auto_create_topics: bool,
) -> krabka_client_core::MetadataScope {
krabka_client_core::MetadataScope::Topics {
allow_auto_topic_creation: allow_auto_create_topics,
}
}
fn is_subscribed_topic(subscribe: &[String], name: &str) -> bool {
subscribe.iter().any(|s| s == name)
}
fn is_group_leader(leader: &str, member_id: &str) -> bool {
leader == member_id
}
async fn leave_startup_member(
client: &Client,
coordinator_id: &AtomicI32,
group_id: &str,
member_id: &str,
group_instance_id: Option<String>,
leave_group_timeout: Time,
) {
if !coordinator::should_send_leave_group(
member_id,
group_instance_id.as_deref(),
GroupMembershipOperation::Default,
) {
return;
}
let broker = client.broker(coordinator_id.load(Ordering::Relaxed));
let send = broker.send(coordinator::build_leave_group_request(
group_id.to_string(),
member_id.to_string(),
None,
));
let _ = tokio::time::timeout(leave_group_timeout.to_std(), send).await;
}
fn has_assigned_partitions(assigned_partitions: &[(String, i32)]) -> bool {
!assigned_partitions.is_empty()
}
pub(crate) fn starting_offset(committed: i64, auto_offset_reset: AutoOffsetReset) -> i64 {
if committed >= 0 {
committed
} else {
reset_starting_offset(auto_offset_reset)
}
}
pub(crate) fn reset_starting_offset(auto_offset_reset: AutoOffsetReset) -> i64 {
match auto_offset_reset {
AutoOffsetReset::Earliest => poll::BEGINNING_SENTINEL,
AutoOffsetReset::Latest | AutoOffsetReset::ByDuration(_) => poll::LATEST_SENTINEL,
AutoOffsetReset::None => poll::NO_OFFSET_SENTINEL,
}
}
pub const DEFAULT_CONSUMER_DEFAULT_API_TIMEOUT: Time = secs(60);
pub const DEFAULT_CONSUMER_MAX_POLL_INTERVAL: Time = minutes(5);
static CONSUMER_CLIENT_ID_SEQUENCE: std::sync::atomic::AtomicU64 =
std::sync::atomic::AtomicU64::new(1);
fn consumer_client_id(
client_id: Option<String>,
group_id: &str,
group_instance_id: Option<&str>,
) -> String {
match client_id {
Some(client_id) if !client_id.is_empty() => client_id,
_ => match group_instance_id {
Some(group_instance_id) => format!("consumer-{group_id}-{group_instance_id}"),
None => format!(
"consumer-{group_id}-{}",
CONSUMER_CLIENT_ID_SEQUENCE.fetch_add(1, Ordering::Relaxed)
),
},
}
}
pub const DEFAULT_CONSUMER_METADATA_MAX_AGE: Time = minutes(5);
pub const DEFAULT_CONSUMER_FETCH_MAX_WAIT: Time = millis(500);
fn validated_fetch_max_wait(wait: Time) -> Result<Time, String> {
let milliseconds = MinMaxI64::<0, { i32::MAX as i64 }>::new(wait.millis_i64())
.map_err(|error| format!("consumer fetch max wait: {error}"))?
.into_value();
if !wait.secs_f64().is_finite() || Time::from_millis(milliseconds) != wait {
return Err("consumer fetch max wait must be a whole number of milliseconds".to_owned());
}
Ok(wait)
}
pub const DEFAULT_CONSUMER_MAX_POLL_RECORDS: usize = 500;
fn validated_max_poll_records(records: usize) -> Result<usize, String> {
if records == 0 || i32::try_from(records).is_err() {
return Err("consumer max poll records must be from 1 to i32::MAX".to_owned());
}
Ok(records)
}
fn validated_max_poll_interval(interval: Time) -> Result<Time, String> {
let milliseconds = MinMaxI64::<1, { i32::MAX as i64 }>::new(interval.millis_i64())
.map_err(|error| format!("consumer max poll interval: {error}"))?
.into_value();
if !interval.secs_f64().is_finite() || Time::from_millis(milliseconds) != interval {
return Err("consumer max poll interval must be a whole number of milliseconds".to_owned());
}
Ok(interval)
}
fn validated_auto_commit_interval(
enable_auto_commit: bool,
interval: Time,
) -> Result<Option<Duration>, String> {
if !enable_auto_commit {
return Ok(None);
}
let milliseconds = MinMaxI64::<0, { i32::MAX as i64 }>::new(interval.millis_i64())
.map_err(|error| format!("consumer auto commit interval: {error}"))?
.into_value();
if !interval.secs_f64().is_finite() || Time::from_millis(milliseconds) != interval {
return Err(
"consumer auto commit interval must be a whole number of milliseconds".to_owned(),
);
}
Ok(Some(interval.to_std()))
}
fn primed_position(committed_epoch: i32) -> position::PartitionPosition {
position::PartitionPosition {
offset_epoch: krabka_ids::LeaderEpoch(committed_epoch),
..Default::default()
}
}
pub(crate) fn protocol_millis_i32(value: Time) -> i32 {
i32::try_from(value.millis_i64_trunc()).unwrap_or(i32::MAX)
}
#[bon::bon]
impl Consumer {
#[builder(start_fn = builder, finish_fn = build)]
#[tracing::instrument(
name = "consumer.start",
level = "info",
skip_all,
fields(group_id = ?group_id, client_id = ?client_id),
err(level = "info")
)]
pub async fn start(
#[builder(into)] bootstrap: String,
#[builder(into)] client_id: Option<String>,
#[builder(into)]
group_id: Option<String>,
#[builder(default = secs(45))] session_timeout: Time,
#[builder(default = DEFAULT_CONSUMER_MAX_POLL_INTERVAL)] max_poll_interval: Time,
#[builder(default = DEFAULT_CONSUMER_MAX_POLL_RECORDS)] max_poll_records: usize,
#[builder(default = secs(3))] heartbeat_interval: Time,
#[builder(default = DEFAULT_CONSUMER_SUBSCRIPTION_METADATA_REFRESH_INTERVAL)]
subscription_metadata_refresh_interval: Time,
#[builder(into, default)] subscribe: Vec<String>,
subscribe_pattern: Option<crate::TopicPattern>,
#[builder(default = true)]
exclude_internal_topics: bool,
#[builder(into)]
subscribe_regex: Option<String>,
#[builder(default)]
group_protocol: crate::GroupProtocol,
#[builder(into)]
group_remote_assignor: Option<String>,
#[builder(into)] group_instance_id: Option<String>,
#[builder(default = AutoOffsetReset::Latest)] auto_offset_reset: AutoOffsetReset,
#[builder(default = IsolationLevel::ReadUncommitted)] isolation_level: IsolationLevel,
#[builder(into, default = Assignor::DEFAULT_LIST.to_vec())] assignors: Vec<Assignor>,
#[builder(default = krabka_client_core::DEFAULT_FETCH_MIN.bytes_i32())]
fetch_min: i32,
#[builder(default = poll::DEFAULT_FETCH_MAX.bytes_i32())]
fetch_max: i32,
#[builder(default = poll::DEFAULT_FETCH_PARTITION_MAX.bytes_i32())]
fetch_partition_max: i32,
#[builder(default = DEFAULT_CONSUMER_FETCH_MAX_WAIT)] fetch_max_wait: Time,
#[builder(default = DEFAULT_CONSUMER_METADATA_MAX_AGE)] metadata_max_age: Time,
#[builder(default = DEFAULT_CONSUMER_DEFAULT_API_TIMEOUT)]
default_api_timeout: Time,
#[builder(default = secs(30))] request_timeout: Time,
#[builder(default = krabka_client_core::DEFAULT_CLIENT_DNS_TIMEOUT)]
dns_timeout: Time,
#[builder(required, default = Some(krabka_client_core::DEFAULT_SEND_BUFFER))]
send_buffer: Option<ByteSize>,
#[builder(required, default = Some(krabka_client_core::DEFAULT_RECEIVE_BUFFER))]
receive_buffer: Option<ByteSize>,
#[builder(default = krabka_client_core::DEFAULT_CONNECTION_DISPATCH_QUEUE_CAPACITY)]
dispatch_queue_capacity: usize,
#[builder(default = krabka_client_core::DEFAULT_CLIENT_FRAME_MAX)] frame_max: ByteSize,
#[builder(default)]
metadata_recovery_strategy: krabka_client_core::MetadataRecoveryStrategy,
#[builder(default = krabka_client_core::DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER)]
metadata_recovery_rebootstrap_trigger: Time,
#[builder(default = DEFAULT_CONSUMER_LEAVE_GROUP_TIMEOUT)] leave_group_timeout: Time,
#[builder(into)] client_rack: Option<String>,
security: Option<security::ClientSecurity>,
#[builder(default)]
client_dns_lookup: krabka_client_core::ClientDnsLookup,
#[builder(default = krabka_client_core::DEFAULT_RECONNECT_BACKOFF)]
reconnect_backoff: Time,
#[builder(default = krabka_client_core::DEFAULT_RECONNECT_BACKOFF_MAX)]
reconnect_backoff_max: Time,
#[builder(default = krabka_client_core::DEFAULT_CONNECTIONS_MAX_IDLE)]
connections_max_idle: Time,
#[builder(default = krabka_client_core::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT)]
socket_connection_setup_timeout: Time,
#[builder(default = krabka_client_core::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX)]
socket_connection_setup_timeout_max: Time,
rebalance_listener: Option<Box<dyn crate::ConsumerRebalanceListener>>,
#[builder(default = DEFAULT_CONSUMER_RETRY_BACKOFF_MAX)]
retry_backoff_max: Time,
retry_policy: Option<ConsumerRetryPolicy>,
enable_auto_commit: Option<bool>,
#[builder(default = secs(5))] auto_commit_interval: Time,
#[builder(default = true)]
allow_auto_create_topics: bool,
#[builder(default = true)]
enable_metrics_push: bool,
) -> Result<Self, ConsumerError> {
let group_id = group_id.unwrap_or_default();
if group_id.is_empty() && !(subscribe.is_empty() && subscribe_pattern.is_none()) {
return Err(ConsumerError::InvalidGroupId);
}
if group_id.is_empty() && enable_auto_commit == Some(true) {
return Err(ConsumerError::InvalidConfig(
"enable.auto.commit cannot be set to true when default group id (null) is used."
.to_owned(),
));
}
let enable_auto_commit = enable_auto_commit.unwrap_or(!group_id.is_empty());
if group_protocol == GroupProtocol::Classic && group_remote_assignor.is_some() {
return Err(ConsumerError::InvalidConfig(
"group.remote.assignor cannot be set when group.protocol=classic".to_owned(),
));
}
if !subscribe.is_empty() && subscribe_pattern.is_some() {
return Err(ConsumerError::InvalidConfig(
"subscribe and subscribe_pattern are mutually exclusive".to_owned(),
));
}
if subscribe_regex.is_some() && (!subscribe.is_empty() || subscribe_pattern.is_some()) {
return Err(ConsumerError::InvalidConfig(
"Subscription to topics, partitions and pattern are mutually exclusive".to_owned(),
));
}
if let Some(regex) = &subscribe_regex {
if regex.is_empty() {
return Err(ConsumerError::InvalidArgument(
"Topic pattern to subscribe to cannot be empty".to_owned(),
));
}
if group_protocol != GroupProtocol::Consumer {
return Err(ConsumerError::InvalidConfig(
consumer_group::RE2J_NEEDS_CONSUMER_PROTOCOL.to_owned(),
));
}
if group_id.is_empty() {
return Err(ConsumerError::InvalidGroupId);
}
}
if subscribe.iter().any(String::is_empty) {
return Err(ConsumerError::InvalidConfig(
"Topic collection to subscribe to cannot contain null or empty topic".to_owned(),
));
}
if group_instance_id.as_deref().is_some_and(str::is_empty) {
return Err(ConsumerError::InvalidConfig(
"group_instance_id must not be empty".into(),
));
}
krabka_client_core::ClientDnsTimeout::new(dns_timeout)
.map_err(ConsumerError::InvalidConfig)?;
let fetch_min = FetchMinBytes::new(fetch_min)
.map_err(ConsumerError::InvalidConfig)?
.size();
let fetch_max = ConsumerFetchMaxBytes::new(fetch_max)
.map_err(ConsumerError::InvalidConfig)?
.size();
if fetch_min.bytes_i32() > fetch_max.bytes_i32() {
return Err(ConsumerError::InvalidConfig(
"consumer fetch min must not exceed consumer fetch max".to_owned(),
));
}
let fetch_partition_max = ConsumerFetchPartitionMaxBytes::new(fetch_partition_max)
.map_err(ConsumerError::InvalidConfig)?
.size();
let leave_group_timeout = ConsumerLeaveGroupTimeout::new(leave_group_timeout.to_std())
.map_err(ConsumerError::InvalidConfig)?;
let subscription_metadata_refresh_interval =
ConsumerSubscriptionMetadataRefreshInterval::new(
subscription_metadata_refresh_interval.to_std(),
)
.map_err(ConsumerError::InvalidConfig)?;
let dispatch_queue_capacity = ConnectionDispatchQueueCapacity::new(dispatch_queue_capacity)
.map_err(ConsumerError::InvalidConfig)?;
let frame_max =
ClientFrameMax::try_from(frame_max).map_err(ConsumerError::InvalidConfig)?;
let metadata_recovery_rebootstrap_trigger =
krabka_client_core::MetadataRecoveryRebootstrapTrigger::new(
metadata_recovery_rebootstrap_trigger,
)
.map_err(ConsumerError::InvalidConfig)?
.time();
let auto_commit_interval =
validated_auto_commit_interval(enable_auto_commit, auto_commit_interval)
.map_err(ConsumerError::InvalidConfig)?;
let max_poll_interval =
validated_max_poll_interval(max_poll_interval).map_err(ConsumerError::InvalidConfig)?;
let max_poll_records =
validated_max_poll_records(max_poll_records).map_err(ConsumerError::InvalidConfig)?;
let fetch_max_wait =
validated_fetch_max_wait(fetch_max_wait).map_err(ConsumerError::InvalidConfig)?;
if metadata_max_age.secs_f64() < 0.0 || !metadata_max_age.secs_f64().is_finite() {
return Err(ConsumerError::InvalidConfig(
"consumer metadata max age must not be negative".to_owned(),
));
}
if default_api_timeout.secs_f64() < 0.0 || !default_api_timeout.secs_f64().is_finite() {
return Err(ConsumerError::InvalidConfig(
"consumer default api timeout must not be negative".to_owned(),
));
}
let retry_policy = match retry_policy {
Some(retry_policy) => retry_policy,
None => ConsumerRetryPolicy::with_retry_backoff_max(retry_backoff_max)
.map_err(ConsumerError::InvalidConfig)?,
};
let rebalance_protocol =
assignor::rebalance_protocol_of(&assignors).map_err(ConsumerError::InvalidConfig)?;
let client_id = consumer_client_id(client_id, &group_id, group_instance_id.as_deref());
let telemetry = enable_metrics_push.then(|| {
telemetry::ClientTelemetryConfig::consumer(
Some(group_id.as_str()),
group_instance_id.as_deref(),
client_rack.as_deref(),
)
});
let config = StartConfig {
bootstrap,
client_id,
group_id,
session_timeout,
max_poll_interval,
max_poll_records,
heartbeat_interval,
subscription_metadata_refresh_interval: Time::from_std(
subscription_metadata_refresh_interval.duration(),
),
subscribe,
subscribe_pattern,
subscribe_regex,
exclude_internal_topics,
group_protocol,
server_assignor: group_remote_assignor,
group_instance_id,
auto_offset_reset,
isolation_level,
assignors,
rebalance_protocol,
fetch_min,
fetch_max,
fetch_partition_max,
fetch_max_wait,
metadata_max_age,
default_api_timeout,
request_timeout,
dns_timeout,
send_buffer,
receive_buffer,
dispatch_queue_capacity,
frame_max,
metadata_recovery_strategy,
metadata_recovery_rebootstrap_trigger,
leave_group_timeout: Time::from_std(leave_group_timeout.duration()),
client_rack,
security,
client_dns_lookup,
reconnect_backoff,
reconnect_backoff_max,
connections_max_idle,
socket_connection_setup_timeout,
socket_connection_setup_timeout_max,
has_rebalance_listener: rebalance_listener.is_some(),
retry_policy,
auto_commit_interval,
allow_auto_create_topics,
telemetry,
};
Self::retry_start_until_success(config, rebalance_listener).await
}
async fn retry_start_until_success(
config: StartConfig,
rebalance_listener: Option<Box<dyn crate::ConsumerRebalanceListener>>,
) -> Result<Self, ConsumerError> {
let started = tokio::time::Instant::now();
let total = config.retry_policy.startup_deadline().to_std();
let mut backoff = config.retry_policy.startup_initial_backoff().to_std();
loop {
let remaining = total.saturating_sub(started.elapsed());
if remaining.is_zero() {
return Err(ConsumerError::Client(ClientError::Timeout(
config.retry_policy.startup_deadline(),
)));
}
match tokio::time::timeout(
config
.retry_policy
.startup_attempt_timeout()
.to_std()
.min(remaining),
Box::pin(Self::start_once(config.clone())),
)
.await
{
Ok(Ok(mut consumer)) => {
consumer.rebalance_listener =
rebalance_listener.map(rebalance_listener::shared);
return Ok(consumer);
}
Ok(Err(error)) => {
if started.elapsed().as_time() < config.retry_policy.startup_deadline()
&& is_retriable_consumer_start_error(&error)
{
tracing::warn!(
group = %config.group_id,
%error,
"consumer startup failed transiently; retrying with a fresh connection"
);
} else {
return Err(error);
}
}
Err(_elapsed) => {
if started.elapsed().as_time() >= config.retry_policy.startup_deadline() {
return Err(ConsumerError::Client(ClientError::Timeout(
config.retry_policy.startup_deadline(),
)));
}
tracing::warn!(
group = %config.group_id,
timeout = ?config.retry_policy.startup_attempt_timeout(),
"consumer startup exceeded attempt timeout \
(likely a cold-boot group-join stall); \
retrying with a fresh connection"
);
}
}
tokio::time::sleep(backoff.min(total.saturating_sub(started.elapsed()))).await;
backoff = backoff
.saturating_mul(2)
.min(config.retry_policy.startup_max_backoff().to_std());
}
}
#[tracing::instrument(
name = "consumer.start_once",
level = "info",
skip_all,
fields(
group_id = %config.group_id,
coordinator_id = tracing::field::Empty,
member_id = tracing::field::Empty,
generation = tracing::field::Empty,
is_leader = tracing::field::Empty,
assigned_partitions = tracing::field::Empty,
),
err(level = "info")
)]
async fn start_once(mut config: StartConfig) -> Result<Self, ConsumerError> {
if config.subscribe.is_empty()
&& config.subscribe_pattern.is_none()
&& config.subscribe_regex.is_none()
{
return start_without_subscription(config).await;
}
if config.subscribe_pattern.is_some() {
config.subscribe = pattern_topics(&config).await?;
}
if config.group_protocol == GroupProtocol::Consumer {
return start_without_subscription(config).await;
}
let finish_config = config.clone();
let telemetry = config.telemetry.clone();
let StartConfig {
bootstrap,
client_id,
group_id,
session_timeout,
max_poll_interval,
subscribe,
group_instance_id,
assignors,
request_timeout,
dns_timeout,
send_buffer,
receive_buffer,
dispatch_queue_capacity,
frame_max,
metadata_recovery_strategy,
metadata_recovery_rebootstrap_trigger,
client_rack,
security,
client_dns_lookup,
reconnect_backoff,
reconnect_backoff_max,
connections_max_idle,
socket_connection_setup_timeout,
socket_connection_setup_timeout_max,
allow_auto_create_topics,
..
} = config;
let client = Client::builder()
.bootstrap(&bootstrap)
.client_id(client_id.clone())
.request_timeout(request_timeout)
.dns_timeout(dns_timeout)
.send_buffer(send_buffer)
.receive_buffer(receive_buffer)
.dispatch_queue_capacity(dispatch_queue_capacity.get())
.frame_max(frame_max.size())
.metadata_recovery_strategy(metadata_recovery_strategy)
.metadata_recovery_rebootstrap_trigger(metadata_recovery_rebootstrap_trigger)
.maybe_security(security.clone())
.client_dns_lookup(client_dns_lookup)
.reconnect_backoff(reconnect_backoff)
.reconnect_backoff_max(reconnect_backoff_max)
.connections_max_idle(connections_max_idle)
.socket_connection_setup_timeout(socket_connection_setup_timeout)
.socket_connection_setup_timeout_max(socket_connection_setup_timeout_max)
.metadata_scope(subscription_metadata_scope(allow_auto_create_topics))
.maybe_telemetry(telemetry)
.build()
.await?;
client.metadata_topics().set(subscribe.iter().cloned());
let session_timeout_ms = protocol_millis_i32(session_timeout);
let rebalance_timeout_ms = protocol_millis_i32(max_poll_interval);
let coordinator_retry = CoordinatorRetryPolicy::from(config.retry_policy);
let coordinator_id = Arc::new(AtomicI32::new(
find_coordinator(&client, &group_id, coordinator_retry).await?,
));
tracing::Span::current().record("coordinator_id", coordinator_id.load(Ordering::Relaxed));
let protocols = initial_join_protocols(&assignors, &subscribe, client_rack.as_deref());
let join_fields = startup_join_fields(
&group_id,
group_instance_id.as_deref(),
protocols.clone(),
(session_timeout_ms, rebalance_timeout_ms),
);
let refind = CoordinatorRefind {
client: &client,
group_id: &group_id,
coordinator_id: &coordinator_id,
};
let r1 = CoordinatorRetryRequest::builder()
.retry(coordinator_retry)
.code(|r: &JoinGroupResponse| r.error_code)
.make(|| refind.send_join_group(join_fields.request(String::new())))
.refind(refind)
.build()
.run_with_refind()
.await?;
let member_id = first_join_member_id(&r1)?;
tracing::Span::current().record("member_id", member_id.as_str());
let cleanup_client = client.clone();
let cleanup_coordinator_id = Arc::clone(&coordinator_id);
let cleanup_group_id = group_id.clone();
let cleanup_member_id = member_id.clone();
let cleanup_group_instance_id = group_instance_id.clone();
let cleanup_leave_group_timeout = finish_config.leave_group_timeout;
let start_result = finish_startup(
finish_config,
client,
coordinator_id,
member_id,
protocols,
(session_timeout_ms, rebalance_timeout_ms),
)
.await;
match start_result {
Ok(consumer) => Ok(consumer),
Err(error) => {
leave_startup_member(
&cleanup_client,
&cleanup_coordinator_id,
&cleanup_group_id,
&cleanup_member_id,
cleanup_group_instance_id,
cleanup_leave_group_timeout,
)
.await;
Err(ConsumerError::StartupAfterJoin(Box::new(error)))
}
}
}
}
async fn finish_startup(
config: StartConfig,
client: Client,
coordinator_id: Arc<AtomicI32>,
member_id: String,
protocols: Vec<(String, bytes::Bytes)>,
timeouts_ms: (i32, i32),
) -> Result<Consumer, ConsumerError> {
let spawn_config = config.clone();
let coordinator_retry = CoordinatorRetryPolicy::from(config.retry_policy);
let StartConfig {
group_id,
group_instance_id,
auto_offset_reset,
..
} = config;
let mut join_fields = startup_join_fields(
&group_id,
group_instance_id.as_deref(),
protocols,
timeouts_ms,
);
join_fields.reason = format!("need to re-join with the given member-id: {member_id}");
let r2 = coordinator::JoinGroupRetry::builder()
.client(&client)
.group_id(&group_id)
.coordinator_id(&coordinator_id)
.retry(coordinator_retry)
.fields(&join_fields)
.member_id(member_id.clone())
.build()
.run()
.await?;
if r2.error_code != 0 {
return Err(ConsumerError::Server(r2.error_code));
}
let InitialAssignment {
assigned_partitions,
group_topics,
topic_ids,
topic_partitions,
} = resolve_initial_assignment(&spawn_config, &client, &coordinator_id, &member_id, &r2)
.await?;
let mut next_offsets: HashMap<(String, i32), i64> = HashMap::new();
let mut positions: HashMap<(String, i32), position::PartitionPosition> = HashMap::new();
if has_assigned_partitions(&assigned_partitions) {
let mut by_topic: HashMap<String, Vec<i32>> = HashMap::new();
for (t, p) in &assigned_partitions {
by_topic.entry(t.clone()).or_default().push(*p);
}
let of = coordinator::send_offset_fetch(
&client,
&group_id,
&coordinator_id,
&offset_wire::build_offset_fetch(&group_id, &by_topic),
coordinator_retry,
)
.await?;
for (name, partition_index, committed, committed_epoch) in
offset_wire::parse_offset_fetch(&of)
{
let starting = starting_offset(committed, auto_offset_reset);
next_offsets.insert((name.clone(), partition_index), starting);
positions.insert((name, partition_index), primed_position(committed_epoch));
}
}
spawn_consumer(
spawn_config,
client,
coordinator_id,
member_id,
StartupState::builder()
.generation_id(r2.generation_id)
.group_topics(group_topics)
.assigned_partitions(assigned_partitions)
.next_offsets(next_offsets)
.positions(positions)
.topic_ids(topic_ids)
.topic_partitions(topic_partitions)
.build(),
)
.await
}
struct InitialAssignment {
assigned_partitions: Vec<(String, i32)>,
group_topics: Vec<String>,
topic_ids: HashMap<String, WireUuid>,
topic_partitions: HashMap<String, i32>,
}
async fn resolve_initial_assignment(
config: &StartConfig,
client: &Client,
coordinator_id: &Arc<AtomicI32>,
member_id: &str,
r2: &JoinGroupResponse,
) -> Result<InitialAssignment, ConsumerError> {
let group_id = &config.group_id;
let subscribe = &config.subscribe;
let group_instance_id = &config.group_instance_id;
let protocol_name = r2.protocol_name.as_deref().unwrap_or_default();
let is_leader = is_group_leader(&r2.leader, member_id);
let group_topics = if is_leader {
coordinator::group_subscription_topics(r2, subscribe)
} else {
subscribe.clone()
};
client.metadata_topics().set(group_topics.iter().cloned());
let md = client.refresh_metadata().await;
client.metadata_topics().set(subscribe.iter().cloned());
let md = md?;
let mut topic_ids: HashMap<String, WireUuid> = HashMap::new();
let mut topic_partitions: HashMap<String, i32> = HashMap::new();
for t in &md.topics {
let Some(name) = &t.name else { continue };
if is_subscribed_topic(&group_topics, name) {
let count = i32::try_from(t.partitions.len()).unwrap_or(i32::MAX);
topic_partitions.insert(name.clone(), count);
topic_ids.insert(name.clone(), t.topic_id);
}
}
tracing::Span::current().record("generation", r2.generation_id);
tracing::Span::current().record("is_leader", is_leader);
let assignments_for_sync: Vec<SyncGroupRequestAssignment> = if is_leader {
coordinator::leader_assignments(&config.assignors, r2, &md)?
} else {
Vec::new()
};
let r3 = CoordinatorRetryRequest::builder()
.retry(config.retry_policy.into())
.code(|r: &SyncGroupResponse| r.error_code)
.make(|| {
let group_id = group_id.clone();
let protocol_name = protocol_name.to_string();
let member_id = member_id.to_string();
let assignments_for_sync = assignments_for_sync.clone();
let generation_id = r2.generation_id;
let group_instance_id = group_instance_id.clone();
let target = coordinator_id.load(Ordering::Relaxed);
async move {
client
.broker(target)
.send(
coordinator::SyncGroupRequestConfig::builder()
.group_id(group_id)
.generation_id(generation_id)
.member_id(member_id)
.maybe_group_instance_id(group_instance_id)
.protocol_name(protocol_name)
.assignments(assignments_for_sync)
.build(),
)
.await
.map_err(ConsumerError::from)
}
})
.refind(CoordinatorRefind {
client,
group_id,
coordinator_id,
})
.build()
.run_with_refind()
.await?;
if r3.error_code != 0 {
return Err(ConsumerError::Server(r3.error_code));
}
let assigned_partitions = decode_assignment(&r3.assignment);
tracing::Span::current().record("assigned_partitions", assigned_partitions.len());
Ok(InitialAssignment {
assigned_partitions,
group_topics,
topic_ids,
topic_partitions,
})
}
#[derive(bon::Builder, Default)]
struct StartupState {
#[builder(default)]
generation_id: i32,
#[builder(default)]
group_topics: Vec<String>,
#[builder(default)]
assigned_partitions: Vec<(String, i32)>,
#[builder(default)]
next_offsets: HashMap<(String, i32), i64>,
#[builder(default)]
positions: HashMap<(String, i32), position::PartitionPosition>,
#[builder(default)]
topic_ids: HashMap<String, WireUuid>,
#[builder(default)]
topic_partitions: HashMap<String, i32>,
}
#[cfg(test)]
#[derive(bon::Builder)]
#[builder(start_fn = builder, finish_fn = build)]
struct TestConsumerStart {
bootstrap: String,
config: StartConfig,
startup: StartupState,
}
#[cfg(test)]
async fn start_test_consumer(start: TestConsumerStart) -> Consumer {
let client = Client::builder()
.bootstrap(start.bootstrap)
.build()
.await
.expect("client");
spawn_consumer(
start.config,
client,
Arc::new(AtomicI32::new(0)),
"member-a".into(),
start.startup,
)
.await
.expect("spawn consumer")
}
async fn start_without_subscription(config: StartConfig) -> Result<Consumer, ConsumerError> {
let client = build_consumer_client(
&config,
ConsumerClientOptions::builder()
.metadata_scope(subscription_metadata_scope(config.allow_auto_create_topics))
.maybe_telemetry(config.telemetry.clone())
.build(),
)
.await?;
let coordinator_id = if config.group_id.is_empty() {
-1
} else {
find_coordinator(&client, &config.group_id, config.retry_policy.into()).await?
};
let member_id = if config.group_protocol == GroupProtocol::Consumer {
uuid::Uuid::new_v4().to_string()
} else {
String::new()
};
client
.metadata_topics()
.set(config.subscribe.iter().cloned());
spawn_consumer(
config,
client,
Arc::new(AtomicI32::new(coordinator_id)),
member_id,
StartupState::builder().generation_id(-1).build(),
)
.await
}
async fn coordinator_task_client(
config: &StartConfig,
group_topics: &[String],
) -> Result<Option<Client>, ConsumerError> {
if config.subscribe.is_empty()
&& config.subscribe_pattern.is_none()
&& config.subscribe_regex.is_none()
{
return Ok(None);
}
let client = build_consumer_client(
config,
ConsumerClientOptions::builder()
.metadata_scope(subscription_metadata_scope(config.allow_auto_create_topics))
.build(),
)
.await?;
let topics = if group_topics.is_empty() {
&config.subscribe
} else {
group_topics
};
client.metadata_topics().set(topics.iter().cloned());
Ok(Some(client))
}
async fn spawn_consumer(
config: StartConfig,
client: Client,
coordinator_id: Arc<AtomicI32>,
member_id: String,
startup: StartupState,
) -> Result<Consumer, ConsumerError> {
let coordinator_client = coordinator_task_client(&config, &startup.group_topics).await?;
let StartConfig {
group_id,
session_timeout,
max_poll_interval,
max_poll_records,
heartbeat_interval,
subscription_metadata_refresh_interval,
subscribe,
subscribe_pattern,
subscribe_regex,
exclude_internal_topics,
group_protocol,
server_assignor,
group_instance_id,
auto_offset_reset,
isolation_level,
assignors,
rebalance_protocol,
fetch_min,
fetch_max,
fetch_partition_max,
fetch_max_wait,
metadata_max_age,
default_api_timeout,
request_timeout,
leave_group_timeout,
client_rack,
has_rebalance_listener,
retry_policy,
auto_commit_interval,
..
} = config;
let subscription = subscription::shared(
subscribe.clone(),
subscribe_pattern,
exclude_internal_topics,
);
if subscribe_regex.is_some() {
subscription.send_modify(|state| state.regex.clone_from(&subscribe_regex));
}
let StartupState {
generation_id,
group_topics: _,
assigned_partitions,
next_offsets,
positions,
topic_ids,
topic_partitions,
} = startup;
let standalone = coordinator_client.is_none();
let coordinator_client = coordinator_client.unwrap_or_else(|| client.clone());
let (identity, next_ownership_id) =
initial_commit_identity(&assigned_partitions, generation_id, &member_id);
let initial_assignment = assigned_partitions.clone();
let assigned = Arc::new(Mutex::new(assigned_partitions));
let assignment_changed = Arc::new(Notify::new());
let commit_identity = Arc::new(Mutex::new(identity));
let commit_serialization = Arc::new(Mutex::new(()));
let commit_async_state = Arc::new(AtomicU8::new(0));
let next_offsets = Arc::new(Mutex::new(next_offsets));
let end_offsets = Arc::new(Mutex::new(HashMap::new()));
let positions = Arc::new(Mutex::new(positions));
let topic_ids = Arc::new(Mutex::new(topic_ids));
let current_generation = Arc::new(AtomicI32::new(generation_id));
let poll_error = coordinator::PollErrorSlot::default();
let poll_signal = coordinator::PollSignal::default();
let rebalance_pending =
tokio::sync::watch::Sender::new(group_protocol == GroupProtocol::Consumer);
let rebalance_pending_receiver = rebalance_pending.subscribe();
let close_operation = tokio::sync::watch::Sender::new(control::CloseRequest::default());
let (listener_sender, listener_calls) = tokio::sync::mpsc::unbounded_channel();
let (enforced_rebalances, enforced_rebalance_reasons) = tokio::sync::mpsc::unbounded_channel();
let assigned_callback_pending = rebalance_listener::AssignedCallbackPending::default();
if has_rebalance_listener && !standalone {
queue_initial_assign_call(&listener_sender, initial_assignment);
}
let auto_commit = auto_commit_interval.map(commit::AutoCommit::new);
let shutdown = CancellationToken::new();
let published_member_id = tokio::sync::watch::Sender::new(member_id.clone());
let live_member_id = published_member_id.subscribe();
let state = CoordinatorState {
client: coordinator_client,
group_id: group_id.clone(),
coordinator_id: Arc::clone(&coordinator_id),
member_id,
published_member_id,
commit_identity: Arc::clone(&commit_identity),
group_instance_id: group_instance_id.clone(),
generation_id,
current_generation: Arc::clone(¤t_generation),
assignors,
rebalance_protocol,
subscription: Arc::clone(&subscription),
subscription_changes: subscription.subscribe(),
joined_topics: subscribe.clone(),
seen_unsubscribes: 0,
group_protocol,
server_assignor,
member_epoch: 0,
sent_heartbeat_fields: consumer_group::SentFields::default(),
assigned: Arc::clone(&assigned),
assignment_changed: Arc::clone(&assignment_changed),
next_ownership_id,
next_offsets: Arc::clone(&next_offsets),
end_offsets: Arc::clone(&end_offsets),
positions: Arc::clone(&positions),
topic_ids: Arc::clone(&topic_ids),
session_timeout,
max_poll_interval,
heartbeat_interval,
subscription_metadata_refresh_interval,
leave_group_timeout,
auto_offset_reset,
client_rack: client_rack.clone(),
initial_subscribed_counts: topic_partitions,
retry_policy: retry_policy.into(),
poll_error: Arc::clone(&poll_error),
auto_commit: auto_commit.clone(),
commit_serialization: Arc::clone(&commit_serialization),
join_prepared: false,
polls: poll_signal.subscribe(),
rebalance_pending,
close_operation: close_operation.subscribe(),
rejoin_reason: String::new(),
listener_calls: has_rebalance_listener.then_some(listener_sender),
enforced_rebalances: enforced_rebalance_reasons,
lost_partitions: Vec::new(),
assigned_callback_pending: Arc::clone(&assigned_callback_pending),
poll_timer: coordinator::PollTimer::new(max_poll_interval),
poll_timeout_in_callback: false,
};
let coord_handle = (!standalone).then(|| match group_protocol {
GroupProtocol::Classic => tokio::spawn(coordinator::run(state, shutdown.clone())),
GroupProtocol::Consumer => tokio::spawn(consumer_group::run(state, shutdown.clone())),
});
Ok(Consumer {
client,
group_id,
coordinator_id,
retry_policy: retry_policy.into(),
member_id: live_member_id,
commit_identity,
commit_serialization,
commit_async_state,
commit_async_callbacks: Arc::default(),
group_instance_id: group_instance_id.clone(),
current_generation,
subscription,
group_protocol,
assigned,
assignment_changed,
next_offsets,
end_offsets,
positions,
topic_ids,
session_timeout,
heartbeat_interval,
rebalance_protocol,
coordinator_shutdown: shutdown,
coordinator_handle: coord_handle,
isolation_level,
fetch_min,
fetch_max,
fetch_partition_max,
fetch_max_wait,
fetches: poll::Fetches::default(),
client_rack,
metadata_max_age,
request_timeout,
wakeup: control::WakeupHandle::default(),
enforced_rebalances,
default_api_timeout,
paused: std::sync::Mutex::default(),
auto_offset_reset,
poll_error,
auto_commit,
poll_signal,
rebalance_pending: rebalance_pending_receiver,
max_poll_records,
fetch_buffer: fetch_buffer::FetchBuffer::default(),
close_operation,
rebalance_listener: None,
listener_calls,
assigned_callback_pending,
})
}
fn is_retriable_consumer_start_error(error: &ConsumerError) -> bool {
matches!(
error,
ConsumerError::Server(14 | 15 | 16 | 22 | 25 | 27 | 79)
| ConsumerError::Client(ClientError::Disconnected)
) || matches!(
error,
ConsumerError::StartupAfterJoin(inner)
if is_retriable_consumer_start_error(inner)
|| matches!(
inner.as_ref(),
ConsumerError::Client(ClientError::Timeout(_))
)
)
}
impl Consumer {
#[must_use]
pub fn group_id(&self) -> &str {
&self.group_id
}
#[must_use]
pub fn member_id(&self) -> String {
self.member_id.borrow().clone()
}
#[must_use]
pub fn generation_id(&self) -> i32 {
self.current_generation.load(Ordering::Relaxed)
}
#[must_use]
pub fn group_metadata(&self) -> ConsumerGroupMetadata {
ConsumerGroupMetadata {
group_id: self.group_id.clone(),
generation_id: self.current_generation.load(Ordering::Relaxed),
member_id: self.member_id(),
group_instance_id: self.group_instance_id.clone(),
}
}
pub async fn assignment(&self) -> Vec<(String, i32)> {
self.assigned.lock().await.clone()
}
pub async fn at_log_end(&self) -> bool {
let assigned = self.assigned.lock().await;
if assigned.is_empty() {
return false;
}
let positions = self.next_offsets.lock().await;
let ends = self.end_offsets.lock().await;
assigned.iter().all(|partition| {
positions
.get(partition)
.filter(|position| !poll::is_reset_sentinel(**position))
.zip(ends.get(partition))
.is_some_and(|(position, end)| position >= end)
})
}
pub async fn close(self) -> Result<(), ConsumerError> {
self.close_with(CloseOptions::builder().build()).await
}
pub async fn client_instance_id(
&self,
timeout: Time,
) -> Result<Option<uuid::Uuid>, ConsumerError> {
client_instance_id(&self.client, timeout).await
}
#[tracing::instrument(
name = "consumer.close",
level = "info",
skip_all,
fields(group_id = %self.group_id, member_id = %self.member_id(), ?options),
err(level = "info")
)]
pub async fn close_with(mut self, options: crate::CloseOptions) -> Result<(), ConsumerError> {
let timeout = control::close_timeout(options)?.to_std();
let started = tokio::time::Instant::now();
let request_timeout = self.request_timeout.to_std();
let coordinator_deadline = started + timeout.min(request_timeout);
self.auto_commit_on_close(coordinator_deadline).await;
self.wait_for_async_commits(coordinator_deadline).await;
let listener_result = self.leave_prepare().await;
self.close_operation.send_replace(control::CloseRequest {
operation: options.group_membership_operation,
});
self.coordinator_shutdown.cancel();
if let Some(h) = self.coordinator_handle.take() {
let _ = h.await;
}
let left = timeout
.saturating_sub(started.elapsed())
.min(request_timeout);
self.close_fetch_sessions(left).await;
let left = timeout.saturating_sub(started.elapsed());
let _ = tokio::time::timeout(left, self.client.close_telemetry()).await;
listener_result
}
pub(crate) async fn leave_prepare(&mut self) -> Result<(), ConsumerError> {
if self.rebalance_listener.is_none() {
return Ok(());
}
let pending = self.run_pending_listener_calls().await;
let owned = self.assigned.lock().await.clone();
if owned.is_empty() {
return pending;
}
let kind = if self.member_id.borrow().is_empty() || *self.rebalance_pending.borrow() {
rebalance_listener::ListenerCallKind::Lost
} else {
rebalance_listener::ListenerCallKind::Revoked
};
let result = self.run_listener(kind, &owned).await;
pending.and(result)
}
}
#[cfg(test)]
impl Consumer {
pub(crate) fn test_builder(
client: krabka_client_core::Client,
) -> test_support::ConsumerTestBuilder {
test_support::ConsumerTest::builder(client)
}
}
#[cfg(test)]
mod protocol_millis_tests {
use assert2::check;
use krabka_units::{Time, convert::TimeExt as _, millis, minutes, secs};
use super::protocol_millis_i32;
#[test]
fn protocol_millis_truncates_rather_than_rounding() {
for (_name, value, expected) in [
("whole default session timeout", secs(45), 45_000),
("whole default rebalance timeout", minutes(1), 60_000),
("whole millisecond", millis(37), 37),
(
"fractional millisecond truncates down",
Time::from_secs_f64(0.0016),
1,
),
(
"fractional millisecond that would round up",
Time::from_secs_f64(1.9999),
1_999,
),
("below one millisecond", Time::from_secs_f64(0.0009), 0),
] {
check!(protocol_millis_i32(value) == expected);
}
}
#[test]
fn protocol_millis_saturates_at_i32_max() {
check!(protocol_millis_i32(Time::from_secs(10_000_000_000)) == i32::MAX);
}
}
#[cfg(test)]
mod consumer_retry_policy_tests {
use assert2::{assert, check};
use krabka_units::{Time, convert::TimeExt as _, millis, minutes, secs};
use super::ConsumerRetryPolicy;
#[test]
fn consumer_retry_policy_defaults_and_overrides() {
let defaults = ConsumerRetryPolicy::default();
check!(defaults.startup_attempt_timeout() == secs(90));
check!(defaults.startup_deadline() == minutes(5));
check!(defaults.startup_initial_backoff() == millis(500));
check!(defaults.startup_max_backoff() == millis(1000));
check!(defaults.coordinator_retry_timeout() == secs(30));
check!(defaults.coordinator_initial_backoff() == millis(100));
check!(defaults.coordinator_max_backoff() == secs(1));
let configured = ConsumerRetryPolicy::builder()
.startup_attempt_timeout(secs(11))
.startup_deadline(secs(12))
.startup_initial_backoff(millis(13))
.startup_max_backoff(millis(14))
.coordinator_retry_timeout(secs(15))
.coordinator_initial_backoff(millis(16))
.coordinator_max_backoff(millis(17))
.build()
.expect("valid retry policy");
check!(configured.startup_attempt_timeout() == secs(11));
check!(configured.startup_deadline() == secs(12));
check!(configured.startup_initial_backoff() == millis(13));
check!(configured.startup_max_backoff() == millis(14));
check!(configured.coordinator_retry_timeout() == secs(15));
check!(configured.coordinator_initial_backoff() == millis(16));
check!(configured.coordinator_max_backoff() == millis(17));
}
#[test]
fn with_retry_backoff_max_caps_startup_and_coordinator_backoff() {
for (name, retry_backoff_max) in [
("Kafka's default", millis(1000)),
("custom, above both fixed initial backoffs", secs(3)),
("below both fixed initial backoffs", millis(100)),
("between the two fixed initial backoffs", millis(250)),
] {
let policy = ConsumerRetryPolicy::with_retry_backoff_max(retry_backoff_max)
.unwrap_or_else(|error| panic!("{name}: {error}"));
check!(policy.startup_max_backoff() == retry_backoff_max);
check!(policy.coordinator_max_backoff() == retry_backoff_max);
check!(
policy.startup_initial_backoff() <= retry_backoff_max,
"{name}: startup initial backoff must not exceed retry_backoff_max"
);
check!(
policy.coordinator_initial_backoff() <= retry_backoff_max,
"{name}: coordinator initial backoff must not exceed retry_backoff_max"
);
}
}
#[test]
fn with_retry_backoff_max_rejects_invalid_values() {
for retry_backoff_max in [
Time::ZERO,
Time::from_secs_f64(0.0005),
Time::from_secs_f64(f64::INFINITY),
] {
assert!(ConsumerRetryPolicy::with_retry_backoff_max(retry_backoff_max).is_err());
}
}
#[test]
fn consumer_retry_policy_rejects_invalid_values() {
for invalid in [
Time::ZERO,
Time::from_secs_f64(0.0005),
Time::from_secs_f64(f64::INFINITY),
] {
assert!(
ConsumerRetryPolicy::builder()
.startup_attempt_timeout(invalid)
.build()
.is_err()
);
assert!(
ConsumerRetryPolicy::builder()
.startup_deadline(invalid)
.build()
.is_err()
);
assert!(
ConsumerRetryPolicy::builder()
.startup_initial_backoff(invalid)
.build()
.is_err()
);
assert!(
ConsumerRetryPolicy::builder()
.startup_max_backoff(invalid)
.build()
.is_err()
);
assert!(
ConsumerRetryPolicy::builder()
.coordinator_retry_timeout(invalid)
.build()
.is_err()
);
assert!(
ConsumerRetryPolicy::builder()
.coordinator_initial_backoff(invalid)
.build()
.is_err()
);
assert!(
ConsumerRetryPolicy::builder()
.coordinator_max_backoff(invalid)
.build()
.is_err()
);
}
let invalid_orderings = [
ConsumerRetryPolicy::builder()
.startup_attempt_timeout(secs(13))
.startup_deadline(secs(12))
.build(),
ConsumerRetryPolicy::builder()
.startup_initial_backoff(millis(3))
.startup_max_backoff(millis(2))
.build(),
ConsumerRetryPolicy::builder()
.coordinator_initial_backoff(millis(3))
.coordinator_max_backoff(millis(2))
.build(),
];
for policy in invalid_orderings {
assert!(policy.is_err());
}
}
}
#[cfg(test)]
mod security_arg_tests {
use assert2::check;
use krabka_client_core::{
ClientError, MockBroker,
security::{ClientSecurity, SaslCredentials},
};
use krabka_protocol::{
Encode,
owned::{
api_versions_request, leave_group_request,
leave_group_response::{self, LeaveGroupResponse},
},
};
use krabka_security::ListenerProtocol;
use super::*;
async fn invalid_consumer_config<F, E>(build: F) -> E
where
F: std::future::Future<Output = Result<Consumer, E>>,
E: std::fmt::Debug,
{
build.await.err().expect("invalid configuration")
}
fn leave_group_response_at(version: i16) -> Vec<u8> {
let mut buf = bytes::BytesMut::new();
if version >= leave_group_response::FLEXIBLE_MIN {
buf.extend_from_slice(&[0x00]);
}
LeaveGroupResponse::default()
.encode(&mut buf, version)
.unwrap();
buf.to_vec()
}
#[test]
fn leave_group_timeout_uses_default_and_valid_override() {
let default = ConsumerLeaveGroupTimeout::default();
assert2::assert!(default.duration() == Duration::from_secs(5));
assert2::assert!(default.milliseconds() == 5_000);
let timeout = ConsumerLeaveGroupTimeout::new(Duration::from_millis(37))
.expect("positive whole milliseconds");
assert2::assert!(timeout.duration() == Duration::from_millis(37));
assert2::assert!(timeout.milliseconds() == 37);
}
#[test]
fn leave_group_timeout_validates_millisecond_boundaries() {
assert2::assert!(ConsumerLeaveGroupTimeout::new(Duration::ZERO).is_err());
assert2::assert!(
ConsumerLeaveGroupTimeout::new(Duration::from_millis(1) + Duration::from_nanos(1))
.is_err()
);
assert2::assert!(ConsumerLeaveGroupTimeout::new(Duration::from_millis(u64::MAX)).is_ok());
assert2::assert!(ConsumerLeaveGroupTimeout::new(Duration::from_secs(u64::MAX)).is_err());
}
#[test]
fn subscription_metadata_refresh_interval_uses_default_and_valid_override() {
let default = ConsumerSubscriptionMetadataRefreshInterval::default();
assert2::assert!(default.duration() == Duration::from_secs(5));
assert2::assert!(default.milliseconds() == 5_000);
let interval = ConsumerSubscriptionMetadataRefreshInterval::new(Duration::from_millis(37))
.expect("positive whole milliseconds");
assert2::assert!(interval.duration() == Duration::from_millis(37));
assert2::assert!(interval.milliseconds() == 37);
}
#[test]
fn subscription_metadata_refresh_interval_validates_millisecond_boundaries() {
assert2::assert!(ConsumerSubscriptionMetadataRefreshInterval::new(Duration::ZERO).is_err());
assert2::assert!(
ConsumerSubscriptionMetadataRefreshInterval::new(
Duration::from_millis(1) + Duration::from_nanos(1)
)
.is_err()
);
assert2::assert!(
ConsumerSubscriptionMetadataRefreshInterval::new(Duration::from_millis(u64::MAX))
.is_ok()
);
assert2::assert!(
ConsumerSubscriptionMetadataRefreshInterval::new(Duration::from_secs(u64::MAX))
.is_err()
);
}
#[tokio::test]
async fn invalid_subscription_metadata_refresh_interval_fails_before_broker_lookup() {
let error = invalid_consumer_config(
Consumer::builder()
.bootstrap("invalid.invalid:9092")
.group_id("metadata-refresh-validation")
.subscribe(["topic".to_owned()])
.subscription_metadata_refresh_interval(krabka_units::secs(0))
.build(),
)
.await;
assert2::assert!(
error
.to_string()
.contains("consumer subscription metadata refresh interval")
);
}
#[tokio::test]
async fn negative_sub_millisecond_times_fail_before_broker_lookup() {
let negative = Time::from_secs_f64(-0.0001);
let mut actual = Vec::new();
for (name, default_api_timeout, metadata_max_age) in [
(
"default api timeout",
negative,
DEFAULT_CONSUMER_METADATA_MAX_AGE,
),
(
"metadata max age",
DEFAULT_CONSUMER_DEFAULT_API_TIMEOUT,
negative,
),
] {
let build = Consumer::builder()
.bootstrap("invalid.invalid:9092")
.group_id("negative-times")
.subscribe(["topic".to_owned()])
.default_api_timeout(default_api_timeout)
.metadata_max_age(metadata_max_age)
.build();
let error = tokio::time::timeout(Duration::from_secs(5), build)
.await
.map(|result| result.err().map(|error| error.to_string()));
actual.push((name, error.ok().flatten()));
}
assert2::assert!(
actual
== vec![
(
"default api timeout",
Some("invalid configuration: consumer default api timeout must not be negative".to_owned())
),
(
"metadata max age",
Some("invalid configuration: consumer metadata max age must not be negative".to_owned())
),
]
);
}
#[tokio::test]
async fn low_retry_backoff_max_clamps_initial_backoffs_instead_of_failing() {
let mock = MockBroker::start(|_api_key, _version, _corr_id, _body| None).await;
let build = Consumer::builder()
.bootstrap(mock.addr.to_string())
.client_id("low-retry-backoff-max-consumer")
.group_id("low-retry-backoff-max-group")
.subscribe(vec!["orders".to_string()])
.request_timeout(krabka_units::millis(50))
.retry_backoff_max(millis(100))
.build();
let res = tokio::time::timeout(Duration::from_millis(500), build).await;
mock.stop();
let Err(err) = res.expect("a valid low retry_backoff_max must not hang the build") else {
panic!("silent broker must time out during build")
};
assert2::assert!(
matches!(err, ConsumerError::Client(ClientError::Timeout(_))),
"build must fail on the connection timeout, not on config validation: {err}"
);
}
#[tokio::test]
async fn invalid_retry_backoff_max_fails_before_broker_lookup() {
let build = Consumer::builder()
.bootstrap("invalid.invalid:9092")
.group_id("retry-backoff-max")
.subscribe(["topic".to_owned()])
.retry_backoff_max(Time::ZERO)
.build();
let error = tokio::time::timeout(Duration::from_secs(5), build)
.await
.expect("build must fail before any broker lookup")
.err()
.expect("a zero retry backoff max must be rejected");
assert2::assert!(error.to_string().contains("backoff"));
}
#[tokio::test]
async fn the_subscription_of_a_build_is_checked_before_broker_lookup() {
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, topics, pattern, expected) in [
(
"a subscription without a group id",
vec!["orders".to_owned()],
None,
"To use the group management or offset commit APIs, you must provide a valid group.id in the consumer configuration.",
),
(
"both",
vec!["orders".to_owned()],
Some(TopicPattern::new(|_| true)),
"invalid configuration: subscribe and subscribe_pattern are mutually exclusive",
),
(
"empty topic name",
vec![String::new()],
None,
"invalid configuration: Topic collection to subscribe to cannot contain null or empty topic",
),
] {
let error = Consumer::builder()
.bootstrap("invalid.invalid:9092")
.maybe_group_id(
(name != "a subscription without a group id")
.then_some("subscription-validation"),
)
.subscribe(topics)
.maybe_subscribe_pattern(pattern)
.build()
.await
.err()
.map(|error| error.to_string());
actual.push((name, error));
wanted.push((name, Some(expected.to_owned())));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn the_consumer_group_protocol_options_are_checked_before_broker_lookup() {
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, protocol, assignor, regex, topics, expected) in [
(
"a remote assignor with the classic protocol",
GroupProtocol::Classic,
Some("uniform"),
None,
Vec::new(),
"invalid configuration: group.remote.assignor cannot be set when group.protocol=classic",
),
(
"a regular expression with the classic protocol",
GroupProtocol::Classic,
None,
Some("orders-.*"),
Vec::new(),
"invalid configuration: Subscribe to RE2/J pattern is not supported when using the CLASSIC protocol defined in config group.protocol",
),
(
"an empty regular expression",
GroupProtocol::Consumer,
None,
Some(""),
Vec::new(),
"invalid argument: Topic pattern to subscribe to cannot be empty",
),
(
"a regular expression and topics",
GroupProtocol::Consumer,
None,
Some("orders-.*"),
vec!["orders".to_owned()],
"invalid configuration: Subscription to topics, partitions and pattern are mutually exclusive",
),
] {
let error = Consumer::builder()
.bootstrap("invalid.invalid:9092")
.group_id("protocol-validation")
.subscribe(topics)
.group_protocol(protocol)
.maybe_group_remote_assignor(assignor)
.maybe_subscribe_regex(regex)
.build()
.await
.err()
.map(|error| error.to_string());
actual.push((name, error));
wanted.push((name, Some(expected.to_owned())));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn invalid_leave_group_timeout_fails_before_broker_lookup() {
let error = invalid_consumer_config(
Consumer::builder()
.bootstrap("invalid.invalid:9092")
.group_id("leave-validation")
.subscribe(["topic".to_owned()])
.leave_group_timeout(krabka_units::secs(0))
.build(),
)
.await;
assert2::assert!(error.to_string().contains("consumer leave-group timeout"));
}
#[tokio::test]
async fn invalid_client_resource_policy_fails_before_broker_lookup() {
let error = Consumer::builder()
.bootstrap("invalid.invalid:9092")
.group_id("client-policy-validation")
.subscribe(["topic".to_owned()])
.dispatch_queue_capacity(0)
.build()
.await
.err()
.expect("invalid configuration");
assert2::assert!(error.to_string().contains("client dispatch queue capacity"));
}
#[tokio::test]
async fn startup_member_cleanup_leaves_the_group_for_a_dynamic_member_only() {
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, group_instance_id, expected) in [
("dynamic member", None, true),
("static member", Some("instance-a".to_owned()), false),
] {
let saw_leave = Arc::new(std::sync::atomic::AtomicBool::new(false));
let saw_leave_in_mock = Arc::clone(&saw_leave);
let mock = MockBroker::start(move |api_key, version, _corr_id, _body| {
if api_key == api_versions_request::API_KEY {
Some(test_support::leave_group_api_versions())
} else if api_key == leave_group_request::API_KEY {
saw_leave_in_mock.store(true, Ordering::SeqCst);
Some(leave_group_response_at(version))
} else {
None
}
})
.await;
let client = Client::builder()
.bootstrap(mock.addr.to_string())
.request_timeout(krabka_units::millis(100))
.build()
.await
.unwrap();
let coordinator_id = AtomicI32::new(0);
leave_startup_member(
&client,
&coordinator_id,
"group-a",
"member-a",
group_instance_id,
krabka_units::millis(37),
)
.await;
mock.stop();
actual.push((name, saw_leave.load(Ordering::SeqCst)));
wanted.push((name, expected));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn startup_member_cleanup_bounds_stalled_leave_with_configured_timeout() {
let fixture = test_support::StalledLeaveGroup::start().await;
tokio::time::timeout(
Duration::from_secs(1),
leave_startup_member(
&fixture.client,
&fixture.coordinator_id,
"group-a",
"member-a",
None,
krabka_units::millis(37),
),
)
.await
.expect("configured leave deadline bounds cleanup");
fixture.broker.stop();
assert2::assert!(fixture.saw_leave.load(Ordering::SeqCst));
}
#[test]
fn initial_join_protocols_follow_the_assignor_list() {
use builder::DecodedSubscription;
let protocols = initial_join_protocols(
&[
Assignor::Range,
Assignor::Sticky,
Assignor::CooperativeSticky,
],
&["topic".into()],
Some("rack-a"),
);
let decoded: Vec<(String, DecodedSubscription)> = protocols
.iter()
.map(|(name, metadata)| (name.clone(), builder::decode_subscription(metadata)))
.collect();
let subscription = |user_data: Option<Bytes>| DecodedSubscription {
topics: vec!["topic".into()],
owned: Vec::new(),
generation_id: -1,
rack_id: Some("rack-a".into()),
user_data,
};
check!(
decoded
== vec![
("range".to_owned(), subscription(None)),
("sticky".to_owned(), subscription(None)),
(
"cooperative-sticky".to_owned(),
subscription(Some(Bytes::from_static(&[0xff, 0xff, 0xff, 0xff]))),
),
]
);
}
#[test]
fn startup_join_fields_preserve_member_timeouts_protocols_and_empty_reason() {
let metadata = bytes::Bytes::from_static(b"metadata");
let fields = startup_join_fields(
"group-a",
Some("instance-a"),
vec![
("range".to_owned(), metadata.clone()),
("cooperative-sticky".to_owned(), metadata.clone()),
],
(45_000, 60_000),
);
assert2::assert!(fields.group_id == "group-a");
assert2::assert!(fields.group_instance_id.as_deref() == Some("instance-a"));
assert2::assert!(fields.session_timeout_ms == 45_000);
assert2::assert!(fields.rebalance_timeout_ms == 60_000);
assert2::assert!(
fields.protocols
== vec![
("range".to_owned(), metadata.clone()),
("cooperative-sticky".to_owned(), metadata),
]
);
assert2::assert!(fields.reason.is_empty());
}
#[test]
fn first_join_member_id_accepts_required_or_success_and_rejects_errors() {
enum Expected<'a> {
Member(&'a str),
Server(i16),
RebalanceFailed,
}
for (_name, error_code, member_id, expected) in [
(
"member id required",
79,
"member-a",
Expected::Member("member-a"),
),
("success", 0, "member-b", Expected::Member("member-b")),
("server error", 42, "member-c", Expected::Server(42)),
("empty member", 0, "", Expected::RebalanceFailed),
] {
let response = JoinGroupResponse {
error_code,
member_id: member_id.into(),
..Default::default()
};
let actual = first_join_member_id(&response);
let matches = match expected {
Expected::Member(expected) => matches!(actual, Ok(actual) if actual == expected),
Expected::Server(expected) => {
matches!(actual, Err(ConsumerError::Server(code)) if code == expected)
}
Expected::RebalanceFailed => {
matches!(actual, Err(ConsumerError::RebalanceFailed(_)))
}
};
assert2::assert!(matches);
}
}
#[test]
fn is_subscribed_topic_matches_exact_topic_names() {
let subscribe = vec!["orders".to_string(), "payments".to_string()];
for (name, expected) in [
("orders", true),
("payments", true),
("shipments", false),
("order", false),
] {
assert2::assert!(is_subscribed_topic(&subscribe, name) == expected);
}
}
#[test]
fn is_group_leader_matches_exact_member_id() {
for (_name, leader, member_id, expected) in [
("exact leader", "member-a", "member-a", true),
("different member", "member-a", "member-b", false),
("empty member", "member-a", "", false),
] {
assert2::assert!(is_group_leader(leader, member_id) == expected);
}
}
#[test]
fn build_sync_assignment_preserves_member_and_assignment_payload() {
let assignment =
coordinator::build_sync_group_assignment("member-a".into(), &[("orders".into(), 3)]);
let decoded = decode_assignment(&assignment.assignment);
assert2::assert!(
(assignment.member_id.as_str(), decoded) == ("member-a", vec![("orders".into(), 3)])
);
}
#[test]
fn offset_prime_helpers_preserve_assignment_presence_offsets_and_epochs() {
for (_name, partitions, expected) in [
("empty assignment", vec![], false),
("assigned partition", vec![("orders".to_string(), 0)], true),
] {
assert2::assert!(has_assigned_partitions(&partitions) == expected);
}
test_support::assert_starting_offsets(starting_offset);
let position = primed_position(9);
assert2::assert!(
position
== position::PartitionPosition {
offset_epoch: krabka_ids::LeaderEpoch(9),
leader_id: -1,
leader_epoch: krabka_ids::LeaderEpoch(-1),
awaiting_validation: false,
}
);
}
#[tokio::test]
async fn consumer_builder_uses_request_timeout_for_initial_connect_handshake() {
let mock = MockBroker::start(|_api_key, _version, _corr_id, _body| None).await;
let build = Consumer::builder()
.bootstrap(mock.addr.to_string())
.client_id("timeout-consumer")
.group_id("timeout-group")
.subscribe(vec!["orders".to_string()])
.request_timeout(krabka_units::millis(100))
.build();
let res = tokio::time::timeout(Duration::from_millis(500), build).await;
mock.stop();
let Err(err) = res.expect("consumer build should not retain the 30s connect timeout")
else {
panic!("silent broker must time out during build")
};
assert2::assert!(matches!(
err,
ConsumerError::Client(ClientError::Timeout(d))
if d == krabka_units::millis(100)
));
}
#[tokio::test]
async fn a_refused_start_records_no_error() {
use krabka_client_core::mock_test_support::recorded_levels;
let (invalid, levels) = recorded_levels(Box::pin(
Consumer::builder()
.bootstrap("127.0.0.1:1")
.group_id("log-level-group")
.subscribe(vec!["orders".to_string()])
.fetch_min(0)
.build(),
))
.await;
assert2::assert!(matches!(invalid, Err(ConsumerError::InvalidConfig(_))));
check!(levels == vec![tracing::Level::INFO]);
let mock = MockBroker::start(|_api_key, _version, _corr_id, _body| None).await;
let (silent, levels) = recorded_levels(Box::pin(
Consumer::builder()
.bootstrap(mock.addr.to_string())
.group_id("log-level-group")
.subscribe(vec!["orders".to_string()])
.request_timeout(millis(100))
.build(),
))
.await;
mock.stop();
assert2::assert!(silent.is_err());
check!(!levels.is_empty());
check!(!levels.contains(&tracing::Level::ERROR), "{levels:?}");
}
#[tokio::test]
async fn consumer_builder_rejects_non_positive_fetch_budgets_before_network_io() {
let min = Consumer::builder()
.bootstrap("127.0.0.1:1")
.group_id("timeout-group")
.subscribe(vec!["orders".to_string()])
.fetch_min(0)
.build()
.await;
assert2::assert!(matches!(min, Err(ConsumerError::InvalidConfig(_))));
let max = Consumer::builder()
.bootstrap("127.0.0.1:1")
.group_id("timeout-group")
.subscribe(vec!["orders".to_string()])
.fetch_max(0)
.build()
.await;
assert2::assert!(matches!(max, Err(ConsumerError::InvalidConfig(_))));
let partition_max = Consumer::builder()
.bootstrap("127.0.0.1:1")
.group_id("timeout-group")
.subscribe(vec!["orders".to_string()])
.fetch_partition_max(0)
.build()
.await;
assert2::assert!(matches!(
partition_max,
Err(ConsumerError::InvalidConfig(_))
));
let inverted = Consumer::builder()
.bootstrap("127.0.0.1:1")
.group_id("timeout-group")
.subscribe(vec!["orders".to_string()])
.fetch_min(2)
.fetch_max(1)
.build()
.await;
assert2::assert!(matches!(
inverted,
Err(ConsumerError::InvalidConfig(message))
if message == "consumer fetch min must not exceed consumer fetch max"
));
}
#[test]
fn consumer_fetch_byte_settings_validate_and_convert_to_uom() {
for value in [1, i32::MAX] {
let size = ByteSize::from_bytes_i64(i64::from(value));
let max = ConsumerFetchMaxBytes::try_from(size).expect("positive fetch maximum");
let partition_max = ConsumerFetchPartitionMaxBytes::try_from(size)
.expect("positive partition fetch maximum");
check!(max.bytes() == value);
check!(partition_max.bytes() == value);
check!(max.size().bytes_i32() == value);
check!(partition_max.size().bytes_i32() == value);
}
for invalid in [
ByteSize::ZERO,
ByteSize::from_bytes_f64(-1.0),
ByteSize::from_bytes_f64(1.5),
ByteSize::from_bytes_f64(f64::from(i32::MAX) + 1.0),
] {
check!(ConsumerFetchMaxBytes::try_from(invalid).is_err());
check!(ConsumerFetchPartitionMaxBytes::try_from(invalid).is_err());
}
}
async fn test_consumer() -> Consumer {
let client = Client::builder()
.bootstrap("127.0.0.1:1")
.client_id("test-client")
.build()
.await
.unwrap();
Consumer::test_builder(client)
.coordinator_id(Arc::new(AtomicI32::new(3)))
.generation_id(7)
.commit_identity(test_support::orders_partition_commit_identity(7))
.group_instance_id("instance-a")
.subscription(subscription::shared(
vec!["orders".into(), "payments".into()],
None,
true,
))
.assigned(vec![("orders".into(), 0)])
.build()
}
#[tokio::test]
async fn accessors_return_consumer_identity_subscription_and_assignment() {
let consumer = test_consumer().await;
check!(
(
consumer.group_id(),
consumer.member_id(),
consumer.generation_id(),
consumer.subscription(),
consumer.assignment().await,
) == (
"group-a",
"member-a".to_owned(),
7,
vec!["orders".to_string(), "payments".to_string()],
vec![("orders".into(), 0)],
)
);
let metadata = consumer.group_metadata();
assert2::assert!(
metadata
== ConsumerGroupMetadata {
group_id: "group-a".into(),
generation_id: 7,
member_id: "member-a".into(),
group_instance_id: Some("instance-a".into()),
}
);
}
#[tokio::test]
async fn log_end_requires_every_assigned_position_to_reach_a_reported_end() {
let consumer = test_consumer().await;
assert2::assert!(!consumer.at_log_end().await);
consumer
.next_offsets
.lock()
.await
.insert(("orders".into(), 0), 12);
consumer
.end_offsets
.lock()
.await
.insert(("orders".into(), 0), 12);
assert2::assert!(consumer.at_log_end().await);
consumer.seek("orders", 0, 0).await.unwrap();
let end = consumer
.end_offsets
.lock()
.await
.get(&("orders".to_owned(), 0))
.copied();
assert2::assert!((consumer.at_log_end().await, end) == (false, Some(12)));
}
#[tokio::test]
async fn generation_tracks_coordinator_rejoins_via_shared_atomic() {
let consumer = test_consumer().await;
assert2::assert!(
(
consumer.generation_id(),
consumer.group_metadata().generation_id
) == (7, 7)
);
consumer.current_generation.store(11, Ordering::Relaxed);
assert2::assert!(
(
consumer.generation_id(),
consumer.group_metadata().generation_id
) == (11, 11)
);
}
#[tokio::test]
async fn close_cancels_shutdown_token_even_without_spawned_handle() {
let consumer = test_consumer().await;
let shutdown = consumer.coordinator_shutdown.clone();
consumer.close().await.unwrap();
assert2::assert!(shutdown.is_cancelled());
}
#[test]
fn retriable_error_classification() {
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
use krabka_client_core::ClientError;
let transient_codes: &[i16] = &[14, 15, 16, 22, 25, 27, 79];
for &code in transient_codes {
assert2::assert!(is_retriable_consumer_start_error(&ConsumerError::Server(
code
)));
}
assert2::assert!(!is_retriable_consumer_start_error(&ConsumerError::Server(
42
)));
for (_name, error, expected) in [
(
"disconnected",
ConsumerError::Client(ClientError::Disconnected),
true,
),
(
"timeout",
ConsumerError::Client(ClientError::Timeout(krabka_units::secs(1))),
false,
),
(
"startup after join",
ConsumerError::StartupAfterJoin(Box::new(ConsumerError::Client(
ClientError::Timeout(krabka_units::secs(1)),
))),
true,
),
(
"connect refused",
ConsumerError::Client(ClientError::Connect {
addr: SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 9092),
source: std::io::Error::new(std::io::ErrorKind::ConnectionRefused, "refused"),
}),
false,
),
("not subscribed", ConsumerError::NotSubscribed, false),
(
"invalid config",
ConsumerError::InvalidConfig("group_id required".into()),
false,
),
(
"incompatible version",
ConsumerError::Client(ClientError::IncompatibleVersion {
api_key: 0,
broker_min: 0,
broker_max: 5,
client_min: 7,
client_max: 10,
}),
false,
),
] {
assert2::assert!(is_retriable_consumer_start_error(&error) == expected);
}
}
#[tokio::test]
async fn consumer_builder_accepts_security() {
let security = ClientSecurity {
protocol: ListenerProtocol::SaslPlaintext,
tls: None,
sasl: Some(SaslCredentials::Plain {
username: "u".into(),
password: "p".into(),
}),
sasl_host: None,
};
let res = Consumer::builder()
.bootstrap("127.0.0.1:1")
.group_id("g")
.subscribe(vec!["t".to_string()])
.security(security)
.build()
.await;
assert2::assert!(res.is_err());
}
}
#[cfg(test)]
mod auto_commit_tests {
use std::{collections::VecDeque, sync::atomic::AtomicI16};
use bytes::Buf as _;
use krabka_client_core::{MockBroker, mock_test_support::offset_commit_response};
use krabka_protocol::{
Decode as _, Encode, UnknownTaggedFields,
owned::{
api_versions_request, fetch_request,
fetch_response::FetchResponse,
find_coordinator_request,
find_coordinator_response::FindCoordinatorResponse,
heartbeat_request,
heartbeat_response::HeartbeatResponse,
join_group_request::{self, JoinGroupRequest},
leave_group_request::{self, LeaveGroupRequest},
leave_group_response::LeaveGroupResponse,
list_offsets_request, metadata_request,
metadata_response::MetadataResponse,
offset_commit_request::{
self, OffsetCommitRequest, OffsetCommitRequestPartition, OffsetCommitRequestTopic,
},
offset_fetch_request,
offset_fetch_response::OffsetFetchResponse,
sync_group_request,
},
};
use super::*;
const GROUP: &str = "group-a";
pub(super) const MEMBER: &str = "member-a";
pub(super) const TOPIC: &str = "orders";
const REBALANCE_IN_PROGRESS: i16 = 27;
const AUTO_COMMIT_INTERVAL: Duration = Duration::from_secs(5);
const SESSION_TIMEOUT: Duration = Duration::from_secs(45);
const REBALANCE_WAIT: Duration = Duration::from_secs(30);
#[derive(Clone, Debug, PartialEq)]
pub(super) enum GroupRequest {
OffsetCommit(OffsetCommitRequest),
JoinGroup,
SyncGroup,
LeaveGroup,
CommitSyncReturned,
SyncGroupLate,
CloseTimedOut,
SessionExpired,
}
#[derive(Clone, Copy, Debug)]
pub(super) enum CommitReply {
Drop,
Error(i16),
}
const API_VERSIONS: [(i16, i16, i16); 11] = [
(api_versions_request::API_KEY, 0, 3),
(metadata_request::API_KEY, 0, 8),
(find_coordinator_request::API_KEY, 0, 2),
(join_group_request::API_KEY, 0, 8),
(sync_group_request::API_KEY, 0, 3),
(heartbeat_request::API_KEY, 0, 3),
(leave_group_request::API_KEY, 0, 5),
(offset_commit_request::API_KEY, 2, 7),
(offset_fetch_request::API_KEY, 1, 5),
(fetch_request::API_KEY, 4, 11),
(list_offsets_request::API_KEY, 1, 5),
];
fn encoded_api_versions() -> Vec<u8> {
encode(
&krabka_client_core::mock_test_support::api_versions_response_ranges(
API_VERSIONS.iter().copied(),
),
0,
)
}
pub(super) struct MockCoordinator {
protocol: &'static str,
requests: std::sync::Mutex<Vec<GroupRequest>>,
pub(super) heartbeat_error: AtomicI16,
assignments: std::sync::Mutex<VecDeque<Vec<(String, i32)>>>,
generation: AtomicI32,
pub(super) commit_replies: std::sync::Mutex<VecDeque<CommitReply>>,
rebalance_on_commit: std::sync::atomic::AtomicBool,
pub(super) last_heartbeat: std::sync::Mutex<tokio::time::Instant>,
drop_heartbeats: std::sync::atomic::AtomicBool,
pub(super) joins: std::sync::Mutex<Vec<JoinGroupRequest>>,
pub(super) leaves: std::sync::Mutex<Vec<LeaveGroupRequest>>,
pub(super) drop_leaves: std::sync::atomic::AtomicBool,
}
pub(super) async fn mock_coordinator_broker(coordinator: &Arc<MockCoordinator>) -> MockBroker {
let in_mock = Arc::clone(coordinator);
MockBroker::start(move |api_key, version, _corr_id, body| {
in_mock.respond(api_key, version, body)
})
.await
}
pub(super) fn single_partition_startup(topic_partition_count: i32) -> StartupState {
StartupState::builder()
.generation_id(1)
.assigned_partitions(vec![partition(0)])
.next_offsets(maplit::hashmap! {partition(0) => 12})
.topic_partitions(maplit::hashmap! {TOPIC.to_owned() => topic_partition_count})
.build()
}
pub(super) async fn start_consumer_with_startup(
config: StartConfig,
startup: StartupState,
) -> Consumer {
start_test_consumer(
TestConsumerStart::builder()
.bootstrap(config.bootstrap.clone())
.config(config)
.startup(startup)
.build(),
)
.await
}
pub(super) async fn start_single_partition_consumer(config: StartConfig) -> Consumer {
start_consumer_with_startup(config, single_partition_startup(1)).await
}
pub(super) fn two_partition_startup() -> StartupState {
StartupState::builder()
.generation_id(1)
.assigned_partitions(vec![partition(0), partition(1)])
.next_offsets(maplit::hashmap! {partition(0) => 12, partition(1) => 7})
.topic_partitions(maplit::hashmap! {TOPIC.to_owned() => 2})
.build()
}
fn encode(response: &impl Encode, version: i16) -> Vec<u8> {
let mut buf = bytes::BytesMut::new();
response.encode(&mut buf, version).expect("encode response");
buf.to_vec()
}
impl MockCoordinator {
pub(super) fn range_partition(index: i32) -> Arc<Self> {
Self::new(Assignor::Range, vec![vec![partition(index)]])
}
pub(super) fn new(assignor: Assignor, assignments: Vec<Vec<(String, i32)>>) -> Arc<Self> {
Arc::new(Self {
protocol: assignor.protocol_name(),
requests: std::sync::Mutex::new(Vec::new()),
heartbeat_error: AtomicI16::new(0),
assignments: std::sync::Mutex::new(assignments.into()),
generation: AtomicI32::new(1),
commit_replies: std::sync::Mutex::new(VecDeque::new()),
rebalance_on_commit: std::sync::atomic::AtomicBool::new(false),
last_heartbeat: std::sync::Mutex::new(tokio::time::Instant::now()),
drop_heartbeats: std::sync::atomic::AtomicBool::new(false),
joins: std::sync::Mutex::new(Vec::new()),
leaves: std::sync::Mutex::new(Vec::new()),
drop_leaves: std::sync::atomic::AtomicBool::new(false),
})
}
pub(super) fn set_heartbeat_error(&self, error_code: i16) {
self.heartbeat_error.store(error_code, Ordering::SeqCst);
}
fn record(&self, request: GroupRequest) {
self.requests.lock().expect("requests lock").push(request);
}
pub(super) fn requests(&self) -> Vec<GroupRequest> {
self.requests.lock().expect("requests lock").clone()
}
pub(super) fn sync_groups(&self) -> usize {
self.requests()
.iter()
.filter(|request| **request == GroupRequest::SyncGroup)
.count()
}
pub(super) fn respond(
&self,
api_key: i16,
version: i16,
mut body: &[u8],
) -> Option<Vec<u8>> {
match api_key {
api_versions_request::API_KEY => Some(encoded_api_versions()),
metadata_request::API_KEY => Some(encode(&MetadataResponse::default(), version)),
heartbeat_request::API_KEY => {
*self.last_heartbeat.lock().expect("last heartbeat lock") =
tokio::time::Instant::now();
let error_code = self.heartbeat_error.swap(0, Ordering::SeqCst);
if error_code == 0 && self.drop_heartbeats.load(Ordering::SeqCst) {
return None;
}
Some(encode(
&HeartbeatResponse {
error_code,
..Default::default()
},
version,
))
}
join_group_request::API_KEY => {
let mut last_heartbeat =
self.last_heartbeat.lock().expect("last heartbeat lock");
if last_heartbeat.elapsed() > SESSION_TIMEOUT {
self.record(GroupRequest::SessionExpired);
}
*last_heartbeat = tokio::time::Instant::now();
drop(last_heartbeat);
self.record(GroupRequest::JoinGroup);
let client_id_len = body.get_i16();
body.advance(usize::try_from(client_id_len.max(0)).expect("client id length"));
let flexible = version >= join_group_request::FLEXIBLE_MIN;
if flexible {
body.advance(1);
}
self.joins
.lock()
.expect("joins lock")
.push(JoinGroupRequest::decode(&mut body, version).expect("decode join"));
let mut response = if flexible { vec![0] } else { Vec::new() };
response.extend(encode(
&JoinGroupResponse {
generation_id: self.generation.fetch_add(1, Ordering::SeqCst) + 1,
protocol_name: Some(self.protocol.to_owned()),
protocol_type: Some("consumer".into()),
leader: "leader".into(),
member_id: MEMBER.into(),
..Default::default()
},
version,
));
Some(response)
}
sync_group_request::API_KEY => {
self.record(GroupRequest::SyncGroup);
let mut assignments = self.assignments.lock().expect("assignments lock");
let assignment = if assignments.len() > 1 {
assignments.pop_front().expect("assignment")
} else {
assignments.front().cloned().expect("assignment")
};
drop(assignments);
Some(encode(
&SyncGroupResponse {
assignment: builder::encode_assignment(&assignment),
..Default::default()
},
version,
))
}
fetch_request::API_KEY => Some(encode(&FetchResponse::default(), version)),
list_offsets_request::API_KEY => {
Some(test_support::list_offsets_mock_response(&mut body, version))
}
offset_fetch_request::API_KEY => {
Some(encode(&OffsetFetchResponse::default(), version))
}
leave_group_request::API_KEY => {
self.record(GroupRequest::LeaveGroup);
let client_id_len = body.get_i16();
body.advance(usize::try_from(client_id_len.max(0)).expect("client id length"));
let flexible = version >= leave_group_request::FLEXIBLE_MIN;
if flexible {
body.advance(1);
}
self.leaves
.lock()
.expect("leaves lock")
.push(LeaveGroupRequest::decode(&mut body, version).expect("decode leave"));
if self.drop_leaves.load(Ordering::SeqCst) {
return None;
}
let mut response = if flexible { vec![0] } else { Vec::new() };
response.extend(encode(&LeaveGroupResponse::default(), version));
Some(response)
}
find_coordinator_request::API_KEY => Some(encode(
&FindCoordinatorResponse {
node_id: 0,
..Default::default()
},
version,
)),
offset_commit_request::API_KEY => self.respond_offset_commit(version, body),
_ => None,
}
}
fn respond_offset_commit(&self, version: i16, mut body: &[u8]) -> Option<Vec<u8>> {
let client_id_len = body.get_i16();
body.advance(usize::try_from(client_id_len.max(0)).expect("client id length"));
let mut request =
OffsetCommitRequest::decode(&mut body, version).expect("decode commit");
request.topics.sort_by(|a, b| a.name.cmp(&b.name));
for topic in &mut request.topics {
topic.partitions.sort_by_key(|p| p.partition_index);
}
if self.rebalance_on_commit.swap(false, Ordering::SeqCst) {
self.heartbeat_error
.store(REBALANCE_IN_PROGRESS, Ordering::SeqCst);
}
let error_code = match self
.commit_replies
.lock()
.expect("commit replies lock")
.pop_front()
{
None => 0,
Some(CommitReply::Error(code)) => code,
Some(CommitReply::Drop) => {
self.record(GroupRequest::OffsetCommit(request));
return None;
}
};
let response = offset_commit_response(&request, |_| error_code);
self.record(GroupRequest::OffsetCommit(request));
Some(encode(&response, version))
}
}
pub(super) fn partition(index: i32) -> (String, i32) {
(TOPIC.to_owned(), index)
}
fn commit(generation: i32, offsets: &[(i32, i64, i32)]) -> GroupRequest {
GroupRequest::OffsetCommit(OffsetCommitRequest {
group_id: GROUP.into(),
generation_id_or_member_epoch: generation,
member_id: MEMBER.into(),
group_instance_id: None,
retention_time_ms: -1,
topics: vec![OffsetCommitRequestTopic {
name: TOPIC.into(),
partitions: offsets
.iter()
.map(
|(partition_index, committed_offset, committed_leader_epoch)| {
OffsetCommitRequestPartition {
partition_index: *partition_index,
committed_offset: *committed_offset,
committed_leader_epoch: *committed_leader_epoch,
committed_metadata: Some(String::new()),
unknown_tagged_fields: UnknownTaggedFields::default(),
}
},
)
.collect(),
..Default::default()
}],
unknown_tagged_fields: UnknownTaggedFields::default(),
})
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn metadata_requests_name_the_subscribed_topics() {
use krabka_protocol::owned::{
find_coordinator_response::FindCoordinatorResponse, metadata_request::MetadataRequest,
};
let named = |topics: &[&str], allow| {
krabka_client_core::topics_request(topics.iter().map(|&t| t.to_owned()), allow)
};
for (name, subscribe, allow_auto_create_topics, expected) in [
(
"subscribe a and b",
vec!["a", "b"],
true,
named(&["a", "b"], true),
),
("auto creation off", vec!["a"], false, named(&["a"], false)),
] {
let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
let handler_requests = Arc::clone(&requests);
let mock = MockBroker::start(move |api_key, version, _corr_id, mut body| {
match api_key {
api_versions_request::API_KEY => Some(encoded_api_versions()),
find_coordinator_request::API_KEY => Some(encode(
&FindCoordinatorResponse {
node_id: 1,
host: "127.0.0.1".into(),
port: 0,
..Default::default()
},
version,
)),
metadata_request::API_KEY => {
let client_id_len = body.get_i16();
body.advance(usize::try_from(client_id_len.max(0)).expect("length"));
let request =
MetadataRequest::decode(&mut body, version).expect("decode Metadata");
handler_requests
.lock()
.expect("requests lock")
.push(request);
Some(encode(&MetadataResponse::default(), version))
}
_ => None,
}
})
.await;
let mut config =
start_config(mock.addr.to_string(), Assignor::Range, false, minutes(1));
config.subscribe = subscribe.into_iter().map(str::to_owned).collect();
config.allow_auto_create_topics = allow_auto_create_topics;
let start = tokio::spawn(Consumer::start_once(config));
let first = test_support::FirstRecorded::builder()
.requests(Arc::clone(&requests))
.timeout(Duration::from_secs(10))
.poll_interval(Duration::from_millis(5))
.build()
.wait()
.await
.expect("the consumer sends Metadata");
start.abort();
mock.stop();
assert2::check!(first == expected, "{name}");
}
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn the_coordinator_task_of_a_leader_build_watches_the_group_topics() {
use krabka_protocol::owned::{
find_coordinator_response::FindCoordinatorResponse,
heartbeat_response::HeartbeatResponse,
};
let requests = Arc::new(std::sync::Mutex::new(Vec::new()));
let handler_requests = Arc::clone(&requests);
let mock = MockBroker::start(move |api_key, version, _corr_id, mut body| match api_key {
api_versions_request::API_KEY => Some(encoded_api_versions()),
find_coordinator_request::API_KEY => {
Some(encode(&FindCoordinatorResponse::default(), version))
}
heartbeat_request::API_KEY => Some(encode(&HeartbeatResponse::default(), version)),
metadata_request::API_KEY => {
let names = test_support::metadata_topic_names(&mut body, version);
handler_requests.lock().expect("requests lock").push(names);
Some(encode(&MetadataResponse::default(), version))
}
_ => None,
})
.await;
let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, minutes(1));
config.heartbeat_interval = millis(50);
config.subscription_metadata_refresh_interval = millis(50);
let consumer = start_test_consumer(
TestConsumerStart::builder()
.bootstrap(mock.addr.to_string())
.config(config)
.startup(
StartupState::builder()
.generation_id(1)
.group_topics(vec![TOPIC.to_owned(), "payments".to_owned()])
.assigned_partitions(vec![partition(0)])
.next_offsets(maplit::hashmap! {partition(0) => 12})
.topic_partitions(maplit::hashmap! {TOPIC.to_owned() => 1})
.build(),
)
.build(),
)
.await;
let first = test_support::FirstRecorded::builder()
.requests(Arc::clone(&requests))
.timeout(Duration::from_secs(5))
.poll_interval(Duration::from_millis(10))
.build()
.wait()
.await;
drop(consumer);
mock.stop();
assert2::assert!(first == Ok(vec![TOPIC.to_owned(), "payments".to_owned()]));
}
pub(super) fn start_config(
bootstrap: String,
assignor: Assignor,
auto_commit: bool,
rebalance_timeout: Time,
) -> StartConfig {
StartConfig {
bootstrap,
client_id: "auto-commit-test".into(),
group_id: GROUP.into(),
session_timeout: Time::from_std(SESSION_TIMEOUT),
max_poll_interval: rebalance_timeout,
max_poll_records: DEFAULT_CONSUMER_MAX_POLL_RECORDS,
heartbeat_interval: secs(3),
subscription_metadata_refresh_interval: minutes(60),
subscribe: vec![TOPIC.into()],
subscribe_pattern: None,
subscribe_regex: None,
exclude_internal_topics: true,
group_protocol: GroupProtocol::Classic,
server_assignor: None,
group_instance_id: None,
auto_offset_reset: AutoOffsetReset::Earliest,
isolation_level: IsolationLevel::ReadUncommitted,
assignors: vec![assignor],
rebalance_protocol: assignor::rebalance_protocol_of(&[assignor])
.expect("assignor protocol"),
fetch_min: krabka_client_core::DEFAULT_FETCH_MIN,
fetch_max: poll::DEFAULT_FETCH_MAX,
fetch_partition_max: poll::DEFAULT_FETCH_PARTITION_MAX,
fetch_max_wait: DEFAULT_CONSUMER_FETCH_MAX_WAIT,
metadata_max_age: DEFAULT_CONSUMER_METADATA_MAX_AGE,
default_api_timeout: DEFAULT_CONSUMER_DEFAULT_API_TIMEOUT,
request_timeout: secs(30),
dns_timeout: krabka_client_core::DEFAULT_CLIENT_DNS_TIMEOUT,
send_buffer: Some(krabka_client_core::DEFAULT_SEND_BUFFER),
receive_buffer: Some(krabka_client_core::DEFAULT_RECEIVE_BUFFER),
dispatch_queue_capacity: ConnectionDispatchQueueCapacity::new(
krabka_client_core::DEFAULT_CONNECTION_DISPATCH_QUEUE_CAPACITY,
)
.expect("dispatch queue capacity"),
frame_max: ClientFrameMax::try_from(krabka_client_core::DEFAULT_CLIENT_FRAME_MAX)
.expect("frame max"),
metadata_recovery_strategy: krabka_client_core::MetadataRecoveryStrategy::default(),
metadata_recovery_rebootstrap_trigger:
krabka_client_core::DEFAULT_METADATA_RECOVERY_REBOOTSTRAP_TRIGGER,
leave_group_timeout: secs(5),
client_rack: None,
security: None,
client_dns_lookup: krabka_client_core::ClientDnsLookup::default(),
reconnect_backoff: krabka_client_core::DEFAULT_RECONNECT_BACKOFF,
reconnect_backoff_max: krabka_client_core::DEFAULT_RECONNECT_BACKOFF_MAX,
connections_max_idle: krabka_client_core::DEFAULT_CONNECTIONS_MAX_IDLE,
socket_connection_setup_timeout:
krabka_client_core::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT,
socket_connection_setup_timeout_max:
krabka_client_core::DEFAULT_SOCKET_CONNECTION_SETUP_TIMEOUT_MAX,
has_rebalance_listener: false,
retry_policy: ConsumerRetryPolicy::default(),
auto_commit_interval: auto_commit.then_some(AUTO_COMMIT_INTERVAL),
allow_auto_create_topics: true,
telemetry: None,
}
}
#[derive(Clone, Copy, Debug)]
enum Step {
Poll,
AdvanceInterval,
ReceiveRecords,
Rebalance { syncs: usize, within: Duration },
CommitSync,
CommitSyncDuringRebalance {
reply: CommitReply,
syncs: usize,
within: Duration,
},
ReplyToNextCommit(CommitReply),
HoldCommitLock,
DropHeartbeats,
Close,
}
pub(super) async fn wait_for_sync_groups(
coordinator: &MockCoordinator,
polls: &coordinator::PollSignal,
syncs: usize,
within: Duration,
) {
let waited = tokio::time::timeout(within, async {
while coordinator.sync_groups() < syncs {
coordinator::note_poll(polls);
tokio::time::sleep(Duration::from_millis(100)).await;
}
})
.await;
if waited.is_err() {
coordinator.record(GroupRequest::SyncGroupLate);
}
tokio::time::sleep(Duration::from_secs(1)).await;
}
async fn run(
auto_commit: bool,
assignor: Assignor,
assignments: Vec<Vec<(String, i32)>>,
steps: &[Step],
) -> Vec<GroupRequest> {
run_with_rebalance_timeout(auto_commit, assignor, assignments, minutes(1), steps).await
}
async fn run_with_rebalance_timeout(
auto_commit: bool,
assignor: Assignor,
assignments: Vec<Vec<(String, i32)>>,
rebalance_timeout: Time,
steps: &[Step],
) -> Vec<GroupRequest> {
let coordinator = MockCoordinator::new(assignor, assignments);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let config = start_config(
mock.addr.to_string(),
assignor,
auto_commit,
rebalance_timeout,
);
let client = Client::builder()
.bootstrap(mock.addr.to_string())
.build()
.await
.expect("client");
let mut consumer = Some(
spawn_consumer(
config,
client,
Arc::new(AtomicI32::new(0)),
MEMBER.into(),
StartupState::builder()
.generation_id(1)
.assigned_partitions(vec![partition(0), partition(1)])
.next_offsets(maplit::hashmap! {partition(0) => 12, partition(1) => 7})
.positions(maplit::hashmap! {partition(0) => primed_position(3)})
.topic_partitions(maplit::hashmap! {TOPIC.to_owned() => 2})
.build(),
)
.await
.expect("spawn consumer"),
);
let mut held_commit_lock = None;
for step in steps {
match step {
Step::Poll => {
let consumer = consumer.as_ref().expect("open consumer");
consumer.maybe_auto_commit_async().await;
drop(consumer.commit_serialization.lock().await);
}
Step::AdvanceInterval => tokio::time::advance(AUTO_COMMIT_INTERVAL).await,
Step::ReceiveRecords => {
let consumer = consumer.as_ref().expect("open consumer");
consumer.next_offsets.lock().await.insert(partition(0), 20);
}
Step::Rebalance { syncs, within } => {
let consumer = consumer.as_ref().expect("open consumer");
coordinator
.heartbeat_error
.store(REBALANCE_IN_PROGRESS, Ordering::SeqCst);
wait_for_sync_groups(&coordinator, &consumer.poll_signal, *syncs, *within)
.await;
}
Step::CommitSync => {
let consumer = consumer.as_ref().expect("open consumer");
consumer.commit_sync().await.expect("commit_sync");
coordinator.record(GroupRequest::CommitSyncReturned);
}
Step::CommitSyncDuringRebalance {
reply,
syncs,
within,
} => {
let consumer = consumer.as_ref().expect("open consumer");
coordinator
.commit_replies
.lock()
.expect("commit replies lock")
.push_back(*reply);
coordinator
.rebalance_on_commit
.store(true, Ordering::SeqCst);
let commit = async {
let _ = consumer.commit_sync().await;
coordinator.record(GroupRequest::CommitSyncReturned);
};
tokio::join!(
commit,
wait_for_sync_groups(&coordinator, &consumer.poll_signal, *syncs, *within)
);
}
Step::ReplyToNextCommit(reply) => {
coordinator
.commit_replies
.lock()
.expect("commit replies lock")
.push_back(*reply);
}
Step::DropHeartbeats => {
coordinator.drop_heartbeats.store(true, Ordering::SeqCst);
}
Step::HoldCommitLock => {
let consumer = consumer.as_ref().expect("open consumer");
held_commit_lock = Some(
Arc::clone(&consumer.commit_serialization)
.lock_owned()
.await,
);
}
Step::Close => {
let close = consumer.take().expect("open consumer").close();
match tokio::time::timeout(Duration::from_mins(2), close).await {
Ok(result) => result.expect("close"),
Err(_) => coordinator.record(GroupRequest::CloseTimedOut),
}
}
}
}
let requests = coordinator.requests();
drop(held_commit_lock);
drop(consumer);
mock.stop();
requests
}
#[derive(bon::Builder)]
struct AutoCommitCase {
name: &'static str,
#[builder(default = true)]
auto_commit: bool,
#[builder(default = Assignor::Range)]
assignor: Assignor,
#[builder(default)]
assignments: Vec<Vec<(String, i32)>>,
steps: Vec<Step>,
expected: Vec<GroupRequest>,
}
fn rebalance_steps(syncs: usize) -> Vec<Step> {
use Step::{Poll, Rebalance, ReceiveRecords};
vec![
Poll,
ReceiveRecords,
Rebalance {
syncs,
within: REBALANCE_WAIT,
},
]
}
#[derive(bon::Builder)]
struct AutoCommitTimeoutCase {
name: &'static str,
#[builder(default = minutes(1))]
rebalance_timeout: Time,
steps: Vec<Step>,
expected: Vec<GroupRequest>,
}
trait NamedGroupRequestCase {
fn name(&self) -> &'static str;
fn expected(&self) -> &[GroupRequest];
}
impl NamedGroupRequestCase for AutoCommitCase {
fn name(&self) -> &'static str {
self.name
}
fn expected(&self) -> &[GroupRequest] {
&self.expected
}
}
impl NamedGroupRequestCase for AutoCommitTimeoutCase {
fn name(&self) -> &'static str {
self.name
}
fn expected(&self) -> &[GroupRequest] {
&self.expected
}
}
#[derive(Debug)]
struct GroupRequestObservation {
name: &'static str,
actual: Vec<GroupRequest>,
expected: Vec<GroupRequest>,
}
async fn assert_group_request_cases<C, F, Fut>(cases: Vec<C>, mut run: F)
where
C: NamedGroupRequestCase,
F: FnMut(C) -> Fut,
Fut: std::future::Future<Output = Vec<GroupRequest>>,
{
let mut observations = Vec::with_capacity(cases.len());
for case in cases {
let name = case.name();
let expected = case.expected().to_vec();
observations.push(GroupRequestObservation {
name,
actual: run(case).await,
expected,
});
}
let failed_cases: Vec<_> = observations
.iter()
.filter(|observation| observation.actual != observation.expected)
.map(|observation| observation.name)
.collect();
assert2::assert!(failed_cases.is_empty());
}
fn rebalance_step(within: Duration) -> Step {
Step::Rebalance { syncs: 1, within }
}
fn commit_sync_then_rebalance_requests(received: &[(i32, i64, i32)]) -> Vec<GroupRequest> {
use GroupRequest::{CommitSyncReturned, JoinGroup, SyncGroup};
vec![
commit(1, received),
CommitSyncReturned,
commit(1, received),
JoinGroup,
SyncGroup,
]
}
#[tokio::test(start_paused = true)]
async fn auto_commit_sends_offset_commit_where_kafka_does() {
use GroupRequest::{JoinGroup, LeaveGroup, SyncGroup};
use Step::{AdvanceInterval, Close, Poll, ReceiveRecords};
let both = [(0, 12, 3), (1, 7, -1)];
let cases = vec![
AutoCommitCase::builder()
.name("on: a poll after the interval commits the positions")
.steps(vec![Poll, AdvanceInterval, Poll])
.expected(vec![commit(1, &both)])
.build(),
AutoCommitCase::builder()
.name("on: a poll before the interval commits nothing")
.steps(vec![Poll, Poll])
.expected(vec![])
.build(),
AutoCommitCase::builder()
.name("off: a poll after the interval commits nothing")
.auto_commit(false)
.steps(vec![Poll, AdvanceInterval, Poll])
.expected(vec![])
.build(),
AutoCommitCase::builder()
.name("on, eager: a rebalance commits the positions of the last poll before JoinGroup")
.assignments(vec![vec![partition(0), partition(1)]])
.steps(rebalance_steps(1))
.expected(vec![commit(1, &both), JoinGroup, SyncGroup])
.build(),
AutoCommitCase::builder()
.name("off, eager: a rebalance commits nothing")
.auto_commit(false)
.assignments(vec![vec![partition(0), partition(1)]])
.steps(rebalance_steps(1))
.expected(vec![JoinGroup, SyncGroup])
.build(),
AutoCommitCase::builder()
.name("on, cooperative: a revoke of partition 1 commits before each JoinGroup")
.assignor(Assignor::CooperativeSticky)
.assignments(vec![vec![partition(0)]])
.steps(rebalance_steps(2))
.expected(vec![
commit(1, &both),
JoinGroup,
SyncGroup,
commit(2, &[(0, 12, 3)]),
JoinGroup,
SyncGroup,
])
.build(),
AutoCommitCase::builder()
.name("off, cooperative: a revoke of partition 1 commits nothing")
.auto_commit(false)
.assignor(Assignor::CooperativeSticky)
.assignments(vec![vec![partition(0)]])
.steps(rebalance_steps(2))
.expected(vec![JoinGroup, SyncGroup, JoinGroup, SyncGroup])
.build(),
AutoCommitCase::builder()
.name("on: close commits the current positions, then leaves the group")
.steps(vec![Poll, ReceiveRecords, Close])
.expected(vec![commit(1, &[(0, 20, 3), (1, 7, -1)]), LeaveGroup])
.build(),
AutoCommitCase::builder()
.name("off: close only leaves the group")
.auto_commit(false)
.assignor(Assignor::CooperativeSticky)
.steps(vec![Poll, ReceiveRecords, Close])
.expected(vec![LeaveGroup])
.build(),
];
assert_group_request_cases(cases, |case| async move {
run(
case.auto_commit,
case.assignor,
case.assignments,
&case.steps,
)
.await
})
.await;
}
#[tokio::test(start_paused = true)]
async fn auto_commit_keeps_the_commit_order_within_the_group_timeouts() {
use CommitReply::{Drop, Error};
use GroupRequest::{CommitSyncReturned, JoinGroup, LeaveGroup, SyncGroup};
use Step::{
Close, CommitSync, CommitSyncDuringRebalance, DropHeartbeats, HoldCommitLock, Poll,
ReceiveRecords, ReplyToNextCommit,
};
let polled = [(0, 12, 3), (1, 7, -1)];
let received = [(0, 20, 3), (1, 7, -1)];
let both_partitions = vec![vec![partition(0), partition(1)]];
let cases = vec![
AutoCommitTimeoutCase::builder()
.name("commit_sync, then a rebalance: the commit before JoinGroup does not go back")
.steps(vec![
Poll,
ReceiveRecords,
CommitSync,
rebalance_step(REBALANCE_WAIT),
])
.expected(commit_sync_then_rebalance_requests(&received))
.build(),
AutoCommitTimeoutCase::builder()
.name(
"a rebalance while commit_sync waits for its response: the commit before \
JoinGroup comes after it",
)
.steps(vec![
Poll,
ReceiveRecords,
CommitSyncDuringRebalance {
reply: Drop,
syncs: 1,
within: Duration::from_mins(1),
},
])
.expected(vec![
commit(1, &received),
commit(1, &received),
CommitSyncReturned,
commit(1, &received),
JoinGroup,
SyncGroup,
])
.build(),
AutoCommitTimeoutCase::builder()
.name(
"a rebalance while commit_sync waits for its response, with a \
REBALANCE_IN_PROGRESS reply: commit_sync fails immediately instead of \
waiting for the rejoin, so the commit before JoinGroup comes after it",
)
.steps(vec![
Poll,
ReceiveRecords,
CommitSyncDuringRebalance {
reply: Error(REBALANCE_IN_PROGRESS),
syncs: 1,
within: REBALANCE_WAIT,
},
])
.expected(commit_sync_then_rebalance_requests(&received))
.build(),
AutoCommitTimeoutCase::builder()
.name(
"another commit holds the commit lock: JoinGroup at the rebalance timeout, \
without a commit",
)
.rebalance_timeout(secs(10))
.steps(vec![Poll, HoldCommitLock, rebalance_step(REBALANCE_WAIT)])
.expected(vec![JoinGroup, SyncGroup])
.build(),
AutoCommitTimeoutCase::builder()
.name(
"another commit holds the commit lock past the session timeout: heartbeats keep \
the member in the group",
)
.steps(vec![
Poll,
HoldCommitLock,
rebalance_step(Duration::from_secs(90)),
])
.expected(vec![JoinGroup, SyncGroup])
.build(),
AutoCommitTimeoutCase::builder()
.name(
"the coordinator does not answer a heartbeat during the commit before JoinGroup: \
JoinGroup at the rebalance timeout",
)
.rebalance_timeout(secs(10))
.steps(vec![
Poll,
HoldCommitLock,
DropHeartbeats,
rebalance_step(Duration::from_secs(20)),
])
.expected(vec![JoinGroup, SyncGroup])
.build(),
AutoCommitTimeoutCase::builder()
.name(
"the coordinator does not answer the commit before JoinGroup: JoinGroup at the \
rebalance timeout",
)
.rebalance_timeout(secs(10))
.steps(vec![Poll, ReplyToNextCommit(Drop), rebalance_step(REBALANCE_WAIT)])
.expected(vec![commit(1, &polled), JoinGroup, SyncGroup])
.build(),
AutoCommitTimeoutCase::builder()
.name(
"another commit holds the commit lock: close leaves the group after the close \
timeout",
)
.steps(vec![Poll, ReceiveRecords, HoldCommitLock, Close])
.expected(vec![LeaveGroup])
.build(),
];
assert_group_request_cases(cases, |case| {
let assignments = both_partitions.clone();
async move {
run_with_rebalance_timeout(
true,
Assignor::Range,
assignments,
case.rebalance_timeout,
&case.steps,
)
.await
}
})
.await;
}
#[test]
fn auto_commit_interval_follows_kafka_config_bounds() {
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, enable, interval, expected) in [
("off", false, Time::from_millis(-1), Ok(None)),
("default", true, secs(5), Ok(Some(Duration::from_secs(5)))),
("zero", true, Time::from_millis(0), Ok(Some(Duration::ZERO))),
(
"i32::MAX milliseconds",
true,
Time::from_millis(i64::from(i32::MAX)),
Ok(Some(Duration::from_millis(u64::from(
i32::MAX.unsigned_abs(),
)))),
),
("negative", true, Time::from_millis(-1), Err(())),
(
"above i32::MAX milliseconds",
true,
Time::from_millis(i64::from(i32::MAX) + 1),
Err(()),
),
(
"half a millisecond",
true,
Time::from_secs_f64(0.0005),
Err(()),
),
(
"i32::MAX milliseconds and a fraction",
true,
Time::from_secs_f64((f64::from(i32::MAX) + 0.25) / 1e3),
Err(()),
),
("not a number", true, Time::from_secs_f64(f64::NAN), Err(())),
(
"infinite",
true,
Time::from_secs_f64(f64::INFINITY),
Err(()),
),
] {
let result = validated_auto_commit_interval(enable, interval).map_err(|_| ());
actual.push((name, result));
wanted.push((name, expected));
}
assert2::assert!(actual == wanted);
}
}
#[cfg(test)]
mod group_rejoin_tests {
use super::{
auto_commit_tests::{
GroupRequest, MockCoordinator, partition, start_config, wait_for_sync_groups,
},
*,
};
const FENCED_INSTANCE_ID: i16 = 82;
#[derive(Debug, PartialEq)]
struct FenceObservation {
first_poll: Result<usize, String>,
commit_after_fence: Result<(), String>,
requests_before_next_poll: Vec<GroupRequest>,
next_poll: Result<usize, String>,
requests_after_next_poll: Vec<GroupRequest>,
assignment: Vec<(String, i32)>,
generation: i32,
}
#[tokio::test(start_paused = true)]
async fn poll_after_a_fenced_instance_id_joins_the_group_again() {
let coordinator = MockCoordinator::range_partition(1);
coordinator.set_heartbeat_error(FENCED_INSTANCE_ID);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, minutes(1));
config.group_instance_id = Some("instance-a".into());
let mut consumer = start_test_consumer(
TestConsumerStart::builder()
.bootstrap(mock.addr.to_string())
.config(config)
.startup(super::auto_commit_tests::two_partition_startup())
.build(),
)
.await;
tokio::time::timeout(Duration::from_secs(30), async {
while !coordinator::poll_error_pending(&consumer.poll_error) {
tokio::time::sleep(Duration::from_millis(100)).await;
}
})
.await
.expect("the coordinator fences the member");
let first_poll = consumer
.poll(millis(100))
.await
.map(|records| records.len())
.map_err(|error| error.to_string());
tokio::time::sleep(Duration::from_secs(20)).await;
let commit_after_fence = consumer
.commit_sync()
.await
.map_err(|error| error.to_string());
let requests_before_next_poll = coordinator.requests();
let next_poll = consumer
.poll(millis(100))
.await
.map(|records| records.len())
.map_err(|error| error.to_string());
wait_for_sync_groups(
&coordinator,
&consumer.poll_signal,
1,
Duration::from_secs(30),
)
.await;
let observation = FenceObservation {
first_poll,
commit_after_fence,
requests_before_next_poll,
next_poll,
requests_after_next_poll: coordinator
.requests()
.into_iter()
.filter(|request| *request != GroupRequest::SessionExpired)
.collect(),
assignment: consumer.assignment().await,
generation: consumer.generation_id(),
};
consumer.close().await.expect("close");
mock.stop();
assert2::assert!(
observation
== FenceObservation {
first_poll: Err(
ConsumerError::FencedInstanceId("instance-a".into()).to_string()
),
commit_after_fence: Err(ConsumerError::CommitFailed.to_string()),
requests_before_next_poll: Vec::new(),
next_poll: Ok(0),
requests_after_next_poll: vec![
GroupRequest::JoinGroup,
GroupRequest::SyncGroup
],
assignment: vec![partition(1)],
generation: 2,
}
);
}
}
#[cfg(test)]
mod poll_interval_tests {
use krabka_protocol::owned::leave_group_request::LeaveGroupRequest;
use super::{
auto_commit_tests::{MEMBER, MockCoordinator, TOPIC, start_config},
*,
};
#[derive(Clone, Copy, Debug)]
enum Polls {
EachTenthOfASecond,
Late,
}
#[derive(Debug, PartialEq)]
struct PollTimeoutObservation {
leaves: Vec<LeaveGroupRequest>,
heartbeats_after_timeout: bool,
commit: Result<(), String>,
member_id_before_next_poll: String,
joins_after_next_poll: Vec<(String, Option<String>)>,
member_id_after_next_poll: String,
}
async fn run_poll_timeout_case(
group_instance_id: Option<&str>,
polls: Polls,
) -> PollTimeoutObservation {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, millis(500));
config.group_instance_id = group_instance_id.map(str::to_owned);
config.heartbeat_interval = millis(100);
let mut consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
let spawned_at = *coordinator
.last_heartbeat
.lock()
.expect("last heartbeat lock");
while *coordinator
.last_heartbeat
.lock()
.expect("last heartbeat lock")
== spawned_at
{
tokio::time::sleep(Duration::from_millis(10)).await;
}
coordinator::note_poll(&consumer.poll_signal);
let poller = matches!(polls, Polls::EachTenthOfASecond).then(|| {
let poll_signal = consumer.poll_signal.clone();
tokio::spawn(async move {
loop {
tokio::time::sleep(Duration::from_millis(100)).await;
coordinator::note_poll(&poll_signal);
}
})
});
tokio::time::sleep(Duration::from_millis(700)).await;
let heartbeat_at_timeout = *coordinator
.last_heartbeat
.lock()
.expect("last heartbeat lock");
tokio::time::sleep(Duration::from_millis(400)).await;
let heartbeats_after_timeout = *coordinator
.last_heartbeat
.lock()
.expect("last heartbeat lock")
!= heartbeat_at_timeout;
let commit = consumer
.commit_sync()
.await
.map_err(|error| error.to_string());
let leaves = coordinator.leaves.lock().expect("leaves lock").clone();
let member_id_before_next_poll = consumer.member_id();
let joins_before = coordinator.joins.lock().expect("joins lock").len();
consumer.poll(millis(0)).await.expect("next poll");
let _ = tokio::time::timeout(Duration::from_secs(5), async {
while consumer.member_id().is_empty() {
coordinator::note_poll(&consumer.poll_signal);
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await;
let joins_after_next_poll = coordinator.joins.lock().expect("joins lock")[joins_before..]
.iter()
.map(|join| (join.member_id.clone(), join.group_instance_id.clone()))
.collect();
let member_id_after_next_poll = consumer.member_id();
if let Some(poller) = poller {
poller.abort();
}
coordinator.set_heartbeat_error(0);
drop(consumer);
mock.stop();
PollTimeoutObservation {
leaves,
heartbeats_after_timeout,
commit,
member_id_before_next_poll,
joins_after_next_poll,
member_id_after_next_poll,
}
}
#[tokio::test]
async fn a_member_leaves_the_group_when_poll_does_not_come_within_max_poll_interval() {
let commit_failed = Err(ConsumerError::CommitFailed.to_string());
let leave =
test_support::dynamic_member_leave_request("consumer poll timeout has expired.");
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, group_instance_id, polls, expected) in [
(
"dynamic member polls in time",
None,
Polls::EachTenthOfASecond,
PollTimeoutObservation {
leaves: Vec::new(),
heartbeats_after_timeout: true,
commit: Ok(()),
member_id_before_next_poll: MEMBER.into(),
joins_after_next_poll: Vec::new(),
member_id_after_next_poll: MEMBER.into(),
},
),
(
"dynamic member polls late",
None,
Polls::Late,
PollTimeoutObservation {
leaves: vec![leave],
heartbeats_after_timeout: false,
commit: commit_failed.clone(),
member_id_before_next_poll: String::new(),
joins_after_next_poll: vec![(String::new(), None)],
member_id_after_next_poll: MEMBER.into(),
},
),
(
"static member polls late",
Some("instance-a"),
Polls::Late,
PollTimeoutObservation {
leaves: Vec::new(),
heartbeats_after_timeout: false,
commit: commit_failed,
member_id_before_next_poll: String::new(),
joins_after_next_poll: vec![(String::new(), Some("instance-a".into()))],
member_id_after_next_poll: MEMBER.into(),
},
),
] {
actual.push((name, run_poll_timeout_case(group_instance_id, polls).await));
wanted.push((name, expected));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn a_poll_during_the_poll_timeout_leave_fetches_nothing_and_starts_the_join() {
let coordinator = MockCoordinator::range_partition(0);
coordinator
.drop_leaves
.store(true, std::sync::atomic::Ordering::SeqCst);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, millis(300));
config.heartbeat_interval = millis(50);
config.leave_group_timeout = secs(1);
let mut consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
tokio::time::timeout(Duration::from_secs(10), async {
while coordinator.leaves.lock().expect("leaves lock").is_empty() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("the member sends LeaveGroup");
let records = consumer.poll(millis(100)).await.expect("poll");
let joined = tokio::time::timeout(Duration::from_secs(5), async {
while coordinator.joins.lock().expect("joins lock").is_empty() {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.is_ok();
let observation = (
records.len(),
joined,
coordinator.leaves.lock().expect("leaves lock").len(),
);
drop(consumer);
mock.stop();
assert2::assert!(observation == (0, true, 1));
}
#[tokio::test]
async fn join_group_sends_max_poll_interval_as_the_rebalance_timeout() {
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, max_poll_interval, expected) in [
("default", None, 300_000),
("configured", Some(secs(10)), 10_000),
] {
let coordinator = MockCoordinator::new(Assignor::Range, vec![Vec::new()]);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let consumer = Consumer::builder()
.bootstrap(mock.addr.to_string())
.group_id("group-a")
.subscribe(vec![TOPIC.to_owned()])
.maybe_max_poll_interval(max_poll_interval)
.build()
.await
.expect("build");
let rebalance_timeouts: Vec<i32> = coordinator
.joins
.lock()
.expect("joins lock")
.iter()
.map(|join| join.rebalance_timeout_ms)
.collect();
drop(consumer);
mock.stop();
actual.push((name, rebalance_timeouts));
wanted.push((name, vec![expected, expected]));
}
assert2::assert!(actual == wanted);
}
#[test]
fn client_id_is_generated_as_kafka_does() {
let first = consumer_client_id(None, "g", None);
let sequence: u64 = first
.strip_prefix("consumer-g-")
.and_then(|n| n.parse().ok())
.expect("generated sequence");
let actual = [
consumer_client_id(None, "g", None),
consumer_client_id(Some(String::new()), "g", None),
consumer_client_id(None, "g", Some("i-1")),
consumer_client_id(Some(String::new()), "g", Some("i-1")),
consumer_client_id(Some("app".into()), "g", Some("i-1")),
];
let next = |step: u64| format!("consumer-g-{}", sequence + step);
assert2::assert!(
actual
== [
next(1),
next(2),
"consumer-g-i-1".to_owned(),
"consumer-g-i-1".to_owned(),
"app".to_owned(),
]
);
}
#[test]
fn max_poll_settings_follow_kafka_config_bounds() {
let interval = |value: Time| validated_max_poll_interval(value).map_err(|_| ());
let records = |value: usize| validated_max_poll_records(value).map_err(|_| ());
let actual = (
[
interval(Time::from_millis(0)),
interval(Time::from_millis(1)),
interval(DEFAULT_CONSUMER_MAX_POLL_INTERVAL),
interval(Time::from_millis(i64::from(i32::MAX))),
interval(Time::from_millis(i64::from(i32::MAX) + 1)),
interval(Time::from_secs_f64(0.0015)),
],
[
records(0),
records(1),
records(DEFAULT_CONSUMER_MAX_POLL_RECORDS),
records(usize::try_from(i32::MAX).expect("fits")),
records(usize::try_from(i32::MAX).expect("fits") + 1),
],
);
let expected = (
[
Err(()),
Ok(Time::from_millis(1)),
Ok(minutes(5)),
Ok(Time::from_millis(i64::from(i32::MAX))),
Err(()),
Err(()),
],
[
Err(()),
Ok(1),
Ok(500),
Ok(usize::try_from(i32::MAX).expect("fits")),
Err(()),
],
);
assert2::assert!(actual == expected);
}
}
#[cfg(test)]
mod group_membership_tests {
use std::sync::{Arc, atomic::Ordering};
use krabka_client_core::MockBroker;
use super::{
auto_commit_tests::{
GroupRequest, MEMBER, MockCoordinator, TOPIC, partition, start_config,
},
*,
};
async fn started_consumer(mock: &MockBroker, group_instance_id: Option<&str>) -> Consumer {
started_consumer_with(mock, group_instance_id, &[Assignor::Range]).await
}
async fn started_consumer_with(
mock: &MockBroker,
group_instance_id: Option<&str>,
assignors: &[Assignor],
) -> Consumer {
let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, minutes(1));
config.assignors = assignors.to_vec();
config.rebalance_protocol =
assignor::rebalance_protocol_of(assignors).expect("assignor protocol");
config.group_instance_id = group_instance_id.map(str::to_owned);
config.heartbeat_interval = millis(50);
let client = Client::builder()
.bootstrap(mock.addr.to_string())
.build()
.await
.expect("client");
spawn_consumer(
config,
client,
Arc::new(AtomicI32::new(0)),
MEMBER.into(),
super::auto_commit_tests::single_partition_startup(1),
)
.await
.expect("spawn consumer")
}
struct StartedRangeConsumer {
coordinator: Arc<MockCoordinator>,
mock: MockBroker,
consumer: Consumer,
}
impl StartedRangeConsumer {
async fn start(group_instance_id: Option<&str>) -> Self {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let consumer = started_consumer(&mock, group_instance_id).await;
Self {
coordinator,
mock,
consumer,
}
}
}
#[tokio::test]
async fn close_leaves_the_group_as_the_membership_operation_says() {
let leave = test_support::dynamic_member_leave_request("the consumer is being closed");
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, group_instance_id, operation, expected) in [
(
"dynamic member, default",
None,
GroupMembershipOperation::Default,
vec![leave.clone()],
),
(
"static member, default",
Some("i-1"),
GroupMembershipOperation::Default,
Vec::new(),
),
(
"static member, leave group",
Some("i-1"),
GroupMembershipOperation::LeaveGroup,
vec![leave.clone()],
),
(
"dynamic member, remain in group",
None,
GroupMembershipOperation::RemainInGroup,
Vec::new(),
),
] {
let fixture = StartedRangeConsumer::start(group_instance_id).await;
fixture
.consumer
.close_with(CloseOptions::group_membership_operation(operation))
.await
.expect("close");
fixture.mock.stop();
actual.push((
name,
fixture
.coordinator
.leaves
.lock()
.expect("leaves lock")
.clone(),
));
wanted.push((name, expected));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn a_join_after_the_member_was_reset_during_preparation_owns_nothing() {
let coordinator =
MockCoordinator::new(Assignor::CooperativeSticky, vec![vec![partition(0)]]);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(
mock.addr.to_string(),
Assignor::CooperativeSticky,
true,
minutes(1),
);
config.heartbeat_interval = millis(50);
let mut consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
consumer.poll(millis(20)).await.expect("first poll");
coordinator
.commit_replies
.lock()
.expect("commit replies lock")
.push_back(super::auto_commit_tests::CommitReply::Drop);
coordinator.set_heartbeat_error(27);
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
let mut reset = false;
while coordinator.joins.lock().expect("joins lock").is_empty()
&& tokio::time::Instant::now() < deadline
{
let _ = consumer.poll(millis(20)).await;
let committing = coordinator
.requests()
.iter()
.any(|request| matches!(request, GroupRequest::OffsetCommit(_)));
if committing && !reset {
coordinator.set_heartbeat_error(25);
reset = true;
}
}
let subscriptions: Vec<(Vec<(String, i32)>, i32)> = coordinator
.joins
.lock()
.expect("joins lock")
.iter()
.flat_map(|join| join.protocols.iter())
.map(|protocol| {
let subscription = builder::decode_subscription(&protocol.metadata);
(subscription.owned, subscription.generation_id)
})
.collect();
drop(consumer);
mock.stop();
assert2::assert!(subscriptions.first() == Some(&(Vec::new(), -1)));
}
#[tokio::test]
async fn close_timeout_bounds_the_auto_commit() {
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, timeout, expected) in [("300 ms", millis(300), (Ok(()), true, 1))] {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(mock.addr.to_string(), Assignor::Range, true, minutes(1));
config.heartbeat_interval = millis(50);
let consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
let held = Arc::clone(&consumer.commit_serialization)
.lock_owned()
.await;
let started = tokio::time::Instant::now();
let closed = tokio::time::timeout(
Duration::from_secs(10),
consumer.close_with(CloseOptions::timeout(timeout)),
)
.await
.expect("close returns")
.map_err(|error| error.to_string());
let fast = started.elapsed() < Duration::from_secs(3);
drop(held);
mock.stop();
let leaves = coordinator.leaves.lock().expect("leaves lock").len();
actual.push((name, (closed, fast, leaves)));
wanted.push((name, expected));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn close_waits_for_a_pending_asynchronous_commit() {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let consumer = started_consumer(&mock, None).await;
let held = Arc::clone(&consumer.commit_serialization)
.lock_owned()
.await;
let release = tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(200)).await;
drop(held);
});
let called = Arc::new(std::sync::atomic::AtomicBool::new(false));
let called_in_callback = Arc::clone(&called);
consumer
.commit_async_with_callback(move |_, _| {
called_in_callback.store(true, Ordering::SeqCst);
})
.expect("commit_async_with_callback");
consumer.close().await.expect("close");
let called_before_close_returned = called.load(Ordering::SeqCst);
release.await.expect("release");
mock.stop();
assert2::assert!(called_before_close_returned);
}
#[derive(Clone, Copy, Debug)]
enum Change {
Subscribe(&'static str),
Unsubscribe,
UnsubscribeThenSubscribe,
PatternWithoutMatch,
}
type SubscriptionOutcome = (
Vec<(Vec<String>, Option<String>)>,
Vec<Option<String>>,
Vec<String>,
);
#[tokio::test]
async fn subscription_changes_follow_kafka() {
let changed = |from: &str, to: &str| {
Some(format!(
"the subscription has changed from [{from}] to [{to}] since the last group join"
))
};
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, change, expected) in [
(
"subscribe to another topic",
Change::Subscribe("payments"),
(
vec![(vec!["payments".to_owned()], changed("orders", "payments"))],
vec![],
vec![],
),
),
(
"subscribe to the same topic",
Change::Subscribe("orders"),
(vec![], vec![], vec![]),
),
(
"unsubscribe",
Change::Unsubscribe,
(
vec![],
vec![Some("the consumer unsubscribed from all topics".to_owned())],
vec!["not subscribed to any topic".to_owned()],
),
),
(
"unsubscribe, then subscribe again",
Change::UnsubscribeThenSubscribe,
(
vec![(
vec!["orders".to_owned()],
Some("consumer pro-actively leaving the group".to_owned()),
)],
vec![Some("the consumer unsubscribed from all topics".to_owned())],
vec!["not subscribed to any topic".to_owned()],
),
),
(
"a pattern that matches no topic",
Change::PatternWithoutMatch,
(vec![(vec![], changed("orders", ""))], vec![], vec![]),
),
] {
let StartedRangeConsumer {
coordinator,
mock,
mut consumer,
} = StartedRangeConsumer::start(None).await;
let mut errors = Vec::new();
match change {
Change::Subscribe(topic) => consumer.subscribe([topic]).await.expect("subscribe"),
Change::Unsubscribe => consumer.unsubscribe().await.expect("unsubscribe"),
Change::UnsubscribeThenSubscribe => {
consumer.unsubscribe().await.expect("unsubscribe");
if let Err(error) = consumer.poll(millis(20)).await {
errors.push(error.to_string());
}
consumer.subscribe(["orders"]).await.expect("subscribe");
}
Change::PatternWithoutMatch => {
consumer
.subscribe_pattern(TopicPattern::new(|_| false))
.expect("subscribe_pattern");
}
}
let deadline = tokio::time::Instant::now() + Duration::from_millis(500);
while tokio::time::Instant::now() < deadline {
if let Err(error) = consumer.poll(millis(20)).await
&& !errors.contains(&error.to_string())
{
errors.push(error.to_string());
}
if !coordinator.joins.lock().expect("joins lock").is_empty() {
break;
}
tokio::time::sleep(Duration::from_millis(10)).await;
}
let joins = coordinator
.joins
.lock()
.expect("joins lock")
.iter()
.map(|join| {
let subscription = builder::decode_subscription(&join.protocols[0].metadata);
(subscription.topics, join.reason.clone())
})
.collect();
let leaves = coordinator
.leaves
.lock()
.expect("leaves lock")
.iter()
.flat_map(|leave| leave.members.iter().map(|member| member.reason.clone()))
.collect();
drop(consumer);
mock.stop();
let outcome: SubscriptionOutcome = (joins, leaves, errors);
actual.push((name, outcome));
wanted.push((name, expected));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn unsubscribe_gives_up_the_assignment_before_it_returns() {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut consumer = started_consumer(&mock, None).await;
let before = consumer.assignment().await;
consumer.coordinator_shutdown.cancel();
if let Some(handle) = consumer.coordinator_handle.take() {
handle.await.expect("coordinator task");
}
consumer.unsubscribe().await.expect("unsubscribe");
let after = consumer.assignment().await;
let owned = consumer.commit_identity.lock().await.ownership_ids.len();
drop(consumer);
mock.stop();
assert2::assert!((before, after, owned) == (vec![partition(0)], vec![], 0));
}
#[tokio::test]
async fn a_manual_assignment_after_unsubscribe_survives_and_blocks_enforce_rebalance() {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut consumer = started_consumer(&mock, None).await;
consumer.unsubscribe().await.expect("unsubscribe");
consumer
.assign(&[(TOPIC.to_owned(), 1)])
.await
.expect("assign");
tokio::time::sleep(Duration::from_millis(300)).await;
let assignment = consumer.assignment().await;
let enforced = consumer
.enforce_rebalance(None)
.map_err(|error| error.to_string());
let rejected = consumer
.assign(&[(TOPIC.to_owned(), -1)])
.await
.map_err(|error| error.to_string());
drop(consumer);
mock.stop();
assert2::assert!(
(assignment, enforced, rejected)
== (
vec![(TOPIC.to_owned(), 1)],
Err("illegal state: Tried to force a rebalance but the consumer has a manual assignment.".to_owned()),
Err("invalid argument: Partition index of orders--1 is negative".to_owned())
)
);
}
#[tokio::test]
async fn enforce_rebalance_joins_on_the_next_poll_with_the_reason() {
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, reason, expected) in [
("no reason", None, "rebalance enforced by user"),
("empty reason", Some(""), "rebalance enforced by user"),
("reason", Some("scale out"), "scale out"),
] {
let StartedRangeConsumer {
coordinator,
mock,
mut consumer,
} = StartedRangeConsumer::start(None).await;
consumer
.enforce_rebalance(reason)
.expect("enforce_rebalance");
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while coordinator.joins.lock().expect("joins lock").is_empty()
&& tokio::time::Instant::now() < deadline
{
consumer.poll(millis(20)).await.expect("poll");
}
let reasons: Vec<Option<String>> = coordinator
.joins
.lock()
.expect("joins lock")
.iter()
.map(|join| join.reason.clone())
.collect();
drop(consumer);
mock.stop();
actual.push((name, reasons));
wanted.push((name, vec![Some(expected.to_owned())]));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn a_poll_that_waits_for_an_eager_join_starts_a_join_requested_after_it_began() {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut consumer = started_consumer(&mock, None).await;
coordinator::note_poll(&consumer.poll_signal);
coordinator.set_heartbeat_error(27);
tokio::time::timeout(Duration::from_secs(5), async {
while !*consumer.rebalance_pending.borrow() {
tokio::time::sleep(Duration::from_millis(5)).await;
}
})
.await
.expect("the heartbeat requests a rebalance");
let started = tokio::time::Instant::now();
let joined = consumer
.wait_for_rebalance(started + Duration::from_secs(5))
.await
.expect("no listener");
let within_timeout = started.elapsed() < Duration::from_secs(4);
let joins = coordinator.joins.lock().expect("joins lock").len();
drop(consumer);
mock.stop();
assert2::assert!((joined, within_timeout, joins) == (true, true, 1));
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn continuous_polls_do_not_stop_the_heartbeats() {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let consumer = started_consumer(&mock, None).await;
let first_heartbeat = *coordinator
.last_heartbeat
.lock()
.expect("last heartbeat lock");
let poll_signal = consumer.poll_signal.clone();
tokio::task::spawn_blocking(move || {
let end = std::time::Instant::now() + Duration::from_millis(500);
while std::time::Instant::now() < end {
coordinator::note_poll(&poll_signal);
}
})
.await
.expect("poller");
let heartbeat_during_polls = *coordinator
.last_heartbeat
.lock()
.expect("last heartbeat lock")
> first_heartbeat + Duration::from_millis(200);
drop(consumer);
mock.stop();
assert2::assert!(heartbeat_during_polls);
}
#[tokio::test]
async fn the_poll_timeout_leave_goes_out_before_the_join_of_the_next_poll() {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(mock.addr.to_string(), Assignor::Range, false, millis(200));
config.heartbeat_interval = millis(50);
let consumer = super::auto_commit_tests::start_single_partition_consumer(config).await;
tokio::time::timeout(Duration::from_secs(5), async {
while !*consumer.rebalance_pending.borrow() {
tokio::task::yield_now().await;
}
})
.await
.expect("the poll timeout expires");
coordinator::note_poll(&consumer.poll_signal);
tokio::time::timeout(Duration::from_secs(5), async {
while coordinator.sync_groups() == 0 {
tokio::time::sleep(Duration::from_millis(10)).await;
}
})
.await
.expect("the member joins again");
let order: Vec<GroupRequest> = coordinator
.requests()
.into_iter()
.filter(|request| matches!(request, GroupRequest::LeaveGroup | GroupRequest::JoinGroup))
.take(2)
.collect();
drop(consumer);
mock.stop();
assert2::assert!(order == vec![GroupRequest::LeaveGroup, GroupRequest::JoinGroup]);
}
#[tokio::test]
async fn rejoin_sends_a_subscription_per_assignor_with_the_owned_partitions_of_the_protocol() {
use builder::{DecodedSubscription, decode_subscription};
let subscription =
|owned: Vec<(String, i32)>, user_data: Option<Bytes>| DecodedSubscription {
topics: vec![TOPIC.to_owned()],
owned,
generation_id: 1,
rack_id: None,
user_data,
};
let generation_one = Some(Bytes::from_static(&[0, 0, 0, 1]));
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, assignors, expected) in [
(
"default list, eager",
Assignor::DEFAULT_LIST.to_vec(),
vec![
("range".to_owned(), subscription(Vec::new(), None)),
(
"cooperative-sticky".to_owned(),
subscription(Vec::new(), generation_one.clone()),
),
],
),
(
"cooperative sticky, cooperative",
vec![Assignor::CooperativeSticky],
vec![(
"cooperative-sticky".to_owned(),
subscription(vec![partition(0)], generation_one),
)],
),
(
"sticky, eager",
vec![Assignor::Sticky],
vec![(
"sticky".to_owned(),
subscription(
Vec::new(),
assignor::subscription_user_data(
Assignor::Sticky,
Some(&[partition(0)]),
1,
),
),
)],
),
] {
let coordinator = MockCoordinator::range_partition(0);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let consumer = started_consumer_with(&mock, None, &assignors).await;
coordinator.set_heartbeat_error(27);
let protocols = tokio::time::timeout(Duration::from_secs(10), async {
loop {
coordinator::note_poll(&consumer.poll_signal);
if let Some(join) = coordinator.joins.lock().expect("joins lock").first() {
break join
.protocols
.iter()
.map(|protocol| {
(
protocol.name.clone(),
decode_subscription(&protocol.metadata),
)
})
.collect::<Vec<_>>();
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("the consumer joins again");
drop(consumer);
mock.stop();
actual.push((name, protocols));
wanted.push((name, expected));
}
assert2::assert!(actual == wanted);
}
#[tokio::test]
async fn join_group_carries_the_rejoin_reason() {
let mut actual = Vec::new();
let mut wanted = Vec::new();
for (name, heartbeat_error, expected) in [
(
"rebalance in progress",
27,
("member-a", "group is already rebalancing"),
),
(
"illegal generation",
22,
(
"member-a",
"encountered ILLEGAL_GENERATION from HEARTBEAT response",
),
),
(
"unknown member id",
25,
("", "encountered UNKNOWN_MEMBER_ID from HEARTBEAT response"),
),
] {
let fixture = StartedRangeConsumer::start(None).await;
fixture
.coordinator
.heartbeat_error
.store(heartbeat_error, Ordering::SeqCst);
let join = tokio::time::timeout(Duration::from_secs(10), async {
loop {
coordinator::note_poll(&fixture.consumer.poll_signal);
if let Some(join) = fixture
.coordinator
.joins
.lock()
.expect("joins lock")
.first()
{
break (join.member_id.clone(), join.reason.clone());
}
tokio::time::sleep(Duration::from_millis(20)).await;
}
})
.await
.expect("the consumer joins again");
drop(fixture.consumer);
fixture.mock.stop();
actual.push((name, join));
wanted.push((name, (expected.0.to_owned(), Some(expected.1.to_owned()))));
}
assert2::assert!(actual == wanted);
}
}
#[cfg(test)]
mod rebalance_listener_tests {
use std::sync::atomic::Ordering;
use super::{
auto_commit_tests::{GroupRequest, MockCoordinator, TOPIC, partition, start_config},
*,
};
use crate::{
ConsumerRebalanceListener, RebalanceListenerError,
test_support::{ListenerStep as Step, RecordingListener},
};
#[derive(Clone, Debug, PartialEq)]
enum Call {
Revoked(Vec<(String, i32)>),
Assigned(Vec<(String, i32)>),
Lost(Vec<(String, i32)>),
CommitInRevoke(Result<(), String>),
AwaitingCallback(Vec<(String, i32)>),
PollError(String),
Position(Option<i64>),
}
struct Recorder {
calls: Arc<std::sync::Mutex<Vec<Call>>>,
commit_on_revoke: bool,
fail_revoke: bool,
seek_on_assign: Option<i64>,
}
impl Recorder {
fn record(&self, call: Call) {
self.calls.lock().expect("calls lock").push(call);
}
}
#[async_trait::async_trait]
impl ConsumerRebalanceListener for Recorder {
async fn on_partitions_revoked(
&mut self,
consumer: &Consumer,
partitions: &[(String, i32)],
) -> Result<(), RebalanceListenerError> {
self.record(Call::Revoked(partitions.to_vec()));
if self.commit_on_revoke {
let result = consumer
.commit_sync()
.await
.map_err(|error| error.to_string());
self.record(Call::CommitInRevoke(result));
}
if self.fail_revoke {
return Err("revoke failed".into());
}
Ok(())
}
async fn on_partitions_assigned(
&mut self,
consumer_for_assignment: &Consumer,
assigned_partitions: &[(String, i32)],
) -> Result<(), RebalanceListenerError> {
self.record(Call::Assigned(assigned_partitions.to_vec()));
let mut awaiting: Vec<_> = consumer_for_assignment
.assigned_callback_pending
.lock()
.expect("pending lock")
.iter()
.cloned()
.collect();
awaiting.sort();
if !assigned_partitions.is_empty() {
self.record(Call::AwaitingCallback(awaiting));
}
if let (Some(offset), Some(partition)) =
(self.seek_on_assign, assigned_partitions.first())
{
consumer_for_assignment
.seek(partition.0.clone(), partition.1, offset)
.await
.expect("seek");
}
Ok(())
}
async fn on_partitions_lost(
&mut self,
_consumer: &Consumer,
partitions: &[(String, i32)],
) -> Result<(), RebalanceListenerError> {
self.record(Call::Lost(partitions.to_vec()));
Ok(())
}
}
#[derive(Clone, Copy, Debug)]
enum Event {
Heartbeat(i16),
Close,
Subscribe(&'static str),
}
struct Case {
assignor: Assignor,
owned: Vec<(String, i32)>,
assignments: Vec<Vec<(String, i32)>>,
event: Event,
commit_on_revoke: bool,
fail_revoke: bool,
seek_on_assign: Option<i64>,
calls: usize,
}
#[derive(Debug, PartialEq)]
struct CaseResult {
calls: Vec<Call>,
commit_generations: Vec<i32>,
}
#[derive(bon::Builder)]
struct ListenerScenario {
name: &'static str,
case: Case,
expected_calls: Vec<Call>,
expected_commit_generations: Vec<i32>,
}
async fn run_case(case: Case) -> CaseResult {
let coordinator = MockCoordinator::new(case.assignor, case.assignments);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(mock.addr.to_string(), case.assignor, false, minutes(1));
config.heartbeat_interval = millis(50);
config.has_rebalance_listener = true;
let mut consumer = super::auto_commit_tests::start_consumer_with_startup(
config,
StartupState::builder()
.generation_id(1)
.assigned_partitions(case.owned.clone())
.next_offsets(case.owned.iter().map(|key| (key.clone(), 12)).collect())
.topic_partitions(maplit::hashmap! {TOPIC.to_owned() => 2})
.build(),
)
.await;
let calls = Arc::new(std::sync::Mutex::new(Vec::new()));
consumer.rebalance_listener = Some(rebalance_listener::shared(Box::new(Recorder {
calls: Arc::clone(&calls),
commit_on_revoke: case.commit_on_revoke,
fail_revoke: case.fail_revoke,
seek_on_assign: case.seek_on_assign,
})));
consumer.poll(millis(20)).await.expect("first poll");
if case.seek_on_assign.is_some() {
let position = consumer
.next_offsets
.lock()
.await
.get(&partition(0))
.copied();
calls
.lock()
.expect("calls lock")
.push(Call::Position(position));
}
match case.event {
event @ (Event::Heartbeat(_) | Event::Subscribe(_)) => {
match event {
Event::Heartbeat(error_code) => coordinator
.heartbeat_error
.store(error_code, Ordering::SeqCst),
Event::Subscribe(topic) => {
consumer.subscribe([topic]).await.expect("subscribe");
}
Event::Close => {}
}
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while calls.lock().expect("calls lock").len() < case.calls
&& tokio::time::Instant::now() < deadline
{
poll_recording(&mut consumer, &calls).await;
}
for _ in 0..10 {
poll_recording(&mut consumer, &calls).await;
}
drop(consumer);
}
Event::Close => consumer.close().await.expect("close"),
}
let generations = coordinator
.requests()
.into_iter()
.filter_map(|request| match request {
GroupRequest::OffsetCommit(commit) => Some(commit.generation_id_or_member_epoch),
_ => None,
})
.collect();
mock.stop();
let calls = calls.lock().expect("calls lock").clone();
CaseResult {
calls,
commit_generations: generations,
}
}
async fn assert_listener_scenarios(scenarios: impl IntoIterator<Item = ListenerScenario>) {
for scenario in scenarios {
let expected = CaseResult {
calls: scenario.expected_calls,
commit_generations: scenario.expected_commit_generations,
};
assert2::assert!(
run_case(scenario.case).await == expected,
"{}",
scenario.name
);
}
}
async fn poll_recording(consumer: &mut Consumer, calls: &std::sync::Mutex<Vec<Call>>) {
if let Err(error) = consumer.poll(millis(20)).await {
calls
.lock()
.expect("calls lock")
.push(Call::PollError(error.to_string()));
}
}
#[tokio::test]
async fn a_cancelled_poll_keeps_the_listener_and_runs_the_callback_again() {
let coordinator = MockCoordinator::new(
Assignor::CooperativeSticky,
vec![vec![partition(0), partition(1)]],
);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(
mock.addr.to_string(),
Assignor::CooperativeSticky,
false,
minutes(1),
);
config.heartbeat_interval = millis(50);
config.has_rebalance_listener = true;
let mut consumer = super::auto_commit_tests::start_consumer_with_startup(
config,
super::auto_commit_tests::single_partition_startup(2),
)
.await;
let steps = Arc::new(std::sync::Mutex::new(Vec::new()));
let listener = RecordingListener::builder()
.steps(Arc::clone(&steps))
.slow_partition(1)
.slow_for(Duration::from_millis(500))
.build();
consumer.rebalance_listener = Some(rebalance_listener::shared(Box::new(listener)));
consumer.poll(millis(20)).await.expect("first poll");
coordinator.set_heartbeat_error(27);
let ended = |steps: &std::sync::Mutex<Vec<Step>>| {
steps
.lock()
.expect("steps lock")
.contains(&Step::Ended("assigned", vec![partition(1)]))
};
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while !ended(&steps) && tokio::time::Instant::now() < deadline {
let _ =
tokio::time::timeout(Duration::from_millis(100), consumer.poll(millis(20))).await;
}
let awaiting = consumer
.assigned_callback_pending
.lock()
.expect("pending lock")
.len();
consumer.close().await.expect("close");
mock.stop();
let steps = steps.lock().expect("steps lock").clone();
assert2::assert!(
(steps, awaiting)
== (
vec![
Step::Started("assigned", vec![partition(0)]),
Step::Ended("assigned", vec![partition(0)]),
Step::Started("assigned", vec![partition(1)]),
Step::Started("assigned", vec![partition(1)]),
Step::Ended("assigned", vec![partition(1)]),
Step::Started("revoked", vec![partition(0), partition(1)]),
Step::Ended("revoked", vec![partition(0), partition(1)]),
],
0
)
);
}
#[tokio::test]
async fn poll_timeout_during_a_callback_leaves_the_group() {
let coordinator = MockCoordinator::new(
Assignor::CooperativeSticky,
vec![vec![partition(0), partition(1)]],
);
let mock = super::auto_commit_tests::mock_coordinator_broker(&coordinator).await;
let mut config = start_config(
mock.addr.to_string(),
Assignor::CooperativeSticky,
false,
millis(400),
);
config.heartbeat_interval = millis(50);
config.has_rebalance_listener = true;
let mut consumer = start_test_consumer(
TestConsumerStart::builder()
.bootstrap(mock.addr.to_string())
.config(config)
.startup(super::auto_commit_tests::single_partition_startup(2))
.build(),
)
.await;
let listener = RecordingListener::builder()
.slow_partition(1)
.slow_for(Duration::from_millis(1200))
.build();
consumer.rebalance_listener = Some(rebalance_listener::shared(Box::new(listener)));
consumer.poll(millis(20)).await.expect("first poll");
coordinator.set_heartbeat_error(27);
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
while coordinator.leaves.lock().expect("leaves lock").is_empty()
&& tokio::time::Instant::now() < deadline
{
let _ = consumer.poll(millis(20)).await;
}
let reasons: Vec<Option<String>> = coordinator
.leaves
.lock()
.expect("leaves lock")
.iter()
.flat_map(|leave| leave.members.iter().map(|member| member.reason.clone()))
.collect();
drop(consumer);
mock.stop();
assert2::assert!(reasons == vec![Some("consumer poll timeout has expired.".to_owned())]);
}
fn case(
assignor: Assignor,
owned: Vec<(String, i32)>,
assignment: Vec<(String, i32)>,
event: Event,
calls: usize,
) -> Case {
Case {
assignor,
owned,
assignments: vec![assignment],
event,
commit_on_revoke: false,
fail_revoke: false,
seek_on_assign: None,
calls,
}
}
fn awaiting(partitions: Vec<(String, i32)>) -> Call {
Call::AwaitingCallback(partitions)
}
fn empty_assignment_call() -> Call {
Call::Assigned(Vec::new())
}
fn range_revoke_calls() -> Vec<Call> {
vec![
Call::Assigned(vec![partition(0)]),
awaiting(vec![]),
Call::Revoked(vec![partition(0)]),
]
}
fn range_revoke_commit_calls() -> Vec<Call> {
let mut calls = range_revoke_calls();
calls.extend([
Call::CommitInRevoke(Ok(())),
Call::Assigned(vec![partition(0)]),
awaiting(vec![partition(0)]),
]);
calls
}
#[tokio::test]
async fn listener_calls_follow_kafkas_order() {
let p = partition;
Box::pin(assert_listener_scenarios([
ListenerScenario::builder()
.name("eager rebalance")
.case(case(
Assignor::Range,
vec![p(0), p(1)],
vec![p(0)],
Event::Heartbeat(27),
5,
))
.expected_calls(vec![
Call::Assigned(vec![p(0), p(1)]),
awaiting(vec![]),
Call::Revoked(vec![p(0), p(1)]),
Call::Assigned(vec![p(0)]),
awaiting(vec![p(0)]),
])
.expected_commit_generations(vec![])
.build(),
ListenerScenario::builder()
.name("cooperative rebalance that revokes")
.case(case(
Assignor::CooperativeSticky,
vec![p(0), p(1)],
vec![p(0)],
Event::Heartbeat(27),
5,
))
.expected_calls(vec![
Call::Assigned(vec![p(0), p(1)]),
awaiting(vec![]),
Call::Revoked(vec![p(1)]),
empty_assignment_call(),
empty_assignment_call(),
])
.expected_commit_generations(vec![])
.build(),
ListenerScenario::builder()
.name("cooperative rebalance that adds")
.case(case(
Assignor::CooperativeSticky,
vec![p(0)],
vec![p(0), p(1)],
Event::Heartbeat(27),
4,
))
.expected_calls(vec![
Call::Assigned(vec![p(0)]),
awaiting(vec![]),
Call::Assigned(vec![p(1)]),
awaiting(vec![p(1)]),
])
.expected_commit_generations(vec![])
.build(),
ListenerScenario::builder()
.name("unknown member id")
.case(case(
Assignor::Range,
vec![p(0)],
vec![p(0)],
Event::Heartbeat(25),
5,
))
.expected_calls(vec![
Call::Assigned(vec![p(0)]),
awaiting(vec![]),
Call::Lost(vec![p(0)]),
Call::Assigned(vec![p(0)]),
awaiting(vec![p(0)]),
])
.expected_commit_generations(vec![])
.build(),
ListenerScenario::builder()
.name("close")
.case(case(
Assignor::Range,
vec![p(0)],
vec![p(0)],
Event::Close,
3,
))
.expected_calls(range_revoke_calls())
.expected_commit_generations(vec![])
.build(),
ListenerScenario::builder()
.name("a commit in the eager revoke callback")
.case(Case {
commit_on_revoke: true,
..case(
Assignor::Range,
vec![p(0)],
vec![p(0)],
Event::Heartbeat(27),
6,
)
})
.expected_calls(range_revoke_commit_calls())
.expected_commit_generations(vec![1])
.build(),
]))
.await;
}
#[tokio::test]
async fn listener_calls_follow_kafka_around_commits_errors_and_seeks() {
let p = partition;
assert_listener_scenarios([
ListenerScenario::builder()
.name("a commit in the cooperative revoke callback carries the new generation")
.case(Case {
commit_on_revoke: true,
..case(
Assignor::CooperativeSticky,
vec![p(0), p(1)],
vec![p(0)],
Event::Heartbeat(27),
6,
)
})
.expected_calls(vec![
Call::Assigned(vec![p(0), p(1)]),
awaiting(vec![]),
Call::Revoked(vec![p(1)]),
Call::CommitInRevoke(Ok(())),
empty_assignment_call(),
empty_assignment_call(),
])
.expected_commit_generations(vec![2])
.build(),
ListenerScenario::builder()
.name("a failed cooperative revoke callback fails the poll after the rebalance")
.case(Case {
fail_revoke: true,
..case(
Assignor::CooperativeSticky,
vec![p(0), p(1)],
vec![p(0)],
Event::Heartbeat(27),
6,
)
})
.expected_calls(vec![
Call::Assigned(vec![p(0), p(1)]),
awaiting(vec![]),
Call::Revoked(vec![p(1)]),
empty_assignment_call(),
empty_assignment_call(),
Call::PollError("rebalance listener failed: revoke failed".into()),
])
.expected_commit_generations(vec![])
.build(),
ListenerScenario::builder()
.name("a cooperative join after a subscription change revokes the partitions of the old topic")
.case(case(
Assignor::CooperativeSticky,
vec![p(0), p(1)],
vec![],
Event::Subscribe("payments"),
4,
))
.expected_calls(vec![
Call::Assigned(vec![p(0), p(1)]),
awaiting(vec![]),
Call::Revoked(vec![p(0), p(1)]),
empty_assignment_call(),
])
.expected_commit_generations(vec![])
.build(),
ListenerScenario::builder()
.name("a seek in the assign callback moves the position of the same poll")
.case(Case {
seek_on_assign: Some(3),
..case(Assignor::Range, vec![p(0)], vec![p(0)], Event::Close, 4)
})
.expected_calls(vec![
Call::Assigned(vec![p(0)]),
awaiting(vec![]),
Call::Position(Some(3)),
Call::Revoked(vec![p(0)]),
])
.expected_commit_generations(vec![])
.build(),
])
.await;
}
}
#[cfg(test)]
mod transport_configuration_tests {
use assert2::assert;
use bytes::BytesMut;
use krabka_client_core::MockBroker;
use krabka_protocol::{
Encode,
owned::{api_versions_request, api_versions_response::ApiVersionsResponse},
};
use krabka_units::{kibibytes, millis};
use super::{ClientError, Consumer};
#[tokio::test]
async fn consumer_retains_dns_and_socket_overrides_in_its_client() {
let mock = MockBroker::start(|api, _, _, _| {
if api != api_versions_request::API_KEY {
return None;
}
let mut body = BytesMut::new();
ApiVersionsResponse::default()
.encode(&mut body, 0)
.expect("encode ApiVersions");
Some(body.to_vec())
})
.await;
let consumer = Consumer::builder()
.bootstrap(mock.addr.to_string())
.request_timeout(millis(10))
.dns_timeout(millis(95))
.send_buffer(None)
.receive_buffer(Some(kibibytes(19)))
.enable_metrics_push(false)
.build()
.await
.expect("build consumer");
let options = consumer.client.connection_options();
assert!(
(
options.dns_timeout.time(),
options.send_buffer,
options.receive_buffer
) == (millis(95), None, Some(kibibytes(19)))
);
assert!(consumer.client.refresh_metadata().await.is_err());
consumer.close().await.expect("close consumer");
mock.stop();
}
#[tokio::test(start_paused = true)]
async fn consumer_total_startup_deadline_bounds_attempts_and_retry_sleeps() {
use krabka_units::{convert::TimeExt as _, secs};
use super::ConsumerRetryPolicy;
for backoff in [millis(100), secs(1), secs(2)] {
let mock = MockBroker::start(|_, _, _, _| None).await;
let policy = ConsumerRetryPolicy::builder()
.startup_attempt_timeout(secs(2))
.startup_deadline(secs(3))
.startup_initial_backoff(backoff)
.startup_max_backoff(backoff)
.coordinator_retry_timeout(secs(1))
.coordinator_initial_backoff(millis(100))
.coordinator_max_backoff(secs(1))
.build()
.expect("policy");
let started = tokio::time::Instant::now();
let result = Consumer::builder()
.bootstrap(mock.addr.to_string())
.group_id("startup-guard")
.subscribe(["startup-topic".to_owned()])
.retry_policy(policy)
.enable_metrics_push(false)
.build()
.await;
assert!(matches!(result, Err(super::ConsumerError::Client(
ClientError::Timeout(deadline)
)) if deadline == secs(3)));
assert!(started.elapsed() == secs(3).to_std());
mock.stop();
}
}
#[tokio::test]
async fn consumer_rejects_invalid_dns_before_retrying_bootstrap() {
let result = Consumer::builder()
.bootstrap("127.0.0.1:1")
.dns_timeout(millis(0))
.build()
.await;
assert!(matches!(
result,
Err(super::ConsumerError::InvalidConfig(_))
));
}
}