laminar_connectors/rabbitmq/sink/
mod.rs

1//! RabbitMQ sink: publisher-confirmed, optionally mandatory publishes to one exchange.
2//!
3//! Publishes are pipelined within `max.pending` messages and `max.pending.bytes`; `flush` (and
4//! therefore every checkpoint's pre-commit) waits until the broker has confirmed each one.
5//! A confirm means the broker accepted the message for every queue it routed to; durable
6//! storage additionally needs persistent delivery mode and durable (for node loss, replicated)
7//! queues, which the operator provisions. Any nack, return, or lost outcome ends this sink
8//! generation; replay from the last checkpoint may then publish duplicates (at-least-once).
9
10mod properties;
11
12use std::collections::VecDeque;
13use std::time::Duration;
14
15use arrow_array::RecordBatch;
16use arrow_schema::SchemaRef;
17use async_trait::async_trait;
18use lapin::options::{BasicPublishOptions, ConfirmSelectOptions};
19use lapin::{BasicProperties, Channel, Confirmation, Connection, PublisherConfirm};
20
21use super::config::{DeliveryMode, RabbitMqSinkConfig};
22use super::connection::{close, connect, prepare_exchange};
23use super::metrics::{saturating, SinkMetrics};
24use crate::config::ConnectorConfig;
25use crate::connector::{
26    SinkConnector, SinkConsistency, SinkContract, SinkInputMode, SinkTopology, WriteResult,
27};
28use crate::error::ConnectorError;
29use crate::serde::RecordSerializer;
30use properties::{base_properties, payload_schema, BoundProperties};
31
32/// RabbitMQ AMQP 0-9-1 exchange sink.
33pub struct RabbitMqSink {
34    schema: SchemaRef,
35    registry: Option<prometheus::Registry>,
36    config: Option<RabbitMqSinkConfig>,
37    writer: Option<Writer>,
38    runtime: Option<Runtime>,
39    failure: Option<String>,
40}
41
42struct Writer {
43    serializer: Box<dyn RecordSerializer>,
44    payload_indices: Vec<usize>,
45    base: BasicProperties,
46}
47
48struct Runtime {
49    connection: Connection,
50    channel: Channel,
51    pending: VecDeque<(PublisherConfirm, usize)>,
52    pending_bytes: usize,
53    metrics: SinkMetrics,
54}
55
56/// Outcome tally of one confirm drain.
57#[derive(Default)]
58struct Drained {
59    confirmed: usize,
60    nacked: usize,
61    returned: Option<String>,
62    ambiguous: Option<String>,
63}
64
65impl Drained {
66    fn into_result(self) -> Result<usize, ConnectorError> {
67        if let Some(cause) = self.ambiguous {
68            return Err(ConnectorError::outcome_unknown(
69                format!(
70                    "RabbitMQ publish outcome unknown after {} confirmed: {cause}",
71                    self.confirmed
72                ),
73                true,
74            ));
75        }
76        if let Some(cause) = self.returned {
77            return Err(ConnectorError::ConfigurationError(format!(
78                "a mandatory publish was returned as unroutable ({cause}); bind a queue to the \
79                 exchange for this routing key"
80            )));
81        }
82        if self.nacked > 0 {
83            return Err(ConnectorError::WriteError(format!(
84                "the broker negatively confirmed {} publish(es)",
85                self.nacked
86            )));
87        }
88        Ok(self.confirmed)
89    }
90}
91
92impl Runtime {
93    /// Await confirms oldest-first until at most `keep` messages and `keep_bytes` remain pending.
94    async fn drain(
95        &mut self,
96        keep: usize,
97        keep_bytes: usize,
98        deadline: tokio::time::Instant,
99    ) -> Result<usize, ConnectorError> {
100        let mut drained = Drained::default();
101        while self.pending.len() > keep || self.pending_bytes > keep_bytes {
102            let Some((confirm, bytes)) = self.pending.pop_front() else {
103                break;
104            };
105            self.pending_bytes -= bytes;
106            match tokio::time::timeout_at(deadline, confirm).await {
107                Ok(Ok(Confirmation::Ack(None))) => drained.confirmed += 1,
108                Ok(Ok(Confirmation::Ack(Some(returned)))) => {
109                    self.metrics.returned.inc();
110                    drained.returned.get_or_insert_with(|| {
111                        format!("{} {}", returned.reply_code, returned.reply_text.as_str())
112                    });
113                }
114                Ok(Ok(Confirmation::Nack(_))) => {
115                    self.metrics.nacked.inc();
116                    drained.nacked += 1;
117                }
118                Ok(Ok(Confirmation::NotRequested)) => {
119                    drained
120                        .ambiguous
121                        .get_or_insert_with(|| "the channel is not in confirm mode".to_string());
122                }
123                Ok(Err(error)) => {
124                    self.metrics.ambiguous.inc();
125                    drained.ambiguous.get_or_insert_with(|| error.to_string());
126                }
127                Err(_) => {
128                    self.metrics.ambiguous.inc();
129                    drained
130                        .ambiguous
131                        .get_or_insert_with(|| "no broker confirm before the deadline".to_string());
132                }
133            }
134        }
135        self.metrics.confirmed.inc_by(drained.confirmed as u64);
136        self.metrics.pending.set(saturating(self.pending.len()));
137        self.metrics
138            .pending_bytes
139            .set(saturating(self.pending_bytes));
140        drained.into_result()
141    }
142}
143
144impl RabbitMqSink {
145    /// Create an unopened sink; metrics register on `registry` at open.
146    #[must_use]
147    pub fn new(schema: SchemaRef, registry: Option<&prometheus::Registry>) -> Self {
148        Self {
149            schema,
150            registry: registry.cloned(),
151            config: None,
152            writer: None,
153            runtime: None,
154            failure: None,
155        }
156    }
157
158    fn fail(&mut self, error: ConnectorError) -> ConnectorError {
159        tracing::error!(%error, "RabbitMQ sink generation ended");
160        self.failure.get_or_insert(error.to_string());
161        error
162    }
163
164    fn live(&mut self) -> Result<(&RabbitMqSinkConfig, &Writer, &mut Runtime), ConnectorError> {
165        if let Some(cause) = &self.failure {
166            return Err(ConnectorError::InvalidState {
167                expected: "a live RabbitMQ sink".into(),
168                actual: cause.clone(),
169            });
170        }
171        match (&self.config, &self.writer, self.runtime.as_mut()) {
172            (Some(config), Some(writer), Some(runtime)) => Ok((config, writer, runtime)),
173            _ => Err(ConnectorError::InvalidState {
174                expected: "an open RabbitMQ sink".into(),
175                actual: "open() has not completed".into(),
176            }),
177        }
178    }
179
180    fn writer(parsed: &RabbitMqSinkConfig, input: &SchemaRef) -> Result<Writer, ConnectorError> {
181        let payload = payload_schema(parsed, input)?;
182        crate::serde::schema_contract::validate_writer(parsed.format, &payload)?;
183        let payload_indices = payload
184            .fields()
185            .iter()
186            .map(|field| input.index_of(field.name()))
187            .collect::<Result<_, _>>()
188            .map_err(|error| ConnectorError::SchemaMismatch(error.to_string()))?;
189        Ok(Writer {
190            serializer: crate::serde::create_serializer(parsed.format)?,
191            base: base_properties(parsed, &payload),
192            payload_indices,
193        })
194    }
195
196    async fn publish(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
197        let (config, writer, runtime) = self.live()?;
198        let rows = BoundProperties::bind(config, batch)?;
199        let payload = batch
200            .project(&writer.payload_indices)
201            .map_err(|error| ConnectorError::SchemaMismatch(error.to_string()))?;
202        let records = writer.serializer.serialize(&payload)?;
203        if records.len() != batch.num_rows() {
204            return Err(ConnectorError::Internal(format!(
205                "RabbitMQ serializer returned {} records for {} rows",
206                records.len(),
207                batch.num_rows()
208            )));
209        }
210        let deadline = tokio::time::Instant::now() + config.confirm_timeout;
211        let options = BasicPublishOptions {
212            mandatory: config.mandatory,
213            ..BasicPublishOptions::default()
214        };
215        let mut bytes = 0_u64;
216        for (row, record) in records.iter().enumerate() {
217            let keep = config.max_pending - 1;
218            let keep_bytes = config.max_pending_bytes.saturating_sub(record.len());
219            runtime.drain(keep, keep_bytes, deadline).await?;
220            let confirm = runtime
221                .channel
222                .basic_publish(
223                    config.exchange.as_str().into(),
224                    rows.routing_key(row).into(),
225                    options,
226                    record,
227                    rows.properties(row, &writer.base),
228                )
229                .await
230                .map_err(|error| {
231                    let prior = row > 0 || !runtime.pending.is_empty();
232                    if prior {
233                        ConnectorError::outcome_unknown(format!("RabbitMQ publish: {error}"), true)
234                    } else {
235                        ConnectorError::WriteError(format!("RabbitMQ publish: {error}"))
236                    }
237                })?;
238            runtime.pending.push_back((confirm, record.len()));
239            runtime.pending_bytes += record.len();
240            bytes += record.len() as u64;
241        }
242        runtime.metrics.published.inc_by(records.len() as u64);
243        runtime.metrics.bytes.inc_by(bytes);
244        runtime
245            .metrics
246            .pending
247            .set(saturating(runtime.pending.len()));
248        runtime
249            .metrics
250            .pending_bytes
251            .set(saturating(runtime.pending_bytes));
252        Ok(WriteResult::new(records.len(), bytes))
253    }
254}
255
256#[async_trait]
257impl SinkConnector for RabbitMqSink {
258    fn contract(&self, config: &ConnectorConfig) -> Result<SinkContract, ConnectorError> {
259        let parsed = RabbitMqSinkConfig::from_config(config)?;
260        let consistency = match parsed.delivery_mode {
261            DeliveryMode::Persistent => SinkConsistency::DurableAtLeastOnce,
262            DeliveryMode::Transient => SinkConsistency::Ephemeral,
263        };
264        Ok(SinkContract::new(
265            consistency,
266            SinkTopology::MultiWriter,
267            SinkInputMode::AppendOnly,
268        ))
269    }
270
271    async fn resolve_schema(
272        &mut self,
273        config: &ConnectorConfig,
274        input: SchemaRef,
275    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
276        use crate::schema::resolution::{logical_binding, SchemaDirection, SchemaOrigin};
277        let parsed = RabbitMqSinkConfig::from_config(config)?;
278        Self::writer(&parsed, &input)?;
279        logical_binding(config, SchemaDirection::Sink, SchemaOrigin::Query, &input)
280    }
281
282    async fn open(&mut self, config: &ConnectorConfig) -> Result<(), ConnectorError> {
283        let parsed = RabbitMqSinkConfig::from_config(config)?;
284        let input = config
285            .arrow_schema()
286            .or_else(|| {
287                config
288                    .schema_binding()
289                    .map(|binding| std::sync::Arc::new(binding.logical.clone()))
290            })
291            .unwrap_or_else(|| self.schema.clone());
292        let writer = Self::writer(&parsed, &input)?;
293        let metrics = SinkMetrics::new(&parsed.connection.connection_name, self.registry.as_ref())?;
294        let connection = connect(&parsed.connection).await?;
295        let channel = connection.create_channel().await.map_err(|error| {
296            ConnectorError::ConnectionFailed(format!("RabbitMQ channel: {error}"))
297        })?;
298        channel
299            .confirm_select(ConfirmSelectOptions::default())
300            .await
301            .map_err(|error| {
302                ConnectorError::ConnectionFailed(format!("confirm.select: {error}"))
303            })?;
304        prepare_exchange(&channel, &parsed.exchange, parsed.exchange_declare.as_ref()).await?;
305        tracing::info!(
306            exchange = %parsed.exchange,
307            mandatory = parsed.mandatory,
308            delivery_mode = ?parsed.delivery_mode,
309            endpoints = ?parsed.connection,
310            "RabbitMQ sink publishing"
311        );
312        self.schema = input;
313        self.writer = Some(writer);
314        self.runtime = Some(Runtime {
315            connection,
316            channel,
317            pending: VecDeque::new(),
318            pending_bytes: 0,
319            metrics,
320        });
321        self.config = Some(parsed);
322        self.failure = None;
323        Ok(())
324    }
325
326    async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
327        match self.publish(batch).await {
328            Ok(result) => Ok(result),
329            Err(error) if self.failure.is_some() => Err(error),
330            Err(error) => Err(self.fail(error)),
331        }
332    }
333
334    fn schema(&self) -> SchemaRef {
335        self.schema.clone()
336    }
337
338    fn suggested_write_timeout(&self) -> Duration {
339        // A write may wait for confirms to free pending capacity, bounded by the confirm timeout.
340        self.config
341            .as_ref()
342            .map_or(Duration::from_secs(30), |config| config.confirm_timeout)
343    }
344
345    async fn flush(&mut self) -> Result<(), ConnectorError> {
346        let (config, _, runtime) = self.live()?;
347        let deadline = tokio::time::Instant::now() + config.confirm_timeout;
348        match runtime.drain(0, 0, deadline).await {
349            Ok(_) => Ok(()),
350            Err(error) => Err(self.fail(error)),
351        }
352    }
353
354    async fn close(&mut self) -> Result<(), ConnectorError> {
355        let drained = if self.failure.is_none() && self.runtime.is_some() {
356            self.flush().await
357        } else {
358            Ok(())
359        };
360        let closed = match self.runtime.take() {
361            Some(runtime) => {
362                close(
363                    &runtime.channel,
364                    &runtime.connection,
365                    Duration::from_secs(5),
366                )
367                .await
368            }
369            None => Ok(()),
370        };
371        drained.and(closed)
372    }
373}
374
375#[cfg(test)]
376mod tests;
laminar_connectors/rabbitmq/sink/mod.rs

laminar_connectors/rabbitmq/sink/
mod.rs
