shared_framework/queue/
mod.rs1use futures::StreamExt;
32use lapin::{options::*, types::FieldTable, BasicProperties, Channel, Connection, ConnectionProperties};
33use serde::{de::DeserializeOwned, Serialize};
34
35#[derive(Debug, Clone)]
39pub struct PublishOptions {
40 pub persistent: bool,
42 pub prefetch_count: u16,
44 pub use_retry_queue: bool,
46 pub use_dead_letter: bool,
48 pub max_retries: usize,
50 pub retry_interval_ms: u64,
52 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
70pub 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 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 pub async fn declare(&self) -> anyhow::Result<()> {
100 let args = FieldTable::default();
101 if self.options.use_retry_queue {
102 }
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 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
180pub struct MessageProducer {
182 channel: Channel,
183}
184
185impl MessageProducer {
186 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 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}