datum_mq/
sink.rs

1use std::sync::{Arc, Mutex};
2
3use datum::{Flow, Keep, NotUsed, Sink, StreamCompletion, StreamError};
4
5use crate::{
6    KafkaMetrics, KafkaProducerSettings, MqError, MqResult, ProducerRecord,
7    native::{NativeKafkaProducerControl, NativeKafkaProducerHandle},
8};
9
10const PRODUCER_BATCH_SIZE: usize = 256;
11
12/// Kafka producer sink entry points.
13pub struct KafkaSink;
14
15impl KafkaSink {
16    /// Creates a producer sink.
17    ///
18    /// The input record carries its target topic. The sink completes each element
19    /// by enqueuing bounded 256-record groups into the native producer owner.
20    /// Delivery failures fail the stream; call
21    /// [`KafkaProducerControl::drain_and_shutdown`] at shutdown to wait for final
22    /// acknowledgements before closing. Production uses record-batch v2 plus
23    /// bounded, idempotent retries by default, preserving the no-duplicate
24    /// guarantee. Its bounded per-broker pipeline serializes each topic-partition
25    /// and halts and settles issued responses before retrying.
26    #[must_use]
27    pub fn plain(settings: KafkaProducerSettings) -> Sink<ProducerRecord, KafkaProducerControl> {
28        Flow::identity()
29            .grouped(PRODUCER_BATCH_SIZE)
30            .to_mat(Self::batched(settings), Keep::right)
31    }
32
33    fn batched(settings: KafkaProducerSettings) -> Sink<Vec<ProducerRecord>, KafkaProducerControl> {
34        Sink::setup(move |_materializer, _attributes| {
35            let metrics = KafkaMetrics::default();
36            let (native, handle) = NativeKafkaProducerControl::start(&settings, metrics.clone())
37                .expect("native Kafka producer configuration must be valid at materialization");
38            let control = KafkaProducerControl::new(native, metrics);
39            Sink::foreach_result(move |records| {
40                let handle: NativeKafkaProducerHandle = handle.clone();
41                handle.send(records).map_err(StreamError::from)
42            })
43            .map_materialized_value(move |completion| {
44                control.attach_completion(completion);
45                control.clone()
46            })
47        })
48    }
49}
50
51/// Materialized control for a Kafka producer sink.
52#[derive(Clone)]
53pub struct KafkaProducerControl {
54    state: Arc<ProducerControlState>,
55}
56
57struct ProducerControlState {
58    native: NativeKafkaProducerControl,
59    metrics: KafkaMetrics,
60    completion: Mutex<Option<StreamCompletion<NotUsed>>>,
61}
62
63impl KafkaProducerControl {
64    fn new(native: NativeKafkaProducerControl, metrics: KafkaMetrics) -> Self {
65        Self {
66            state: Arc::new(ProducerControlState {
67                native,
68                metrics,
69                completion: Mutex::new(None),
70            }),
71        }
72    }
73
74    fn attach_completion(&self, completion: StreamCompletion<NotUsed>) {
75        if let Ok(mut slot) = self.state.completion.lock() {
76            *slot = Some(completion);
77        }
78    }
79
80    #[must_use]
81    pub fn metrics(&self) -> KafkaMetrics {
82        self.state.metrics.clone()
83    }
84
85    /// Waits for accepted input records to receive delivery acknowledgements, then flushes.
86    pub fn drain_and_shutdown(&self) -> MqResult<()> {
87        let completion = self
88            .state
89            .completion
90            .lock()
91            .map_err(|_| MqError::Failed("Kafka producer completion lock poisoned".to_owned()))?
92            .take();
93        let completion_result = completion.map_or(Ok(()), |completion| {
94            completion
95                .wait()
96                .map(|_| ())
97                .map_err(|error| MqError::Failed(error.to_string()))
98        });
99        completion_result.and(self.state.native.finish())
100    }
101
102    pub fn flush(&self) -> MqResult<()> {
103        self.state.native.flush()
104    }
105}
datum_mq/sink.rs

datum_mq/
sink.rs
