#![allow(non_snake_case)]
use async_trait::async_trait;
use rdkafka::admin::{AdminClient, AdminOptions, NewTopic, TopicReplication};
use rdkafka::client::DefaultClientContext;
use rdkafka::config::{ClientConfig, RDKafkaLogLevel};
use rdkafka::consumer::{Consumer, DefaultConsumerContext, StreamConsumer};
use rdkafka::message::{BorrowedHeaders, Headers, Message};
use rdkafka::producer::{FutureProducer, FutureRecord};
use rdkafka::topic_partition_list::TopicPartitionList;
use std::collections::HashMap as FxHashMap;
use std::sync::Arc;
use std::time::Duration;
use tokio::sync::Mutex;
use crate::core::{RiError, RiResult};
use crate::queue::{RiQueue, RiQueueMessage, RiQueueProducer, RiQueueConsumer, RiQueueStats};
type KafkaConsumer = StreamConsumer<DefaultConsumerContext>;
#[allow(dead_code)]
#[derive(Clone)]
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiKafkaQueue {
brokers: String,
topic: String,
producer: Arc<FutureProducer>,
consumer: Arc<KafkaConsumer>,
admin_client: Arc<AdminClient<DefaultClientContext>>,
}
impl RiKafkaQueue {
pub async fn new(brokers: &str, topic: &str) -> RiResult<Self> {
let mut config = ClientConfig::new();
config.set("bootstrap.servers", brokers);
config.set("message.timeout.ms", "30000");
config.set("request.timeout.ms", "10000");
config.set("session.timeout.ms", "30000");
config.set("enable.auto.commit", "false");
config.set("auto.offset.reset", "earliest");
config.set_log_level(RDKafkaLogLevel::Warning);
let producer: FutureProducer = config
.create::<FutureProducer>()
.map_err(|e| RiError::Queue(format!("Failed to create Kafka producer: {}", e)))?;
let consumer: KafkaConsumer = config
.create::<KafkaConsumer>()
.map_err(|e| RiError::Queue(format!("Failed to create Kafka consumer: {}", e)))?;
let admin_client: AdminClient<DefaultClientContext> = config
.create::<AdminClient<DefaultClientContext>>()
.map_err(|e| RiError::Queue(format!("Failed to create Kafka admin client: {}", e)))?;
let queue = Self {
brokers: brokers.to_string(),
topic: topic.to_string(),
producer: Arc::new(producer),
consumer: Arc::new(consumer),
admin_client: Arc::new(admin_client),
};
queue.ensure_topic_exists().await?;
Ok(queue)
}
async fn ensure_topic_exists(&self) -> RiResult<()> {
let metadata = self.consumer.fetch_metadata(None, Duration::from_secs(5))
.map_err(|e| RiError::Queue(format!("Failed to get Kafka metadata: {}", e)))?;
let topic_exists = metadata.topics().iter().any(|t| t.name() == self.topic);
if !topic_exists {
let new_topic = NewTopic::new(&self.topic, 1, TopicReplication::Fixed(1));
let admin_options = AdminOptions::new();
self.admin_client.create_topics(std::slice::from_ref(&new_topic), &admin_options).await
.map_err(|e| RiError::Queue(format!("Failed to create Kafka topic: {}", e)))?;
tokio::time::sleep(Duration::from_secs(1)).await;
}
Ok(())
}
async fn get_topic_metadata(&self) -> RiResult<i32> {
let metadata = self.consumer.fetch_metadata(Some(&self.topic), Duration::from_secs(5))
.map_err(|e| RiError::Queue(format!("Failed to get Kafka metadata: {}", e)))?;
if let Some(topic_meta) = metadata.topics().first() {
Ok(topic_meta.partitions().len() as i32)
} else {
Ok(0)
}
}
}
#[async_trait]
impl RiQueue for RiKafkaQueue {
async fn create_producer(&self) -> RiResult<Box<dyn RiQueueProducer>> {
Ok(Box::new(KafkaQueueProducer {
producer: self.producer.clone(),
topic: self.topic.clone(),
}))
}
async fn create_consumer(&self, consumer_group: &str) -> RiResult<Box<dyn RiQueueConsumer>> {
let consumer = self.consumer.clone();
let mut partition_list = TopicPartitionList::new();
let partition_count = self.get_topic_metadata().await?;
for i in 0..partition_count.max(1) {
partition_list.add_partition(&self.topic, i);
}
consumer.assign(&partition_list)
.map_err(|e| RiError::Queue(format!("Failed to assign partitions: {}", e)))?;
Ok(Box::new(KafkaQueueConsumer {
consumer,
topic: self.topic.clone(),
consumer_group: consumer_group.to_string(),
paused: Arc::new(Mutex::new(false)),
}))
}
async fn get_stats(&self) -> RiResult<RiQueueStats> {
let _partition_count = self.get_topic_metadata().await?;
let topic = self.topic.clone();
Ok(RiQueueStats {
queue_name: topic.clone(),
message_count: 0,
consumer_count: 1,
producer_count: 1,
processed_messages: 0,
failed_messages: 0,
avg_processing_time_ms: 0.0,
total_bytes_sent: 0,
total_bytes_received: 0,
last_message_time: 0,
})
}
async fn purge(&self) -> RiResult<()> {
let admin_options = AdminOptions::new();
self.admin_client.delete_topics(std::slice::from_ref(&self.topic.as_str()), &admin_options).await
.map_err(|e| RiError::Queue(format!("Failed to purge Kafka topic: {}", e)))?;
tokio::time::sleep(Duration::from_secs(1)).await;
self.ensure_topic_exists().await?;
Ok(())
}
async fn delete(&self) -> RiResult<()> {
self.purge().await
}
}
pub struct KafkaQueueProducer {
producer: Arc<FutureProducer>,
topic: String,
}
#[async_trait]
impl RiQueueProducer for KafkaQueueProducer {
async fn send(&self, message: RiQueueMessage) -> RiResult<()> {
let payload = if message.payload.is_empty() {
vec![]
} else {
message.payload
};
let key = message.id.as_bytes();
let future_record = FutureRecord::to(&self.topic)
.key(key)
.payload(&payload);
self.producer.send(future_record, Duration::from_secs(10)).await
.map_err(|(e, _)| RiError::Queue(format!("Failed to send message to Kafka: {}", e)))?;
Ok(())
}
async fn send_batch(&self, messages: Vec<RiQueueMessage>) -> RiResult<()> {
for message in messages {
self.send(message).await?;
}
Ok(())
}
}
#[allow(dead_code)]
pub struct KafkaQueueConsumer {
consumer: Arc<KafkaConsumer>,
topic: String,
consumer_group: String,
paused: Arc<Mutex<bool>>,
}
#[async_trait]
impl RiQueueConsumer for KafkaQueueConsumer {
async fn receive(&self) -> RiResult<Option<RiQueueMessage>> {
let paused = *self.paused.lock().await;
if paused {
return Ok(None);
}
let message = tokio::time::timeout(Duration::from_secs(5), self.consumer.recv()).await;
match message {
Ok(Ok(msg)) => {
let payload = msg.payload().unwrap_or(&[]).to_vec();
let key = msg.key().map(|k| String::from_utf8_lossy(k).to_string()).unwrap_or_default();
let timestamp = msg.timestamp().to_millis().unwrap_or(0) as u64;
let headers: FxHashMap<String, String> = msg.headers()
.map(|h: &BorrowedHeaders| {
h.iter().filter_map(|header| {
header.value.map(|v| (header.key.to_string(), String::from_utf8_lossy(v).to_string()))
}).collect()
})
.unwrap_or_default();
let message = RiQueueMessage {
id: key,
payload,
headers,
timestamp: std::time::UNIX_EPOCH + Duration::from_millis(timestamp),
retry_count: 0,
max_retries: 3,
};
Ok(Some(message))
}
Ok(Err(e)) => Err(RiError::Queue(format!("Kafka receive error: {}", e))),
Err(_) => Ok(None),
}
}
async fn ack(&self, _message_id: &str) -> RiResult<()> {
self.consumer.commit_consumer_state(rdkafka::consumer::CommitMode::Sync)
.map_err(|e| RiError::Queue(format!("Failed to commit offset: {}", e)))?;
Ok(())
}
async fn nack(&self, message_id: &str) -> RiResult<()> {
log::info!("Message negatively acknowledged: {}", message_id);
Ok(())
}
async fn pause(&self) -> RiResult<()> {
let mut paused = self.paused.lock().await;
*paused = true;
Ok(())
}
async fn resume(&self) -> RiResult<()> {
let mut paused = self.paused.lock().await;
*paused = false;
Ok(())
}
}