queuerious_lapin/

consumer.rs

//! Tracked consumer wrapping lapin::Consumer.

use std::sync::Arc;

use futures_util::StreamExt;
use queuerious::{Backend, JobEvent, QueuriousClient};

use crate::config::LapinAdapterConfig;
use crate::delivery::TrackedDelivery;

#[cfg(feature = "metrics")]
use std::collections::HashMap;

#[cfg(feature = "metrics")]
use crate::metrics::SlidingWindowCounters;

/// A wrapped `lapin::Consumer` that automatically reports job lifecycle events.
///
/// Each delivery is wrapped in a [`TrackedDelivery`] that reports `job_started`
/// when received (via a direct API call to get the server-assigned `execution_id`),
/// and `job_completed`/`job_failed` when acked/nacked.
///
/// Use [`next_tracked()`](Self::next_tracked) to consume deliveries with full
/// lifecycle tracking.
///
/// # Example
///
/// ```no_run
/// use queuerious::QueuriousClient;
/// use queuerious_lapin::TrackedConsumer;
///
/// # async fn example(channel: &lapin::Channel) -> Result<(), Box<dyn std::error::Error>> {
/// let client = std::sync::Arc::new(
///     QueuriousClient::builder()
///         .api_key("qk_your_key")
///         .build()?
/// );
///
/// let consumer = channel.basic_consume(
///     "my-queue", "my-consumer",
///     lapin::options::BasicConsumeOptions::default(),
///     lapin::types::FieldTable::default(),
/// ).await?;
///
/// let mut tracked = TrackedConsumer::new(consumer, "my-queue", client);
///
/// while let Some(delivery) = tracked.next_tracked().await {
///     let delivery = delivery?;
///     // Process message...
///     delivery.ack(lapin::options::BasicAckOptions::default()).await?;
/// }
/// # Ok(())
/// # }
/// ```
pub struct TrackedConsumer {
    inner: lapin::Consumer,
    client: Arc<QueuriousClient>,
    queue_name: String,
    config: LapinAdapterConfig,
    #[cfg(feature = "metrics")]
    counters: Option<Arc<HashMap<String, SlidingWindowCounters>>>,
}

impl TrackedConsumer {
    /// Wrap a lapin consumer for automatic event tracking.
    pub fn new(
        consumer: lapin::Consumer,
        queue_name: impl Into<String>,
        client: Arc<QueuriousClient>,
    ) -> Self {
        Self::with_config(consumer, queue_name, client, LapinAdapterConfig::default())
    }

    /// Wrap a lapin consumer with custom configuration.
    pub fn with_config(
        consumer: lapin::Consumer,
        queue_name: impl Into<String>,
        client: Arc<QueuriousClient>,
        config: LapinAdapterConfig,
    ) -> Self {
        Self {
            inner: consumer,
            client,
            queue_name: queue_name.into(),
            config,
            #[cfg(feature = "metrics")]
            counters: None,
        }
    }

    /// Attach shared metrics counters for Tier 0 data collection.
    ///
    /// When set, each [`TrackedDelivery`] produced by this consumer will
    /// record processing duration and success/failure to the counters map
    /// under this consumer's queue name.
    ///
    /// The same counters map should be shared with
    /// [`RabbitMqMetricsCollector`](crate::metrics::RabbitMqMetricsCollector).
    #[cfg(feature = "metrics")]
    pub fn with_counters(mut self, counters: Arc<HashMap<String, SlidingWindowCounters>>) -> Self {
        self.counters = Some(counters);
        self
    }

    /// Get a reference to the shared counters, if set.
    #[cfg(feature = "metrics")]
    pub fn counters(&self) -> Option<&Arc<HashMap<String, SlidingWindowCounters>>> {
        self.counters.as_ref()
    }

    /// Process the next delivery with full lifecycle tracking.
    ///
    /// Reports `job_started` via a direct API call and obtains the server-assigned
    /// `execution_id`, enabling accurate `job_completed`/`job_failed` tracking
    /// when the returned [`TrackedDelivery`] is acked or nacked.
    ///
    /// Returns `None` when the consumer stream ends.
    pub async fn next_tracked(&mut self) -> Option<Result<TrackedDelivery, TrackedConsumerError>> {
        let delivery = match self.inner.next().await {
            Some(Ok(d)) => d,
            Some(Err(e)) => return Some(Err(TrackedConsumerError::Lapin(e))),
            None => return None,
        };

        // Parse the JSON payload once — shared across job_id, correlation_id,
        // and event payload extraction to avoid redundant deserialization.
        let parsed_payload: Option<serde_json::Value> = serde_json::from_slice(&delivery.data).ok();

        let job_id = self
            .config
            .job_id_extractor
            .as_ref()
            .map(|f| f(&delivery))
            .unwrap_or_else(|| crate::config::extract_job_id(&delivery, parsed_payload.as_ref()));

        let job_type = self
            .config
            .job_type_extractor
            .as_ref()
            .map(|f| f(&delivery))
            .unwrap_or_else(|| {
                crate::config::extract_job_type(
                    &delivery.properties,
                    &delivery.routing_key,
                    &self.config.default_job_type,
                )
            });

        let payload = if self.config.include_payload {
            if delivery.data.len() > self.config.max_payload_size {
                serde_json::json!({
                    "_queuerious": "payload_too_large",
                    "message": format!(
                        "Payload not stored ({} bytes exceeds {} byte limit)",
                        delivery.data.len(),
                        self.config.max_payload_size
                    ),
                    "original_size_bytes": delivery.data.len()
                })
            } else {
                parsed_payload.clone().unwrap_or(serde_json::Value::Null)
            }
        } else {
            serde_json::Value::Object(serde_json::Map::new())
        };

        // Extract metadata from AMQP properties
        let metadata = self
            .config
            .metadata_extractor
            .as_ref()
            .map(|f| f(&delivery))
            .or_else(|| {
                if self.config.include_default_metadata {
                    let meta = crate::config::extract_default_metadata(&delivery);
                    // Only include if there's actual metadata
                    if meta.as_object().is_none_or(|m| m.is_empty()) {
                        None
                    } else {
                        Some(meta)
                    }
                } else {
                    None
                }
            });

        let mut builder =
            JobEvent::started(&self.queue_name, Backend::RabbitMQ, &job_id, &job_type)
                .payload(payload);

        if let Some(meta) = metadata {
            builder = builder.metadata(meta);
        }

        // Extract correlation_id: custom extractor -> default chain
        let correlation_id = self
            .config
            .correlation_id_extractor
            .as_ref()
            .and_then(|f| f(&delivery))
            .or_else(|| crate::config::extract_correlation_id(&delivery, parsed_payload.as_ref()));

        if let Some(cid) = correlation_id {
            builder = builder.correlation_id(cid);
        }

        if let Some(max) = self.config.max_attempts {
            builder = builder.max_attempts(max);
        }

        // Propagate retry count from AMQP headers as attempt_number so the
        // server records the correct attempt instead of always starting at 1.
        if let Some(retry_count) = crate::config::extract_retry_count(&delivery) {
            builder = builder.attempt_number(retry_count + 1);
        }

        let event = builder.build();

        // Build the TrackedDelivery with optional metrics support.
        let make_tracked = |delivery, execution_id, created_at| {
            let tracked = TrackedDelivery::new(
                delivery,
                self.client.clone(),
                self.queue_name.clone(),
                execution_id,
                created_at,
            );

            #[cfg(feature = "metrics")]
            let tracked = if let Some(ref counters) = self.counters {
                tracked.with_counters(counters.clone())
            } else {
                tracked
            };

            tracked
        };

        // Send directly and get execution_id from the response.
        match self.client.report_single(event).await {
            Ok(response) => {
                let (execution_id, created_at) = if let Some(result) = response.results.first() {
                    (result.execution_id, result.created_at)
                } else {
                    tracing::warn!(
                        "Ingest response missing results; lifecycle tracking unavailable"
                    );
                    (uuid::Uuid::nil(), chrono::Utc::now())
                };

                Some(Ok(make_tracked(delivery, execution_id, created_at)))
            }
            Err(e) => {
                tracing::warn!(error = %e, "Failed to report job_started; proceeding without tracking");
                Some(Ok(make_tracked(
                    delivery,
                    uuid::Uuid::nil(),
                    chrono::Utc::now(),
                )))
            }
        }
    }
}

/// Errors from the tracked consumer.
#[derive(Debug, thiserror::Error)]
pub enum TrackedConsumerError {
    /// Lapin error.
    #[error("Lapin error: {0}")]
    Lapin(#[from] lapin::Error),
}