#![allow(non_snake_case)]
#![cfg(feature = "rabbitmq")]
use async_trait::async_trait;
use lapin::{Connection, ConnectionProperties, Channel, Queue, Consumer};
use lapin::options::{QueueDeclareOptions, BasicConsumeOptions, BasicPublishOptions, BasicAckOptions};
use lapin::types::FieldTable;
use std::collections::HashMap as FxHashMap;
use std::sync::atomic::AtomicU64;
use std::sync::Arc;
use tokio::sync::Mutex;
use futures::StreamExt;
use serde::{Deserialize, Serialize};
#[cfg(feature = "http_client")]
use reqwest;
#[cfg(feature = "http_client")]
use urlencoding;
use crate::core::RiResult;
use crate::queue::{RiQueue, RiQueueMessage, RiQueueProducer, RiQueueConsumer, RiQueueStats};
pub struct RiRabbitMQQueue {
name: String,
#[allow(dead_code)]
connection: Arc<Connection>,
channel: Arc<Channel>,
#[allow(dead_code)]
queue: Arc<Queue>,
#[allow(dead_code)]
management_url: Option<String>,
#[allow(dead_code)]
management_username: Option<String>,
#[allow(dead_code)]
management_password: Option<String>,
}
#[derive(Debug, Serialize, Deserialize)]
#[allow(dead_code)]
struct RabbitMQQueueInfo {
name: String,
messages: u64,
consumers: u64,
message_stats: Option<RabbitMQMessageStats>,
}
#[derive(Debug, Serialize, Deserialize)]
#[allow(dead_code)]
struct RabbitMQMessageStats {
publish: Option<u64>,
deliver_no_ack: Option<u64>,
get_no_ack: Option<u64>,
redeliver: Option<u64>,
deliver: Option<u64>,
get: Option<u64>,
}
impl RiRabbitMQQueue {
pub async fn new(name: &str, connection_string: &str) -> RiResult<Self> {
let connection = Connection::connect(connection_string, ConnectionProperties::default()).await?;
Self::new_with_connection(name, connection).await
}
pub async fn new_with_connection(name: &str, connection: lapin::Connection) -> RiResult<Self> {
let channel = connection.create_channel().await?;
let queue = channel
.queue_declare(
name,
QueueDeclareOptions {
durable: true,
..Default::default()
},
FieldTable::default(),
)
.await?;
Ok(Self {
name: name.to_string(),
connection: Arc::new(connection),
channel: Arc::new(channel),
queue: Arc::new(queue),
management_url: None,
management_username: None,
management_password: None,
})
}
pub async fn new_with_management(
name: &str,
connection: lapin::Connection,
management_url: &str,
management_username: &str,
management_password: &str,
) -> RiResult<Self> {
let channel = connection.create_channel().await?;
let queue = channel
.queue_declare(
name,
QueueDeclareOptions {
durable: true,
..Default::default()
},
FieldTable::default(),
)
.await?;
Ok(Self {
name: name.to_string(),
connection: Arc::new(connection),
channel: Arc::new(channel),
queue: Arc::new(queue),
management_url: Some(management_url.to_string()),
management_username: Some(management_username.to_string()),
management_password: Some(management_password.to_string()),
})
}
#[cfg(feature = "http_client")]
async fn fetch_rabbitmq_stats(&self) -> RiResult<RiQueueStats> {
let (Some(management_url), Some(username), Some(password)) = (
&self.management_url,
&self.management_username,
&self.management_password,
) else {
return Err(crate::core::RiError::Other(
"Management API not configured. Use new_with_management() to enable.".to_string(),
));
};
let client = reqwest::Client::new();
let url = format!(
"{}/api/queues/%2f/{}",
management_url.trim_end_matches('/'),
urlencoding::encode(&self.name)
);
let response = client
.get(&url)
.basic_auth(username, Some(password))
.send()
.await
.map_err(|e| crate::core::RiError::Other(format!("Failed to connect to RabbitMQ Management API: {}", e)))?;
if !response.status().is_success() {
return Err(crate::core::RiError::Other(format!(
"RabbitMQ Management API returned error: {}",
response.status()
)));
}
let queue_info: RabbitMQQueueInfo = response
.json()
.await
.map_err(|e| crate::core::RiError::Other(format!("Failed to parse RabbitMQ response: {}", e)))?;
let processed_messages = queue_info
.message_stats
.as_ref()
.and_then(|s| s.publish.or(s.deliver).or(s.get))
.unwrap_or(0);
let failed_messages = queue_info
.message_stats
.as_ref()
.and_then(|s| s.redeliver)
.unwrap_or(0);
Ok(RiQueueStats {
queue_name: queue_info.name,
message_count: queue_info.messages,
consumer_count: queue_info.consumers as u32,
producer_count: 0,
processed_messages,
failed_messages,
avg_processing_time_ms: 0.0,
total_bytes_sent: 0,
total_bytes_received: 0,
last_message_time: 0,
})
}
async fn get_basic_stats(&self) -> RiResult<RiQueueStats> {
let queue_info = self.channel
.queue_declare(
&self.name,
lapin::options::QueueDeclareOptions {
passive: true,
..Default::default()
},
lapin::types::FieldTable::default(),
)
.await?;
Ok(RiQueueStats {
queue_name: self.name.clone(),
message_count: queue_info.message_count() as u64,
consumer_count: queue_info.consumer_count() as u32,
producer_count: 0,
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_trait]
impl RiQueue for RiRabbitMQQueue {
async fn create_producer(&self) -> RiResult<Box<dyn RiQueueProducer>> {
Ok(Box::new(RabbitMQProducer {
channel: self.channel.clone(),
queue_name: self.name.clone(),
}))
}
async fn create_consumer(&self, consumer_group: &str) -> RiResult<Box<dyn RiQueueConsumer>> {
let consumer = self.channel
.basic_consume(
&self.name,
consumer_group,
BasicConsumeOptions::default(),
FieldTable::default(),
)
.await?;
Ok(Box::new(RabbitMQConsumer::new(self.channel.clone(), consumer)))
}
async fn get_stats(&self) -> RiResult<RiQueueStats> {
#[cfg(feature = "http_client")]
match self.fetch_rabbitmq_stats().await {
Ok(detailed_stats) => Ok(detailed_stats),
Err(_) => {
self.get_basic_stats().await
}
}
#[cfg(not(feature = "http_client"))]
self.get_basic_stats().await
}
async fn purge(&self) -> RiResult<()> {
self.channel.queue_purge(&self.name, Default::default()).await?;
Ok(())
}
async fn delete(&self) -> RiResult<()> {
self.channel.queue_delete(&self.name, Default::default()).await?;
Ok(())
}
}
struct RabbitMQProducer {
channel: Arc<Channel>,
queue_name: String,
}
#[async_trait]
impl RiQueueProducer for RabbitMQProducer {
async fn send(&self, message: RiQueueMessage) -> RiResult<()> {
let payload = serde_json::to_vec(&message)?;
self.channel
.basic_publish(
"",
&self.queue_name,
BasicPublishOptions::default(),
&payload,
Default::default(),
)
.await?;
Ok(())
}
async fn send_batch(&self, messages: Vec<RiQueueMessage>) -> RiResult<()> {
for message in messages {
self.send(message).await?;
}
Ok(())
}
}
#[allow(dead_code)]
struct RabbitMQConsumer {
channel: Arc<Channel>,
consumer: Arc<Mutex<Consumer>>,
paused: Arc<Mutex<bool>>,
delivery_tags: Arc<Mutex<FxHashMap<u64, String>>>,
next_delivery_tag: Arc<AtomicU64>,
}
impl RabbitMQConsumer {
fn new(channel: Arc<Channel>, consumer: Consumer) -> Self {
Self {
channel,
consumer: Arc::new(Mutex::new(consumer)),
paused: Arc::new(Mutex::new(false)),
delivery_tags: Arc::new(Mutex::new(FxHashMap::default())),
next_delivery_tag: Arc::new(AtomicU64::new(1)),
}
}
}
#[async_trait]
impl RiQueueConsumer for RabbitMQConsumer {
async fn receive(&self) -> RiResult<Option<RiQueueMessage>> {
let paused = *self.paused.lock().await;
if paused {
return Ok(None);
}
let mut consumer = self.consumer.lock().await;
match consumer.next().await {
Some(delivery_result) => {
let delivery = delivery_result.map_err(|e| crate::core::RiError::Other(format!("Consumer error: {e}")))?;
const MAX_MESSAGE_SIZE: usize = 10 * 1024 * 1024; if delivery.data.len() > MAX_MESSAGE_SIZE {
log::warn!(
"[Ri.RabbitMQ] Message too large: {} bytes (max {} bytes), rejecting",
delivery.data.len(), MAX_MESSAGE_SIZE
);
return Err(crate::core::RiError::Other(format!(
"Message too large: {} bytes (max {} bytes)",
delivery.data.len(), MAX_MESSAGE_SIZE
)));
}
let _message_id = {
let delivery_tag = delivery.delivery_tag;
let message_id = format!("msg_{}", uuid::Uuid::new_v4());
let mut tags = self.delivery_tags.lock().await;
tags.insert(delivery_tag, message_id.clone());
message_id
};
let message: RiQueueMessage = serde_json::from_slice(&delivery.data)
.map_err(|e| {
log::warn!("[Ri.RabbitMQ] Failed to deserialize message: {}", e);
crate::core::RiError::Other(format!("Failed to deserialize message: {}", e))
})?;
Ok(Some(message))
},
None => Ok(None)
}
}
async fn ack(&self, message_id: &str) -> RiResult<()> {
log::debug!("Acknowledging message: {}", message_id);
let delivery_tag = {
let tags = self.delivery_tags.lock().await;
tags.iter()
.find(|(_, id)| *id == message_id)
.map(|(tag, _)| *tag)
};
if let Some(tag) = delivery_tag {
let channel = self.channel.clone();
channel.basic_ack(tag, BasicAckOptions { multiple: false }).await
.map_err(|e| crate::core::RiError::Other(format!("Failed to ack message: {e}")))?;
let mut tags = self.delivery_tags.lock().await;
tags.remove(&tag);
log::debug!("Message {} acknowledged successfully", message_id);
} else {
log::warn!("Message ID not found for acknowledgment: {}", message_id);
}
Ok(())
}
async fn nack(&self, message_id: &str) -> RiResult<()> {
log::debug!("Negatively acknowledging message: {}", message_id);
let delivery_tag = {
let tags = self.delivery_tags.lock().await;
tags.iter()
.find(|(_, id)| *id == message_id)
.map(|(tag, _)| *tag)
};
if let Some(tag) = delivery_tag {
let channel = self.channel.clone();
use lapin::options::BasicNackOptions;
channel.basic_nack(tag, BasicNackOptions { multiple: false, requeue: true }).await
.map_err(|e| crate::core::RiError::Other(format!("Failed to nack message: {e}")))?;
let mut tags = self.delivery_tags.lock().await;
tags.remove(&tag);
log::debug!("Message {} negatively acknowledged (will be requeued)", message_id);
} else {
log::warn!("Message ID not found for negative acknowledgment: {}", 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(())
}
}