#![allow(non_snake_case)]
use async_trait::async_trait;
use serde::{Serialize, Deserialize};
use std::collections::HashMap as FxHashMap;
use std::time::SystemTime;
use crate::core::RiResult;
#[derive(Debug, Clone)]
pub enum RiQueueError {
BackendError(String),
ConfigError(String),
ConnectionError(String),
MessageNotFound(String),
ConsumerGroupError(String),
SerializationError(String),
}
impl std::fmt::Display for RiQueueError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
RiQueueError::BackendError(msg) => write!(f, "Queue backend error: {}", msg),
RiQueueError::ConfigError(msg) => write!(f, "Queue configuration error: {}", msg),
RiQueueError::ConnectionError(msg) => write!(f, "Queue connection error: {}", msg),
RiQueueError::MessageNotFound(msg) => write!(f, "Message not found: {}", msg),
RiQueueError::ConsumerGroupError(msg) => write!(f, "Consumer group error: {}", msg),
RiQueueError::SerializationError(msg) => write!(f, "Serialization error: {}", msg),
}
}
}
impl std::error::Error for RiQueueError {}
impl From<RiQueueError> for crate::core::RiError {
fn from(error: RiQueueError) -> Self {
crate::core::RiError::Queue(error.to_string())
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RiQueueMessage {
pub id: String,
pub payload: Vec<u8>,
pub headers: FxHashMap<String, String>,
pub timestamp: SystemTime,
pub retry_count: u32,
pub max_retries: u32,
}
impl RiQueueMessage {
pub const MAX_PAYLOAD_SIZE: usize = 10 * 1024 * 1024;
pub const MAX_HEADER_KEY_LEN: usize = 256;
pub const MAX_HEADER_VALUE_LEN: usize = 4096;
pub const MAX_HEADERS: usize = 64;
pub fn new(payload: Vec<u8>) -> Self {
let payload = if payload.len() > Self::MAX_PAYLOAD_SIZE {
log::warn!(
"[Ri.Queue] Payload size {} exceeds maximum {}, truncating",
payload.len(),
Self::MAX_PAYLOAD_SIZE
);
payload[..Self::MAX_PAYLOAD_SIZE].to_vec()
} else {
payload
};
Self {
id: uuid::Uuid::new_v4().to_string(),
payload,
headers: FxHashMap::default(),
timestamp: SystemTime::now(),
retry_count: 0,
max_retries: 3,
}
}
pub fn new_validated(payload: Vec<u8>) -> RiResult<Self> {
if payload.len() > Self::MAX_PAYLOAD_SIZE {
return Err(RiQueueError::ConfigError(format!(
"Payload size {} exceeds maximum allowed {}",
payload.len(),
Self::MAX_PAYLOAD_SIZE
)).into());
}
Ok(Self {
id: uuid::Uuid::new_v4().to_string(),
payload,
headers: FxHashMap::default(),
timestamp: SystemTime::now(),
retry_count: 0,
max_retries: 3,
})
}
pub fn with_headers(mut self, headers: FxHashMap<String, String>) -> Self {
let mut safe_headers = FxHashMap::with_capacity(
headers.len().min(Self::MAX_HEADERS)
);
for (key, value) in headers {
if safe_headers.len() >= Self::MAX_HEADERS {
log::warn!("[Ri.Queue] Maximum header count {} reached, ignoring remaining headers", Self::MAX_HEADERS);
break;
}
let safe_key = if key.len() > Self::MAX_HEADER_KEY_LEN {
log::warn!("[Ri.Queue] Header key '{}' truncated to {} characters", key, Self::MAX_HEADER_KEY_LEN);
key[..Self::MAX_HEADER_KEY_LEN].to_string()
} else {
key
};
let safe_value = if value.len() > Self::MAX_HEADER_VALUE_LEN {
log::warn!("[Ri.Queue] Header value for key '{}' truncated to {} characters", safe_key, Self::MAX_HEADER_VALUE_LEN);
value[..Self::MAX_HEADER_VALUE_LEN].to_string()
} else {
value
};
if safe_key.chars().any(|c| c.is_control()) {
log::warn!("[Ri.Queue] Header key contains control characters, skipping");
continue;
}
safe_headers.insert(safe_key, safe_value);
}
self.headers = safe_headers;
self
}
pub fn with_max_retries(mut self, max_retries: u32) -> Self {
self.max_retries = max_retries;
self
}
pub fn increment_retry(&mut self) {
self.retry_count += 1;
}
pub fn can_retry(&self) -> bool {
self.retry_count < self.max_retries
}
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiQueueMessage {
#[new]
fn py_new(payload: Vec<u8>) -> Self {
Self::new(payload)
}
#[staticmethod]
fn py_new_with_string(payload: String) -> Self {
Self::new(payload.into_bytes())
}
fn get_payload_string(&self) -> String {
String::from_utf8_lossy(&self.payload).to_string()
}
fn get_id(&self) -> String {
self.id.clone()
}
}
#[cfg_attr(feature = "pyo3", pyo3::prelude::pyclass)]
#[derive(Debug, Clone)]
pub struct RiQueueStats {
pub queue_name: String,
pub message_count: u64,
pub consumer_count: u32,
pub producer_count: u32,
pub processed_messages: u64,
pub failed_messages: u64,
pub avg_processing_time_ms: f64,
pub total_bytes_sent: u64,
pub total_bytes_received: u64,
pub last_message_time: u64,
}
#[cfg(feature = "pyo3")]
#[pyo3::prelude::pymethods]
impl RiQueueStats {
#[new]
fn py_new(queue_name: String) -> Self {
Self {
queue_name,
message_count: 0,
consumer_count: 0,
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]
pub trait RiQueueProducer: Send + Sync {
async fn send(&self, message: RiQueueMessage) -> RiResult<()>;
async fn send_batch(&self, messages: Vec<RiQueueMessage>) -> RiResult<()>;
async fn send_multi(&self, messages: &[RiQueueMessage]) -> RiResult<()> {
for message in messages {
self.send(message.clone()).await?;
}
Ok(())
}
}
#[async_trait]
pub trait RiQueueConsumer: Send + Sync {
async fn receive(&self) -> RiResult<Option<RiQueueMessage>>;
async fn ack(&self, message_id: &str) -> RiResult<()>;
async fn nack(&self, message_id: &str) -> RiResult<()>;
async fn pause(&self) -> RiResult<()>;
async fn resume(&self) -> RiResult<()>;
async fn receive_multi(&self, count: usize) -> RiResult<Vec<Option<RiQueueMessage>>> {
let mut messages = Vec::with_capacity(count);
for _ in 0..count {
messages.push(self.receive().await?);
}
Ok(messages)
}
async fn ack_multi(&self, message_ids: &[String]) -> RiResult<()> {
for id in message_ids {
self.ack(id).await?;
}
Ok(())
}
}
#[async_trait]
pub trait RiQueue: Send + Sync {
async fn create_producer(&self) -> RiResult<Box<dyn RiQueueProducer>>;
async fn create_consumer(&self, consumer_group: &str) -> RiResult<Box<dyn RiQueueConsumer>>;
async fn get_stats(&self) -> RiResult<RiQueueStats>;
async fn purge(&self) -> RiResult<()>;
async fn delete(&self) -> RiResult<()>;
}