feature: added management api for redis

Author: densumesh

broccoli-queue/Cargo.toml

@@ -32,6 +32,7 @@ redis = { version = "0.27.5", features = [
     "aio",
 ], optional = true }
 dashmap = "6.1.0"
+derive_more = { version = "2.0.1", features = ["display"], optional = true }
 
 [dev-dependencies]
 chrono = { version = "0.4.39", features = ["serde"] }
@@ -52,7 +53,7 @@ rabbitmq = ["dep:lapin", "dep:deadpool", "dep:deadpool-lapin"]
 test-fairness = []
 
 # Allows for access to the management API for the queue.
-management = []
+management = ["dep:derive_more"]
 
 [[bench]]
 name = "queue_benchmark"

broccoli-queue/examples/consumer.rs

@@ -23,7 +23,7 @@ async fn process_job(job: JobPayload) -> Result<(), BroccoliError> {
     // Simulate some work
     tokio::time::sleep(Duration::from_secs(1)).await;
 
-    Ok(())
+    Err(BroccoliError::Consume("Failed to process job".into()))
 }
 
 async fn success_handler(msg: JobPayload) -> Result<(), BroccoliError> {

broccoli-queue/src/brokers/management.rs

@@ -1,41 +1,42 @@
+use derive_more::Display;
+
 use crate::error::BroccoliError;
 
 use super::broker::Broker;
 
 #[async_trait::async_trait]
 /// Trait for managing queues.
 pub trait QueueManagement {
-    /// Retries all messages in the queue.
+    /// Retries messages in the queue.
     async fn retry_queue(
         &self,
-        queue_name: &str,
-        disambiguator: Option<String>,
+        queue_name: String,
         source_type: QueueType,
     ) -> Result<usize, BroccoliError>;
-    /// Gets the size of the queue.
-    async fn get_queue_size(
-        &self,
-        queue_name: &str,
-        queue_type: QueueType,
-    ) -> Result<usize, BroccoliError>;
     /// Gets the status of specific or all queues. If `queue_name` is `None`, returns the status of all queues.
     async fn get_queue_status(
         &self,
-        queue_name: Option<&str>,
+        queue_name: Option<String>,
     ) -> Result<Vec<QueueStatus>, BroccoliError>;
 }
 
 pub(crate) trait BrokerWithManagement: Broker + QueueManagement {}
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq, Eq, Display)]
 /// Enum representing the type of queue.
 pub enum QueueType {
     /// Failed queue.
+    #[display("failed")]
     Failed,
     /// Processing queue.
+    #[display("processing")]
     Processing,
     /// Main queue.
+    #[display("main")]
     Main,
+    /// Fairness queue.
+    #[display("fairness")]
+    Fairness,
 }
 
 #[derive(Debug, Clone)]
@@ -51,4 +52,6 @@ pub struct QueueStatus {
     pub processing: usize,
     /// Number of messages that failed to be processed.
     pub failed: usize,
+    /// If the queue is a fairness queue, the number of disambiguators.
+    pub disambiguator_count: Option<usize>,
 }

broccoli-queue/src/brokers/rabbitmq/management.rs

@@ -1,6 +1,5 @@
 use lapin::{
-    options::{BasicAckOptions, BasicGetOptions, BasicPublishOptions, QueueDeclareOptions},
-    types::FieldTable,
+    options::{BasicAckOptions, BasicGetOptions, BasicPublishOptions},
     BasicProperties,
 };
 
@@ -15,8 +14,7 @@ use super::RabbitMQBroker;
 impl QueueManagement for RabbitMQBroker {
     async fn retry_queue(
         &self,
-        queue_name: &str,
-        _disambiguator: Option<String>,
+        queue_name: String,
         source_type: QueueType,
     ) -> Result<usize, BroccoliError> {
         let pool = self.ensure_pool()?;
@@ -34,6 +32,11 @@ impl QueueManagement for RabbitMQBroker {
                     "Cannot retry from ingestion queue".into(),
                 ))
             }
+            QueueType::Fairness => {
+                return Err(BroccoliError::InvalidOperation(
+                    "Cannot retry from fairness queue".into(),
+                ))
+            }
         };
 
         let mut count = 0;
@@ -44,7 +47,7 @@ impl QueueManagement for RabbitMQBroker {
             channel
                 .basic_publish(
                     "broccoli",
-                    queue_name,
+                    &queue_name,
                     BasicPublishOptions::default(),
                     &delivery.data,
                     BasicProperties::default(),
@@ -60,35 +63,10 @@ impl QueueManagement for RabbitMQBroker {
         Ok(count)
     }
 
-    async fn get_queue_size(
+    async fn get_queue_status(
         &self,
-        queue_name: &str,
-        queue_type: QueueType,
-    ) -> Result<usize, BroccoliError> {
-        let pool = self.ensure_pool()?;
-        let conn = pool
-            .get()
-            .await
-            .map_err(|e| BroccoliError::Consume(format!("Failed to consume message: {e:?}")))?;
-        let channel = conn.create_channel().await?;
-
-        let queue = match queue_type {
-            QueueType::Failed => format!("{queue_name}_failed"),
-            QueueType::Processing => format!("{queue_name}_processing"),
-            QueueType::Main => queue_name.to_string(),
-        };
-
-        let queue_info = channel
-            .queue_declare(
-                &queue,
-                QueueDeclareOptions::default(),
-                FieldTable::default(),
-            )
-            .await?;
-        Ok(queue_info.message_count() as usize)
-    }
-
-    async fn get_queue_status(&self) -> Result<Vec<QueueStatus>, BroccoliError> {
+        queue_name: Option<String>,
+    ) -> Result<Vec<QueueStatus>, BroccoliError> {
         let pool = self.ensure_pool()?;
         let conn = pool
             .get()
@@ -98,7 +76,7 @@ impl QueueManagement for RabbitMQBroker {
 
         // List queues through management API or channel operations
         // This is a simplified version - in practice you'd want to use the RabbitMQ Management API
-        let mut statuses = Vec::new();
+        let statuses = Vec::new();
 
         // Implementation note: RabbitMQ doesn't provide memory usage through regular AMQP
         // You would need to use the HTTP Management API to get this information

broccoli-queue/src/brokers/redis/broker.rs

@@ -10,7 +10,7 @@ use crate::{
 
 use super::utils::OptionalInternalBrokerMessage;
 
-pub(crate) type RedisPool = bb8_redis::bb8::Pool<bb8_redis::RedisConnectionManager>;
+pub type RedisPool = bb8_redis::bb8::Pool<bb8_redis::RedisConnectionManager>;
 
 #[derive(Default)]
 /// A message broker implementation for Redis.