rust_rabbit/

consumer.rs

use crate::{
    connection::Connection, 
    error::RustRabbitError, 
    message::{ErrorType, MessageEnvelope},
    retry::RetryConfig
};
use futures_lite::stream::StreamExt;
use lapin::{
    options::{BasicAckOptions, BasicConsumeOptions, QueueDeclareOptions},
    types::FieldTable,
    Channel,
};
use serde::de::DeserializeOwned;
use std::future::Future;
use std::sync::Arc;
use tokio::sync::Semaphore;
use tracing::{debug, error, warn};

/// Message wrapper with retry tracking
#[derive(Debug)]
pub struct Message<T>
where
    T: Clone,
{
    pub data: T,
    pub retry_attempt: u32,
    tag: u64,
    channel: Arc<Channel>,
}

impl<T> Clone for Message<T>
where
    T: Clone,
{
    fn clone(&self) -> Self {
        Self {
            data: self.data.clone(),
            retry_attempt: self.retry_attempt,
            tag: self.tag,
            channel: Arc::clone(&self.channel),
        }
    }
}

impl<T> Message<T>
where
    T: Clone,
{
    /// Acknowledge the message
    pub async fn ack(&self) -> Result<(), RustRabbitError> {
        self.channel
            .basic_ack(self.tag, BasicAckOptions::default())
            .await
            .map_err(RustRabbitError::from)
    }

    /// Reject and requeue the message
    pub async fn nack(&self, requeue: bool) -> Result<(), RustRabbitError> {
        self.channel
            .basic_nack(
                self.tag,
                lapin::options::BasicNackOptions {
                    multiple: false,
                    requeue,
                },
            )
            .await
            .map_err(RustRabbitError::from)
    }
}

/// Consumer configuration builder
pub struct ConsumerBuilder {
    connection: Arc<Connection>,
    queue_name: String,
    exchange_name: Option<String>,
    routing_key: Option<String>,
    retry_config: Option<RetryConfig>,
    prefetch_count: Option<u16>,
    auto_ack: bool,
}

impl ConsumerBuilder {
    pub fn new(connection: Arc<Connection>, queue_name: impl Into<String>) -> Self {
        Self {
            connection,
            queue_name: queue_name.into(),
            exchange_name: None,
            routing_key: None,
            retry_config: None,
            prefetch_count: Some(10),
            auto_ack: true,
        }
    }

    /// Bind to an exchange with routing key
    pub fn bind_to_exchange(
        mut self,
        exchange: impl Into<String>,
        routing_key: impl Into<String>,
    ) -> Self {
        self.exchange_name = Some(exchange.into());
        self.routing_key = Some(routing_key.into());
        self
    }

    /// Set routing key (for use with bind_to_exchange)
    pub fn routing_key(mut self, routing_key: impl Into<String>) -> Self {
        self.routing_key = Some(routing_key.into());
        self
    }

    /// Set concurrency level (same as prefetch count)
    pub fn concurrency(mut self, count: u16) -> Self {
        self.prefetch_count = Some(count);
        self
    }

    /// Configure retry behavior
    pub fn with_retry(mut self, retry_config: RetryConfig) -> Self {
        self.retry_config = Some(retry_config);
        self
    }

    /// Set prefetch count
    pub fn with_prefetch(mut self, count: u16) -> Self {
        self.prefetch_count = Some(count);
        self
    }

    /// Disable auto-acknowledge (manual ack required)
    pub fn manual_ack(mut self) -> Self {
        self.auto_ack = false;
        self
    }

    /// Build the consumer
    pub fn build(self) -> Consumer {
        Consumer {
            connection: self.connection,
            queue_name: self.queue_name,
            exchange_name: self.exchange_name,
            routing_key: self.routing_key,
            retry_config: self.retry_config,
            prefetch_count: self.prefetch_count.unwrap_or(10),
            auto_ack: self.auto_ack,
        }
    }
}

/// Simplified Consumer for message consumption
pub struct Consumer {
    connection: Arc<Connection>,
    queue_name: String,
    exchange_name: Option<String>,
    routing_key: Option<String>,
    #[allow(dead_code)]
    retry_config: Option<RetryConfig>,
    prefetch_count: u16,
    auto_ack: bool,
}

impl Consumer {
    /// Create a new consumer builder
    pub fn builder(connection: Arc<Connection>, queue_name: impl Into<String>) -> ConsumerBuilder {
        ConsumerBuilder::new(connection, queue_name)
    }

    /// Start consuming messages
    pub async fn consume<T, H, Fut>(&self, handler: H) -> Result<(), RustRabbitError>
    where
        T: DeserializeOwned + Send + Clone + Sync + 'static,
        H: Fn(Message<T>) -> Fut + Send + Sync + Clone + 'static,
        Fut: Future<Output = Result<(), Box<dyn std::error::Error + Send + Sync>>> + Send,
    {
        let channel = self.connection.create_channel().await?;

        // Set prefetch count
        channel
            .basic_qos(
                self.prefetch_count,
                lapin::options::BasicQosOptions::default(),
            )
            .await?;

        // Setup infrastructure (queues, exchanges)
        self.setup_infrastructure(&channel).await?;

        // Start consuming
        let mut consumer = channel
            .basic_consume(
                &self.queue_name,
                "",
                BasicConsumeOptions::default(),
                FieldTable::default(),
            )
            .await?;

        let semaphore = Arc::new(Semaphore::new(self.prefetch_count as usize));

        debug!("Started consuming from queue: {}", self.queue_name);

        // Process messages (simplified - no retry for now)
        while let Some(delivery_result) = consumer.next().await {
            let delivery = delivery_result?;
            let permit = semaphore.clone().acquire_owned().await.unwrap();
            let handler_clone = handler.clone();
            let auto_ack = self.auto_ack;
            let channel_clone = Arc::new(channel.clone());

            tokio::spawn(async move {
                let _permit = permit;

                // Deserialize message
                match serde_json::from_slice::<T>(&delivery.data) {
                    Ok(data) => {
                        let message = Message {
                            data,
                            retry_attempt: 0, // Simplified for now
                            tag: delivery.delivery_tag,
                            channel: channel_clone.clone(),
                        };

                        // Process message
                        match handler_clone(message.clone()).await {
                            Ok(()) => {
                                if auto_ack {
                                    if let Err(e) = message.ack().await {
                                        error!("Failed to ack message: {}", e);
                                    }
                                }
                                debug!("Message processed successfully");
                            }
                            Err(e) => {
                                error!("Handler error: {}", e);
                                if auto_ack {
                                    // Simple reject without retry for now
                                    if let Err(e) = message.nack(false).await {
                                        error!("Failed to nack message: {}", e);
                                    }
                                }
                            }
                        }
                    }
                    Err(e) => {
                        error!("Failed to deserialize message: {}", e);
                        if auto_ack {
                            // Reject malformed messages
                            if let Err(e) = channel_clone
                                .basic_nack(
                                    delivery.delivery_tag,
                                    lapin::options::BasicNackOptions {
                                        multiple: false,
                                        requeue: false,
                                    },
                                )
                                .await
                            {
                                error!("Failed to nack malformed message: {}", e);
                            }
                        }
                    }
                }
            });
        }

        Ok(())
    }



    /// Setup queue and exchange infrastructure
    async fn setup_infrastructure(&self, channel: &Channel) -> Result<(), RustRabbitError> {
        // Declare queue
        channel
            .queue_declare(
                &self.queue_name,
                QueueDeclareOptions {
                    durable: true,
                    ..Default::default()
                },
                FieldTable::default(),
            )
            .await?;

        // Bind to exchange if specified
        if let (Some(exchange), Some(routing_key)) = (&self.exchange_name, &self.routing_key) {
            channel
                .queue_bind(
                    &self.queue_name,
                    exchange,
                    routing_key,
                    lapin::options::QueueBindOptions::default(),
                    FieldTable::default(),
                )
                .await?;
        }

        Ok(())
    }

    /// Start consuming message envelopes with full retry support
    pub async fn consume_envelopes<T, H, Fut>(&self, handler: H) -> Result<(), RustRabbitError>
    where
        T: DeserializeOwned + Send + Clone + Sync + 'static + serde::Serialize,
        H: Fn(MessageEnvelope<T>) -> Fut + Send + Sync + Clone + 'static,
        Fut: Future<Output = Result<(), Box<dyn std::error::Error + Send + Sync>>> + Send,
    {
        let channel = self.connection.create_channel().await?;
        let retry_config = self.retry_config.clone();

        // Set prefetch count
        channel
            .basic_qos(
                self.prefetch_count,
                lapin::options::BasicQosOptions::default(),
            )
            .await?;

        // Setup queue and exchange
        self.setup_infrastructure(&channel).await?;

        // Create consumer
        let mut consumer = channel
            .basic_consume(
                &self.queue_name,
                "rust-rabbit-envelope-consumer",
                BasicConsumeOptions::default(),
                FieldTable::default(),
            )
            .await?;

        let semaphore = Arc::new(Semaphore::new(self.prefetch_count as usize));

        debug!("Started consuming envelopes from queue: {}", self.queue_name);

        // Process message envelopes with retry support
        while let Some(delivery_result) = consumer.next().await {
            let delivery = delivery_result?;
            let permit = semaphore.clone().acquire_owned().await.unwrap();
            let handler_clone = handler.clone();
            let auto_ack = self.auto_ack;
            let channel_clone = Arc::new(channel.clone());
            let retry_config_clone = retry_config.clone();
            let queue_name = self.queue_name.clone();
            let connection = self.connection.clone();

            tokio::spawn(async move {
                let _permit = permit;

                // Try to deserialize as MessageEnvelope
                match serde_json::from_slice::<MessageEnvelope<T>>(&delivery.data) {
                    Ok(mut envelope) => {
                        debug!(
                            "Processing envelope {} (attempt {}/{})",
                            envelope.metadata.message_id,
                            envelope.metadata.retry_attempt + 1,
                            envelope.metadata.max_retries + 1
                        );

                        // Process message
                        match handler_clone(envelope.clone()).await {
                            Ok(()) => {
                                if auto_ack {
                                    if let Err(e) = channel_clone
                                        .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
                                        .await
                                    {
                                        error!("Failed to ack message: {}", e);
                                    }
                                }
                                debug!("Envelope {} processed successfully", envelope.metadata.message_id);
                            }
                            Err(e) => {
                                error!("Handler error for envelope {}: {}", envelope.metadata.message_id, e);
                                
                                // Determine error type (simplified classification)
                                let error_type = classify_error(e.as_ref());
                                
                                // Add error to envelope
                                envelope = envelope.with_error(
                                    &e.to_string(),
                                    error_type,
                                    Some(&format!("Queue: {}", queue_name))
                                );

                                if auto_ack {
                                    // Check if we should retry
                                    if let Some(retry_cfg) = &retry_config_clone {
                                        if !envelope.is_retry_exhausted() {
                                            // Calculate delay and schedule retry
                                            if let Some(_delay) = retry_cfg.calculate_delay(envelope.metadata.retry_attempt - 1) {
                                                warn!(
                                                    "Scheduling retry {} for envelope {} (simple requeue for now)",
                                                    envelope.metadata.retry_attempt,
                                                    envelope.metadata.message_id,
                                                );
                                                
                                                // TODO: Implement proper delay-based retry scheduling
                                                // For now, just nack and requeue
                                                if let Err(e) = channel_clone
                                                    .basic_nack(
                                                        delivery.delivery_tag,
                                                        lapin::options::BasicNackOptions {
                                                            multiple: false,
                                                            requeue: true, // Simple requeue for now
                                                        },
                                                    )
                                                    .await
                                                {
                                                    error!("Failed to nack message for retry: {}", e);
                                                }
                                            } else {
                                                // No more retries, send to DLQ
                                                Self::send_to_dlq(&envelope, retry_cfg, &connection, &queue_name).await;
                                                
                                                // ACK original message
                                                if let Err(e) = channel_clone
                                                    .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
                                                    .await
                                                {
                                                    error!("Failed to ack message after DLQ: {}", e);
                                                }
                                            }
                                        } else {
                                            // Retry exhausted, send to DLQ
                                            warn!("Retry exhausted for envelope {}", envelope.metadata.message_id);
                                            Self::send_to_dlq(&envelope, retry_cfg, &connection, &queue_name).await;
                                            
                                            // ACK original message
                                            if let Err(e) = channel_clone
                                                .basic_ack(delivery.delivery_tag, BasicAckOptions::default())
                                                .await
                                            {
                                                error!("Failed to ack message after DLQ: {}", e);
                                            }
                                        }
                                    } else {
                                        // No retry config, just nack
                                        if let Err(e) = channel_clone
                                            .basic_nack(
                                                delivery.delivery_tag,
                                                lapin::options::BasicNackOptions {
                                                    multiple: false,
                                                    requeue: false,
                                                },
                                            )
                                            .await
                                        {
                                            error!("Failed to nack message: {}", e);
                                        }
                                    }
                                }
                            }
                        }
                    }
                    Err(e) => {
                        error!("Failed to deserialize message envelope: {}", e);
                        if auto_ack {
                            // Reject malformed messages
                            if let Err(e) = channel_clone
                                .basic_nack(
                                    delivery.delivery_tag,
                                    lapin::options::BasicNackOptions {
                                        multiple: false,
                                        requeue: false,
                                    },
                                )
                                .await
                            {
                                error!("Failed to nack malformed envelope: {}", e);
                            }
                        }
                    }
                }
            });
        }

        Ok(())
    }

    /// Send failed message to Dead Letter Queue
    async fn send_to_dlq<T>(
        envelope: &MessageEnvelope<T>, 
        retry_config: &RetryConfig,
        connection: &Arc<Connection>,
        queue_name: &str,
    ) where
        T: serde::Serialize,
    {
        match connection.create_channel().await {
            Ok(dlq_channel) => {
                let dlq_name = retry_config.get_dead_letter_queue(queue_name);
                
                // Declare DLQ
                if let Err(e) = dlq_channel
                    .queue_declare(
                        &dlq_name,
                        QueueDeclareOptions {
                            durable: true,
                            ..Default::default()
                        },
                        FieldTable::default(),
                    )
                    .await
                {
                    error!("Failed to declare DLQ {}: {}", dlq_name, e);
                    return;
                }

                // Publish to DLQ with failure summary
                let failure_summary = envelope.get_failure_summary();
                let dlq_payload = serde_json::json!({
                    "envelope": envelope,
                    "failure_summary": failure_summary,
                    "sent_to_dlq_at": chrono::Utc::now(),
                });

                if let Ok(payload_bytes) = serde_json::to_vec(&dlq_payload) {
                    if let Err(e) = dlq_channel
                        .basic_publish(
                            "",
                            &dlq_name,
                            lapin::options::BasicPublishOptions::default(),
                            &payload_bytes,
                            lapin::BasicProperties::default(),
                        )
                        .await
                    {
                        error!("Failed to publish to DLQ {}: {}", dlq_name, e);
                    } else {
                        warn!("Sent envelope {} to DLQ: {}", envelope.metadata.message_id, failure_summary);
                    }
                }
            }
            Err(e) => {
                error!("Failed to create DLQ channel: {}", e);
            }
        }
    }
}

/// Classify error type based on error message (simplified heuristics)
fn classify_error(error: &(dyn std::error::Error + Send + Sync)) -> ErrorType {
    let error_msg = error.to_string().to_lowercase();
    
    if error_msg.contains("timeout") 
        || error_msg.contains("connection") 
        || error_msg.contains("network") 
        || error_msg.contains("temporary") {
        ErrorType::Transient
    } else if error_msg.contains("rate limit") 
        || error_msg.contains("quota") 
        || error_msg.contains("resource") {
        ErrorType::Resource
    } else if error_msg.contains("validation") 
        || error_msg.contains("authentication") 
        || error_msg.contains("authorization") 
        || error_msg.contains("invalid") 
        || error_msg.contains("bad request") {
        ErrorType::Permanent
    } else {
        ErrorType::Unknown
    }
}