#![allow(non_snake_case)]
use std::sync::Arc;
use std::collections::HashMap as FxHashMap;
use tokio::sync::RwLock;
use async_trait::async_trait;
use crate::core::{RiResult, AsyncServiceModule, RiServiceContext};
use crate::queue::{RiQueue, RiQueueConfig, RiQueueBackendType};
#[cfg(feature = "pyo3")]
use pyo3::PyResult;
struct QueueConnectionPool {
backend_type: RiQueueBackendType,
connections: Arc<RwLock<Vec<Arc<dyn std::any::Any + Send + Sync>>>>,
max_connections: usize,
#[allow(dead_code)]
connection_string: String,
}
impl QueueConnectionPool {
fn new(backend_type: RiQueueBackendType, connection_string: String, max_connections: usize) -> Self {
Self {
backend_type,
connections: Arc::new(RwLock::new(Vec::new())),
max_connections,
connection_string,
}
}
#[allow(dead_code)]
async fn get_connection(&self) -> RiResult<Arc<dyn std::any::Any + Send + Sync>> {
let mut connections = self.connections.write().await;
if let Some(conn) = connections.pop() {
return Ok(conn);
}
if connections.len() < self.max_connections {
let conn = self.create_connection().await?;
return Ok(conn);
}
Err(crate::core::RiError::Other("Connection pool exhausted".to_string()))
}
async fn return_connection(&self, connection: Arc<dyn std::any::Any + Send + Sync>) {
let mut connections = self.connections.write().await;
if connections.len() < self.max_connections {
connections.push(connection);
}
}
async fn create_connection(&self) -> RiResult<Arc<dyn std::any::Any + Send + Sync>> {
match self.backend_type {
#[cfg(feature = "rabbitmq")]
RiQueueBackendType::RabbitMQ => {
use lapin::{Connection, ConnectionProperties};
let conn = Connection::connect(&self.connection_string, ConnectionProperties::default()).await?;
Ok(Arc::new(conn))
}
#[cfg(not(feature = "rabbitmq"))]
RiQueueBackendType::RabbitMQ => {
Err(crate::core::RiError::Config(
"RabbitMQ support is disabled. Enable the 'rabbitmq' feature to use RabbitMQ backend.".to_string(),
))
}
#[cfg(feature = "redis")]
RiQueueBackendType::Redis => {
use redis::Client;
let client = Client::open(self.connection_string.as_str())?;
let conn = client.get_multiplexed_async_connection().await?;
Ok(Arc::new(conn))
}
#[cfg(not(feature = "redis"))]
RiQueueBackendType::Redis => {
Err(crate::core::RiError::Config(
"Redis support is disabled. Enable the 'redis' feature to use Redis backend.".to_string(),
))
}
#[cfg(all(feature = "kafka", not(windows)))]
RiQueueBackendType::Kafka => {
Ok(Arc::new(()))
}
#[cfg(any(not(feature = "kafka"), windows))]
RiQueueBackendType::Kafka => {
Err(crate::core::RiError::Config(
"Kafka support is disabled. Enable the 'kafka' feature to use Kafka backend.".to_string(),
))
}
RiQueueBackendType::Memory => {
Ok(Arc::new(()))
}
}
}
async fn close_all(&self) {
let mut connections = self.connections.write().await;
connections.clear();
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiQueueModule {
queue_manager: Arc<RiQueueManager>,
}
impl RiQueueModule {
pub async fn new(config: RiQueueConfig) -> RiResult<Self> {
let queue_manager = Arc::new(RiQueueManager::new(config).await?);
Ok(Self { queue_manager })
}
pub fn with_config(_config: RiQueueConfig) -> RiResult<Self> {
let queue_manager = Arc::new(RiQueueManager::default());
Ok(Self { queue_manager })
}
pub fn queue_manager(&self) -> Arc<RiQueueManager> {
self.queue_manager.clone()
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiQueueModule {
#[new]
fn py_new() -> PyResult<Self> {
use crate::queue::RiQueueConfig;
use crate::queue::RiQueueBackendType;
let config = RiQueueConfig {
enabled: true,
backend_type: RiQueueBackendType::Memory,
connection_string: "memory://localhost".to_string(),
max_connections: 10,
message_max_size: 1024 * 1024,
consumer_timeout_ms: 30000,
producer_timeout_ms: 30000,
retry_policy: crate::queue::config::RiRetryPolicy::default(),
dead_letter_config: None,
};
let runtime = tokio::runtime::Handle::current();
let result = runtime.block_on(async {
Self::new(config).await
});
match result {
Ok(module) => Ok(module),
Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create queue module: {e}"))),
}
}
}
#[async_trait]
impl AsyncServiceModule for RiQueueModule {
async fn init(&mut self, _ctx: &mut RiServiceContext) -> RiResult<()> {
self.queue_manager.init().await
}
async fn after_shutdown(&mut self, _ctx: &mut RiServiceContext) -> RiResult<()> {
self.queue_manager.shutdown().await
}
fn name(&self) -> &str {
"dms-queue"
}
fn is_critical(&self) -> bool {
false
}
}
impl crate::core::ServiceModule for RiQueueModule {
fn name(&self) -> &str {
"Ri.Queue"
}
fn is_critical(&self) -> bool {
false
}
fn priority(&self) -> i32 {
15
}
fn dependencies(&self) -> Vec<&str> {
vec![]
}
fn init(&mut self, _ctx: &mut crate::core::RiServiceContext) -> crate::core::RiResult<()> {
Ok(())
}
fn start(&mut self, _ctx: &mut crate::core::RiServiceContext) -> crate::core::RiResult<()> {
Ok(())
}
fn shutdown(&mut self, _ctx: &mut crate::core::RiServiceContext) -> crate::core::RiResult<()> {
Ok(())
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
pub struct RiQueueManager {
config: RiQueueConfig,
queues: Arc<RwLock<FxHashMap<String, Arc<dyn RiQueue>>>>,
connection_pool: Option<Arc<QueueConnectionPool>>,
}
impl RiQueueManager {
pub async fn new(config: RiQueueConfig) -> RiResult<Self> {
let backend_type = config.backend_type.clone();
let connection_string = config.connection_string.clone();
let connection_pool = match backend_type {
RiQueueBackendType::RabbitMQ | RiQueueBackendType::Redis => {
Some(Arc::new(QueueConnectionPool::new(
backend_type,
connection_string,
10, )))
}
#[cfg(all(feature = "kafka", not(windows)))]
RiQueueBackendType::Kafka => {
Some(Arc::new(QueueConnectionPool::new(
backend_type,
connection_string,
10, )))
}
#[cfg(any(not(feature = "kafka"), windows))]
RiQueueBackendType::Kafka => {
None
}
_ => None,
};
Ok(Self {
config: RiQueueConfig {
enabled: config.enabled,
backend_type: config.backend_type.clone(),
connection_string: config.connection_string.clone(),
max_connections: config.max_connections,
message_max_size: config.message_max_size,
consumer_timeout_ms: config.consumer_timeout_ms,
producer_timeout_ms: config.producer_timeout_ms,
retry_policy: config.retry_policy.clone(),
dead_letter_config: config.dead_letter_config.clone(),
},
queues: Arc::new(RwLock::new(FxHashMap::default())),
connection_pool,
})
}
}
impl Default for RiQueueManager {
fn default() -> Self {
Self {
config: RiQueueConfig {
enabled: true,
backend_type: RiQueueBackendType::Memory,
connection_string: "memory://localhost".to_string(),
max_connections: 10,
message_max_size: 1024 * 1024,
consumer_timeout_ms: 30000,
producer_timeout_ms: 30000,
retry_policy: crate::queue::config::RiRetryPolicy::default(),
dead_letter_config: None,
},
queues: Arc::new(RwLock::new(FxHashMap::default())),
connection_pool: None,
}
}
}
impl RiQueueManager {
pub async fn init(&self) -> RiResult<()> {
if let Some(ref pool) = self.connection_pool {
for _ in 0..3 {
if let Ok(conn) = pool.create_connection().await {
let _ = pool.return_connection(conn).await;
}
}
}
Ok(())
}
pub async fn create_queue(&self, name: &str) -> RiResult<Arc<dyn RiQueue>> {
let mut queues = self.queues.write().await;
if let Some(queue) = queues.get(name) {
return Ok(queue.clone());
}
let queue = self.create_backend_queue(name).await?;
queues.insert(name.to_string(), queue.clone());
Ok(queue)
}
async fn create_backend_queue(&self, name: &str) -> RiResult<Arc<dyn RiQueue>> {
match self.config.backend_type {
RiQueueBackendType::Memory => {
Ok(Arc::new(crate::queue::backends::RiMemoryQueue::new(name)))
}
#[cfg(feature = "rabbitmq")]
RiQueueBackendType::RabbitMQ => {
if let Some(ref pool) = self.connection_pool {
let conn = pool.get_connection().await?;
if let Some(_lapin_conn) = conn.downcast_ref::<lapin::Connection>() {
let queue = crate::queue::backends::RiRabbitMQQueue::new(
name,
&self.config.connection_string,
)
.await?;
let _ = pool.return_connection(conn).await;
return Ok(Arc::new(queue));
}
}
Ok(Arc::new(
crate::queue::backends::RiRabbitMQQueue::new(
name,
&self.config.connection_string,
)
.await?,
))
}
#[cfg(not(feature = "rabbitmq"))]
RiQueueBackendType::RabbitMQ => {
Err(crate::core::RiError::Config(
"RabbitMQ support is disabled. Enable the 'rabbitmq' feature to use RabbitMQ backend.".to_string(),
))
}
#[cfg(all(feature = "kafka", not(windows)))]
RiQueueBackendType::Kafka => {
Ok(Arc::new(
crate::queue::backends::RiKafkaQueue::new(
name,
&self.config.connection_string,
)
.await?,
))
}
#[cfg(any(not(feature = "kafka"), windows))]
RiQueueBackendType::Kafka => {
Err(crate::core::RiError::Config(
"Kafka support is disabled. Enable the 'kafka' feature to use Kafka backend.".to_string(),
))
}
#[cfg(feature = "redis")]
RiQueueBackendType::Redis => {
Ok(Arc::new(
crate::queue::backends::RiRedisQueue::new(
name,
&self.config.connection_string,
)
.await?,
))
}
#[cfg(not(feature = "redis"))]
RiQueueBackendType::Redis => {
Err(crate::core::RiError::Config(
"Redis support is disabled. Enable the 'redis' feature to use Redis backend.".to_string(),
))
}
}
}
pub async fn get_queue(&self, name: &str) -> Option<Arc<dyn RiQueue>> {
let queues = self.queues.read().await;
queues.get(name).cloned()
}
pub async fn list_queues(&self) -> Vec<String> {
let queues = self.queues.read().await;
queues.keys().cloned().collect()
}
pub async fn delete_queue(&self, name: &str) -> RiResult<()> {
let mut queues = self.queues.write().await;
if let Some(queue) = queues.remove(name) {
queue.delete().await?;
}
Ok(())
}
pub async fn shutdown(&self) -> RiResult<()> {
let mut queues = self.queues.write().await;
for (_, queue) in queues.drain() {
let _ = queue.purge().await;
}
if let Some(ref pool) = self.connection_pool {
pool.close_all().await;
}
Ok(())
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiQueueManager {
#[pyo3(name = "create_queue")]
fn create_queue_impl(&self, name: String) -> PyResult<String> {
let rt = tokio::runtime::Runtime::new().map_err(|e| {
pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create runtime: {}", e))
})?;
rt.block_on(async {
self.create_queue(&name)
.await
.map(|_| name)
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create queue: {e}")))
})
}
#[pyo3(name = "get_queue")]
fn get_queue_impl(&self, name: String) -> PyResult<Option<()>> {
let rt = tokio::runtime::Runtime::new().map_err(|e| {
pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create runtime: {}", e))
})?;
Ok(rt.block_on(async {
self.get_queue(&name).await.map(|_| ())
}))
}
#[pyo3(name = "list_queues")]
fn list_queues_impl(&self) -> PyResult<Vec<String>> {
let rt = tokio::runtime::Runtime::new().map_err(|e| {
pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create runtime: {}", e))
})?;
Ok(rt.block_on(async {
self.list_queues().await
}))
}
#[pyo3(name = "delete_queue")]
fn delete_queue_impl(&self, name: String) -> PyResult<()> {
let rt = tokio::runtime::Runtime::new().map_err(|e| {
pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to create runtime: {}", e))
})?;
rt.block_on(async {
self.delete_queue(&name)
.await
.map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(format!("Failed to delete queue: {e}")))
})
}
}