use std::{
collections::{BTreeMap, BTreeSet, HashMap, HashSet},
time::{Duration, Instant},
};
use crate::native::{
KafkaClientError, KafkaClientResult, NativeSecurityConfig, VERSION,
connection::BrokerConnection,
model::{
KafkaPayloadBatch, LocalCommitState, NativeKafkaMetrics, OffsetCommit, PayloadBatchBuilder,
StartOffset, TopicPartition, TopicPartitionAssignment,
},
profile::{self, ProfileBucket},
protocol::{
API_KEY_API_VERSIONS, API_KEY_FETCH, API_KEY_FIND_COORDINATOR, API_KEY_HEARTBEAT,
API_KEY_INIT_PRODUCER_ID, API_KEY_JOIN_GROUP, API_KEY_LEAVE_GROUP, API_KEY_LIST_OFFSETS,
API_KEY_METADATA, API_KEY_OFFSET_COMMIT, API_KEY_OFFSET_FETCH, API_KEY_PRODUCE,
API_KEY_SASL_AUTHENTICATE, API_KEY_SASL_HANDSHAKE, API_KEY_SYNC_GROUP,
API_VERSION_API_VERSIONS, API_VERSION_FETCH, API_VERSION_FIND_COORDINATOR,
API_VERSION_HEARTBEAT, API_VERSION_INIT_PRODUCER_ID, API_VERSION_JOIN_GROUP,
API_VERSION_LEAVE_GROUP, API_VERSION_LIST_OFFSETS, API_VERSION_METADATA,
API_VERSION_OFFSET_COMMIT, API_VERSION_OFFSET_FETCH, API_VERSION_PRODUCE,
API_VERSION_SYNC_GROUP, ApiVersionsResponse, ConsumerGroupProtocol, FetchBodyDecoder,
FetchPartitionOutcome, FetchPartitionRequest, FetchTopicRequest, FindCoordinatorResponse,
HeartbeatResponse, InitProducerIdResponse, JoinGroupMember, JoinGroupResponse,
LeaveGroupResponse, ListOffsetPartitionRequest, ListOffsetTopicRequest, MetadataResponse,
OffsetCommitPartitionRequest, OffsetCommitTopicRequest, OffsetFetchResponse,
OffsetFetchResult, OffsetFetchTopicRequest, SyncGroupAssignment, SyncGroupResponse,
decode_consumer_protocol_assignment, decode_consumer_protocol_subscription,
decode_list_offsets_response, decode_offset_commit_response, encode_api_versions_request,
encode_consumer_protocol_assignment, encode_consumer_protocol_subscription,
encode_fetch_request, encode_find_coordinator_request, encode_heartbeat_request,
encode_init_producer_id_request, encode_join_group_request, encode_leave_group_request,
encode_list_offsets_request, encode_metadata_request, encode_offset_commit_request,
encode_offset_fetch_request, encode_sync_group_request, request_header_version,
response_header_version,
},
};
const OFFSET_OUT_OF_RANGE: i16 = 1;
const ILLEGAL_GENERATION: i16 = 22;
const UNKNOWN_MEMBER_ID: i16 = 25;
const REBALANCE_IN_PROGRESS: i16 = 27;
const MEMBER_ID_REQUIRED: i16 = 79;
const EARLIEST_TIMESTAMP: i64 = -2;
const LATEST_TIMESTAMP: i64 = -1;
const UNKNOWN_MEMBER_ID_VALUE: &str = "";
#[derive(Debug, Clone)]
pub struct NativeKafkaClientConfig {
pub bootstrap_servers: String,
pub client_id: String,
pub max_in_flight: usize,
pub request_timeout: Duration,
pub broker_floor: &'static str,
pub(crate) security: NativeSecurityConfig,
}
impl NativeKafkaClientConfig {
#[must_use]
pub fn new(bootstrap_servers: impl Into<String>) -> Self {
Self {
bootstrap_servers: bootstrap_servers.into(),
client_id: "datum-kafka-native".to_owned(),
max_in_flight: 4,
request_timeout: Duration::from_secs(30),
broker_floor: "Kafka 2.8+",
security: NativeSecurityConfig::default(),
}
}
#[must_use]
pub fn with_client_id(mut self, client_id: impl Into<String>) -> Self {
self.client_id = client_id.into();
self
}
}
#[derive(Debug, Clone)]
pub struct FetchTuning {
pub max_wait_ms: i32,
pub min_bytes: i32,
pub max_bytes: i32,
pub partition_max_bytes: i32,
pub output_records: usize,
pub output_payload_bytes: usize,
pub include_timestamps: bool,
pub high_watermark: usize,
pub low_watermark: usize,
pub max_retries: usize,
}
impl Default for FetchTuning {
fn default() -> Self {
Self {
max_wait_ms: 500,
min_bytes: 1,
max_bytes: 50 * 1024 * 1024,
partition_max_bytes: 1024 * 1024,
output_records: 10_000,
output_payload_bytes: 10_000 * 256,
include_timestamps: false,
high_watermark: 20_000,
low_watermark: 10_000,
max_retries: 3,
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum NativeCommitPolicy {
Manual,
External,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum AutoOffsetReset {
Earliest,
Latest,
Error,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Subscription {
Assignment(Vec<TopicPartitionAssignment>),
Topics(Vec<String>),
}
impl Subscription {
#[must_use]
pub fn assignment(assignments: Vec<TopicPartitionAssignment>) -> Self {
Self::Assignment(assignments)
}
#[must_use]
pub fn topics(topics: impl IntoIterator<Item = impl Into<String>>) -> Self {
Self::Topics(topics.into_iter().map(Into::into).collect())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum GroupAssignor {
CooperativeSticky,
Range,
}
impl GroupAssignor {
fn protocol_name(self) -> &'static str {
match self {
Self::CooperativeSticky => "cooperative-sticky",
Self::Range => "range",
}
}
fn cooperative(self) -> bool {
matches!(self, Self::CooperativeSticky)
}
}
#[derive(Debug, Clone)]
pub struct NativeKafkaConsumerConfig {
pub client: NativeKafkaClientConfig,
pub fetch: FetchTuning,
pub assignments: Vec<TopicPartitionAssignment>,
pub subscription: Subscription,
pub group_id: String,
pub group_instance_id: Option<String>,
pub group_assignor: GroupAssignor,
pub session_timeout: Duration,
pub rebalance_timeout: Duration,
pub heartbeat_interval: Duration,
pub commit_policy: NativeCommitPolicy,
pub commit_batch_size: usize,
pub commit_interval: Duration,
pub auto_offset_reset: AutoOffsetReset,
pub metrics: NativeKafkaMetrics,
}
impl NativeKafkaConsumerConfig {
#[must_use]
pub fn new(
bootstrap_servers: impl Into<String>,
assignments: Vec<TopicPartitionAssignment>,
) -> Self {
Self {
client: NativeKafkaClientConfig::new(bootstrap_servers),
fetch: FetchTuning::default(),
subscription: Subscription::Assignment(assignments.clone()),
assignments,
group_id: "datum-kafka-native".to_owned(),
group_instance_id: None,
group_assignor: GroupAssignor::CooperativeSticky,
session_timeout: Duration::from_secs(45),
rebalance_timeout: Duration::from_secs(300),
heartbeat_interval: Duration::from_secs(3),
commit_policy: NativeCommitPolicy::Manual,
commit_batch_size: 10_000,
commit_interval: Duration::from_millis(100),
auto_offset_reset: AutoOffsetReset::Latest,
metrics: NativeKafkaMetrics::default(),
}
}
#[must_use]
pub fn subscribe(bootstrap_servers: impl Into<String>, topics: Vec<String>) -> Self {
let mut config = Self::new(bootstrap_servers, Vec::new());
config.subscription = Subscription::Topics(topics);
config
}
#[must_use]
pub fn with_subscription(mut self, subscription: Subscription) -> Self {
if let Subscription::Assignment(assignments) = &subscription {
self.assignments = assignments.clone();
}
self.subscription = subscription;
self
}
#[must_use]
pub fn with_group_id(mut self, group_id: impl Into<String>) -> Self {
self.group_id = group_id.into();
self
}
#[must_use]
pub fn with_group_instance_id(mut self, group_instance_id: impl Into<String>) -> Self {
self.group_instance_id = Some(group_instance_id.into());
self
}
#[must_use]
pub fn with_group_assignor(mut self, assignor: GroupAssignor) -> Self {
self.group_assignor = assignor;
self
}
#[must_use]
pub fn with_commit_policy(mut self, policy: NativeCommitPolicy) -> Self {
self.commit_policy = policy;
self
}
#[must_use]
pub fn with_auto_offset_reset(mut self, reset: AutoOffsetReset) -> Self {
self.auto_offset_reset = reset;
self
}
}
#[derive(Debug)]
pub struct NativeKafkaClient {
config: NativeKafkaClientConfig,
bootstrap: BrokerConnection,
brokers: HashMap<i32, BrokerInfo>,
connections: HashMap<i32, BrokerConnection>,
coordinators: HashMap<String, i32>,
cluster_id: Option<String>,
}
#[derive(Debug, Clone)]
struct BrokerInfo {
host: String,
port: i32,
}
impl BrokerInfo {
fn addr(&self) -> String {
format!("{}:{}", self.host, self.port)
}
}
#[derive(Debug, Clone, Copy)]
struct ApiVersionRange {
min: i16,
max: i16,
}
#[derive(Debug)]
pub struct NativeKafkaConsumer {
client: NativeKafkaClient,
pub(crate) fetch: FetchTuning,
group_id: String,
group_instance_id: Option<String>,
commit_policy: NativeCommitPolicy,
commit_batch_size: usize,
commit_interval: Duration,
auto_offset_reset: AutoOffsetReset,
metrics: NativeKafkaMetrics,
partitions: Vec<PartitionState>,
active_partitions: Vec<TopicPartition>,
paused_partitions: BTreeSet<TopicPartition>,
group: Option<GroupState>,
}
#[derive(Debug, Clone)]
struct PartitionState {
topic: String,
partition: i32,
leader_id: i32,
leader_epoch: i32,
fetch_offset: i64,
log_start_offset: i64,
high_watermark: i64,
}
impl PartitionState {
fn topic_partition(&self) -> TopicPartition {
TopicPartition::new(self.topic.clone(), self.partition)
}
}
#[derive(Debug, Clone)]
struct GroupState {
topics: Vec<String>,
assignor: GroupAssignor,
member_id: String,
generation_id: i32,
protocol_name: String,
last_heartbeat: Instant,
heartbeat_interval: Duration,
session_timeout: Duration,
rebalance_timeout: Duration,
}
#[derive(Debug, Clone)]
struct MemberSubscription {
member_id: String,
group_instance_id: Option<String>,
topics: Vec<String>,
owned_partitions: Vec<TopicPartition>,
generation_id: i32,
}
impl NativeKafkaClient {
pub async fn connect(config: NativeKafkaClientConfig) -> KafkaClientResult<Self> {
let bootstrap_addr = first_bootstrap(&config.bootstrap_servers)?;
let bootstrap_host = broker_host(&bootstrap_addr)?;
let bootstrap = BrokerConnection::connect(
bootstrap_addr,
bootstrap_host,
config.client_id.clone(),
config.max_in_flight,
config.request_timeout,
&config.security,
)
.await?;
let api_versions = negotiate_api_versions(&bootstrap, &config).await?;
require_client_versions(&api_versions, config.broker_floor)?;
initialize_security(&bootstrap, &config, &api_versions).await?;
let mut client = Self {
config,
bootstrap,
brokers: HashMap::new(),
connections: HashMap::new(),
coordinators: HashMap::new(),
cluster_id: None,
};
client.refresh_metadata(&[]).await?;
Ok(client)
}
pub async fn consumer(
mut self,
config: NativeKafkaConsumerConfig,
) -> KafkaClientResult<NativeKafkaConsumer> {
if config.commit_policy == NativeCommitPolicy::Manual && config.group_id.is_empty() {
return Err(KafkaClientError::InvalidConfig(
"native Kafka manual commits require a non-empty group id".to_owned(),
));
}
let (partitions, active_partitions, group) = match &config.subscription {
Subscription::Assignment(assignments) => {
if assignments.is_empty() {
return Err(KafkaClientError::InvalidConfig(
"native Kafka consumer requires at least one manual assignment".to_owned(),
));
}
let topics = assignments
.iter()
.map(|assignment| assignment.topic.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let metadata = self.refresh_metadata(&topics).await?;
let mut partitions = Vec::with_capacity(assignments.len());
for assignment in assignments {
partitions.push(partition_from_metadata(&metadata, assignment)?);
}
self.resolve_start_offsets(
&mut partitions,
assignments,
&config.group_id,
config.auto_offset_reset,
)
.await?;
let active_partitions = assignments
.iter()
.map(TopicPartitionAssignment::topic_partition)
.collect();
config.metrics.add_assigned(assignments.len() as u64);
(partitions, active_partitions, None)
}
Subscription::Topics(topics) => {
let topics = normalized_topics(topics)?;
let metadata = self.refresh_metadata(&topics).await?;
let group = GroupState {
topics,
assignor: config.group_assignor,
member_id: UNKNOWN_MEMBER_ID_VALUE.to_owned(),
generation_id: -1,
protocol_name: config.group_assignor.protocol_name().to_owned(),
last_heartbeat: Instant::now(),
heartbeat_interval: config.heartbeat_interval,
session_timeout: config.session_timeout,
rebalance_timeout: config.rebalance_timeout,
};
let mut consumer = NativeKafkaConsumer {
client: self,
fetch: config.fetch,
group_id: config.group_id,
group_instance_id: config.group_instance_id,
commit_policy: config.commit_policy,
commit_batch_size: config.commit_batch_size,
commit_interval: config.commit_interval,
auto_offset_reset: config.auto_offset_reset,
metrics: config.metrics,
partitions: Vec::new(),
active_partitions: Vec::new(),
paused_partitions: BTreeSet::new(),
group: Some(group),
};
consumer.join_group_and_apply(None, Some(metadata)).await?;
return Ok(consumer);
}
};
Ok(NativeKafkaConsumer {
client: self,
fetch: config.fetch,
group_id: config.group_id,
group_instance_id: config.group_instance_id,
commit_policy: config.commit_policy,
commit_batch_size: config.commit_batch_size,
commit_interval: config.commit_interval,
auto_offset_reset: config.auto_offset_reset,
metrics: config.metrics,
partitions,
active_partitions,
paused_partitions: BTreeSet::new(),
group,
})
}
async fn list_offsets_for_broker(
&mut self,
broker_id: i32,
requests: &[ListOffsetTopicRequest],
) -> KafkaClientResult<Vec<crate::native::protocol::ListOffsetResult>> {
let body = encode_list_offsets_request(requests)?;
let connection = self.connection_for_broker(broker_id).await?;
let response = connection
.request(
API_KEY_LIST_OFFSETS,
API_VERSION_LIST_OFFSETS,
request_header_version(API_VERSION_LIST_OFFSETS),
response_header_version(API_KEY_LIST_OFFSETS, API_VERSION_LIST_OFFSETS),
body,
)
.await?;
decode_list_offsets_response(API_VERSION_LIST_OFFSETS, &response)
}
pub(crate) async fn refresh_metadata(
&mut self,
topics: &[String],
) -> KafkaClientResult<MetadataResponse> {
let body = encode_metadata_request(topics)?;
let response = self
.bootstrap
.request(
API_KEY_METADATA,
API_VERSION_METADATA,
request_header_version(API_VERSION_METADATA),
response_header_version(API_KEY_METADATA, API_VERSION_METADATA),
body,
)
.await?;
let metadata = MetadataResponse::decode(API_VERSION_METADATA, &response)?;
for broker in &metadata.brokers {
self.brokers.insert(
broker.node_id,
BrokerInfo {
host: broker.host.clone(),
port: broker.port,
},
);
}
self.cluster_id = metadata.cluster_id.clone();
Ok(metadata)
}
pub(crate) async fn init_producer_id(&mut self) -> KafkaClientResult<(i64, i16)> {
let body = encode_init_producer_id_request(None, -1, -1, -1)?;
let response = self
.bootstrap
.request(
API_KEY_INIT_PRODUCER_ID,
API_VERSION_INIT_PRODUCER_ID,
request_header_version(API_VERSION_INIT_PRODUCER_ID),
response_header_version(API_KEY_INIT_PRODUCER_ID, API_VERSION_INIT_PRODUCER_ID),
body,
)
.await?;
let response = InitProducerIdResponse::decode(API_VERSION_INIT_PRODUCER_ID, &response)?;
if response.error_code != 0 {
return Err(KafkaClientError::broker(
"InitProducerId",
response.error_code,
));
}
Ok((response.producer_id, response.producer_epoch))
}
pub(crate) async fn connection_for_broker(
&mut self,
node_id: i32,
) -> KafkaClientResult<BrokerConnection> {
if let Some(connection) = self.connections.get(&node_id) {
return Ok(connection.clone());
}
let broker = self
.brokers
.get(&node_id)
.ok_or_else(|| KafkaClientError::protocol(format!("unknown broker {node_id}")))?
.clone();
let connection = BrokerConnection::connect(
broker.addr(),
broker.host.clone(),
self.config.client_id.clone(),
self.config.max_in_flight,
self.config.request_timeout,
&self.config.security,
)
.await?;
let api_versions = negotiate_api_versions(&connection, &self.config).await?;
require_client_versions(&api_versions, self.config.broker_floor)?;
initialize_security(&connection, &self.config, &api_versions).await?;
self.connections.insert(node_id, connection.clone());
Ok(connection)
}
pub(crate) fn invalidate_broker_connection(&mut self, node_id: i32) {
self.connections.remove(&node_id);
}
async fn connection_for_coordinator(
&mut self,
group_id: &str,
) -> KafkaClientResult<BrokerConnection> {
if let Some(node_id) = self.coordinators.get(group_id).copied() {
return self.connection_for_broker(node_id).await;
}
let mut coordinator = None;
let mut last_error = None;
for attempt in 0..10 {
match self.find_coordinator(group_id).await {
Ok(found) => {
coordinator = Some(found);
break;
}
Err(error) if error.retriable() && attempt < 9 => {
last_error = Some(error);
tokio::time::sleep(Duration::from_millis(100)).await;
}
Err(error) => return Err(error),
}
}
let coordinator = coordinator.ok_or_else(|| {
last_error.unwrap_or_else(|| {
KafkaClientError::protocol("FindCoordinator retry loop exhausted")
})
})?;
self.brokers.insert(
coordinator.node_id,
BrokerInfo {
host: coordinator.host,
port: coordinator.port,
},
);
self.coordinators
.insert(group_id.to_owned(), coordinator.node_id);
self.connection_for_broker(coordinator.node_id).await
}
async fn find_coordinator(
&mut self,
group_id: &str,
) -> KafkaClientResult<FindCoordinatorResponse> {
let body = encode_find_coordinator_request(group_id)?;
let response = self
.bootstrap
.request(
API_KEY_FIND_COORDINATOR,
API_VERSION_FIND_COORDINATOR,
request_header_version(API_VERSION_FIND_COORDINATOR),
response_header_version(API_KEY_FIND_COORDINATOR, API_VERSION_FIND_COORDINATOR),
body,
)
.await?;
let response = FindCoordinatorResponse::decode(API_VERSION_FIND_COORDINATOR, &response)?;
if response.error_code == 0 {
Ok(response)
} else {
Err(KafkaClientError::broker(
"FindCoordinator",
response.error_code,
))
}
}
async fn offset_commit(
&mut self,
group_id: &str,
generation_id: i32,
member_id: &str,
group_instance_id: Option<&str>,
commits: &[OffsetCommit],
) -> KafkaClientResult<()> {
if commits.is_empty() {
return Ok(());
}
let topics = offset_commit_topics(commits);
for attempt in 0..2 {
let body = encode_offset_commit_request(
group_id,
generation_id,
member_id,
group_instance_id,
&topics,
)?;
let connection = self.connection_for_coordinator(group_id).await?;
let response = connection
.request(
API_KEY_OFFSET_COMMIT,
API_VERSION_OFFSET_COMMIT,
request_header_version(API_VERSION_OFFSET_COMMIT),
response_header_version(API_KEY_OFFSET_COMMIT, API_VERSION_OFFSET_COMMIT),
body,
)
.await?;
let results = decode_offset_commit_response(API_VERSION_OFFSET_COMMIT, &response)?;
if let Some(error) = results.iter().find(|result| result.error_code != 0) {
if attempt == 0 && coordinator_error(error.error_code) {
self.coordinators.remove(group_id);
continue;
}
return Err(KafkaClientError::broker("OffsetCommit", error.error_code));
}
return Ok(());
}
Err(KafkaClientError::protocol(
"OffsetCommit coordinator retry loop exhausted",
))
}
async fn join_group(
&mut self,
group_id: &str,
state: &GroupState,
group_instance_id: Option<&str>,
protocols: &[ConsumerGroupProtocol],
) -> KafkaClientResult<JoinGroupResponse> {
for attempt in 0..3 {
let body = encode_join_group_request(
group_id,
millis_i32(state.session_timeout, "session.timeout.ms")?,
millis_i32(state.rebalance_timeout, "rebalance.timeout.ms")?,
&state.member_id,
group_instance_id,
protocols,
)?;
let connection = self.connection_for_coordinator(group_id).await?;
let response = connection
.request(
API_KEY_JOIN_GROUP,
API_VERSION_JOIN_GROUP,
request_header_version(API_VERSION_JOIN_GROUP),
response_header_version(API_KEY_JOIN_GROUP, API_VERSION_JOIN_GROUP),
body,
)
.await?;
let response = JoinGroupResponse::decode(API_VERSION_JOIN_GROUP, &response)?;
if response.error_code == 0 || response.error_code == MEMBER_ID_REQUIRED {
return Ok(response);
}
if attempt < 2 && coordinator_error(response.error_code) {
self.coordinators.remove(group_id);
continue;
}
if attempt < 2 && response.error_code == REBALANCE_IN_PROGRESS {
tokio::time::sleep(Duration::from_millis(50)).await;
continue;
}
return Err(KafkaClientError::broker("JoinGroup", response.error_code));
}
Err(KafkaClientError::protocol(
"JoinGroup coordinator retry loop exhausted",
))
}
async fn sync_group(
&mut self,
group_id: &str,
generation_id: i32,
member_id: &str,
group_instance_id: Option<&str>,
assignments: &[SyncGroupAssignment],
) -> KafkaClientResult<SyncGroupResponse> {
for attempt in 0..2 {
let body = encode_sync_group_request(
group_id,
generation_id,
member_id,
group_instance_id,
assignments,
)?;
let connection = self.connection_for_coordinator(group_id).await?;
let response = connection
.request(
API_KEY_SYNC_GROUP,
API_VERSION_SYNC_GROUP,
request_header_version(API_VERSION_SYNC_GROUP),
response_header_version(API_KEY_SYNC_GROUP, API_VERSION_SYNC_GROUP),
body,
)
.await?;
let response = SyncGroupResponse::decode(API_VERSION_SYNC_GROUP, &response)?;
if response.error_code == 0 || response.error_code == REBALANCE_IN_PROGRESS {
return Ok(response);
}
if attempt == 0 && coordinator_error(response.error_code) {
self.coordinators.remove(group_id);
continue;
}
return Err(KafkaClientError::broker("SyncGroup", response.error_code));
}
Err(KafkaClientError::protocol(
"SyncGroup coordinator retry loop exhausted",
))
}
async fn heartbeat(
&mut self,
group_id: &str,
generation_id: i32,
member_id: &str,
group_instance_id: Option<&str>,
) -> KafkaClientResult<HeartbeatResponse> {
let body = encode_heartbeat_request(group_id, generation_id, member_id, group_instance_id)?;
let connection = self.connection_for_coordinator(group_id).await?;
let response = connection
.request(
API_KEY_HEARTBEAT,
API_VERSION_HEARTBEAT,
request_header_version(API_VERSION_HEARTBEAT),
response_header_version(API_KEY_HEARTBEAT, API_VERSION_HEARTBEAT),
body,
)
.await?;
HeartbeatResponse::decode(API_VERSION_HEARTBEAT, &response)
}
async fn leave_group(
&mut self,
group_id: &str,
member_id: &str,
group_instance_id: Option<&str>,
) -> KafkaClientResult<()> {
if member_id.is_empty() {
return Ok(());
}
let body = encode_leave_group_request(group_id, member_id, group_instance_id)?;
let connection = self.connection_for_coordinator(group_id).await?;
let response = connection
.request(
API_KEY_LEAVE_GROUP,
API_VERSION_LEAVE_GROUP,
request_header_version(API_VERSION_LEAVE_GROUP),
response_header_version(API_KEY_LEAVE_GROUP, API_VERSION_LEAVE_GROUP),
body,
)
.await?;
let response = LeaveGroupResponse::decode(API_VERSION_LEAVE_GROUP, &response)?;
if response.error_code != 0 {
return Err(KafkaClientError::broker("LeaveGroup", response.error_code));
}
if let Some(error) = response
.members
.iter()
.find(|member| member.error_code != 0)
{
return Err(KafkaClientError::broker(
"LeaveGroup member",
error.error_code,
));
}
Ok(())
}
async fn offset_fetch(
&mut self,
group_id: &str,
partitions: &[TopicPartition],
) -> KafkaClientResult<Vec<OffsetFetchResult>> {
let topics = offset_fetch_topics(partitions);
for attempt in 0..2 {
let body = encode_offset_fetch_request(group_id, &topics)?;
let connection = self.connection_for_coordinator(group_id).await?;
let response = connection
.request(
API_KEY_OFFSET_FETCH,
API_VERSION_OFFSET_FETCH,
request_header_version(API_VERSION_OFFSET_FETCH),
response_header_version(API_KEY_OFFSET_FETCH, API_VERSION_OFFSET_FETCH),
body,
)
.await?;
let response = OffsetFetchResponse::decode(API_VERSION_OFFSET_FETCH, &response)?;
if response.error_code != 0 {
if attempt == 0 && coordinator_error(response.error_code) {
self.coordinators.remove(group_id);
continue;
}
return Err(KafkaClientError::broker("OffsetFetch", response.error_code));
}
if let Some(error) = response
.offsets
.iter()
.find(|result| result.error_code != 0)
{
return Err(KafkaClientError::broker(
"OffsetFetch partition",
error.error_code,
));
}
return Ok(response.offsets);
}
Err(KafkaClientError::protocol(
"OffsetFetch coordinator retry loop exhausted",
))
}
async fn resolve_start_offsets(
&mut self,
partitions: &mut [PartitionState],
assignments: &[TopicPartitionAssignment],
group_id: &str,
auto_offset_reset: AutoOffsetReset,
) -> KafkaClientResult<()> {
let mut resolved = BTreeMap::<TopicPartition, i64>::new();
let committed_assignments = assignments
.iter()
.filter(|assignment| assignment.start_offset == StartOffset::Committed)
.map(TopicPartitionAssignment::topic_partition)
.collect::<Vec<_>>();
if !committed_assignments.is_empty() {
let offsets = self.offset_fetch(group_id, &committed_assignments).await?;
for assignment in assignments
.iter()
.filter(|assignment| assignment.start_offset == StartOffset::Committed)
{
let topic_partition = assignment.topic_partition();
if let Some(offset) = offsets.iter().find(|offset| {
offset.topic == assignment.topic
&& offset.partition_index == assignment.partition
&& offset.committed_offset >= 0
}) {
resolved.insert(topic_partition, offset.committed_offset);
}
}
}
let mut requests =
BTreeMap::<i32, BTreeMap<String, Vec<ListOffsetPartitionRequest>>>::new();
for assignment in assignments {
let timestamp = match assignment.start_offset {
StartOffset::Beginning => EARLIEST_TIMESTAMP,
StartOffset::End => LATEST_TIMESTAMP,
StartOffset::Timestamp(timestamp) => timestamp,
StartOffset::Committed if resolved.contains_key(&assignment.topic_partition()) => {
continue;
}
StartOffset::Committed => match reset_timestamp(auto_offset_reset)? {
Some(timestamp) => timestamp,
None => {
return Err(KafkaClientError::InvalidConfig(format!(
"no committed offset for {}:{} and auto-offset-reset=error",
assignment.topic, assignment.partition
)));
}
},
StartOffset::Offset(_) => continue,
};
let leader_epoch = partitions
.iter()
.find(|partition| {
partition.topic == assignment.topic
&& partition.partition == assignment.partition
})
.map_or(-1, |partition| partition.leader_epoch);
let leader_id = partitions
.iter()
.find(|partition| {
partition.topic == assignment.topic
&& partition.partition == assignment.partition
})
.map_or(-1, |partition| partition.leader_id);
requests
.entry(leader_id)
.or_default()
.entry(assignment.topic.clone())
.or_default()
.push(ListOffsetPartitionRequest {
partition_index: assignment.partition,
current_leader_epoch: leader_epoch,
timestamp,
});
}
let mut offsets = BTreeMap::<TopicPartition, i64>::new();
for (leader_id, topics) in requests {
if leader_id < 0 {
return Err(KafkaClientError::protocol(
"cannot list offsets before partition leader is known",
));
}
let topics = topics
.into_iter()
.map(|(name, partitions)| ListOffsetTopicRequest { name, partitions })
.collect::<Vec<_>>();
for offset in self.list_offsets_for_broker(leader_id, &topics).await? {
if offset.error_code == 0 {
offsets.insert(
TopicPartition::new(offset.topic, offset.partition_index),
offset.offset,
);
} else {
return Err(KafkaClientError::broker("ListOffsets", offset.error_code));
}
}
}
for assignment in assignments {
let offset = match assignment.start_offset {
StartOffset::Offset(offset) => offset,
StartOffset::Beginning
| StartOffset::End
| StartOffset::Timestamp(_)
| StartOffset::Committed => resolved
.get(&assignment.topic_partition())
.copied()
.or_else(|| offsets.get(&assignment.topic_partition()).copied())
.ok_or_else(|| {
KafkaClientError::protocol(format!(
"ListOffsets response omitted {}:{}",
assignment.topic, assignment.partition
))
})?,
};
if let Some(partition) = partitions.iter_mut().find(|partition| {
partition.topic == assignment.topic && partition.partition == assignment.partition
}) {
partition.fetch_offset = offset;
}
}
Ok(())
}
}
impl NativeKafkaConsumer {
pub async fn connect(config: NativeKafkaConsumerConfig) -> KafkaClientResult<Self> {
let client = NativeKafkaClient::connect(config.client.clone()).await?;
client.consumer(config).await
}
pub async fn fetch_batch(&mut self) -> KafkaClientResult<Option<KafkaPayloadBatch>> {
let mut retries = 0_usize;
loop {
match self.fetch_once().await {
Ok(batch) => return Ok(batch),
Err(error) if error.retriable() && retries < self.fetch.max_retries => {
retries += 1;
self.client.connections.clear();
let topics = self
.partitions
.iter()
.map(|partition| partition.topic.clone())
.collect::<std::collections::BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let metadata = self.client.refresh_metadata(&topics).await?;
for partition in &mut self.partitions {
if let Some(next) = partition_from_topic_metadata(&metadata, partition) {
partition.leader_id = next.leader_id;
partition.leader_epoch = next.leader_epoch;
}
}
}
Err(error) => return Err(error),
}
}
}
pub fn active_partitions(&self) -> &[TopicPartition] {
&self.active_partitions
}
pub(crate) fn refresh_pauses(&mut self, commit_state: &LocalCommitState) {
let outstanding = commit_state.outstanding_by_partition();
for partition in self.partitions.iter().map(PartitionState::topic_partition) {
let count = outstanding.get(&partition).copied().unwrap_or(0) as usize;
if self.paused_partitions.contains(&partition) {
if count <= self.fetch.low_watermark {
self.paused_partitions.remove(&partition);
}
} else if count >= self.fetch.high_watermark {
self.paused_partitions.insert(partition);
}
}
}
pub(crate) fn group_enabled(&self) -> bool {
self.group.is_some()
}
pub(crate) async fn poll_group(
&mut self,
commit_state: &LocalCommitState,
) -> KafkaClientResult<()> {
let Some(group) = self.group.as_ref() else {
return Ok(());
};
if group.last_heartbeat.elapsed() < group.heartbeat_interval {
return Ok(());
}
let group_id = self.group_id.clone();
let generation_id = group.generation_id;
let member_id = group.member_id.clone();
let group_instance_id = self.group_instance_id.clone();
let response = self
.client
.heartbeat(
&group_id,
generation_id,
&member_id,
group_instance_id.as_deref(),
)
.await?;
if let Some(group) = self.group.as_mut() {
group.last_heartbeat = Instant::now();
}
match response.error_code {
0 => Ok(()),
REBALANCE_IN_PROGRESS => self.join_group_and_apply(Some(commit_state), None).await,
ILLEGAL_GENERATION | UNKNOWN_MEMBER_ID => {
self.mark_current_assignment_lost(commit_state);
Err(KafkaClientError::broker("Heartbeat", response.error_code))
}
code if coordinator_error(code) => {
self.client.coordinators.remove(&group_id);
self.join_group_and_apply(Some(commit_state), None).await
}
code => Err(KafkaClientError::broker("Heartbeat", code)),
}
}
pub(crate) async fn leave_group(&mut self) -> KafkaClientResult<()> {
let Some(group) = self.group.as_ref() else {
return Ok(());
};
let group_id = self.group_id.clone();
let member_id = group.member_id.clone();
let group_instance_id = self.group_instance_id.clone();
self.client
.leave_group(&group_id, &member_id, group_instance_id.as_deref())
.await
}
pub(crate) fn commit_policy(&self) -> NativeCommitPolicy {
self.commit_policy
}
pub(crate) fn commit_batch_size(&self) -> usize {
self.commit_batch_size
}
pub(crate) fn commit_interval(&self) -> Duration {
self.commit_interval
}
pub(crate) async fn commit_offsets(
&mut self,
commits: &[OffsetCommit],
) -> KafkaClientResult<()> {
self.metrics.broker_commit_request();
let (generation_id, member_id, group_instance_id) = self
.group
.as_ref()
.map(|group| {
(
group.generation_id,
group.member_id.as_str(),
self.group_instance_id.as_deref(),
)
})
.unwrap_or((-1, "", None));
match self
.client
.offset_commit(
&self.group_id.clone(),
generation_id,
member_id,
group_instance_id,
commits,
)
.await
{
Ok(()) => Ok(()),
Err(error) => {
self.metrics.broker_commit_failed();
Err(error)
}
}
}
async fn join_group_and_apply(
&mut self,
commit_state: Option<&LocalCommitState>,
metadata: Option<MetadataResponse>,
) -> KafkaClientResult<()> {
let mut metadata = metadata;
for _ in 0..4 {
let assignment = match self.join_group_once(metadata.take()).await {
Ok(assignment) => assignment,
Err(
error @ KafkaClientError::Broker {
code: ILLEGAL_GENERATION | UNKNOWN_MEMBER_ID,
..
},
) => {
if let Some(commit_state) = commit_state {
self.mark_current_assignment_lost(commit_state);
}
return Err(error);
}
Err(error) => return Err(error),
};
let revoked = self.revoked_partitions(&assignment);
if !revoked.is_empty()
&& let Some(commit_state) = commit_state
{
self.force_commit_partitions(commit_state, &revoked).await?;
self.metrics.add_revoked(revoked.len() as u64);
}
let assigned = assignment.iter().cloned().collect::<BTreeSet<_>>();
self.apply_group_assignment(assigned).await?;
if revoked.is_empty()
|| !self
.group
.as_ref()
.is_some_and(|group| group.assignor.cooperative())
{
return Ok(());
}
}
Err(KafkaClientError::protocol(
"cooperative rebalance did not settle after repeated rejoins",
))
}
async fn join_group_once(
&mut self,
metadata: Option<MetadataResponse>,
) -> KafkaClientResult<Vec<TopicPartition>> {
let mut metadata = metadata;
loop {
let group = self
.group
.as_ref()
.ok_or_else(|| {
KafkaClientError::protocol("group join requested for manual consumer")
})?
.clone();
let group_id = self.group_id.clone();
let group_instance_id = self.group_instance_id.clone();
let owned = self.active_partitions.clone();
let protocols = [ConsumerGroupProtocol {
name: group.assignor.protocol_name().to_owned(),
metadata: encode_consumer_protocol_subscription(
&group.topics,
&owned,
group.generation_id,
)?,
}];
let response = self
.client
.join_group(&group_id, &group, group_instance_id.as_deref(), &protocols)
.await?;
if let Some(group) = self.group.as_mut() {
group.member_id = response.member_id.clone();
group.last_heartbeat = Instant::now();
}
if response.error_code == MEMBER_ID_REQUIRED {
continue;
}
if response.error_code != 0 {
return Err(KafkaClientError::broker("JoinGroup", response.error_code));
}
let protocol_name = response.protocol_name.clone();
let assignor = assignor_for_protocol(&protocol_name)?;
let is_leader = response.leader == response.member_id;
let assignments = if is_leader {
let members = decode_member_subscriptions(&response.members)?;
let member_topics = members
.iter()
.flat_map(|member| member.topics.iter().cloned())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let metadata = match metadata.take() {
Some(metadata) if metadata_covers(&metadata, &member_topics) => metadata,
_ => self.client.refresh_metadata(&member_topics).await?,
};
let plan = assign_partitions(assignor, &members, &metadata);
encode_sync_assignments(plan)?
} else {
Vec::new()
};
let sync = self
.client
.sync_group(
&group_id,
response.generation_id,
&response.member_id,
group_instance_id.as_deref(),
&assignments,
)
.await?;
if sync.error_code == REBALANCE_IN_PROGRESS {
continue;
}
if matches!(sync.error_code, ILLEGAL_GENERATION | UNKNOWN_MEMBER_ID) {
return Err(KafkaClientError::broker("SyncGroup", sync.error_code));
}
if sync.error_code != 0 {
return Err(KafkaClientError::broker("SyncGroup", sync.error_code));
}
if let Some(group) = self.group.as_mut() {
group.generation_id = response.generation_id;
group.protocol_name = protocol_name;
group.assignor = assignor;
group.last_heartbeat = Instant::now();
}
return decode_consumer_protocol_assignment(&sync.assignment);
}
}
fn revoked_partitions(&self, assignment: &[TopicPartition]) -> BTreeSet<TopicPartition> {
let assigned = assignment.iter().cloned().collect::<BTreeSet<_>>();
self.active_partitions
.iter()
.filter(|partition| !assigned.contains(*partition))
.cloned()
.collect()
}
async fn force_commit_partitions(
&mut self,
commit_state: &LocalCommitState,
partitions: &BTreeSet<TopicPartition>,
) -> KafkaClientResult<()> {
let commits = commit_state.due_commits_for(partitions);
if commits.is_empty() {
return Ok(());
}
if self.commit_policy == NativeCommitPolicy::Manual {
self.commit_offsets(&commits).await?;
}
commit_state.mark_committed(&commits);
Ok(())
}
async fn apply_group_assignment(
&mut self,
assignment: BTreeSet<TopicPartition>,
) -> KafkaClientResult<()> {
let previous_assignment = self
.active_partitions
.iter()
.cloned()
.collect::<BTreeSet<_>>();
let newly_assigned = assignment
.iter()
.filter(|partition| !previous_assignment.contains(*partition))
.count();
let assigned_topics = assignment
.iter()
.map(|partition| partition.topic.clone())
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
let metadata = self.client.refresh_metadata(&assigned_topics).await?;
let mut previous = self
.partitions
.iter()
.map(|partition| (partition.topic_partition(), partition.clone()))
.collect::<BTreeMap<_, _>>();
let mut partitions = Vec::with_capacity(assignment.len());
let mut new_assignments = Vec::new();
for topic_partition in &assignment {
if let Some(existing) = previous.remove(topic_partition) {
partitions.push(existing);
} else {
let assignment = TopicPartitionAssignment::committed(
topic_partition.topic.clone(),
topic_partition.partition,
);
partitions.push(partition_from_metadata(&metadata, &assignment)?);
new_assignments.push(assignment);
}
}
if !new_assignments.is_empty() {
self.client
.resolve_start_offsets(
&mut partitions,
&new_assignments,
&self.group_id,
self.auto_offset_reset,
)
.await?;
}
self.active_partitions = assignment.into_iter().collect();
self.paused_partitions
.retain(|partition| self.active_partitions.contains(partition));
self.partitions = partitions;
self.metrics.add_assigned(newly_assigned as u64);
Ok(())
}
fn mark_current_assignment_lost(&mut self, commit_state: &LocalCommitState) {
let lost = self
.active_partitions
.iter()
.cloned()
.collect::<BTreeSet<_>>();
self.metrics.add_lost(lost.len() as u64);
commit_state.mark_lost(&lost);
self.partitions.clear();
self.active_partitions.clear();
self.paused_partitions.clear();
}
async fn fetch_once(&mut self) -> KafkaClientResult<Option<KafkaPayloadBatch>> {
let requests = self.fetch_requests_by_leader();
let mut builder = PayloadBatchBuilder::with_capacity(
self.fetch.output_records,
self.fetch.output_payload_bytes,
self.fetch.include_timestamps,
);
for (leader_id, topics) in requests {
let offsets = fetch_offset_lookup(&topics);
let body = encode_fetch_request(
self.fetch.max_wait_ms,
self.fetch.min_bytes,
self.fetch.max_bytes,
&topics,
)?;
let connection = self.client.connection_for_broker(leader_id).await?;
let decoder = FetchBodyDecoder::new(API_VERSION_FETCH, offsets, builder)?;
let decoder = profile::measure_async(
ProfileBucket::SocketFetch,
connection.request_decoded(
API_KEY_FETCH,
API_VERSION_FETCH,
request_header_version(API_VERSION_FETCH),
response_header_version(API_KEY_FETCH, API_VERSION_FETCH),
body,
decoder,
),
)
.await?;
let (returned_builder, error_code, outcomes) = decoder.into_parts();
builder = returned_builder;
if error_code != 0 {
return Err(KafkaClientError::broker("Fetch", error_code));
}
self.apply_fetch_outcomes(outcomes).await?;
if builder.len() >= self.fetch.output_records {
break;
}
}
Ok(builder.finish(self.active_partitions.clone()))
}
fn fetch_requests_by_leader(&self) -> BTreeMap<i32, Vec<FetchTopicRequest>> {
let mut leaders = BTreeMap::<i32, BTreeMap<String, Vec<FetchPartitionRequest>>>::new();
let has_paused_partitions = !self.paused_partitions.is_empty();
for partition in &self.partitions {
if has_paused_partitions
&& self
.paused_partitions
.contains(&partition.topic_partition())
{
continue;
}
leaders
.entry(partition.leader_id)
.or_default()
.entry(partition.topic.clone())
.or_default()
.push(FetchPartitionRequest {
partition: partition.partition,
current_leader_epoch: partition.leader_epoch,
fetch_offset: partition.fetch_offset,
last_fetched_epoch: -1,
log_start_offset: partition.log_start_offset,
partition_max_bytes: self.fetch.partition_max_bytes,
});
}
leaders
.into_iter()
.map(|(leader, topics)| {
let topics = topics
.into_iter()
.map(|(topic, partitions)| FetchTopicRequest { topic, partitions })
.collect();
(leader, topics)
})
.collect()
}
async fn apply_fetch_outcomes(
&mut self,
outcomes: Vec<FetchPartitionOutcome>,
) -> KafkaClientResult<()> {
for outcome in outcomes {
if outcome.error_code == OFFSET_OUT_OF_RANGE {
self.recover_out_of_range(&outcome.topic, outcome.partition_index)
.await?;
continue;
}
if outcome.error_code != 0 {
return Err(KafkaClientError::broker(
"Fetch partition",
outcome.error_code,
));
}
let Some(partition_state) = self.partitions.iter_mut().find(|partition| {
partition.topic == outcome.topic && partition.partition == outcome.partition_index
}) else {
continue;
};
profile::measure(ProfileBucket::OffsetBookkeeping, || {
partition_state.high_watermark = outcome.high_watermark;
partition_state.log_start_offset = outcome.log_start_offset;
if let Some(max_offset) = outcome.max_offset {
partition_state.fetch_offset = max_offset + 1;
}
if partition_state.fetch_offset < outcome.log_start_offset {
partition_state.fetch_offset = outcome.log_start_offset;
}
});
}
Ok(())
}
async fn recover_out_of_range(&mut self, topic: &str, partition: i32) -> KafkaClientResult<()> {
let Some(timestamp) = reset_timestamp(self.auto_offset_reset)? else {
return Err(KafkaClientError::broker(
"Fetch partition",
OFFSET_OUT_OF_RANGE,
));
};
let Some(partition_state) = self
.partitions
.iter()
.find(|state| state.topic == topic && state.partition == partition)
.cloned()
else {
return Ok(());
};
let topics = [ListOffsetTopicRequest {
name: topic.to_owned(),
partitions: vec![ListOffsetPartitionRequest {
partition_index: partition,
current_leader_epoch: partition_state.leader_epoch,
timestamp,
}],
}];
let offsets = self
.client
.list_offsets_for_broker(partition_state.leader_id, &topics)
.await?;
let Some(offset) = offsets
.into_iter()
.find(|offset| offset.topic == topic && offset.partition_index == partition)
else {
return Err(KafkaClientError::protocol(format!(
"ListOffsets response omitted {topic}:{partition} during out-of-range recovery"
)));
};
if offset.error_code != 0 {
return Err(KafkaClientError::broker("ListOffsets", offset.error_code));
}
if let Some(state) = self
.partitions
.iter_mut()
.find(|state| state.topic == topic && state.partition == partition)
{
profile::measure(ProfileBucket::OffsetBookkeeping, || {
state.fetch_offset = offset.offset;
});
}
self.metrics.offset_out_of_range_reset();
Ok(())
}
}
fn fetch_offset_lookup(topics: &[FetchTopicRequest]) -> BTreeMap<(String, i32), i64> {
let mut offsets = BTreeMap::new();
for topic in topics {
for partition in &topic.partitions {
offsets.insert(
(topic.topic.clone(), partition.partition),
partition.fetch_offset,
);
}
}
offsets
}
async fn negotiate_api_versions(
connection: &BrokerConnection,
config: &NativeKafkaClientConfig,
) -> KafkaClientResult<HashMap<i16, ApiVersionRange>> {
let body = encode_api_versions_request("datum-mq-native", VERSION)?;
let response = connection
.request_before_auth(
API_KEY_API_VERSIONS,
API_VERSION_API_VERSIONS,
request_header_version(API_VERSION_API_VERSIONS),
response_header_version(API_KEY_API_VERSIONS, API_VERSION_API_VERSIONS),
body,
)
.await?;
let response = ApiVersionsResponse::decode(API_VERSION_API_VERSIONS, &response)?;
if response.error_code != 0 {
return Err(KafkaClientError::broker("ApiVersions", response.error_code));
}
let mut versions = HashMap::with_capacity(response.api_keys.len());
for api in response.api_keys {
versions.insert(
api.api_key,
ApiVersionRange {
min: api.min_version,
max: api.max_version,
},
);
}
require_version(
&versions,
API_KEY_API_VERSIONS,
API_VERSION_API_VERSIONS,
config.broker_floor,
)?;
Ok(versions)
}
async fn initialize_security(
connection: &BrokerConnection,
config: &NativeKafkaClientConfig,
versions: &HashMap<i16, ApiVersionRange>,
) -> KafkaClientResult<()> {
if config.security.sasl.is_none() {
return Ok(());
}
require_version(
versions,
API_KEY_SASL_HANDSHAKE,
1,
"Kafka 2.8+ SASL handshake v1",
)?;
require_version(
versions,
API_KEY_SASL_AUTHENTICATE,
1,
"Kafka 2.8+ SASL authentication v1",
)?;
connection.initialize_sasl(1, 1).await
}
fn require_version(
versions: &HashMap<i16, ApiVersionRange>,
api_key: i16,
version: i16,
floor: &str,
) -> KafkaClientResult<()> {
let Some(range) = versions.get(&api_key) else {
return Err(KafkaClientError::unsupported(format!(
"broker did not advertise API key {api_key}; KC-1 requires {floor}"
)));
};
if range.min <= version && range.max >= version {
Ok(())
} else {
Err(KafkaClientError::unsupported(format!(
"broker API key {api_key} supports v{}..v{}, but KC-1 requires v{version} ({floor})",
range.min, range.max
)))
}
}
fn require_client_versions(
versions: &HashMap<i16, ApiVersionRange>,
floor: &str,
) -> KafkaClientResult<()> {
for (api_key, version) in [
(API_KEY_PRODUCE, API_VERSION_PRODUCE),
(API_KEY_METADATA, API_VERSION_METADATA),
(API_KEY_LIST_OFFSETS, API_VERSION_LIST_OFFSETS),
(API_KEY_FETCH, API_VERSION_FETCH),
(API_KEY_FIND_COORDINATOR, API_VERSION_FIND_COORDINATOR),
(API_KEY_OFFSET_COMMIT, API_VERSION_OFFSET_COMMIT),
(API_KEY_OFFSET_FETCH, API_VERSION_OFFSET_FETCH),
(API_KEY_JOIN_GROUP, API_VERSION_JOIN_GROUP),
(API_KEY_HEARTBEAT, API_VERSION_HEARTBEAT),
(API_KEY_LEAVE_GROUP, API_VERSION_LEAVE_GROUP),
(API_KEY_SYNC_GROUP, API_VERSION_SYNC_GROUP),
] {
require_version(versions, api_key, version, floor)?;
}
Ok(())
}
fn normalized_topics(topics: &[String]) -> KafkaClientResult<Vec<String>> {
let topics = topics
.iter()
.map(|topic| topic.trim())
.filter(|topic| !topic.is_empty())
.map(str::to_owned)
.collect::<BTreeSet<_>>()
.into_iter()
.collect::<Vec<_>>();
if topics.is_empty() {
return Err(KafkaClientError::InvalidConfig(
"native Kafka topic subscription requires at least one topic".to_owned(),
));
}
Ok(topics)
}
fn millis_i32(value: Duration, name: &'static str) -> KafkaClientResult<i32> {
i32::try_from(value.as_millis()).map_err(|_| {
KafkaClientError::InvalidConfig(format!("{name} exceeds Kafka int32 millisecond range"))
})
}
fn assignor_for_protocol(protocol_name: &str) -> KafkaClientResult<GroupAssignor> {
match protocol_name {
"cooperative-sticky" => Ok(GroupAssignor::CooperativeSticky),
"range" => Ok(GroupAssignor::Range),
other => Err(KafkaClientError::unsupported(format!(
"group selected unsupported assignor {other}"
))),
}
}
fn metadata_covers(metadata: &MetadataResponse, topics: &[String]) -> bool {
topics.iter().all(|topic| {
metadata
.topics
.iter()
.any(|metadata_topic| metadata_topic.name == *topic)
})
}
fn decode_member_subscriptions(
members: &[JoinGroupMember],
) -> KafkaClientResult<Vec<MemberSubscription>> {
members
.iter()
.map(|member| {
let subscription = decode_consumer_protocol_subscription(&member.metadata)?;
Ok(MemberSubscription {
member_id: member.member_id.clone(),
group_instance_id: member.group_instance_id.clone(),
topics: subscription.topics,
owned_partitions: subscription.owned_partitions,
generation_id: subscription.generation_id,
})
})
.collect()
}
fn encode_sync_assignments(
assignments: BTreeMap<String, Vec<TopicPartition>>,
) -> KafkaClientResult<Vec<SyncGroupAssignment>> {
assignments
.into_iter()
.map(|(member_id, partitions)| {
Ok(SyncGroupAssignment {
member_id,
assignment: encode_consumer_protocol_assignment(&partitions)?,
})
})
.collect()
}
fn assign_partitions(
assignor: GroupAssignor,
members: &[MemberSubscription],
metadata: &MetadataResponse,
) -> BTreeMap<String, Vec<TopicPartition>> {
let target = match assignor {
GroupAssignor::Range => range_assignment(members, metadata),
GroupAssignor::CooperativeSticky => sticky_target_assignment(members, metadata),
};
if assignor.cooperative() {
cooperative_adjustment(target, members)
} else {
target
}
}
fn range_assignment(
members: &[MemberSubscription],
metadata: &MetadataResponse,
) -> BTreeMap<String, Vec<TopicPartition>> {
let mut assignment = empty_assignment(members);
let topic_partitions = partitions_by_topic(metadata);
for (topic, partitions) in topic_partitions {
let mut subscribers = members
.iter()
.filter(|member| member.topics.contains(&topic))
.collect::<Vec<_>>();
subscribers.sort_by(|left, right| member_order_key(left).cmp(&member_order_key(right)));
if subscribers.is_empty() {
continue;
}
let base = partitions.len() / subscribers.len();
let extra = partitions.len() % subscribers.len();
let mut cursor = 0_usize;
for (index, member) in subscribers.iter().enumerate() {
let count = base + usize::from(index < extra);
let end = cursor + count;
if count > 0 {
assignment
.entry(member.member_id.clone())
.or_default()
.extend(
partitions[cursor..end]
.iter()
.map(|partition| TopicPartition::new(topic.clone(), *partition)),
);
}
cursor = end;
}
}
sort_assignment(&mut assignment);
assignment
}
fn sticky_target_assignment(
members: &[MemberSubscription],
metadata: &MetadataResponse,
) -> BTreeMap<String, Vec<TopicPartition>> {
let mut assignment = empty_assignment(members);
if members.is_empty() {
return assignment;
}
let all_partitions = all_topic_partitions(metadata);
let quotas = member_quotas(members, all_partitions.len());
let mut claimed = HashSet::<TopicPartition>::new();
let mut ordered_members = members.iter().collect::<Vec<_>>();
ordered_members.sort_by(|left, right| member_order_key(left).cmp(&member_order_key(right)));
for member in &ordered_members {
let quota = quotas.get(&member.member_id).copied().unwrap_or(0);
let mut owned = member
.owned_partitions
.iter()
.filter(|partition| {
member.topics.contains(&partition.topic)
&& all_partitions.contains(*partition)
&& !claimed.contains(*partition)
})
.cloned()
.collect::<Vec<_>>();
owned.sort();
for partition in owned.into_iter().take(quota) {
claimed.insert(partition.clone());
assignment
.entry(member.member_id.clone())
.or_default()
.push(partition);
}
}
for partition in all_partitions {
if claimed.contains(&partition) {
continue;
}
let Some(member) = ordered_members
.iter()
.filter(|member| member.topics.contains(&partition.topic))
.min_by_key(|member| {
let current = assignment.get(&member.member_id).map_or(0, Vec::len);
let quota = quotas.get(&member.member_id).copied().unwrap_or(0);
(
usize::from(current >= quota),
current,
member_order_key(member),
)
})
else {
continue;
};
assignment
.entry(member.member_id.clone())
.or_default()
.push(partition);
}
sort_assignment(&mut assignment);
assignment
}
fn cooperative_adjustment(
mut assignment: BTreeMap<String, Vec<TopicPartition>>,
members: &[MemberSubscription],
) -> BTreeMap<String, Vec<TopicPartition>> {
let mut previous_owner = BTreeMap::<TopicPartition, (&str, i32)>::new();
for member in members {
for partition in &member.owned_partitions {
match previous_owner.get_mut(partition) {
Some((owner, generation)) if member.generation_id > *generation => {
*owner = &member.member_id;
*generation = member.generation_id;
}
Some(_) => {}
None => {
previous_owner.insert(
partition.clone(),
(member.member_id.as_str(), member.generation_id),
);
}
}
}
}
for (member_id, partitions) in &mut assignment {
partitions.retain(|partition| {
previous_owner
.get(partition)
.is_none_or(|(owner, _)| *owner == member_id)
});
}
assignment
}
fn empty_assignment(members: &[MemberSubscription]) -> BTreeMap<String, Vec<TopicPartition>> {
members
.iter()
.map(|member| (member.member_id.clone(), Vec::new()))
.collect()
}
fn member_quotas(
members: &[MemberSubscription],
partition_count: usize,
) -> BTreeMap<String, usize> {
if members.is_empty() {
return BTreeMap::new();
}
let mut ordered = members.iter().collect::<Vec<_>>();
ordered.sort_by(|left, right| member_order_key(left).cmp(&member_order_key(right)));
let base = partition_count / members.len();
let extra = partition_count % members.len();
ordered
.into_iter()
.enumerate()
.map(|(index, member)| (member.member_id.clone(), base + usize::from(index < extra)))
.collect()
}
fn member_order_key(member: &MemberSubscription) -> (&str, &str) {
(
member
.group_instance_id
.as_deref()
.unwrap_or(member.member_id.as_str()),
member.member_id.as_str(),
)
}
fn partitions_by_topic(metadata: &MetadataResponse) -> BTreeMap<String, Vec<i32>> {
let mut out = BTreeMap::new();
for topic in &metadata.topics {
if topic.error_code != 0 {
continue;
}
let mut partitions = topic
.partitions
.iter()
.filter(|partition| partition.error_code == 0)
.map(|partition| partition.partition_index)
.collect::<Vec<_>>();
partitions.sort_unstable();
out.insert(topic.name.clone(), partitions);
}
out
}
fn all_topic_partitions(metadata: &MetadataResponse) -> Vec<TopicPartition> {
partitions_by_topic(metadata)
.into_iter()
.flat_map(|(topic, partitions)| {
partitions
.into_iter()
.map(move |partition| TopicPartition::new(topic.clone(), partition))
})
.collect()
}
fn sort_assignment(assignment: &mut BTreeMap<String, Vec<TopicPartition>>) {
for partitions in assignment.values_mut() {
partitions.sort();
partitions.dedup();
}
}
fn reset_timestamp(reset: AutoOffsetReset) -> KafkaClientResult<Option<i64>> {
match reset {
AutoOffsetReset::Earliest => Ok(Some(EARLIEST_TIMESTAMP)),
AutoOffsetReset::Latest => Ok(Some(LATEST_TIMESTAMP)),
AutoOffsetReset::Error => Ok(None),
}
}
fn offset_commit_topics(commits: &[OffsetCommit]) -> Vec<OffsetCommitTopicRequest> {
let mut grouped = BTreeMap::<String, Vec<OffsetCommitPartitionRequest>>::new();
for commit in commits {
grouped
.entry(commit.topic.clone())
.or_default()
.push(OffsetCommitPartitionRequest {
partition_index: commit.partition,
committed_offset: commit.offset,
committed_leader_epoch: -1,
committed_metadata: None,
});
}
grouped
.into_iter()
.map(|(name, partitions)| OffsetCommitTopicRequest { name, partitions })
.collect()
}
fn offset_fetch_topics(partitions: &[TopicPartition]) -> Vec<OffsetFetchTopicRequest> {
let mut grouped = BTreeMap::<String, Vec<i32>>::new();
for partition in partitions {
grouped
.entry(partition.topic.clone())
.or_default()
.push(partition.partition);
}
grouped
.into_iter()
.map(|(name, partitions)| OffsetFetchTopicRequest { name, partitions })
.collect()
}
fn coordinator_error(code: i16) -> bool {
matches!(code, 14..=16)
}
fn first_bootstrap(bootstrap_servers: &str) -> KafkaClientResult<String> {
bootstrap_servers
.split(',')
.map(str::trim)
.find(|server| !server.is_empty())
.map(str::to_owned)
.ok_or_else(|| {
KafkaClientError::InvalidConfig(
"bootstrap.servers must include at least one host:port".to_owned(),
)
})
}
fn broker_host(address: &str) -> KafkaClientResult<String> {
let (host, port) = address.rsplit_once(':').ok_or_else(|| {
KafkaClientError::InvalidConfig(format!("Kafka broker address {address} must be host:port"))
})?;
if host.is_empty() || port.parse::<u16>().is_err() {
return Err(KafkaClientError::InvalidConfig(format!(
"Kafka broker address {address} must include a valid host and port"
)));
}
Ok(host.trim_matches(['[', ']']).to_owned())
}
fn partition_from_metadata(
metadata: &MetadataResponse,
assignment: &TopicPartitionAssignment,
) -> KafkaClientResult<PartitionState> {
let topic = metadata
.topics
.iter()
.find(|topic| topic.name == assignment.topic)
.ok_or_else(|| {
KafkaClientError::protocol(format!(
"metadata response omitted topic {}",
assignment.topic
))
})?;
if topic.error_code != 0 {
return Err(KafkaClientError::broker("Metadata topic", topic.error_code));
}
let partition = topic
.partitions
.iter()
.find(|partition| partition.partition_index == assignment.partition)
.ok_or_else(|| {
KafkaClientError::protocol(format!(
"metadata response omitted {}:{}",
assignment.topic, assignment.partition
))
})?;
if partition.error_code != 0 {
return Err(KafkaClientError::broker(
"Metadata partition",
partition.error_code,
));
}
Ok(PartitionState {
topic: assignment.topic.clone(),
partition: assignment.partition,
leader_id: partition.leader_id,
leader_epoch: partition.leader_epoch,
fetch_offset: 0,
log_start_offset: 0,
high_watermark: 0,
})
}
fn partition_from_topic_metadata(
metadata: &MetadataResponse,
current: &PartitionState,
) -> Option<PartitionState> {
let assignment = TopicPartitionAssignment {
topic: current.topic.clone(),
partition: current.partition,
start_offset: StartOffset::Offset(current.fetch_offset),
};
partition_from_metadata(metadata, &assignment)
.ok()
.map(|mut next| {
next.fetch_offset = current.fetch_offset;
next.log_start_offset = current.log_start_offset;
next.high_watermark = current.high_watermark;
next
})
}
#[cfg(test)]
mod tests {
use super::*;
use crate::native::protocol::{MetadataPartition, MetadataTopic};
fn metadata(topic: &str, partitions: i32) -> MetadataResponse {
MetadataResponse {
brokers: Vec::new(),
cluster_id: None,
topics: vec![MetadataTopic {
error_code: 0,
name: topic.to_owned(),
partitions: (0..partitions)
.map(|partition| MetadataPartition {
error_code: 0,
partition_index: partition,
leader_id: 1,
leader_epoch: 0,
})
.collect(),
}],
}
}
fn member(id: &str, owned: Vec<i32>) -> MemberSubscription {
MemberSubscription {
member_id: id.to_owned(),
group_instance_id: None,
topics: vec!["topic-a".to_owned()],
owned_partitions: owned
.into_iter()
.map(|partition| TopicPartition::new("topic-a", partition))
.collect(),
generation_id: 1,
}
}
#[test]
fn range_assignor_is_per_topic_oracle() {
let members = vec![member("a", Vec::new()), member("b", Vec::new())];
let assignment = assign_partitions(GroupAssignor::Range, &members, &metadata("topic-a", 3));
assert_eq!(
assignment.get("a").expect("member a"),
&vec![
TopicPartition::new("topic-a", 0),
TopicPartition::new("topic-a", 1)
]
);
assert_eq!(
assignment.get("b").expect("member b"),
&vec![TopicPartition::new("topic-a", 2)]
);
}
#[test]
fn cooperative_sticky_removes_transfers_until_owner_releases() {
let first_members = vec![member("a", vec![0, 1, 2, 3]), member("b", Vec::new())];
let first = assign_partitions(
GroupAssignor::CooperativeSticky,
&first_members,
&metadata("topic-a", 4),
);
assert_eq!(
first.get("a").expect("member a"),
&vec![
TopicPartition::new("topic-a", 0),
TopicPartition::new("topic-a", 1)
]
);
assert!(first.get("b").expect("member b").is_empty());
let second_members = vec![member("a", vec![0, 1]), member("b", Vec::new())];
let second = assign_partitions(
GroupAssignor::CooperativeSticky,
&second_members,
&metadata("topic-a", 4),
);
assert_eq!(
second.get("b").expect("member b"),
&vec![
TopicPartition::new("topic-a", 2),
TopicPartition::new("topic-a", 3)
]
);
}
}