Skip to main content

krabka_client_streams/runtime/
eos.rs

1//! Exactly-once (EOS v2 / KIP-447) primitives.
2//!
3//! This module holds the processing-guarantee config, the
4//! transactional-producer I/O seam, and the streams group metadata that
5//! `send_offsets_to_transaction` uses. The runtime wires these in at T2/T3.
6
7use async_trait::async_trait;
8
9use crate::error::StreamsClientError;
10
11/// Delivery guarantee for the runtime (`processing.guarantee`).
12#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
13pub enum ProcessingGuarantee {
14    /// Produce-then-commit. A crash mid-cycle can cause a replay. This is the default.
15    #[default]
16    AtLeastOnce,
17    /// Transactional: produce + offset-commit atomically (KIP-447).
18    ExactlyOnceV2,
19}
20
21/// Streams group metadata for `send_offsets_to_transaction`.
22///
23/// This type maps to the native `krabka_client_consumer::ConsumerGroupMetadata`.
24#[derive(Debug, Clone)]
25pub struct StreamsGroupMeta {
26    pub group: String,
27    /// The member epoch (next-gen "generation").
28    pub generation: i32,
29    pub member: String,
30    pub group_instance: Option<String>,
31}
32
33/// EOS-v2 transactional producer seam.
34///
35/// The runtime injects the implementation. Production uses
36/// `BrokerTransactionalProducer`, and tests use `MockTransactionalProducer`. An
37/// implementor must also implement `RecordProducer` for `send`.
38#[async_trait]
39pub trait TransactionalProducer: crate::runtime::io::RecordProducer {
40    async fn init_transactions(&self) -> Result<(), StreamsClientError>;
41    async fn begin_transaction(&self) -> Result<(), StreamsClientError>;
42    async fn send_offsets_to_transaction(
43        &self,
44        offsets: &[(String, i32, i64)],
45        group_meta: &StreamsGroupMeta,
46    ) -> Result<(), StreamsClientError>;
47    async fn commit_transaction(&self) -> Result<(), StreamsClientError>;
48    async fn abort_transaction(&self) -> Result<(), StreamsClientError>;
49}
50
51/// KIP-447 transactional id, stable for each application and thread pair.
52///
53/// A restart fences a zombie with the producer-epoch bump in
54/// `init_transactions`.
55#[must_use]
56pub fn transactional_id(application_id: &str, thread_idx: usize) -> String {
57    format!("{application_id}-{thread_idx}")
58}
59
60#[cfg(test)]
61pub(crate) mod mock {
62    use std::sync::Mutex;
63
64    use bytes::Bytes;
65
66    use super::*;
67
68    type SentRecord = (String, Option<i32>, Option<Bytes>, Option<Bytes>);
69
70    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
71    pub enum Step {
72        Init,
73        Begin,
74        Send,
75        SendOffsets,
76        Commit,
77        Abort,
78    }
79
80    /// Records the call sequence.
81    ///
82    /// `fail_at` makes the FIRST call that reaches that step return an error.
83    /// This drives the abort and rollback tests.
84    #[derive(Default)]
85    pub struct MockTransactionalProducer {
86        pub calls: Mutex<Vec<Step>>,
87        /// `(topic, partition, key, value)` of each `send`. A test-only record log.
88        pub sent: Mutex<Vec<SentRecord>>,
89        pub fail_at: Mutex<Option<Step>>,
90        /// Return the fatal producer-fenced error the first time this step runs.
91        pub fence_at: Mutex<Option<Step>>,
92    }
93    impl MockTransactionalProducer {
94        fn record(&self, s: Step) -> Result<(), StreamsClientError> {
95            self.calls.lock().unwrap().push(s);
96            let mut f = self.fail_at.lock().unwrap();
97            if *f == Some(s) {
98                *f = None;
99                return Err(StreamsClientError::Runtime(format!("mock fail at {s:?}")));
100            }
101            drop(f);
102            let mut fenced = self.fence_at.lock().unwrap();
103            if *fenced == Some(s) {
104                *fenced = None;
105                return Err(krabka_client_producer::ProducerError::FencedProducer.into());
106            }
107            Ok(())
108        }
109    }
110    #[async_trait]
111    impl crate::runtime::io::RecordProducer for MockTransactionalProducer {
112        async fn send(
113            &self,
114            topic: &str,
115            partition: Option<i32>,
116            key: Option<Bytes>,
117            value: Option<Bytes>,
118        ) -> Result<(), StreamsClientError> {
119            self.sent
120                .lock()
121                .unwrap()
122                .push((topic.to_string(), partition, key, value));
123            self.record(Step::Send)
124        }
125        async fn flush(&self) -> Result<(), StreamsClientError> {
126            Ok(())
127        }
128    }
129    #[async_trait]
130    impl TransactionalProducer for MockTransactionalProducer {
131        async fn init_transactions(&self) -> Result<(), StreamsClientError> {
132            self.record(Step::Init)
133        }
134        async fn begin_transaction(&self) -> Result<(), StreamsClientError> {
135            self.record(Step::Begin)
136        }
137        async fn send_offsets_to_transaction(
138            &self,
139            _o: &[(String, i32, i64)],
140            _m: &StreamsGroupMeta,
141        ) -> Result<(), StreamsClientError> {
142            self.record(Step::SendOffsets)
143        }
144        async fn commit_transaction(&self) -> Result<(), StreamsClientError> {
145            self.record(Step::Commit)
146        }
147        async fn abort_transaction(&self) -> Result<(), StreamsClientError> {
148            self.record(Step::Abort)
149        }
150    }
151}
152
153#[cfg(test)]
154mod tests {
155    use assert2::check;
156
157    use super::*;
158
159    #[test]
160    fn transactional_id_is_stable_per_thread() {
161        check!(transactional_id("word-count", 0) == "word-count-0");
162        check!(transactional_id("word-count", 1) == "word-count-1");
163    }
164
165    #[tokio::test]
166    async fn mock_records_calls_and_can_fail() {
167        use mock::{MockTransactionalProducer, Step};
168
169        use crate::runtime::io::RecordProducer;
170        let p = MockTransactionalProducer {
171            fail_at: std::sync::Mutex::new(Some(Step::Commit)),
172            ..Default::default()
173        };
174        p.begin_transaction().await.unwrap();
175        p.send("out", None, None, None).await.unwrap();
176        check!(p.commit_transaction().await.is_err());
177        check!(*p.calls.lock().unwrap() == vec![Step::Begin, Step::Send, Step::Commit]);
178    }
179}