camel_component_redis/

consumer.rs

use async_trait::async_trait;
use camel_component_api::{Body, CamelError, Exchange, Message};
use camel_component_api::{ConcurrencyModel, Consumer, ConsumerContext};
use futures_util::StreamExt;
use redis::Msg;
use std::time::Duration;
use tokio::task::JoinHandle;
use tokio_util::sync::CancellationToken;
use tracing::{error, info, warn};

use crate::config::{RedisCommand, RedisEndpointConfig};

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum QueuePopCommand {
    Blpop,
    Brpop,
}

fn queue_command_name(pop_command: QueuePopCommand) -> &'static str {
    match pop_command {
        QueuePopCommand::Blpop => "BLPOP",
        QueuePopCommand::Brpop => "BRPOP",
    }
}

/// Mode of operation for the Redis consumer.
#[derive(Debug, Clone)]
pub enum RedisConsumerMode {
    /// Pub/Sub mode for real-time message streams.
    PubSub {
        /// Channels to subscribe to (SUBSCRIBE)
        channels: Vec<String>,
        /// Patterns to subscribe to (PSUBSCRIBE)
        patterns: Vec<String>,
    },
    /// Queue mode for blocking list operations.
    Queue {
        /// Key to watch for items
        key: String,
        /// Timeout in seconds for blocking pop
        timeout: u64,
        /// Blocking pop command to use (left or right)
        pop_command: QueuePopCommand,
    },
}

/// Redis consumer implementation supporting both Pub/Sub and Queue modes.
pub struct RedisConsumer {
    config: RedisEndpointConfig,
    mode: RedisConsumerMode,
    /// Cancellation token for graceful shutdown
    cancel_token: Option<CancellationToken>,
    /// Handle to the spawned consumer task
    task_handle: Option<JoinHandle<Result<(), CamelError>>>,
}

impl RedisConsumer {
    /// Creates a new RedisConsumer with the given configuration.
    ///
    /// The mode is automatically determined from the command type in the config:
    /// - SUBSCRIBE → PubSub with channels
    /// - PSUBSCRIBE → PubSub with patterns
    /// - BLPOP/BRPOP → Queue mode
    pub fn new(config: RedisEndpointConfig) -> Self {
        let mode = match &config.command {
            RedisCommand::Subscribe => RedisConsumerMode::PubSub {
                channels: config.channels.clone(),
                patterns: vec![],
            },
            RedisCommand::Psubscribe => RedisConsumerMode::PubSub {
                channels: vec![],
                patterns: config.channels.clone(),
            },
            RedisCommand::Blpop | RedisCommand::Brpop => {
                let key = config.key.clone().unwrap_or_else(|| "queue".to_string());
                let pop_command = if config.command == RedisCommand::Brpop {
                    QueuePopCommand::Brpop
                } else {
                    QueuePopCommand::Blpop
                };
                RedisConsumerMode::Queue {
                    key,
                    timeout: config.timeout,
                    pop_command,
                }
            }
            _ => {
                warn!(
                    "Invalid consumer command: {:?}, defaulting to BLPOP",
                    config.command
                );
                RedisConsumerMode::Queue {
                    key: config.key.clone().unwrap_or_else(|| "queue".to_string()),
                    timeout: config.timeout,
                    pop_command: QueuePopCommand::Blpop,
                }
            }
        };

        Self {
            config,
            mode,
            cancel_token: None,
            task_handle: None,
        }
    }
}

#[async_trait]
impl Consumer for RedisConsumer {
    async fn start(&mut self, ctx: ConsumerContext) -> Result<(), CamelError> {
        // Create cancellation token for this consumer
        let cancel_token = CancellationToken::new();
        self.cancel_token = Some(cancel_token.clone());

        // Clone config and mode for the spawned task
        let config = self.config.clone();
        let mode = self.mode.clone();

        info!("Starting Redis consumer in {:?} mode", mode);

        // Spawn the appropriate consumer loop based on mode
        let handle =
            match mode {
                RedisConsumerMode::PubSub { channels, patterns } => tokio::spawn(
                    run_pubsub_consumer(config, channels, patterns, ctx, cancel_token),
                ),
                RedisConsumerMode::Queue {
                    key,
                    timeout,
                    pop_command,
                } => tokio::spawn(run_queue_consumer(
                    config,
                    key,
                    timeout,
                    pop_command,
                    ctx,
                    cancel_token,
                )),
            };

        self.task_handle = Some(handle);
        Ok(())
    }

    async fn stop(&mut self) -> Result<(), CamelError> {
        info!("Stopping Redis consumer");

        // Cancel the token to signal shutdown
        if let Some(token) = &self.cancel_token {
            token.cancel();
        }

        // Wait for the task to complete
        if let Some(handle) = self.task_handle.take() {
            match handle.await {
                Ok(result) => {
                    if let Err(e) = result {
                        error!("Consumer task exited with error: {}", e);
                    }
                }
                Err(e) => {
                    error!("Failed to join consumer task: {}", e);
                }
            }
        }

        self.cancel_token = None;
        info!("Redis consumer stopped");
        Ok(())
    }

    /// Redis consumers are sequential by default to maintain message order.
    ///
    /// This default is chosen for the following reasons:
    /// - **Pub/Sub**: Messages often need ordering (e.g., event streams, notifications)
    /// - **Queue (BLPOP)**: Queue items should be processed in order
    /// - **Backpressure**: Sequential processing naturally applies backpressure
    ///   when the consumer is slower than the producer
    ///
    /// Users can override this with `.concurrent(n)` in the route DSL if they
    /// want parallel processing and ordering is not a concern.
    fn concurrency_model(&self) -> ConcurrencyModel {
        ConcurrencyModel::Sequential
    }
}

/// Runs a Pub/Sub consumer loop.
///
/// Creates a dedicated PubSub connection and subscribes to the specified
/// channels and/or patterns. Messages are converted to Exchanges and sent
/// through the consumer context.
async fn run_pubsub_consumer(
    config: RedisEndpointConfig,
    channels: Vec<String>,
    patterns: Vec<String>,
    ctx: ConsumerContext,
    cancel_token: CancellationToken,
) -> Result<(), CamelError> {
    info!("PubSub consumer connecting to {}", config.redis_url());

    // Create dedicated PubSub connection
    let client = redis::Client::open(config.redis_url())
        .map_err(|e| CamelError::ProcessorError(format!("Failed to create Redis client: {}", e)))?;

    let mut pubsub = client.get_async_pubsub().await.map_err(|e| {
        CamelError::ProcessorError(format!("Failed to create PubSub connection: {}", e))
    })?;

    // Subscribe to channels
    for channel in &channels {
        info!("Subscribing to channel: {}", channel);
        pubsub.subscribe(channel).await.map_err(|e| {
            CamelError::ProcessorError(format!("Failed to subscribe to channel {}: {}", channel, e))
        })?;
    }

    // Subscribe to patterns
    for pattern in &patterns {
        info!("Subscribing to pattern: {}", pattern);
        pubsub.psubscribe(pattern).await.map_err(|e| {
            CamelError::ProcessorError(format!("Failed to subscribe to pattern {}: {}", pattern, e))
        })?;
    }

    info!("PubSub consumer started, waiting for messages");

    // Message loop
    let mut stream = pubsub.on_message();
    loop {
        tokio::select! {
            _ = cancel_token.cancelled() => {
                info!("PubSub consumer received shutdown signal");
                break;
            }
            msg = stream.next() => {
                if let Some(msg) = msg {
                    let exchange = build_exchange_from_pubsub(msg);
                    if let Err(e) = ctx.send(exchange).await {
                        error!("Failed to send exchange to pipeline: {}", e);
                        // Don't break - continue processing messages
                    }
                } else {
                    // Stream ended
                    warn!("PubSub stream ended");
                    break;
                }
            }
        }
    }

    Ok(())
}

/// Runs a Queue consumer loop using BLPOP or BRPOP.
///
/// Creates a dedicated connection and performs blocking list pop operations.
/// Items are converted to Exchanges and sent through the consumer context.
async fn run_queue_consumer(
    config: RedisEndpointConfig,
    key: String,
    timeout: u64,
    pop_command: QueuePopCommand,
    ctx: ConsumerContext,
    cancel_token: CancellationToken,
) -> Result<(), CamelError> {
    info!(
        "Queue consumer connecting to {} for key '{}' with {} timeout {}s",
        config.redis_url(),
        key,
        queue_command_name(pop_command),
        timeout
    );

    // Create dedicated multiplexed connection
    let client = redis::Client::open(config.redis_url())
        .map_err(|e| CamelError::ProcessorError(format!("Failed to create Redis client: {}", e)))?;

    let mut conn = client
        .get_multiplexed_async_connection()
        .await
        .map_err(|e| CamelError::ProcessorError(format!("Failed to create connection: {}", e)))?;

    info!("Queue consumer started, waiting for items");

    // Blocking pop loop (BLPOP/BRPOP)
    let queue_cmd = queue_command_name(pop_command);
    loop {
        tokio::select! {
            _ = cancel_token.cancelled() => {
                info!("Queue consumer received shutdown signal");
                break;
            }
            result = async {
                let cmd = redis::cmd(queue_cmd)
                    .arg(&key)
                    .arg(timeout)
                    .to_owned();
                cmd.query_async::<Option<(String, String)>>(&mut conn).await
            } =>
            {
                match result {
                    Ok(Some((key, value))) => {
                        let exchange = build_exchange_from_blpop(key, value);
                        if let Err(e) = ctx.send(exchange).await {
                            error!("Failed to send exchange to pipeline: {}", e);
                            // Don't break - continue processing items
                        }
                    }
                    Ok(None) => {
                        // Timeout - continue loop
                        // This is normal for blocking POP with timeout
                    }
                    Err(e) => {
                        if e.is_timeout() {
                            // Timeout - continue loop silently
                        } else {
                            error!("{} error: {}", queue_cmd, e);
                            tokio::time::sleep(Duration::from_millis(100)).await;
                        }
                    }
                }
            }
        }
    }

    Ok(())
}

fn build_pubsub_exchange(payload: String, channel: String, pattern: Option<String>) -> Exchange {
    let mut exchange = Exchange::new(Message::new(Body::Text(payload)));
    exchange
        .input
        .set_header("CamelRedis.Channel", serde_json::Value::String(channel));

    if let Some(pattern) = pattern {
        exchange
            .input
            .set_header("CamelRedis.Pattern", serde_json::Value::String(pattern));
    }

    exchange
}

/// Builds an Exchange from a Pub/Sub message.
///
/// Sets the following headers:
/// - `CamelRedis.Channel`: The channel the message was published to
/// - `CamelRedis.Pattern`: The pattern matched (if applicable, for PSUBSCRIBE)
fn build_exchange_from_pubsub(msg: Msg) -> Exchange {
    let payload: String = msg
        .get_payload()
        .unwrap_or_else(|_| "<error decoding payload>".to_string());
    let channel = msg.get_channel_name().to_string();
    let pattern = if msg.from_pattern() {
        msg.get_pattern::<String>().ok()
    } else {
        None
    };

    build_pubsub_exchange(payload, channel, pattern)
}

/// Builds an Exchange from a BLPOP result.
///
/// Sets the following headers:
/// - `CamelRedis.Key`: The list key the item was popped from
fn build_exchange_from_blpop(key: String, value: String) -> Exchange {
    let mut exchange = Exchange::new(Message::new(Body::Text(value)));

    // Set key header
    exchange
        .input
        .set_header("CamelRedis.Key", serde_json::Value::String(key));

    exchange
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::sync::mpsc;

    fn create_test_config(command: RedisCommand) -> RedisEndpointConfig {
        RedisEndpointConfig {
            host: Some("localhost".to_string()),
            port: Some(6379),
            command,
            channels: vec!["test".to_string()],
            key: Some("test-queue".to_string()),
            timeout: 1,
            password: None,
            db: 0,
        }
    }

    #[test]
    fn test_consumer_new_subscribe() {
        let config = create_test_config(RedisCommand::Subscribe);
        let consumer = RedisConsumer::new(config);

        match consumer.mode {
            RedisConsumerMode::PubSub { channels, patterns } => {
                assert_eq!(channels, vec!["test".to_string()]);
                assert!(patterns.is_empty());
            }
            _ => panic!("Expected PubSub mode"),
        }
    }

    #[test]
    fn test_consumer_new_psubscribe() {
        let config = create_test_config(RedisCommand::Psubscribe);
        let consumer = RedisConsumer::new(config);

        match consumer.mode {
            RedisConsumerMode::PubSub { channels, patterns } => {
                assert!(channels.is_empty());
                assert_eq!(patterns, vec!["test".to_string()]);
            }
            _ => panic!("Expected PubSub mode"),
        }
    }

    #[test]
    fn test_consumer_new_blpop() {
        let config = create_test_config(RedisCommand::Blpop);
        let consumer = RedisConsumer::new(config);

        match consumer.mode {
            RedisConsumerMode::Queue {
                key,
                timeout,
                pop_command,
            } => {
                assert_eq!(key, "test-queue");
                assert_eq!(timeout, 1);
                assert_eq!(pop_command, QueuePopCommand::Blpop);
            }
            _ => panic!("Expected Queue mode"),
        }
    }

    #[test]
    fn test_consumer_new_brpop_uses_right_pop_command() {
        let config = create_test_config(RedisCommand::Brpop);
        let consumer = RedisConsumer::new(config);

        match consumer.mode {
            RedisConsumerMode::Queue { pop_command, .. } => {
                assert_eq!(pop_command, QueuePopCommand::Brpop);
            }
            _ => panic!("Expected Queue mode"),
        }
    }

    #[test]
    fn test_consumer_new_blpop_default_key() {
        let mut config = create_test_config(RedisCommand::Blpop);
        config.key = None;
        let consumer = RedisConsumer::new(config);

        match consumer.mode {
            RedisConsumerMode::Queue {
                key, pop_command, ..
            } => {
                assert_eq!(key, "queue");
                assert_eq!(pop_command, QueuePopCommand::Blpop);
            }
            _ => panic!("Expected Queue mode"),
        }
    }

    #[test]
    fn test_consumer_new_non_consumer_command_defaults_to_queue_mode() {
        let config = create_test_config(RedisCommand::Set);
        let consumer = RedisConsumer::new(config);

        match consumer.mode {
            RedisConsumerMode::Queue {
                key,
                timeout,
                pop_command,
            } => {
                assert_eq!(key, "test-queue");
                assert_eq!(timeout, 1);
                assert_eq!(pop_command, QueuePopCommand::Blpop);
            }
            _ => panic!("Expected Queue mode"),
        }
    }

    #[test]
    fn test_queue_command_name_matches_pop_side() {
        assert_eq!(queue_command_name(QueuePopCommand::Blpop), "BLPOP");
        assert_eq!(queue_command_name(QueuePopCommand::Brpop), "BRPOP");
    }

    #[test]
    fn test_consumer_concurrency_model_is_sequential() {
        let config = create_test_config(RedisCommand::Subscribe);
        let consumer = RedisConsumer::new(config);
        assert_eq!(consumer.concurrency_model(), ConcurrencyModel::Sequential);
    }

    #[test]
    fn test_build_exchange_from_blpop() {
        let exchange = build_exchange_from_blpop("mykey".to_string(), "myvalue".to_string());

        assert_eq!(exchange.input.body.as_text(), Some("myvalue"));

        let header = exchange.input.header("CamelRedis.Key");
        assert_eq!(
            header,
            Some(&serde_json::Value::String("mykey".to_string()))
        );
    }

    #[test]
    fn test_build_pubsub_exchange_without_pattern() {
        let exchange = build_pubsub_exchange("hello".to_string(), "news".to_string(), None);

        assert_eq!(exchange.input.body.as_text(), Some("hello"));
        assert_eq!(
            exchange.input.header("CamelRedis.Channel"),
            Some(&serde_json::json!("news"))
        );
        assert!(exchange.input.header("CamelRedis.Pattern").is_none());
    }

    #[test]
    fn test_build_pubsub_exchange_with_pattern() {
        let exchange = build_pubsub_exchange(
            "hello".to_string(),
            "news.eu".to_string(),
            Some("news.*".to_string()),
        );

        assert_eq!(
            exchange.input.header("CamelRedis.Pattern"),
            Some(&serde_json::json!("news.*"))
        );
    }

    #[tokio::test]
    async fn test_consumer_stops_gracefully() {
        let config = create_test_config(RedisCommand::Blpop);
        let mut consumer = RedisConsumer::new(config);

        // Create a mock context (won't actually be used in this test)
        let (tx, _rx) = mpsc::channel(16);
        let cancel_token = CancellationToken::new();
        let ctx = ConsumerContext::new(tx, cancel_token.clone());

        // Start should succeed
        let start_result = consumer.start(ctx).await;
        assert!(start_result.is_ok());

        // Give task a moment to start
        tokio::time::sleep(Duration::from_millis(10)).await;

        // Stop should succeed
        let stop_result = consumer.stop().await;
        assert!(stop_result.is_ok());
    }
}