Skip to main content

ruststream_lapin/
broker.rs

1//! The broker handle: connection lifecycle, subscriptions, and publisher constructors.
2
3use 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/// The live connection plus the shared fire-and-forget publish channel.
22#[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
38/// The connection cell shared by the broker and everything it hands out, so publishers obtained
39/// before `Broker::connect` resolve the connection on first use.
40pub(crate) type SharedConn = Arc<OnceCell<ConnState>>;
41
42/// A `RabbitMQ` broker backed by [`lapin`](https://docs.rs/lapin).
43///
44/// Follows the `RustStream` lazy startup contract: [`new`](Self::new) is synchronous and does no
45/// I/O; the network work happens in the idempotent async `Broker::connect`, which the runtime
46/// calls once at startup. Publishers handed out earlier share the connection cell and resolve it
47/// on first use.
48///
49/// By default the broker never creates infrastructure: descriptors describe the EXPECTED
50/// topology, and a missing queue is a subscribe error. Opt into declaration with
51/// [`declare_topology(true)`](Self::declare_topology).
52///
53/// # Examples
54///
55/// ```no_run
56/// use ruststream_lapin::{LapinBroker, QueueType};
57///
58/// let broker = LapinBroker::new("amqp://localhost:5672")
59///     .prefetch(64)
60///     .default_queue_type(QueueType::Quorum);
61/// # let _ = broker;
62/// ```
63#[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    /// Records the connection URI; no I/O happens until `Broker::connect`.
75    ///
76    /// The URI carries credentials, virtual host, and TLS scheme:
77    /// `amqp://user:pass@host:5672/vhost` (or `amqps://` with a TLS feature enabled).
78    #[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    /// Connects eagerly: [`new`](Self::new) followed by `Broker::connect`.
91    ///
92    /// # Errors
93    ///
94    /// Returns [`AmqpError::Connect`] when the connection cannot be established.
95    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    /// A connection name shown in the `RabbitMQ` management UI.
102    #[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    /// Caps unacknowledged deliveries in flight per subscription (`basic.qos`).
109    ///
110    /// This is the back-pressure window for subscriber streams; individual queue descriptors
111    /// can override it. Without it the server imposes no prefetch limit.
112    #[must_use]
113    pub fn prefetch(mut self, prefetch: u16) -> Self {
114        self.prefetch = Some(prefetch);
115        self
116    }
117
118    /// Whether subscribing declares the descriptor's expected topology first. Defaults to
119    /// `false`: managing infrastructure is the user's job, so creation is a deliberate opt-in.
120    ///
121    /// When enabled, subscribing declares the bound exchanges (except the built-in `amq.*`
122    /// ones and the default exchange), the queue, and the bindings.
123    #[must_use]
124    pub fn declare_topology(mut self, declare: bool) -> Self {
125        self.declare = declare;
126        self
127    }
128
129    /// The queue type declared for descriptors that do not set one.
130    ///
131    /// Only consulted when [`declare_topology`](Self::declare_topology) is enabled. Without a
132    /// broker default or a per-queue type, no `x-queue-type` argument is sent and the server
133    /// default applies.
134    #[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    /// Opens a subscription for `def`, declaring its topology first when the broker opted in.
145    ///
146    /// # Errors
147    ///
148    /// Returns [`AmqpError::NotConnected`] before `Broker::connect`, [`AmqpError::Declare`] when
149    /// opted-in declaration fails, [`AmqpError::InvalidOptions`] for contradictory descriptor
150    /// options, and [`AmqpError::Subscribe`] when the channel or consumer cannot be opened (for
151    /// example the queue does not exist and declaration was not opted into).
152    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    /// A fire-and-forget publisher on the shared publish channel.
185    ///
186    /// Upgrade with [`confirms`](LapinPublisher::confirms) or
187    /// [`server_tx`](LapinPublisher::server_tx) for transactional publishing.
188    #[must_use]
189    pub fn publisher(&self) -> LapinPublisher {
190        LapinPublisher::new(Arc::clone(&self.conn))
191    }
192
193    /// A request/reply client over `RabbitMQ` direct reply-to.
194    #[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    /// Establishes the connection and the shared publish channel; idempotent.
204    ///
205    /// # Errors
206    ///
207    /// Returns [`AmqpError::Connect`] when the URI cannot be parsed or the connection fails.
208    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    /// Closes the connection; further operations fail with [`AmqpError::NotConnected`] or a
232    /// channel error. Idempotent: closing an already-closed connection succeeds.
233    ///
234    /// # Errors
235    ///
236    /// Returns [`AmqpError::Connect`] when the close handshake fails.
237    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// `Self::subscribe` inside this impl would resolve to the trait method and recurse; the type
252// name is the only way to reach the inherent one.
253#[allow(clippy::use_self)]
254impl Subscribe for LapinBroker {
255    type Subscriber = LapinSubscriber;
256
257    /// Subscribes to the queue `name` with descriptor defaults (durable, shared).
258    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
269/// Extracts the `host[:port]` part of an AMQP URI for `AsyncAPI` metadata; never fails, because
270/// metadata must not block startup on a URI the connection itself will reject anyway.
271fn 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        // The default exchange and the amq.* built-ins exist on every broker and must not be
287        // redeclared.
288        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}