shared-framework 0.0.17

//! RabbitMQ messaging over lapin: durable queues with optional retry and dead-letter queues.
//!
//! [`QueueConsumer`] receives JSON-encoded messages of type `T`, while
//! [`MessageProducer`] publishes them. [`PublishOptions`] controls durability,
//! prefetch, retry, and dead-letter behavior.
//!
//! Queue names: the main queue is `<name>`, the retry queue is `<name>.retry`
//! (dead-letters back after the TTL), and the dead-letter queue is
//! `<name>.dlq`. Failed messages go to the retry queue while the retry count
//! is below `max_retries`, then to the dead-letter queue when enabled.
//!
//! Use this module when work should be processed asynchronously by another task.
//!
//! ```ignore
//! # use crate::queue::{MessageProducer, PublishOptions, QueueConsumer};
//! # async fn example() -> anyhow::Result<()> {
//! let producer = MessageProducer::new("amqp://guest:guest@localhost:5672").await?;
//! producer.publish("jobs", &serde_json::json!({"id": 1})).await?;
//!
//! let consumer = QueueConsumer::<serde_json::Value>::connect(
//!     "amqp://guest:guest@localhost:5672",
//!     "jobs",
//!     PublishOptions::default(),
//! )
//! .await?;
//! consumer.declare().await?;
//! # Ok(())
//! # }
//! ```

use futures::StreamExt;
use lapin::{options::*, types::FieldTable, BasicProperties, Channel, Connection, ConnectionProperties};
use serde::{de::DeserializeOwned, Serialize};

/// Options controlling queue declaration and failure handling.
///
/// `T` elsewhere is the JSON message type; these options apply to any message type.
#[derive(Debug, Clone)]
pub struct PublishOptions {
    /// Declare queues as durable. Defaults to `true`.
    pub persistent: bool,
    /// Maximum unacknowledged messages per consumer. Defaults to `16`.
    pub prefetch_count: u16,
    /// Declare and use the `<name>.retry` queue. Defaults to `false`.
    pub use_retry_queue: bool,
    /// Declare and use the `<name>.dlq` queue. Defaults to `false`.
    pub use_dead_letter: bool,
    /// Times a message is resent to the retry queue before dead-lettering. Defaults to `5`.
    pub max_retries: usize,
    /// Delay before a retry-queue message is redelivered. Defaults to `30_000` ms.
    pub retry_interval_ms: u64,
    /// Per-message handler timeout; `None` means no timeout. Defaults to `None`.
    pub processing_timeout_ms: Option<u64>,
}

impl Default for PublishOptions {
    fn default() -> Self {
        Self {
            persistent: true,
            prefetch_count: 16,
            use_retry_queue: false,
            use_dead_letter: false,
            max_retries: 5,
            retry_interval_ms: 30_000,
            processing_timeout_ms: None,
        }
    }
}

/// Lapin-backed consumer for JSON messages of type `T`.
///
/// `T` is the message payload deserialized from each delivery body. The
/// consumer holds one channel bound to `queue_name` with the given options.
pub struct QueueConsumer<T> {
    channel: Channel,
    queue_name: String,
    options: PublishOptions,
    _marker: std::marker::PhantomData<T>,
}

impl<T> QueueConsumer<T>
where
    T: DeserializeOwned + Send + Sync + 'static,
{
    /// Connects to the broker, opens a channel, and applies the prefetch count.
    ///
    /// Returns an error if the connection, channel, or QoS setup fails.
    pub async fn connect(addr: &str, queue_name: &str, options: PublishOptions) -> anyhow::Result<Self> {
        let conn = Connection::connect(addr, ConnectionProperties::default()).await?;
        let channel = conn.create_channel().await?;
        channel.basic_qos(options.prefetch_count, BasicQosOptions::default()).await?;
        Ok(Self { channel, queue_name: queue_name.to_string(), options, _marker: std::marker::PhantomData })
    }

    /// Declares the durable main queue plus the retry and dead-letter queues when enabled.
    ///
    /// The retry queues dead-letters back to the main queue after
    /// `retry_interval_ms`. Returns an error if any declaration fails.
    pub async fn declare(&self) -> anyhow::Result<()> {
        let args = FieldTable::default();
        if self.options.use_retry_queue {
            // dead-letter back to main queue via retry queue ttl
        }
        self.channel
            .queue_declare(self.queue_name.clone().into(), QueueDeclareOptions { durable: true, ..Default::default() }, args)
            .await?;
        if self.options.use_retry_queue {
            let retry = format!("{}.retry", self.queue_name);
            let mut retry_args = FieldTable::default();
            retry_args.insert("x-dead-letter-exchange".into(), lapin::types::AMQPValue::LongString("".into()));
            retry_args.insert("x-dead-letter-routing-key".into(), lapin::types::AMQPValue::LongString(self.queue_name.clone().into()));
            retry_args.insert("x-message-ttl".into(), lapin::types::AMQPValue::LongLongInt(self.options.retry_interval_ms as i64));
            self.channel.queue_declare(retry.into(), QueueDeclareOptions { durable: true, ..Default::default() }, retry_args).await?;
        }
        if self.options.use_dead_letter {
            let dlq = format!("{}.dlq", self.queue_name);
            self.channel.queue_declare(dlq.into(), QueueDeclareOptions { durable: true, ..Default::default() }, FieldTable::default()).await?;
        }
        Ok(())
    }

    /// Consumes deliveries in a loop, acknowledging successes and rerouting failures.
    ///
    /// `handler` receives each decoded message plus its `redelivered` flag and
    /// returns success or failure. Successful handlers are acknowledged;
    /// handler errors, timeouts, and undecodable bodies go through the
    /// retry/dead-letter policy. `F` is the handler closure type and `Fut` its
    /// returned future. Runs until the broker stream ends.
    pub async fn consume<F, Fut>(&self, mut handler: F) -> anyhow::Result<()>
    where
        F: FnMut(T, bool) -> Fut + Send + 'static,
        Fut: Future<Output = anyhow::Result<()>> + Send,
    {
        let mut consumer = self.channel.basic_consume(self.queue_name.clone().into(), format!("consumer-{}", uuid::Uuid::new_v4()).into(), BasicConsumeOptions::default(), FieldTable::default()).await?;
        while let Some(delivery) = consumer.next().await {
            if let Ok(delivery) = delivery {
                let data: Result<T, _> = serde_json::from_slice(&delivery.data);
                match data {
                    Ok(msg) => {
                        let redelivered = delivery.redelivered;
                        let res = if let Some(timeout) = self.options.processing_timeout_ms {
                            tokio::time::timeout(std::time::Duration::from_millis(timeout), handler(msg, redelivered)).await
                        } else {
                            Ok(handler(msg, redelivered).await)
                        };
                        match res {
                            Ok(Ok(_)) => { let _ = delivery.ack(BasicAckOptions::default()).await; }
                            _ => { let _ = self.handle_nack(&delivery).await; }
                        }
                    }
                    Err(_) => { let _ = self.handle_nack(&delivery).await; }
                }
            }
        }
        Ok(())
    }

    async fn handle_nack(&self, delivery: &lapin::message::Delivery) -> anyhow::Result<()> {
        let retry_count: i64 = delivery.properties.headers().as_ref()
            .and_then(|h| h.inner().get("x-retry-count"))
            .and_then(|v| match v { lapin::types::AMQPValue::LongLongInt(i) => Some(*i), _ => None })
            .unwrap_or(0);
        if self.options.use_retry_queue && (retry_count as usize) < self.options.max_retries {
            let mut headers = FieldTable::default();
            headers.insert("x-retry-count".into(), lapin::types::AMQPValue::LongLongInt(retry_count + 1));
            let props = BasicProperties::default().with_headers(headers).with_delivery_mode(2);
            self.channel.basic_publish("".into(), format!("{}.retry", self.queue_name).into(), BasicPublishOptions::default(), &delivery.data, props).await?;
            delivery.ack(BasicAckOptions::default()).await?;
        } else if self.options.use_dead_letter {
            let props = BasicProperties::default().with_delivery_mode(2);
            self.channel.basic_publish("".into(), format!("{}.dlq", self.queue_name).into(), BasicPublishOptions::default(), &delivery.data, props).await?;
            delivery.ack(BasicAckOptions::default()).await?;
        } else {
            delivery.nack(BasicNackOptions { requeue: false, multiple: false }).await?;
        }
        Ok(())
    }
}

/// Publishes JSON-encoded messages to a broker over one lapin channel.
pub struct MessageProducer {
    channel: Channel,
}

impl MessageProducer {
    /// Connects to the broker and opens a channel for publishing.
    ///
    /// Returns an error if the connection or channel cannot be created.
    pub async fn new(addr: &str) -> anyhow::Result<Self> {
        let conn = Connection::connect(addr, ConnectionProperties::default()).await?;
        Ok(Self { channel: conn.create_channel().await? })
    }

    /// Publishes a JSON-encoded payload to `routing_key` as a persistent message.
    ///
    /// `T` is the payload type serialized to JSON. Returns an error if
    /// serialization or publishing fails.
    pub async fn publish<T: Serialize>(&self, routing_key: &str, payload: &T) -> anyhow::Result<()> {
        let body = serde_json::to_vec(payload)?;
        self.channel.basic_publish("".into(), routing_key.to_string().into(), BasicPublishOptions::default(), &body, BasicProperties::default().with_delivery_mode(2)).await?;
        Ok(())
    }
}