Skip to main content

shared_framework/queue/
mod.rs

1//! RabbitMQ messaging over lapin: durable queues with optional retry and dead-letter queues.
2//!
3//! [`QueueConsumer`] receives JSON-encoded messages of type `T`, while
4//! [`MessageProducer`] publishes them. [`PublishOptions`] controls durability,
5//! prefetch, retry, and dead-letter behavior.
6//!
7//! Queue names: the main queue is `<name>`, the retry queue is `<name>.retry`
8//! (dead-letters back after the TTL), and the dead-letter queue is
9//! `<name>.dlq`. Failed messages go to the retry queue while the retry count
10//! is below `max_retries`, then to the dead-letter queue when enabled.
11//!
12//! Use this module when work should be processed asynchronously by another task.
13//!
14//! ```ignore
15//! # use crate::queue::{MessageProducer, PublishOptions, QueueConsumer};
16//! # async fn example() -> anyhow::Result<()> {
17//! let producer = MessageProducer::new("amqp://guest:guest@localhost:5672").await?;
18//! producer.publish("jobs", &serde_json::json!({"id": 1})).await?;
19//!
20//! let consumer = QueueConsumer::<serde_json::Value>::connect(
21//!     "amqp://guest:guest@localhost:5672",
22//!     "jobs",
23//!     PublishOptions::default(),
24//! )
25//! .await?;
26//! consumer.declare().await?;
27//! # Ok(())
28//! # }
29//! ```
30
31use futures::StreamExt;
32use lapin::{options::*, types::FieldTable, BasicProperties, Channel, Connection, ConnectionProperties};
33use serde::{de::DeserializeOwned, Serialize};
34
35/// Options controlling queue declaration and failure handling.
36///
37/// `T` elsewhere is the JSON message type; these options apply to any message type.
38#[derive(Debug, Clone)]
39pub struct PublishOptions {
40    /// Declare queues as durable. Defaults to `true`.
41    pub persistent: bool,
42    /// Maximum unacknowledged messages per consumer. Defaults to `16`.
43    pub prefetch_count: u16,
44    /// Declare and use the `<name>.retry` queue. Defaults to `false`.
45    pub use_retry_queue: bool,
46    /// Declare and use the `<name>.dlq` queue. Defaults to `false`.
47    pub use_dead_letter: bool,
48    /// Times a message is resent to the retry queue before dead-lettering. Defaults to `5`.
49    pub max_retries: usize,
50    /// Delay before a retry-queue message is redelivered. Defaults to `30_000` ms.
51    pub retry_interval_ms: u64,
52    /// Per-message handler timeout; `None` means no timeout. Defaults to `None`.
53    pub processing_timeout_ms: Option<u64>,
54}
55
56impl Default for PublishOptions {
57    fn default() -> Self {
58        Self {
59            persistent: true,
60            prefetch_count: 16,
61            use_retry_queue: false,
62            use_dead_letter: false,
63            max_retries: 5,
64            retry_interval_ms: 30_000,
65            processing_timeout_ms: None,
66        }
67    }
68}
69
70/// Lapin-backed consumer for JSON messages of type `T`.
71///
72/// `T` is the message payload deserialized from each delivery body. The
73/// consumer holds one channel bound to `queue_name` with the given options.
74pub struct QueueConsumer<T> {
75    channel: Channel,
76    queue_name: String,
77    options: PublishOptions,
78    _marker: std::marker::PhantomData<T>,
79}
80
81impl<T> QueueConsumer<T>
82where
83    T: DeserializeOwned + Send + Sync + 'static,
84{
85    /// Connects to the broker, opens a channel, and applies the prefetch count.
86    ///
87    /// Returns an error if the connection, channel, or QoS setup fails.
88    pub async fn connect(addr: &str, queue_name: &str, options: PublishOptions) -> anyhow::Result<Self> {
89        let conn = Connection::connect(addr, ConnectionProperties::default()).await?;
90        let channel = conn.create_channel().await?;
91        channel.basic_qos(options.prefetch_count, BasicQosOptions::default()).await?;
92        Ok(Self { channel, queue_name: queue_name.to_string(), options, _marker: std::marker::PhantomData })
93    }
94
95    /// Declares the durable main queue plus the retry and dead-letter queues when enabled.
96    ///
97    /// The retry queues dead-letters back to the main queue after
98    /// `retry_interval_ms`. Returns an error if any declaration fails.
99    pub async fn declare(&self) -> anyhow::Result<()> {
100        let args = FieldTable::default();
101        if self.options.use_retry_queue {
102            // dead-letter back to main queue via retry queue ttl
103        }
104        self.channel
105            .queue_declare(self.queue_name.clone().into(), QueueDeclareOptions { durable: true, ..Default::default() }, args)
106            .await?;
107        if self.options.use_retry_queue {
108            let retry = format!("{}.retry", self.queue_name);
109            let mut retry_args = FieldTable::default();
110            retry_args.insert("x-dead-letter-exchange".into(), lapin::types::AMQPValue::LongString("".into()));
111            retry_args.insert("x-dead-letter-routing-key".into(), lapin::types::AMQPValue::LongString(self.queue_name.clone().into()));
112            retry_args.insert("x-message-ttl".into(), lapin::types::AMQPValue::LongLongInt(self.options.retry_interval_ms as i64));
113            self.channel.queue_declare(retry.into(), QueueDeclareOptions { durable: true, ..Default::default() }, retry_args).await?;
114        }
115        if self.options.use_dead_letter {
116            let dlq = format!("{}.dlq", self.queue_name);
117            self.channel.queue_declare(dlq.into(), QueueDeclareOptions { durable: true, ..Default::default() }, FieldTable::default()).await?;
118        }
119        Ok(())
120    }
121
122    /// Consumes deliveries in a loop, acknowledging successes and rerouting failures.
123    ///
124    /// `handler` receives each decoded message plus its `redelivered` flag and
125    /// returns success or failure. Successful handlers are acknowledged;
126    /// handler errors, timeouts, and undecodable bodies go through the
127    /// retry/dead-letter policy. `F` is the handler closure type and `Fut` its
128    /// returned future. Runs until the broker stream ends.
129    pub async fn consume<F, Fut>(&self, mut handler: F) -> anyhow::Result<()>
130    where
131        F: FnMut(T, bool) -> Fut + Send + 'static,
132        Fut: Future<Output = anyhow::Result<()>> + Send,
133    {
134        let mut consumer = self.channel.basic_consume(self.queue_name.clone().into(), format!("consumer-{}", uuid::Uuid::new_v4()).into(), BasicConsumeOptions::default(), FieldTable::default()).await?;
135        while let Some(delivery) = consumer.next().await {
136            if let Ok(delivery) = delivery {
137                let data: Result<T, _> = serde_json::from_slice(&delivery.data);
138                match data {
139                    Ok(msg) => {
140                        let redelivered = delivery.redelivered;
141                        let res = if let Some(timeout) = self.options.processing_timeout_ms {
142                            tokio::time::timeout(std::time::Duration::from_millis(timeout), handler(msg, redelivered)).await
143                        } else {
144                            Ok(handler(msg, redelivered).await)
145                        };
146                        match res {
147                            Ok(Ok(_)) => { let _ = delivery.ack(BasicAckOptions::default()).await; }
148                            _ => { let _ = self.handle_nack(&delivery).await; }
149                        }
150                    }
151                    Err(_) => { let _ = self.handle_nack(&delivery).await; }
152                }
153            }
154        }
155        Ok(())
156    }
157
158    async fn handle_nack(&self, delivery: &lapin::message::Delivery) -> anyhow::Result<()> {
159        let retry_count: i64 = delivery.properties.headers().as_ref()
160            .and_then(|h| h.inner().get("x-retry-count"))
161            .and_then(|v| match v { lapin::types::AMQPValue::LongLongInt(i) => Some(*i), _ => None })
162            .unwrap_or(0);
163        if self.options.use_retry_queue && (retry_count as usize) < self.options.max_retries {
164            let mut headers = FieldTable::default();
165            headers.insert("x-retry-count".into(), lapin::types::AMQPValue::LongLongInt(retry_count + 1));
166            let props = BasicProperties::default().with_headers(headers).with_delivery_mode(2);
167            self.channel.basic_publish("".into(), format!("{}.retry", self.queue_name).into(), BasicPublishOptions::default(), &delivery.data, props).await?;
168            delivery.ack(BasicAckOptions::default()).await?;
169        } else if self.options.use_dead_letter {
170            let props = BasicProperties::default().with_delivery_mode(2);
171            self.channel.basic_publish("".into(), format!("{}.dlq", self.queue_name).into(), BasicPublishOptions::default(), &delivery.data, props).await?;
172            delivery.ack(BasicAckOptions::default()).await?;
173        } else {
174            delivery.nack(BasicNackOptions { requeue: false, multiple: false }).await?;
175        }
176        Ok(())
177    }
178}
179
180/// Publishes JSON-encoded messages to a broker over one lapin channel.
181pub struct MessageProducer {
182    channel: Channel,
183}
184
185impl MessageProducer {
186    /// Connects to the broker and opens a channel for publishing.
187    ///
188    /// Returns an error if the connection or channel cannot be created.
189    pub async fn new(addr: &str) -> anyhow::Result<Self> {
190        let conn = Connection::connect(addr, ConnectionProperties::default()).await?;
191        Ok(Self { channel: conn.create_channel().await? })
192    }
193
194    /// Publishes a JSON-encoded payload to `routing_key` as a persistent message.
195    ///
196    /// `T` is the payload type serialized to JSON. Returns an error if
197    /// serialization or publishing fails.
198    pub async fn publish<T: Serialize>(&self, routing_key: &str, payload: &T) -> anyhow::Result<()> {
199        let body = serde_json::to_vec(payload)?;
200        self.channel.basic_publish("".into(), routing_key.to_string().into(), BasicPublishOptions::default(), &body, BasicProperties::default().with_delivery_mode(2)).await?;
201        Ok(())
202    }
203}