faucet_core/driver/
message.rs

1//! The messaging family (RFC 0014): produce with confirmations, transactions
2//! for exactly-once, consume with acknowledgements.
3//!
4//! A [`MessageDriver`] encodes one record into a message and publishes a
5//! group of messages; [`MessageSinkDriver`] drives everything around that —
6//! grouping, the per-group deadline, which rows a timeout or a cancel leaves
7//! unconfirmed, per-row outcomes, and (with [`MessageTransactions`]) the
8//! exactly-once commit. A [`MessageConsumer`] receives and acknowledges;
9//! [`MessageSourceDriver`] acknowledges a page only once the pipeline has
10//! written it.
11
12use super::sink::{CommitTokenDriver, PageWrite, RowPolicy, SinkDriver};
13use super::source::{PageStream, SourceDriver};
14use crate::accounting::WriteReport;
15use crate::check::{CheckContext, CheckReport, Probe};
16use crate::dlq::BatchAtomicity;
17use crate::error::FaucetError;
18use crate::io::{CallContext, CallSlot};
19use crate::pipeline::StreamPage;
20use crate::traits::RowOutcome;
21use async_trait::async_trait;
22use futures::StreamExt;
23use serde_json::Value;
24use std::collections::HashMap;
25use std::sync::Mutex;
26use std::time::Duration;
27
28/// One message to publish.
29#[derive(Debug, Clone, PartialEq, Eq)]
30#[non_exhaustive]
31pub struct OutboundMessage {
32    /// Where it goes: a subject, topic, queue or routing key.
33    pub destination: String,
34    /// The partition / ordering key, when the backend has one.
35    pub key: Option<Vec<u8>>,
36    /// The body.
37    pub payload: Vec<u8>,
38    /// Headers, in order.
39    pub headers: Vec<(String, String)>,
40}
41
42impl OutboundMessage {
43    /// A message with a body and no key or headers.
44    pub fn new(destination: impl Into<String>, payload: impl Into<Vec<u8>>) -> Self {
45        Self {
46            destination: destination.into(),
47            key: None,
48            payload: payload.into(),
49            headers: Vec::new(),
50        }
51    }
52
53    /// The same message with an ordering key.
54    pub fn with_key(mut self, key: impl Into<Vec<u8>>) -> Self {
55        self.key = Some(key.into());
56        self
57    }
58
59    /// The same message with one more header.
60    pub fn with_header(mut self, name: impl Into<String>, value: impl Into<String>) -> Self {
61        self.headers.push((name.into(), value.into()));
62        self
63    }
64}
65
66/// A producer's primitives (RFC 0014). Wrap it in a [`MessageSinkDriver`],
67/// then a [`DriverSink`](super::DriverSink).
68#[async_trait]
69pub trait MessageDriver: Send + Sync {
70    /// The connector's name.
71    fn connector_name(&self) -> &'static str;
72
73    /// The JSON Schema of the connector's config.
74    fn config_schema(&self) -> Value {
75        serde_json::json!({"type": "object", "properties": {}})
76    }
77
78    /// The destination, for lineage (credential-free).
79    fn dataset_uri(&self) -> String {
80        format!("{}://unknown", self.connector_name())
81    }
82
83    /// The connector's own call defaults.
84    fn call_context(&self) -> CallContext {
85        CallContext::new().without_timeout()
86    }
87
88    /// A non-mutating preflight probe.
89    async fn check(
90        &self,
91        calls: &CallSlot,
92        ctx: &CheckContext,
93    ) -> Result<CheckReport, FaucetError> {
94        let _ = (calls, ctx);
95        Ok(CheckReport::not_implemented())
96    }
97
98    /// What a failed write leaves behind.
99    fn batch_atomicity(&self) -> BatchAtomicity {
100        BatchAtomicity::BestEffort
101    }
102
103    /// The message for `record`. An error fails that row alone.
104    fn encode(&self, record: &Value) -> Result<OutboundMessage, FaucetError>;
105
106    /// Records per publish group; `0` publishes the whole page as one.
107    fn group_size(&self) -> usize {
108        0
109    }
110
111    /// How long one group may take. On expiry its unconfirmed rows, and
112    /// every row after them, fail with a retriable timeout. `None`: only the
113    /// run's call deadline applies.
114    fn group_timeout(&self) -> Option<Duration> {
115        None
116    }
117
118    /// Open the connection, once; called before the first group of every
119    /// page and before a flush. Default: nothing to open.
120    async fn connect(&self, calls: &CallSlot) -> Result<(), FaucetError> {
121        let _ = calls;
122        Ok(())
123    }
124
125    /// Publish a group, returning one outcome per message (`Ok` once the
126    /// backend confirmed it). An `Err` is a transport failure of the whole
127    /// page.
128    async fn publish(
129        &self,
130        calls: &CallSlot,
131        group: Vec<OutboundMessage>,
132    ) -> Result<Vec<RowOutcome>, FaucetError>;
133
134    /// Make published messages durable (wait for the client's buffer).
135    async fn flush(&self, calls: &CallSlot) -> Result<(), FaucetError> {
136        let _ = calls;
137        Ok(())
138    }
139
140    /// The transaction primitives, for exactly-once.
141    fn transactions(&self) -> Option<&dyn MessageTransactions> {
142        None
143    }
144
145    /// Whether the adapter runs each [`publish`](Self::publish) and
146    /// [`flush`](Self::flush) as one backend call (op `publish` / `flush`,
147    /// bounded by [`group_timeout`](Self::group_timeout) and the run's call
148    /// context). Return `false` for a driver whose `publish` makes its own
149    /// calls through the slot — its own op labels, request chunking, retries
150    /// and bounds; the adapter then awaits it directly. Default: `true`.
151    fn publish_is_one_call(&self) -> bool {
152        true
153    }
154
155    /// The error a whole-page write (the plain and reported paths) fails
156    /// with when rows of `offered` failed, given `(index, error)` for each
157    /// in index order. `None` (the default): the first failed row's own
158    /// error.
159    fn page_error(&self, offered: usize, failed: &[(usize, FaucetError)]) -> Option<FaucetError> {
160        let _ = (offered, failed);
161        None
162    }
163}
164
165/// A producer's transaction primitives. With them, [`MessageSinkDriver`]
166/// commits a page's messages and its commit token in one transaction.
167#[async_trait]
168pub trait MessageTransactions: Send + Sync {
169    /// Start a transaction for `scope`.
170    async fn begin(&self, calls: &CallSlot, scope: &str) -> Result<(), FaucetError>;
171
172    /// Record `token` for `scope` inside the open transaction and commit
173    /// everything published since [`begin`](Self::begin).
174    async fn commit(&self, calls: &CallSlot, scope: &str, token: &str) -> Result<(), FaucetError>;
175
176    /// Abandon the open transaction; nothing published in it becomes visible.
177    async fn abort(&self, calls: &CallSlot, scope: &str);
178
179    /// The last token committed for `scope`.
180    async fn last_token(
181        &self,
182        calls: &CallSlot,
183        scope: &str,
184    ) -> Result<Option<String>, FaucetError>;
185}
186
187/// A [`MessageDriver`] as a [`SinkDriver`].
188pub struct MessageSinkDriver<D> {
189    driver: D,
190}
191
192impl<D: MessageDriver> MessageSinkDriver<D> {
193    /// The sink driver over `driver`.
194    pub fn new(driver: D) -> Self {
195        Self { driver }
196    }
197
198    /// The message driver.
199    pub fn driver(&self) -> &D {
200        &self.driver
201    }
202
203    /// The message driver, mutably.
204    pub fn driver_mut(&mut self) -> &mut D {
205        &mut self.driver
206    }
207
208    fn timeout(&self, what: &str, after: Duration) -> FaucetError {
209        FaucetError::Timeout {
210            operation: format!("{} {what}", self.driver.connector_name()),
211            after,
212        }
213    }
214
215    /// Run `fut` as one backend call named `what`, bounded by the group
216    /// timeout. The outer error is the call's own (its timeout, the run's
217    /// deadline or cancel); the inner one is the driver's.
218    async fn bounded<T: Send>(
219        &self,
220        calls: &CallSlot,
221        what: &'static str,
222        fut: impl std::future::Future<Output = Result<T, FaucetError>> + Send,
223    ) -> Result<Result<T, FaucetError>, FaucetError> {
224        let timeout = self.driver.group_timeout();
225        if !self.driver.publish_is_one_call() {
226            return match timeout {
227                Some(after) => tokio::time::timeout(after, fut)
228                    .await
229                    .map_err(|_| self.timeout(what, after)),
230                None => Ok(fut.await),
231            };
232        }
233        calls
234            .call_once(what, async {
235                match timeout {
236                    Some(after) => tokio::time::timeout(after, fut)
237                        .await
238                        .map_err(|_| self.timeout(what, after)),
239                    None => Ok(fut.await),
240                }
241            })
242            .await
243    }
244
245    /// Publish every record in groups; one outcome per record.
246    async fn publish_all(
247        &self,
248        calls: &CallSlot,
249        rows: &[Value],
250    ) -> Result<Vec<RowOutcome>, FaucetError> {
251        let mut outcomes: Vec<RowOutcome> = rows.iter().map(|_| Ok(())).collect();
252        if rows.is_empty() {
253            return Ok(outcomes);
254        }
255        self.driver.connect(calls).await?;
256        let size = match self.driver.group_size() {
257            0 => rows.len(),
258            n => n,
259        };
260        for (g, chunk) in rows.chunks(size).enumerate() {
261            let base = g * size;
262            let mut positions = Vec::with_capacity(chunk.len());
263            let mut group = Vec::with_capacity(chunk.len());
264            for (i, record) in chunk.iter().enumerate() {
265                match self.driver.encode(record) {
266                    Ok(m) => {
267                        positions.push(base + i);
268                        group.push(m);
269                    }
270                    Err(e) => outcomes[base + i] = Err(e),
271                }
272            }
273            if group.is_empty() {
274                continue;
275            }
276            let sent = group.len();
277            match self
278                .bounded(calls, "publish", self.driver.publish(calls, group))
279                .await
280            {
281                Ok(Ok(results)) if results.len() == sent => {
282                    for (pos, r) in positions.into_iter().zip(results) {
283                        outcomes[pos] = r;
284                    }
285                }
286                Ok(Ok(results)) => {
287                    return Err(FaucetError::Sink(format!(
288                        "{}: {} outcomes for a group of {sent} messages",
289                        self.driver.connector_name(),
290                        results.len()
291                    )));
292                }
293                Ok(Err(e)) if g == 0 => return Err(e),
294                Ok(Err(e)) => {
295                    let name = self.driver.connector_name();
296                    fail_open(&mut outcomes[base..], || cut(name, &e));
297                    return Ok(outcomes);
298                }
299                Err(FaucetError::Timeout { after, .. }) => {
300                    let after = self.driver.group_timeout().unwrap_or(after);
301                    fail_open(&mut outcomes[base..], || self.timeout("publish", after));
302                    return Ok(outcomes);
303                }
304                Err(cause) => {
305                    let name = self.driver.connector_name();
306                    fail_open(&mut outcomes[base..], || cut(name, &cause));
307                    return Ok(outcomes);
308                }
309            }
310        }
311        Ok(outcomes)
312    }
313}
314
315/// The error a row gets when its group's call was cut short by `cause` (a
316/// cancel, or the call itself failing).
317fn cut(connector: &str, cause: &FaucetError) -> FaucetError {
318    match cause {
319        FaucetError::Cancelled { operation } => FaucetError::Cancelled {
320            operation: operation.clone(),
321        },
322        other => FaucetError::Sink(format!("{connector} publish: {other}")),
323    }
324}
325
326/// Fail every row of `outcomes` not already decided.
327fn fail_open(outcomes: &mut [RowOutcome], err: impl Fn() -> FaucetError) {
328    for o in outcomes.iter_mut().filter(|o| o.is_ok()) {
329        *o = Err(err());
330    }
331}
332
333#[async_trait]
334impl<D: MessageDriver> SinkDriver for MessageSinkDriver<D> {
335    fn connector_name(&self) -> &'static str {
336        self.driver.connector_name()
337    }
338
339    fn config_schema(&self) -> Value {
340        self.driver.config_schema()
341    }
342
343    fn dataset_uri(&self) -> String {
344        self.driver.dataset_uri()
345    }
346
347    fn call_context(&self) -> CallContext {
348        self.driver.call_context()
349    }
350
351    async fn check(
352        &self,
353        calls: &CallSlot,
354        ctx: &CheckContext,
355    ) -> Result<CheckReport, FaucetError> {
356        self.driver.check(calls, ctx).await
357    }
358
359    fn batch_atomicity(&self) -> BatchAtomicity {
360        self.driver.batch_atomicity()
361    }
362
363    async fn write(
364        &self,
365        calls: &CallSlot,
366        rows: &[Value],
367        policy: RowPolicy,
368    ) -> Result<PageWrite, FaucetError> {
369        let write = PageWrite::from_outcomes(self.publish_all(calls, rows).await?);
370        if policy == RowPolicy::WholePage
371            && !write.failed.is_empty()
372            && let Some(e) = self.driver.page_error(rows.len(), &write.failed)
373        {
374            return Err(e);
375        }
376        Ok(write)
377    }
378
379    async fn flush(&self, calls: &CallSlot) -> Result<(), FaucetError> {
380        self.driver.connect(calls).await?;
381        self.bounded(calls, "flush", self.driver.flush(calls))
382            .await?
383    }
384
385    fn commit_tokens(&self) -> Option<&dyn CommitTokenDriver> {
386        self.driver
387            .transactions()
388            .map(|_| self as &dyn CommitTokenDriver)
389    }
390}
391
392#[async_trait]
393impl<D: MessageDriver> CommitTokenDriver for MessageSinkDriver<D> {
394    async fn last_token(
395        &self,
396        calls: &CallSlot,
397        scope: &str,
398    ) -> Result<Option<String>, FaucetError> {
399        match self.driver.transactions() {
400            Some(t) => t.last_token(calls, scope).await,
401            None => Ok(None),
402        }
403    }
404
405    async fn write_with_token(
406        &self,
407        calls: &CallSlot,
408        rows: &[Value],
409        scope: &str,
410        token: &str,
411    ) -> Result<WriteReport, FaucetError> {
412        let tx = self.driver.transactions().ok_or_else(|| {
413            FaucetError::Sink(format!(
414                "{}: exactly-once needs a transactional producer",
415                self.driver.connector_name()
416            ))
417        })?;
418        tx.begin(calls, scope).await?;
419        let outcomes = match self.publish_all(calls, rows).await {
420            Ok(o) => o,
421            Err(e) => {
422                tx.abort(calls, scope).await;
423                return Err(e);
424            }
425        };
426        if let Some(e) = outcomes.into_iter().find_map(Result::err) {
427            tx.abort(calls, scope).await;
428            return Err(e);
429        }
430        if let Err(e) = tx.commit(calls, scope, token).await {
431            tx.abort(calls, scope).await;
432            return Err(e);
433        }
434        Ok(WriteReport::accepted(rows.len()))
435    }
436}
437
438/// An opaque handle a consumer acknowledges a message by.
439#[derive(Debug, Clone, PartialEq, Eq, Hash)]
440#[non_exhaustive]
441pub struct AckToken(pub String);
442
443/// One received message.
444#[derive(Debug, Clone, PartialEq, Eq)]
445#[non_exhaustive]
446pub struct InboundMessage {
447    /// Where it came from: a subject, topic or queue.
448    pub source: String,
449    /// The body.
450    pub payload: Vec<u8>,
451    /// How to acknowledge it.
452    pub ack: AckToken,
453}
454
455impl InboundMessage {
456    /// A received message.
457    pub fn new(source: impl Into<String>, payload: impl Into<Vec<u8>>, ack: AckToken) -> Self {
458        Self {
459            source: source.into(),
460            payload: payload.into(),
461            ack,
462        }
463    }
464}
465
466/// A consumer's primitives (RFC 0014). Wrap it in a [`MessageSourceDriver`],
467/// then a [`DriverSource`](super::DriverSource).
468#[async_trait]
469pub trait MessageConsumer: Send + Sync {
470    /// The connector's name.
471    fn connector_name(&self) -> &'static str;
472
473    /// The JSON Schema of the connector's config.
474    fn config_schema(&self) -> Value {
475        serde_json::json!({"type": "object", "properties": {}})
476    }
477
478    /// The source, for lineage (credential-free).
479    fn dataset_uri(&self) -> String {
480        format!("{}://unknown", self.connector_name())
481    }
482
483    /// The connector's own call defaults.
484    fn call_context(&self) -> CallContext {
485        CallContext::new()
486    }
487
488    /// Receive up to `max` messages, waiting at most the consumer's own idle
489    /// wait. An empty result means nothing arrived in that wait.
490    async fn receive(
491        &self,
492        calls: &CallSlot,
493        max: usize,
494    ) -> Result<Vec<InboundMessage>, FaucetError>;
495
496    /// Acknowledge messages so the broker does not redeliver them.
497    async fn ack(&self, calls: &CallSlot, tokens: Vec<AckToken>) -> Result<(), FaucetError>;
498
499    /// The record for a received message.
500    fn decode(&self, message: &InboundMessage) -> Result<Value, FaucetError>;
501
502    /// Stop after this many messages; `None`: only idleness stops the read.
503    fn max_messages(&self) -> Option<usize> {
504        None
505    }
506
507    /// Consecutive empty receives that end the read.
508    fn idle_receives_to_stop(&self) -> u32 {
509        1
510    }
511
512    /// A preflight probe that consumes nothing (connect, look up the queue).
513    /// `None` (the default): the probe is skipped, because reading a page
514    /// would take messages off the queue.
515    async fn check(
516        &self,
517        calls: &CallSlot,
518        ctx: &CheckContext,
519    ) -> Result<Option<CheckReport>, FaucetError> {
520        let _ = (calls, ctx);
521        Ok(None)
522    }
523}
524
525/// A [`MessageConsumer`] as a [`SourceDriver`].
526///
527/// Each page carries a bookmark, so the pipeline flushes the sink after
528/// writing it; the page's messages are acknowledged when the next page is
529/// requested (or the stream is polled to its end) — that is, only once the
530/// pipeline has written and flushed them. A run that fails before then
531/// acknowledges nothing it did not write, and the broker redelivers it.
532pub struct MessageSourceDriver<C> {
533    consumer: C,
534    pending: Mutex<Vec<AckToken>>,
535}
536
537impl<C: MessageConsumer> MessageSourceDriver<C> {
538    /// The source driver over `consumer`.
539    pub fn new(consumer: C) -> Self {
540        Self {
541            consumer,
542            pending: Mutex::new(Vec::new()),
543        }
544    }
545
546    /// The consumer.
547    pub fn consumer(&self) -> &C {
548        &self.consumer
549    }
550
551    async fn ack_pending(&self, calls: &CallSlot) -> Result<(), FaucetError> {
552        let tokens = std::mem::take(&mut *self.pending.lock().unwrap_or_else(|e| e.into_inner()));
553        if tokens.is_empty() {
554            return Ok(());
555        }
556        self.consumer.ack(calls, tokens).await
557    }
558}
559
560#[async_trait]
561impl<C: MessageConsumer> SourceDriver for MessageSourceDriver<C> {
562    fn connector_name(&self) -> &'static str {
563        self.consumer.connector_name()
564    }
565
566    fn config_schema(&self) -> Value {
567        self.consumer.config_schema()
568    }
569
570    fn dataset_uri(&self) -> String {
571        self.consumer.dataset_uri()
572    }
573
574    fn call_context(&self) -> CallContext {
575        self.consumer.call_context()
576    }
577
578    fn consumes_destructively(&self) -> bool {
579        true
580    }
581
582    fn read<'a>(
583        &'a self,
584        calls: &'a CallSlot,
585        _context: &'a HashMap<String, Value>,
586        batch_size: usize,
587    ) -> PageStream<'a> {
588        self.pages(calls, batch_size, true, self.consumer.max_messages())
589    }
590
591    async fn fetch(
592        &self,
593        calls: &CallSlot,
594        _context: &HashMap<String, Value>,
595    ) -> Option<Result<(Vec<Value>, Option<Value>), FaucetError>> {
596        let cap = self
597            .consumer
598            .max_messages()
599            .unwrap_or(crate::pipeline::DEFAULT_BATCH_SIZE);
600        let mut pages = self.pages(calls, 0, false, Some(cap));
601        let mut rows = Vec::new();
602        let mut bookmark = None;
603        while let Some(page) = pages.next().await {
604            match page {
605                Ok(page) => {
606                    rows.extend(page.records);
607                    bookmark = page.bookmark;
608                }
609                Err(e) => return Some(Err(e)),
610            }
611        }
612        Some(Ok((rows, bookmark)))
613    }
614
615    async fn check(
616        &self,
617        calls: &CallSlot,
618        ctx: &CheckContext,
619    ) -> Result<Option<CheckReport>, FaucetError> {
620        if let Some(report) = self.consumer.check(calls, ctx).await? {
621            return Ok(Some(report));
622        }
623        Ok(Some(CheckReport::single(Probe::skip(
624            "read",
625            "reading a page would consume messages; no consumer probe is defined",
626        ))))
627    }
628}
629
630impl<C: MessageConsumer> MessageSourceDriver<C> {
631    /// The page stream, ending after `limit` messages. A queue has no end, so
632    /// `batch_size` 0 pages by the default size rather than meaning one page.
633    /// `ack: false` never acknowledges, so a buffered read leaves every message
634    /// for the broker to redeliver; its `limit` keeps a redelivered message
635    /// from extending the read forever.
636    fn pages<'a>(
637        &'a self,
638        calls: &'a CallSlot,
639        batch_size: usize,
640        ack: bool,
641        limit: Option<usize>,
642    ) -> PageStream<'a> {
643        let per_page = if batch_size == 0 {
644            crate::pipeline::DEFAULT_BATCH_SIZE
645        } else {
646            batch_size
647        };
648        Box::pin(async_stream::try_stream! {
649            let mut consumed = 0usize;
650            let mut idle = 0u32;
651            loop {
652                if ack {
653                    self.ack_pending(calls).await?;
654                }
655                let room = match limit {
656                    Some(max) if consumed >= max => break,
657                    Some(max) => per_page.min(max - consumed),
658                    None => per_page,
659                };
660                let messages = self.consumer.receive(calls, room).await?;
661                if messages.is_empty() {
662                    idle += 1;
663                    if idle >= self.consumer.idle_receives_to_stop() {
664                        break;
665                    }
666                    continue;
667                }
668                idle = 0;
669                let mut records = Vec::with_capacity(messages.len());
670                let mut tokens = Vec::with_capacity(messages.len());
671                for m in &messages {
672                    records.push(self.consumer.decode(m)?);
673                    tokens.push(m.ack.clone());
674                }
675                consumed += records.len();
676                if ack {
677                    *self.pending.lock().unwrap_or_else(|e| e.into_inner()) = tokens;
678                }
679                yield StreamPage {
680                    records,
681                    bookmark: Some(serde_json::json!({ "consumed": consumed })),
682                };
683            }
684        })
685    }
686}
687
688#[cfg(test)]
689mod tests {
690    use super::*;
691
692    #[test]
693    fn a_cut_short_group_keeps_a_cancel_and_words_anything_else() {
694        let cancel = FaucetError::Cancelled {
695            operation: "publish call".into(),
696        };
697        assert!(matches!(
698            cut("broker", &cancel),
699            FaucetError::Cancelled { operation } if operation == "publish call"
700        ));
701        let other = cut("broker", &FaucetError::Source("x".into()));
702        assert!(
703            matches!(&other, FaucetError::Sink(m) if m.starts_with("broker publish: ")),
704            "{other}"
705        );
706        let mut outcomes: Vec<RowOutcome> = vec![Err(FaucetError::Sink("own".into())), Ok(())];
707        fail_open(&mut outcomes, || cut("broker", &cancel));
708        assert!(matches!(&outcomes[0], Err(FaucetError::Sink(m)) if m == "own"));
709        assert!(matches!(&outcomes[1], Err(FaucetError::Cancelled { .. })));
710    }
711}
faucet_core/driver/message.rs

faucet_core/driver/
message.rs
