use std::collections::HashMap;
use std::time::Duration;
use anyhow::{Context, Result as AnyResult, bail};
use kafka_protocol::error::ParseResponseErrorCode;
use kafka_protocol::messages::find_coordinator_response::FindCoordinatorResponse;
use kafka_protocol::messages::share_group_heartbeat_response::Assignment;
use kafka_protocol::messages::{
FindCoordinatorRequest, ShareGroupHeartbeatRequest, ShareGroupHeartbeatResponse,
};
use kafka_protocol::protocol::StrBytes;
use tokio::time::sleep;
use tracing::trace;
use crate::constants::{
FIND_COORDINATOR_GROUP_KEY_TYPE, FIND_COORDINATOR_VERSION_CAP,
SHARE_GROUP_HEARTBEAT_VERSION_CAP,
};
use crate::metadata::{BrokerAddress, MetadataRefresh, refresh_metadata};
use crate::network::{BrokerConnection, connect_to_any_bootstrap};
use crate::types::TopicPartitionKey;
use super::types::ShareAssignment;
use super::{KafkaShareConsumer, SHARE_COORDINATOR_RETRY_ATTEMPTS, SHARE_MEMBER_LEAVE_EPOCH};
impl KafkaShareConsumer {
pub(super) async fn heartbeat_with_retries(
&mut self,
include_subscription: bool,
) -> AnyResult<()> {
let mut last_error = None;
let attempts = self.share_coordinator_retry_attempts();
for attempt in 0..attempts {
match self.heartbeat(include_subscription).await {
Ok(()) => return Ok(()),
Err(error) => {
let is_coordinator_error = is_coordinator_error(&error);
if is_coordinator_error {
self.coordinator = None;
}
last_error = Some(error);
if !is_coordinator_error {
break;
}
if attempt + 1 < attempts {
sleep(self.config.retry_backoff).await;
}
}
}
}
Err(last_error.expect("share heartbeat attempt count is at least one"))
}
pub(super) async fn leave_group(&mut self) -> AnyResult<()> {
self.ensure_coordinator().await?;
let version = self
.coordinator
.as_mut()
.context("share group coordinator is not connected")?
.version_with_cap::<ShareGroupHeartbeatRequest>(SHARE_GROUP_HEARTBEAT_VERSION_CAP)?;
let request = ShareGroupHeartbeatRequest::default()
.with_group_id(StrBytes::from_string(self.config.group_id.clone()).into())
.with_member_id(StrBytes::from_string(self.member_id.clone()))
.with_member_epoch(SHARE_MEMBER_LEAVE_EPOCH)
.with_subscribed_topic_names(Some(Vec::new()));
let client_id = self.config.client_id.clone();
let _ = self
.coordinator
.as_mut()
.expect("coordinator checked above")
.send_request::<ShareGroupHeartbeatRequest>(&client_id, version, &request)
.await?;
Ok(())
}
pub(super) fn assignments_by_leader(&self) -> HashMap<i32, Vec<ShareAssignment>> {
let mut grouped = HashMap::<i32, Vec<ShareAssignment>>::new();
for assignment in self.assignments.values() {
grouped
.entry(assignment.leader_id)
.or_default()
.push(assignment.clone());
}
grouped
}
pub(super) async fn refresh_metadata_if_needed(&mut self) -> AnyResult<()> {
if self
.metadata
.needs_any_refresh(self.subscriptions.clone(), self.config.metadata_max_age)
{
refresh_metadata(MetadataRefresh {
bootstrap_servers: &self.config.bootstrap_servers,
client_id: &self.config.client_id,
request_timeout: self.config.request_timeout,
security_protocol: self.config.security_protocol,
tls: &self.config.tls,
sasl: &self.config.sasl,
metadata: &mut self.metadata,
topics: &self.subscriptions,
})
.await?;
}
Ok(())
}
pub(super) async fn refresh_metadata(&mut self) -> AnyResult<()> {
refresh_metadata(MetadataRefresh {
bootstrap_servers: &self.config.bootstrap_servers,
client_id: &self.config.client_id,
request_timeout: self.config.request_timeout,
security_protocol: self.config.security_protocol,
tls: &self.config.tls,
sasl: &self.config.sasl,
metadata: &mut self.metadata,
topics: &self.subscriptions,
})
.await?;
Ok(())
}
pub(super) async fn ensure_coordinator_with_retries(&mut self) -> AnyResult<()> {
let mut last_error = None;
let attempts = self.share_coordinator_retry_attempts();
for attempt in 0..attempts {
match self.ensure_coordinator().await {
Ok(()) => return Ok(()),
Err(error) => {
let is_coordinator_error = is_coordinator_error(&error);
last_error = Some(error);
if !is_coordinator_error {
break;
}
if attempt + 1 < attempts {
sleep(self.config.retry_backoff).await;
}
}
}
}
Err(last_error.expect("share coordinator lookup attempt count is at least one"))
}
pub(super) async fn leader_connection(
&mut self,
leader_id: i32,
) -> AnyResult<&mut BrokerConnection> {
if !self.leader_connections.contains_key(&leader_id) {
if self.metadata.broker(leader_id).is_none() {
self.refresh_metadata().await?;
}
let broker = self
.metadata
.broker(leader_id)
.with_context(|| format!("metadata missing broker {leader_id}"))?;
let connection = BrokerConnection::connect_with_transport(
&broker.address(),
&self.config.client_id,
self.config.request_timeout,
self.config.security_protocol,
&self.config.tls,
&self.config.sasl,
)
.await?;
self.leader_connections.insert(leader_id, connection);
}
Ok(self
.leader_connections
.get_mut(&leader_id)
.expect("leader connection inserted above"))
}
async fn heartbeat(&mut self, include_subscription: bool) -> AnyResult<()> {
self.ensure_coordinator().await?;
let version = self
.coordinator
.as_mut()
.context("share group coordinator is not connected")?
.version_with_cap::<ShareGroupHeartbeatRequest>(SHARE_GROUP_HEARTBEAT_VERSION_CAP)?;
let subscribed_topic_names = include_subscription.then(|| {
self.subscriptions
.iter()
.cloned()
.map(StrBytes::from_string)
.map(Into::into)
.collect()
});
let request = ShareGroupHeartbeatRequest::default()
.with_group_id(StrBytes::from_string(self.config.group_id.clone()).into())
.with_member_id(StrBytes::from_string(self.member_id.clone()))
.with_member_epoch(self.member_epoch)
.with_rack_id(self.config.rack_id.clone().map(StrBytes::from_string))
.with_subscribed_topic_names(subscribed_topic_names);
let client_id = self.config.client_id.clone();
let response: ShareGroupHeartbeatResponse = self
.coordinator
.as_mut()
.expect("coordinator checked above")
.send_request::<ShareGroupHeartbeatRequest>(&client_id, version, &request)
.await?;
self.handle_heartbeat_response(response)?;
Ok(())
}
fn handle_heartbeat_response(
&mut self,
response: ShareGroupHeartbeatResponse,
) -> AnyResult<()> {
if let Some(error) = response.error_code.err() {
bail!("share group heartbeat failed: {error}");
}
if let Some(member_id) = response.member_id {
self.member_id = member_id.to_string();
}
self.member_epoch = response.member_epoch;
if response.heartbeat_interval_ms > 0 {
self.heartbeat_interval = Duration::from_millis(response.heartbeat_interval_ms as u64);
}
if let Some(assignment) = response.assignment {
self.apply_assignment(assignment)?;
}
Ok(())
}
fn apply_assignment(&mut self, assignment: Assignment) -> AnyResult<()> {
let mut assignments = HashMap::new();
for topic in assignment.topic_partitions {
let topic_name = self
.metadata
.topic_name(&topic.topic_id)
.cloned()
.with_context(|| format!("metadata missing topic id {}", topic.topic_id))?;
for partition in topic.partitions {
let partition_metadata = self
.metadata
.partition(&topic_name, partition)
.with_context(|| format!("metadata missing {topic_name}:{partition}"))?;
assignments.insert(
TopicPartitionKey::new(topic_name.clone(), partition),
ShareAssignment {
topic_id: topic.topic_id,
topic: topic_name.clone(),
partition,
leader_id: partition_metadata.leader_id,
leader_epoch: partition_metadata.leader_epoch,
},
);
}
}
self.assignments = assignments;
Ok(())
}
async fn ensure_coordinator(&mut self) -> AnyResult<()> {
if self.coordinator.is_some() {
return Ok(());
}
let mut bootstrap = connect_to_any_bootstrap(
&self.config.bootstrap_servers,
&self.config.client_id,
self.config.request_timeout,
self.config.security_protocol,
&self.config.tls,
&self.config.sasl,
)
.await?;
let version =
bootstrap.version_with_cap::<FindCoordinatorRequest>(FIND_COORDINATOR_VERSION_CAP)?;
let request = if version >= 4 {
FindCoordinatorRequest::default()
.with_key_type(FIND_COORDINATOR_GROUP_KEY_TYPE)
.with_coordinator_keys(vec![StrBytes::from_string(self.config.group_id.clone())])
} else {
FindCoordinatorRequest::default()
.with_key(StrBytes::from_string(self.config.group_id.clone()))
.with_key_type(FIND_COORDINATOR_GROUP_KEY_TYPE)
};
let response = bootstrap
.send_request::<FindCoordinatorRequest>(&self.config.client_id, version, &request)
.await?;
let coordinator = parse_find_coordinator_response(response, version)?;
let connection = BrokerConnection::connect_with_transport(
&coordinator.address(),
&self.config.client_id,
self.config.request_timeout,
self.config.security_protocol,
&self.config.tls,
&self.config.sasl,
)
.await?;
trace!(coordinator = %coordinator.address(), "connected to share group coordinator");
self.coordinator = Some(connection);
Ok(())
}
fn share_coordinator_retry_attempts(&self) -> usize {
self.config
.max_retries
.max(SHARE_COORDINATOR_RETRY_ATTEMPTS)
.max(1)
}
}
fn parse_find_coordinator_response(
response: FindCoordinatorResponse,
version: i16,
) -> AnyResult<BrokerAddress> {
if version >= 4 {
let coordinator = response
.coordinators
.into_iter()
.next()
.context("find coordinator returned no coordinator entries")?;
if let Some(error) = coordinator.error_code.err() {
bail!("find coordinator failed: {error}");
}
return Ok(BrokerAddress::new(
coordinator.host.to_string(),
u16::try_from(coordinator.port)
.with_context(|| format!("invalid coordinator port {}", coordinator.port))?,
));
}
if let Some(error) = response.error_code.err() {
bail!("find coordinator failed: {error}");
}
Ok(BrokerAddress::new(
response.host.to_string(),
u16::try_from(response.port)
.with_context(|| format!("invalid coordinator port {}", response.port))?,
))
}
fn is_coordinator_error(error: &anyhow::Error) -> bool {
let message = error.to_string();
message.contains("NotCoordinator")
|| message.contains("CoordinatorNotAvailable")
|| message.contains("CoordinatorLoadInProgress")
}