1use std::sync::Arc;
4
5use lapin::options::{
6 BasicConsumeOptions, BasicQosOptions, ExchangeDeclareOptions, QueueBindOptions,
7 QueueDeclareOptions,
8};
9use lapin::types::{AMQPValue, FieldTable, ShortString};
10use lapin::{Channel, Connection, ConnectionProperties};
11use ruststream::{Broker, DescribeServer, ServerSpec, Subscribe};
12use tokio::sync::OnceCell;
13
14use crate::convert;
15use crate::error::AmqpError;
16use crate::publisher::LapinPublisher;
17use crate::queue::{QueueType, RabbitQueue};
18use crate::requester::LapinRequester;
19use crate::subscriber::LapinSubscriber;
20
21#[derive(Debug)]
23pub(crate) struct ConnState {
24 connection: Connection,
25 publish_channel: Channel,
26}
27
28impl ConnState {
29 pub(crate) fn connection(&self) -> &Connection {
30 &self.connection
31 }
32
33 pub(crate) fn publish_channel(&self) -> &Channel {
34 &self.publish_channel
35 }
36}
37
38pub(crate) type SharedConn = Arc<OnceCell<ConnState>>;
41
42#[derive(Debug, Clone)]
64pub struct LapinBroker {
65 conn: SharedConn,
66 uri: String,
67 connection_name: Option<String>,
68 prefetch: Option<u16>,
69 declare: bool,
70 default_queue_type: Option<QueueType>,
71}
72
73impl LapinBroker {
74 #[must_use]
79 pub fn new(uri: impl Into<String>) -> Self {
80 Self {
81 conn: Arc::new(OnceCell::new()),
82 uri: uri.into(),
83 connection_name: None,
84 prefetch: None,
85 declare: false,
86 default_queue_type: None,
87 }
88 }
89
90 pub async fn connect(uri: impl Into<String>) -> Result<Self, AmqpError> {
96 let broker = Self::new(uri);
97 Broker::connect(&broker).await?;
98 Ok(broker)
99 }
100
101 #[must_use]
103 pub fn connection_name(mut self, name: impl Into<String>) -> Self {
104 self.connection_name = Some(name.into());
105 self
106 }
107
108 #[must_use]
113 pub fn prefetch(mut self, prefetch: u16) -> Self {
114 self.prefetch = Some(prefetch);
115 self
116 }
117
118 #[must_use]
124 pub fn declare_topology(mut self, declare: bool) -> Self {
125 self.declare = declare;
126 self
127 }
128
129 #[must_use]
135 pub fn default_queue_type(mut self, queue_type: QueueType) -> Self {
136 self.default_queue_type = Some(queue_type);
137 self
138 }
139
140 fn connected(&self) -> Result<&ConnState, AmqpError> {
141 self.conn.get().ok_or(AmqpError::NotConnected)
142 }
143
144 pub async fn subscribe(&self, def: RabbitQueue) -> Result<LapinSubscriber, AmqpError> {
153 let state = self.connected()?;
154 let channel = state
155 .connection
156 .create_channel()
157 .await
158 .map_err(AmqpError::subscribe)?;
159
160 if self.declare {
161 declare_topology(&channel, &def, self.default_queue_type).await?;
162 }
163 if let Some(prefetch) = def.prefetch_or(self.prefetch) {
164 channel
165 .basic_qos(prefetch, BasicQosOptions::default())
166 .await
167 .map_err(AmqpError::subscribe)?;
168 }
169
170 let queue = def.name().to_owned();
171 let consumer = channel
172 .basic_consume(
173 convert::short(&queue, "queue name")?,
174 ShortString::default(),
175 BasicConsumeOptions::default(),
176 FieldTable::default(),
177 )
178 .await
179 .map_err(AmqpError::subscribe)?;
180
181 Ok(LapinSubscriber::new(channel, consumer, queue))
182 }
183
184 #[must_use]
189 pub fn publisher(&self) -> LapinPublisher {
190 LapinPublisher::new(Arc::clone(&self.conn))
191 }
192
193 #[must_use]
195 pub fn requester(&self) -> LapinRequester {
196 LapinRequester::new(Arc::clone(&self.conn))
197 }
198}
199
200impl Broker for LapinBroker {
201 type Error = AmqpError;
202
203 async fn connect(&self) -> Result<(), Self::Error> {
209 self.conn
210 .get_or_try_init(|| async {
211 let mut properties = ConnectionProperties::default();
212 if let Some(name) = &self.connection_name {
213 properties = properties.with_connection_name(name.as_str().into());
214 }
215 let connection = Connection::connect(&self.uri, properties)
216 .await
217 .map_err(AmqpError::connect)?;
218 let publish_channel = connection
219 .create_channel()
220 .await
221 .map_err(AmqpError::connect)?;
222 Ok(ConnState {
223 connection,
224 publish_channel,
225 })
226 })
227 .await?;
228 Ok(())
229 }
230
231 async fn shutdown(&self) -> Result<(), Self::Error> {
238 if let Some(state) = self.conn.get()
239 && state.connection.status().connected()
240 {
241 state
242 .connection
243 .close(200, ShortString::from("OK"))
244 .await
245 .map_err(AmqpError::connect)?;
246 }
247 Ok(())
248 }
249}
250
251#[allow(clippy::use_self)]
254impl Subscribe for LapinBroker {
255 type Subscriber = LapinSubscriber;
256
257 async fn subscribe(&self, name: &str) -> Result<Self::Subscriber, Self::Error> {
259 LapinBroker::subscribe(self, RabbitQueue::new(name)).await
260 }
261}
262
263impl DescribeServer for LapinBroker {
264 fn describe_server(&self) -> ServerSpec {
265 ServerSpec::new(host_of(&self.uri), "amqp")
266 }
267}
268
269fn host_of(uri: &str) -> String {
272 let after_scheme = uri.split_once("://").map_or(uri, |(_, rest)| rest);
273 let after_auth = after_scheme
274 .rsplit_once('@')
275 .map_or(after_scheme, |(_, rest)| rest);
276 let host = after_auth.split(['/', '?']).next().unwrap_or(after_auth);
277 host.to_owned()
278}
279
280async fn declare_topology(
281 channel: &Channel,
282 def: &RabbitQueue,
283 broker_default: Option<QueueType>,
284) -> Result<(), AmqpError> {
285 for (exchange, _) in def.bindings() {
286 if exchange.name().is_empty() || exchange.name().starts_with("amq.") {
289 continue;
290 }
291 channel
292 .exchange_declare(
293 convert::short(exchange.name(), "exchange name")?,
294 exchange.kind().clone(),
295 ExchangeDeclareOptions {
296 durable: exchange.is_durable(),
297 auto_delete: exchange.is_auto_delete(),
298 ..ExchangeDeclareOptions::default()
299 },
300 FieldTable::default(),
301 )
302 .await
303 .map_err(AmqpError::declare)?;
304 }
305
306 let queue_type = def.queue_type_or(broker_default);
307 if queue_type == Some(QueueType::Quorum) && !def.is_durable() {
308 return Err(AmqpError::InvalidOptions(format!(
309 "queue {:?} is a quorum queue and must stay durable; drop `.durable(false)` or pick \
310 `QueueType::Classic`",
311 def.name(),
312 )));
313 }
314
315 let mut arguments = def.declare_arguments().clone();
316 if let Some(queue_type) = queue_type {
317 arguments.insert(
318 ShortString::from("x-queue-type"),
319 AMQPValue::LongString(queue_type.as_str().into()),
320 );
321 }
322 channel
323 .queue_declare(
324 convert::short(def.name(), "queue name")?,
325 QueueDeclareOptions {
326 durable: def.is_durable(),
327 exclusive: def.is_exclusive(),
328 auto_delete: def.is_auto_delete(),
329 ..QueueDeclareOptions::default()
330 },
331 arguments,
332 )
333 .await
334 .map_err(AmqpError::declare)?;
335
336 for (exchange, routing_key) in def.bindings() {
337 channel
338 .queue_bind(
339 convert::short(def.name(), "queue name")?,
340 convert::short(exchange.name(), "exchange name")?,
341 convert::short(routing_key, "routing key")?,
342 QueueBindOptions::default(),
343 FieldTable::default(),
344 )
345 .await
346 .map_err(AmqpError::declare)?;
347 }
348
349 Ok(())
350}
351
352#[cfg(test)]
353mod tests {
354 use super::host_of;
355
356 #[test]
357 fn host_extraction_handles_auth_vhost_and_bare_forms() {
358 assert_eq!(host_of("amqp://localhost:5672"), "localhost:5672");
359 assert_eq!(host_of("amqp://user:pass@rabbit:5672/prod"), "rabbit:5672");
360 assert_eq!(host_of("amqps://rabbit/vhost"), "rabbit");
361 assert_eq!(host_of("rabbit:5672"), "rabbit:5672");
362 }
363}