use std::net::SocketAddr;
use std::sync::Arc;
use tokio::net;
use crate::cluster::ClusterClient;
use crate::error::{KafkaError, KafkaErrorCode, Result};
use crate::protocol::{
CreateTopicsRequest, CreateTopicsResponse, DeleteGroupsRequest, DeleteGroupsResponse,
DeleteTopicsRequest, DeleteTopicsResponse, DescribeGroupsRequest, DescribeGroupsResponse,
FindCoordinatorRequest, FindCoordinatorResponse, ListGroupsRequest, ListGroupsResponse,
ListOffsetsPartition, ListOffsetsRequest, ListOffsetsResponse, ListOffsetsTopic,
MetadataRequest, MetadataResponse, OffsetCommitRequest, OffsetCommitResponse,
OffsetFetchRequest, OffsetFetchRequestGroup, OffsetFetchResponse,
create_topics_request::{CreatableReplicaAssignment, CreatableTopic, CreatableTopicConfig},
delete_topics_request::DeleteTopicState,
offset_commit_request::{OffsetCommitRequestPartition, OffsetCommitRequestTopic},
};
#[derive(Debug, Clone)]
pub struct NewTopic {
pub name: String,
pub num_partitions: i32,
pub replication_factor: i16,
pub replica_assignments: Option<Vec<Vec<i32>>>,
pub configs: Vec<(String, String)>,
}
impl NewTopic {
pub fn new(name: impl Into<String>, num_partitions: i32, replication_factor: i16) -> Self {
Self {
name: name.into(),
num_partitions,
replication_factor,
replica_assignments: None,
configs: Vec::new(),
}
}
pub fn with_config(mut self, key: impl Into<String>, value: impl Into<String>) -> Self {
self.configs.push((key.into(), value.into()));
self
}
pub fn with_replica_assignments(mut self, assignments: Vec<Vec<i32>>) -> Self {
self.replica_assignments = Some(assignments);
self
}
}
#[derive(Debug, Clone)]
pub struct AdminTopicResult {
pub name: String,
pub error_code: KafkaErrorCode,
pub error_message: Option<String>,
}
impl AdminTopicResult {
pub fn is_success(&self) -> bool {
self.error_code.is_ok()
}
pub fn already_exists(&self) -> bool {
self.error_code == KafkaErrorCode::TOPIC_ALREADY_EXISTS
}
}
#[derive(Debug, Clone)]
pub struct AdminTopic {
pub name: String,
pub internal: bool,
pub partitions: usize,
}
#[derive(Debug, Clone)]
pub struct AdminPartitionInfo {
pub partition: i32,
pub leader_id: i32,
pub replicas: Vec<i32>,
pub isr: Vec<i32>,
}
#[derive(Debug, Clone)]
pub struct AdminTopicDescription {
pub name: String,
pub internal: bool,
pub partitions: Vec<AdminPartitionInfo>,
}
#[derive(Debug, Clone)]
pub struct AdminBroker {
pub id: i32,
pub host: String,
pub port: i32,
pub addr: Option<SocketAddr>,
}
#[derive(Debug, Clone)]
pub struct AdminClusterInfo {
pub cluster_id: Option<String>,
pub controller_id: Option<i32>,
pub brokers: Vec<AdminBroker>,
}
#[derive(Debug, Clone)]
pub struct AdminGroup {
pub group_id: String,
pub protocol_type: String,
}
#[derive(Debug, Clone)]
pub struct AdminGroupMember {
pub member_id: String,
pub client_id: String,
pub client_host: String,
}
#[derive(Debug, Clone)]
pub struct GroupOffset {
pub topic: String,
pub partition: i32,
pub committed_offset: i64,
pub log_end_offset: i64,
pub lag: i64,
pub metadata: String,
}
#[derive(Debug, Clone)]
pub struct AdminGroupDescription {
pub group_id: String,
pub state: String,
pub protocol_type: String,
pub members: Vec<AdminGroupMember>,
}
#[derive(Debug, Clone)]
pub struct OffsetCommitSpec {
pub topic: String,
pub partition: i32,
pub offset: i64,
pub metadata: Option<String>,
}
pub struct AdminClient {
cluster: Arc<ClusterClient>,
}
impl AdminClient {
pub(crate) fn new(cluster: Arc<ClusterClient>) -> Self {
Self { cluster }
}
pub async fn create_topics(&self, topics: &[NewTopic]) -> Result<Vec<AdminTopicResult>> {
let creatable: Vec<CreatableTopic> = topics
.iter()
.map(|t| {
let assignments: Vec<CreatableReplicaAssignment> = t
.replica_assignments
.as_ref()
.map(|a| {
a.iter()
.map(|ids| CreatableReplicaAssignment {
partition_index: -1,
broker_ids: ids.clone(),
})
.collect()
})
.unwrap_or_default();
let configs: Vec<CreatableTopicConfig> = t
.configs
.iter()
.map(|(k, v)| CreatableTopicConfig {
name: k.clone(),
value: Some(v.clone()),
})
.collect();
CreatableTopic {
name: t.name.clone(),
num_partitions: t.num_partitions,
replication_factor: t.replication_factor,
assignments,
configs,
}
})
.collect();
let request = CreateTopicsRequest {
topics: creatable,
timeout_ms: 30_000,
validate_only: false,
};
let response: CreateTopicsResponse = self.cluster.send_to_any_broker(&request).await?;
let results = response
.topics
.into_iter()
.map(|t| AdminTopicResult {
name: t.name,
error_code: KafkaErrorCode::from_i16(t.error_code),
error_message: t.error_message,
})
.collect();
Ok(results)
}
pub async fn create_topic(&self, topic: &NewTopic) -> Result<AdminTopicResult> {
let mut results = self.create_topics(std::slice::from_ref(topic)).await?;
results
.pop()
.ok_or_else(|| KafkaError::InvalidConfiguration("no result returned".into()))
}
pub async fn delete_topics(
&self,
topic_names: &[impl AsRef<str>],
) -> Result<Vec<AdminTopicResult>> {
let topics: Vec<DeleteTopicState> = topic_names
.iter()
.map(|n| DeleteTopicState {
name: Some(n.as_ref().to_string()),
topic_id: uuid::Uuid::nil(),
})
.collect();
let topic_names_vec: Vec<String> =
topic_names.iter().map(|n| n.as_ref().to_string()).collect();
let request = DeleteTopicsRequest {
topics: topics.clone(),
topic_names: topic_names_vec,
timeout_ms: 30_000,
};
let response: DeleteTopicsResponse = self.cluster.send_to_any_broker(&request).await?;
let results = response
.responses
.into_iter()
.map(|r| AdminTopicResult {
name: r.name.unwrap_or_default(),
error_code: KafkaErrorCode::from_i16(r.error_code),
error_message: r.error_message,
})
.collect();
Ok(results)
}
pub async fn delete_topic(&self, name: &str) -> Result<AdminTopicResult> {
let mut results = self.delete_topics(&[name]).await?;
results
.pop()
.ok_or_else(|| KafkaError::InvalidConfiguration("no result returned".into()))
}
pub async fn list_topics(&self) -> Result<Vec<AdminTopic>> {
self.cluster.refresh_metadata().await?;
let topics = self.cluster.metadata().get_all_topics().await;
Ok(topics
.into_iter()
.filter(|t| !t.is_internal) .map(|t| AdminTopic {
name: t.name.unwrap_or_default(),
internal: t.is_internal,
partitions: t.partitions.len(),
})
.collect())
}
pub async fn describe_topics(
&self,
topic_names: &[impl AsRef<str>],
) -> Result<Vec<AdminTopicDescription>> {
let name_list: Vec<String> = topic_names.iter().map(|s| s.as_ref().to_string()).collect();
let request_topics: Vec<crate::protocol::MetadataRequestTopic> = name_list
.iter()
.map(|name| crate::protocol::MetadataRequestTopic {
topic_id: uuid::Uuid::nil(),
name: Some(name.clone()),
})
.collect();
let request = MetadataRequest {
topics: Some(request_topics),
allow_auto_topic_creation: false,
include_cluster_authorized_operations: false,
include_topic_authorized_operations: false,
};
let response: MetadataResponse = self.cluster.send_to_any_broker(&request).await?;
let descriptions = response
.topics
.into_iter()
.map(|t| {
let partitions = t
.partitions
.iter()
.map(|p| AdminPartitionInfo {
partition: p.partition_index,
leader_id: p.leader_id,
replicas: p.replica_nodes.clone(),
isr: p.isr_nodes.clone(),
})
.collect();
AdminTopicDescription {
name: t.name.unwrap_or_default(),
internal: t.is_internal,
partitions,
}
})
.collect();
Ok(descriptions)
}
pub async fn describe_cluster(&self) -> Result<AdminClusterInfo> {
self.cluster.refresh_metadata().await?;
let metadata = self.cluster.metadata();
let brokers: Vec<AdminBroker> = metadata
.get_all_brokers()
.await
.into_iter()
.map(|b| {
let host = b.host;
let port = b.port;
let addr_str = format!("{}:{}", host, port);
AdminBroker {
id: b.node_id,
host: host.clone(),
port,
addr: addr_str.parse().ok(),
}
})
.collect();
Ok(AdminClusterInfo {
cluster_id: metadata.get_cluster_id().await,
controller_id: metadata.get_controller_id().await,
brokers,
})
}
pub async fn list_groups(&self) -> Result<Vec<AdminGroup>> {
let request = ListGroupsRequest {
states_filter: vec![],
types_filter: vec![],
};
let response: ListGroupsResponse = self.cluster.send_to_any_broker(&request).await?;
let groups = response
.groups
.into_iter()
.map(|g| AdminGroup {
group_id: g.group_id,
protocol_type: g.protocol_type,
})
.collect();
Ok(groups)
}
pub async fn describe_groups(
&self,
group_ids: &[impl AsRef<str>],
) -> Result<Vec<AdminGroupDescription>> {
let ids: Vec<String> = group_ids.iter().map(|s| s.as_ref().to_string()).collect();
let request = DescribeGroupsRequest {
groups: ids.clone(),
include_authorized_operations: false,
};
let response: DescribeGroupsResponse = self.cluster.send_to_any_broker(&request).await?;
let descriptions = response
.groups
.into_iter()
.map(|g| {
let members = g
.members
.into_iter()
.map(|m| AdminGroupMember {
member_id: m.member_id,
client_id: m.client_id,
client_host: m.client_host,
})
.collect();
AdminGroupDescription {
group_id: g.group_id,
state: g.group_state,
protocol_type: g.protocol_type,
members,
}
})
.collect();
Ok(descriptions)
}
pub async fn delete_group(&self, group_id: &str) -> Result<()> {
let request = DeleteGroupsRequest {
groups_names: vec![group_id.to_string()],
};
let _response: DeleteGroupsResponse = self.cluster.send_to_any_broker(&request).await?;
Ok(())
}
pub async fn fetch_group_offsets(&self, group_id: &str) -> Result<Vec<GroupOffset>> {
let coord = self.find_group_coordinator(group_id).await?;
let request = OffsetFetchRequest {
group_id: group_id.to_string(),
topics: None,
groups: vec![OffsetFetchRequestGroup {
group_id: group_id.to_string(),
member_id: None,
member_epoch: -1,
topics: None,
}],
require_stable: false,
};
let response: OffsetFetchResponse = self.cluster.send_to_broker(coord, &request).await?;
struct RawPartition {
topic: String,
partition_index: i32,
committed_offset: i64,
metadata: Option<String>,
}
let mut raw = Vec::<RawPartition>::new();
if !response.groups.is_empty() {
for grp in response.groups {
if grp.group_id != group_id {
continue;
}
if grp.error_code != 0 {
return Err(KafkaError::NoCoordinator);
}
for t in grp.topics {
let name = t.name;
if name.is_empty() {
continue;
}
for p in t.partitions {
if p.error_code != 0 {
continue;
}
raw.push(RawPartition {
topic: name.clone(),
partition_index: p.partition_index,
committed_offset: p.committed_offset,
metadata: p.metadata.clone(),
});
}
}
}
} else {
if response.error_code != 0 {
return Err(KafkaError::NoCoordinator);
}
for t in &response.topics {
let name = &t.name;
if name.is_empty() {
continue;
}
for p in &t.partitions {
if p.error_code != 0 {
continue;
}
raw.push(RawPartition {
topic: name.clone(),
partition_index: p.partition_index,
committed_offset: p.committed_offset,
metadata: p.metadata.clone(),
});
}
}
}
let mut offsets = Vec::with_capacity(raw.len());
for rp in raw {
let log_end = self
.fetch_log_end_offset(&rp.topic, rp.partition_index)
.await
.unwrap_or(-1);
let lag = if log_end >= 0 && rp.committed_offset >= 0 {
(log_end - rp.committed_offset).max(0)
} else {
-1
};
offsets.push(GroupOffset {
topic: rp.topic,
partition: rp.partition_index,
committed_offset: rp.committed_offset,
log_end_offset: log_end,
lag,
metadata: rp.metadata.unwrap_or_default(),
});
}
Ok(offsets)
}
async fn find_group_coordinator(&self, group_id: &str) -> Result<SocketAddr> {
let request = FindCoordinatorRequest {
key: group_id.to_string(),
key_type: 0,
coordinator_keys: vec![group_id.to_string()],
};
let response: FindCoordinatorResponse = self.cluster.send_to_any_broker(&request).await?;
if response.error_code != 0 {
return Err(KafkaError::NoCoordinator);
}
let (host, port) = if !response.host.is_empty() {
(response.host.clone(), response.port)
} else if let Some(coord) = response.coordinators.first() {
if coord.error_code != 0 {
return Err(KafkaError::NoCoordinator);
}
(coord.host.clone(), coord.port)
} else {
return Err(KafkaError::NoCoordinator);
};
net::lookup_host(format!("{}:{}", host, port))
.await
.map_err(|_| KafkaError::NoCoordinator)?
.next()
.ok_or(KafkaError::NoCoordinator)
}
async fn fetch_log_end_offset(&self, topic: &str, partition: i32) -> Result<i64> {
let leader_addr = self
.cluster
.metadata()
.get_partition_leader(topic, partition)
.await
.ok_or_else(|| KafkaError::PartitionNotFound(topic.to_string(), partition))?;
let request = ListOffsetsRequest {
replica_id: -1,
isolation_level: 0,
topics: vec![ListOffsetsTopic {
name: topic.to_string(),
partitions: vec![ListOffsetsPartition {
partition_index: partition,
current_leader_epoch: -1,
timestamp: -1,
}],
}],
timeout_ms: 5000,
};
let response: ListOffsetsResponse = self
.cluster
.send_to_broker::<ListOffsetsRequest, ListOffsetsResponse>(leader_addr, &request)
.await?;
for t in &response.topics {
if t.name == topic
&& let Some(p) = t.partitions.iter().find(|p| p.partition_index == partition)
{
if p.error_code != 0 {
break;
}
return Ok(p.offset);
}
}
Err(KafkaError::PartitionNotFound(topic.to_string(), partition))
}
pub async fn commit_offsets(&self, group_id: &str, offsets: &[OffsetCommitSpec]) -> Result<()> {
let mut topics: std::collections::HashMap<String, Vec<OffsetCommitRequestPartition>> =
std::collections::HashMap::new();
for spec in offsets {
topics
.entry(spec.topic.clone())
.or_default()
.push(OffsetCommitRequestPartition {
partition_index: spec.partition,
committed_offset: spec.offset,
committed_leader_epoch: -1,
committed_metadata: spec.metadata.clone(),
});
}
let request = OffsetCommitRequest {
group_id: group_id.to_string(),
generation_id_or_member_epoch: -1,
member_id: String::new(),
group_instance_id: None,
retention_time_ms: -1,
topics: topics
.into_iter()
.map(|(name, partitions)| OffsetCommitRequestTopic {
name,
topic_id: uuid::Uuid::nil(),
partitions,
})
.collect(),
};
let _response: OffsetCommitResponse = self.cluster.send_to_any_broker(&request).await?;
Ok(())
}
pub async fn refresh_metadata(&self) -> Result<()> {
self.cluster.refresh_metadata().await
}
pub async fn get_broker_config(&self, key: &str) -> Option<usize> {
self.cluster.query_broker_config(key).await
}
}