krabka_client_streams/runtime/
eos.rs1use async_trait::async_trait;
8
9use crate::error::StreamsClientError;
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
13pub enum ProcessingGuarantee {
14 #[default]
16 AtLeastOnce,
17 ExactlyOnceV2,
19}
20
21#[derive(Debug, Clone)]
25pub struct StreamsGroupMeta {
26 pub group: String,
27 pub generation: i32,
29 pub member: String,
30 pub group_instance: Option<String>,
31}
32
33#[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#[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 #[derive(Default)]
85 pub struct MockTransactionalProducer {
86 pub calls: Mutex<Vec<Step>>,
87 pub sent: Mutex<Vec<SentRecord>>,
89 pub fail_at: Mutex<Option<Step>>,
90 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}