Skip to main content

laminar_connectors/ibmmq/sink/
mod.rs

1//! IBM MQ queue sink: `MQPUT` under syncpoint and `MQCMIT` per unit of work.
2//!
3//! A batch is put in units of at most `max.uow.messages` messages; each unit commits before the
4//! write returns, so a successful write is durable on the queue manager for persistent messages.
5//! A failed put backs its unit out. A lost commit outcome, or a failure after an earlier unit of
6//! the same batch committed, ends this sink generation with an unknown outcome; replay from the
7//! last checkpoint may then put duplicates (at-least-once).
8
9mod message;
10
11use std::time::Duration;
12
13use arrow_array::RecordBatch;
14use arrow_schema::SchemaRef;
15use async_trait::async_trait;
16use mqi::prelude::*;
17use mqi::types::{MessageFormat, MQCMHO, MQPMO, MQSMPO};
18use mqi::{constants, structs, Properties};
19
20use super::config::{IbmMqSinkConfig, Persistence};
21use super::library;
22use super::metrics::SinkMetrics;
23use super::session::{classify, Session};
24use super::worker::Worker;
25use crate::config::ConnectorConfig;
26use crate::connector::{
27    ConnectorTaskOwner, ConnectorTaskTracker, SinkConnector, SinkConsistency, SinkContract,
28    SinkInputMode, SinkTopology, WriteResult,
29};
30use crate::error::ConnectorError;
31use crate::serde::RecordSerializer;
32use message::{base_descriptor, payload_schema, Outgoing};
33
34const CLOSE_TIMEOUT: Duration = Duration::from_secs(30);
35
36/// IBM MQ queue sink.
37pub struct IbmMqSink {
38    schema: SchemaRef,
39    registry: Option<prometheus::Registry>,
40    config: Option<IbmMqSinkConfig>,
41    writer: Option<Writer>,
42    worker: Option<Worker<Session>>,
43    metrics: Option<SinkMetrics>,
44    failure: Option<String>,
45    task_owner: ConnectorTaskOwner,
46    task_tracker: ConnectorTaskTracker,
47}
48
49struct Writer {
50    serializer: Box<dyn RecordSerializer>,
51    payload_indices: Vec<usize>,
52    base: libmqm_sys::MQMD,
53}
54
55/// Why a put batch stopped, and whether any of it was already committed.
56struct PutFailure {
57    committed: usize,
58    error: ConnectorError,
59}
60
61/// Put `messages` in units of `unit` messages, committing each unit.
62fn put_all(session: &Session, messages: &[Outgoing], unit: usize) -> Result<usize, PutFailure> {
63    let mut committed = 0;
64    for part in messages.chunks(unit) {
65        for outgoing in part {
66            if let Err(error) = put_one(session, outgoing) {
67                let backed_out = session.backout();
68                let error = match backed_out {
69                    Ok(()) => error,
70                    Err(backout) => ConnectorError::outcome_unknown(
71                        format!("{error}; MQBACK also failed: {backout}"),
72                        true,
73                    ),
74                };
75                return Err(PutFailure { committed, error });
76            }
77        }
78        session
79            .commit()
80            .map_err(|error| PutFailure { committed, error })?;
81        committed += part.len();
82    }
83    Ok(committed)
84}
85
86fn put_one(session: &Session, outgoing: &Outgoing) -> Result<(), ConnectorError> {
87    let handle = Properties::new(session.connection(), MQCMHO::default())
88        .map_err(|error| classify("MQCRTMH", &error))?;
89    for (name, value) in &outgoing.properties {
90        handle
91            .set_property(name.as_str(), value, MQSMPO::default())
92            .warn_as_error()
93            .map_err(|error| classify(&format!("MQSETMP '{name}'"), &error))?;
94    }
95    let mut options: MQPMO = constants::MQPMO_SYNCPOINT | constants::MQPMO_FAIL_IF_QUIESCING;
96    if outgoing.new_message_id {
97        options |= constants::MQPMO_NEW_MSG_ID;
98    }
99    let md = structs::MQMD::new(outgoing.md);
100    let format = MessageFormat::from_mqmd(&md);
101    let mut handle = handle;
102    session
103        .queue
104        .put_message(
105            &(options, md, &mut handle),
106            &(outgoing.data.as_slice(), format),
107        )
108        .warn_as_error()
109        .map_err(|error| match error.2 {
110            constants::MQRC_MSG_TOO_BIG_FOR_Q | constants::MQRC_MSG_TOO_BIG_FOR_Q_MGR => {
111                ConnectorError::ConfigurationError(format!(
112                    "MQPUT: a {}-byte message exceeds the queue's MAXMSGL: {error}",
113                    outgoing.data.len()
114                ))
115            }
116            _ => classify("MQPUT", &error),
117        })
118}
119
120impl IbmMqSink {
121    /// Create an unopened sink; metrics register on `registry` at open.
122    #[must_use]
123    pub fn new(schema: SchemaRef, registry: Option<&prometheus::Registry>) -> Self {
124        let (task_owner, task_tracker) = ConnectorTaskOwner::new();
125        Self {
126            schema,
127            registry: registry.cloned(),
128            config: None,
129            writer: None,
130            worker: None,
131            metrics: None,
132            failure: None,
133            task_owner,
134            task_tracker,
135        }
136    }
137
138    fn writer(parsed: &IbmMqSinkConfig, input: &SchemaRef) -> Result<Writer, ConnectorError> {
139        let payload = payload_schema(parsed, input)?;
140        crate::serde::schema_contract::validate_writer(parsed.format, &payload)?;
141        let payload_indices = payload
142            .fields()
143            .iter()
144            .map(|field| input.index_of(field.name()))
145            .collect::<Result<_, _>>()
146            .map_err(|error| ConnectorError::SchemaMismatch(error.to_string()))?;
147        Ok(Writer {
148            serializer: crate::serde::create_serializer(parsed.format)?,
149            base: base_descriptor(parsed, &payload)?,
150            payload_indices,
151        })
152    }
153
154    fn fail(&mut self, error: ConnectorError) -> ConnectorError {
155        tracing::error!(%error, "IBM MQ sink generation ended");
156        self.worker = None;
157        self.failure.get_or_insert(error.to_string());
158        error
159    }
160
161    async fn put(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
162        if let Some(cause) = &self.failure {
163            return Err(ConnectorError::InvalidState {
164                expected: "a live IBM MQ sink".into(),
165                actual: cause.clone(),
166            });
167        }
168        let (Some(config), Some(writer), Some(worker), Some(metrics)) =
169            (&self.config, &self.writer, &self.worker, &self.metrics)
170        else {
171            return Err(ConnectorError::InvalidState {
172                expected: "an open IBM MQ sink".into(),
173                actual: "open() has not completed".into(),
174            });
175        };
176        let payload = batch
177            .project(&writer.payload_indices)
178            .map_err(|error| ConnectorError::SchemaMismatch(error.to_string()))?;
179        let records = writer.serializer.serialize(&payload)?;
180        if records.len() != batch.num_rows() {
181            return Err(ConnectorError::Internal(format!(
182                "IBM MQ serializer returned {} records for {} rows",
183                records.len(),
184                batch.num_rows()
185            )));
186        }
187        let outgoing = message::messages(config, &writer.base, batch, records)?;
188        let bytes: u64 = outgoing
189            .iter()
190            .map(|message| message.data.len() as u64)
191            .sum();
192        let count = outgoing.len();
193        let unit = config.max_uow_messages;
194        let result = worker
195            .call(
196                "MQPUT/MQCMIT",
197                config.connection.operation_timeout,
198                move |session| Ok(put_all(session, &outgoing, unit)),
199            )
200            .await?;
201        match result {
202            Ok(committed) => {
203                metrics.put.inc_by(committed as u64);
204                metrics.bytes.inc_by(bytes);
205                metrics.commits.inc_by(committed.div_ceil(unit) as u64);
206                Ok(WriteResult::new(count, bytes))
207            }
208            Err(PutFailure { committed, error }) => {
209                metrics.backouts.inc();
210                if committed > 0 || error.is_outcome_unknown() {
211                    metrics.ambiguous.inc();
212                    Err(ConnectorError::outcome_unknown(
213                        format!(
214                            "IBM MQ batch partially committed ({committed} of {count}): {error}"
215                        ),
216                        true,
217                    ))
218                } else {
219                    Err(error)
220                }
221            }
222        }
223    }
224}
225
226#[async_trait]
227impl SinkConnector for IbmMqSink {
228    fn terminal_task_tracker(&self) -> Option<ConnectorTaskTracker> {
229        Some(self.task_tracker.clone())
230    }
231
232    fn contract(&self, config: &ConnectorConfig) -> Result<SinkContract, ConnectorError> {
233        let parsed = IbmMqSinkConfig::from_config(config)?;
234        // Only an explicitly persistent message is durable at MQCMIT; DEFPSIST is not verified.
235        let consistency = match parsed.persistence {
236            Persistence::Persistent => SinkConsistency::DurableAtLeastOnce,
237            Persistence::NotPersistent | Persistence::AsQueueDefinition => {
238                SinkConsistency::Ephemeral
239            }
240        };
241        Ok(SinkContract::new(
242            consistency,
243            SinkTopology::MultiWriter,
244            SinkInputMode::AppendOnly,
245        ))
246    }
247
248    async fn resolve_schema(
249        &mut self,
250        config: &ConnectorConfig,
251        input: SchemaRef,
252    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
253        use crate::schema::resolution::{logical_binding, SchemaDirection, SchemaOrigin};
254        let parsed = IbmMqSinkConfig::from_config(config)?;
255        Self::writer(&parsed, &input)?;
256        logical_binding(config, SchemaDirection::Sink, SchemaOrigin::Query, &input)
257    }
258
259    async fn open(&mut self, config: &ConnectorConfig) -> Result<(), ConnectorError> {
260        let parsed = IbmMqSinkConfig::from_config(config)?;
261        let input = config
262            .arrow_schema()
263            .or_else(|| {
264                config
265                    .schema_binding()
266                    .map(|binding| std::sync::Arc::new(binding.logical.clone()))
267            })
268            .unwrap_or_else(|| self.schema.clone());
269        let writer = Self::writer(&parsed, &input)?;
270        let library = library::load(parsed.connection.client_library.as_deref())?;
271        let metrics = SinkMetrics::new(&parsed.queue, self.registry.as_ref())?;
272        let guard = self.task_owner.track().ok_or_else(|| {
273            ConnectorError::Internal("IBM MQ sink task generation is retired".into())
274        })?;
275        let session_config = parsed.clone();
276        let worker = Worker::start(
277            &format!("ibmmq-sink-{}", parsed.queue),
278            guard,
279            parsed.connection.operation_timeout,
280            move || {
281                Session::open(
282                    &session_config.connection,
283                    library,
284                    &session_config.queue,
285                    session_config.target_queue_manager.as_deref(),
286                    constants::MQOO_OUTPUT,
287                )
288            },
289        )
290        .await?;
291        tracing::info!(
292            queue = %parsed.queue,
293            persistence = ?parsed.persistence,
294            connection = ?parsed.connection,
295            "IBM MQ sink putting under syncpoint"
296        );
297        self.schema = input;
298        self.writer = Some(writer);
299        self.worker = Some(worker);
300        self.metrics = Some(metrics);
301        self.config = Some(parsed);
302        self.failure = None;
303        Ok(())
304    }
305
306    async fn write_batch(&mut self, batch: &RecordBatch) -> Result<WriteResult, ConnectorError> {
307        match self.put(batch).await {
308            Ok(result) => Ok(result),
309            Err(error) if self.failure.is_some() => Err(error),
310            Err(error) => Err(self.fail(error)),
311        }
312    }
313
314    fn schema(&self) -> SchemaRef {
315        self.schema.clone()
316    }
317
318    fn suggested_write_timeout(&self) -> Duration {
319        self.config
320            .as_ref()
321            .map_or(Duration::from_secs(30), |config| {
322                config.connection.operation_timeout
323            })
324    }
325
326    async fn close(&mut self) -> Result<(), ConnectorError> {
327        match self.worker.take() {
328            Some(worker) => worker.stop(CLOSE_TIMEOUT).await,
329            None => Ok(()),
330        }
331    }
332}
333
334#[cfg(test)]
335mod tests;