crabka_replicator/
worker.rs1use std::sync::Arc;
6use std::time::{Duration, SystemTime, UNIX_EPOCH};
7
8use crabka_connect::{ConnectorRuntime, RuntimeState};
9use tracing::warn;
10
11use crate::checkpoint_store::InternalTopicCheckpointStore;
12use crate::config::{NamingPolicy, PolicyConfig};
13use crate::record::ReplicatedRecord;
14use crate::selector::Selector;
15use crate::sink::{SinkParams, TargetSink};
16use crate::source::SourceConsumer;
17use crate::tasks::checkpoint::{CheckpointParams, CheckpointTask};
18use crate::tasks::heartbeat::{HeartbeatParams, HeartbeatTask};
19
20const MAX_BUILD_ELAPSED: Duration = Duration::from_secs(30);
22const INITIAL_BACKOFF: Duration = Duration::from_millis(250);
24const MAX_BACKOFF: Duration = Duration::from_secs(8);
26
27pub struct FlowWorkerParams {
31 pub flow_name: String,
34 pub source_bootstrap: String,
36 pub target_bootstrap: String,
38 pub source_alias: String,
40 pub target_alias: String,
42 pub naming: NamingPolicy,
44 pub topics: Vec<String>,
46 pub target_zones: Vec<String>,
48 pub policies: Vec<PolicyConfig>,
50 pub group_selector: Selector,
52 pub security_source: Option<crabka_client_core::security::ClientSecurity>,
54 pub security_target: Option<crabka_client_core::security::ClientSecurity>,
56}
57
58pub struct FlowWorker {
61 runtime: crabka_connect::ConnectorHandle,
62 heartbeat: HeartbeatTask,
63 checkpoint: CheckpointTask,
64}
65
66fn now_ms() -> i64 {
68 i64::try_from(
69 SystemTime::now()
70 .duration_since(UNIX_EPOCH)
71 .unwrap_or_default()
72 .as_millis(),
73 )
74 .unwrap_or(i64::MAX)
75}
76
77fn next_backoff(current: Duration) -> Duration {
80 (current * 2).min(MAX_BACKOFF)
81}
82
83impl FlowWorker {
84 pub async fn start(p: FlowWorkerParams) -> crate::Result<Self> {
95 let mut backoff = INITIAL_BACKOFF;
96 let mut elapsed = Duration::ZERO;
97
98 loop {
99 match Self::build(&p).await {
100 Ok(worker) => return Ok(worker),
101 Err(e) => {
102 if elapsed >= MAX_BUILD_ELAPSED {
103 return Err(e);
104 }
105 warn!(
106 flow = %p.flow_name,
107 error = %e,
108 backoff_ms = backoff.as_millis(),
109 "flow worker build failed; retrying after backoff"
110 );
111 tokio::time::sleep(backoff).await;
112 elapsed += backoff;
113 backoff = next_backoff(backoff);
114 }
115 }
116 }
117 }
118
119 async fn build(p: &FlowWorkerParams) -> crate::Result<Self> {
122 let group_id = format!("crabka-replicator-{}", p.flow_name);
123
124 let source = SourceConsumer::start(
125 &p.source_bootstrap,
126 &group_id,
127 &p.topics,
128 p.security_source.clone(),
129 )
130 .await?;
131
132 let sink = TargetSink::start(SinkParams {
133 target_bootstrap: p.target_bootstrap.clone(),
134 source_alias: p.source_alias.clone(),
135 naming: p.naming,
136 target_zones: p.target_zones.clone(),
137 policies: p.policies.clone(),
138 security: p.security_target.clone(),
139 })
140 .await?;
141
142 let store = InternalTopicCheckpointStore::start(
143 &p.target_bootstrap,
144 &p.flow_name,
145 p.security_target.clone(),
146 )
147 .await?;
148
149 let runtime = ConnectorRuntime::<(), ReplicatedRecord>::new()
150 .add_source(source)
151 .add_sink(sink)
152 .checkpoint_store(Arc::new(store))
153 .commit_interval(Duration::from_millis(500))
154 .max_batch(500)
155 .run()?;
156
157 let heartbeat = HeartbeatTask::start(HeartbeatParams {
158 target_bootstrap: p.target_bootstrap.clone(),
159 source_alias: p.source_alias.clone(),
160 target_alias: p.target_alias.clone(),
161 interval: Duration::from_secs(1),
162 now_ms,
163 security: p.security_target.clone(),
164 })
165 .await?;
166
167 let checkpoint = CheckpointTask::start(
168 CheckpointParams {
169 source_bootstrap: p.source_bootstrap.clone(),
170 target_bootstrap: p.target_bootstrap.clone(),
171 source_alias: p.source_alias.clone(),
172 naming: p.naming,
173 group_selector: p.group_selector.clone(),
174 security: p.security_target.clone(),
175 },
176 Duration::from_secs(5),
177 )
178 .await?;
179
180 Ok(Self {
181 runtime,
182 heartbeat,
183 checkpoint,
184 })
185 }
186
187 #[must_use]
190 pub fn state(&self) -> RuntimeState {
191 self.runtime.state()
192 }
193
194 pub async fn shutdown(self) {
197 let _ = self.runtime.shutdown().await;
198 self.heartbeat.shutdown().await;
199 self.checkpoint.shutdown().await;
200 }
201}
202
203#[cfg(test)]
204mod tests {
205 use assert2::assert;
206
207 use super::*;
208
209 #[test]
210 fn now_ms_is_a_recent_epoch_millis() {
211 assert!(super::now_ms() > 1_700_000_000_000);
214 }
215
216 #[test]
217 fn next_backoff_doubles_and_caps() {
218 assert!(super::next_backoff(Duration::from_millis(250)) == Duration::from_millis(500));
220 assert!(super::next_backoff(MAX_BACKOFF) == MAX_BACKOFF);
222 }
223
224 #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
225 async fn worker_replicates_one_flow() {
226 let s_dir = tempfile::TempDir::new().unwrap();
227 let t_dir = tempfile::TempDir::new().unwrap();
228 let source = crabka_broker::Broker::start(crabka_broker::BrokerConfig::for_tests(
229 s_dir.path().to_path_buf(),
230 ))
231 .await
232 .unwrap();
233 let target = crabka_broker::Broker::start(crabka_broker::BrokerConfig::for_tests(
234 t_dir.path().to_path_buf(),
235 ))
236 .await
237 .unwrap();
238 let sb = source.listen_addr().to_string();
239 let tb = target.listen_addr().to_string();
240
241 crate::test_util::create_topic(&sb, "orders", 1).await;
242 crate::test_util::produce(&sb, "orders", b"k", b"v").await;
243
244 let worker = FlowWorker::start(FlowWorkerParams {
245 flow_name: "us-east__eu-west".into(),
246 source_bootstrap: sb,
247 target_bootstrap: tb.clone(),
248 source_alias: "us-east".into(),
249 target_alias: "eu-west".into(),
250 naming: crate::config::NamingPolicy::Default,
251 topics: vec!["orders".to_string()],
252 target_zones: vec!["us".into()],
253 policies: vec![],
254 group_selector: crate::selector::Selector::compile(&[], &[]).unwrap(),
255 security_source: None,
256 security_target: None,
257 })
258 .await
259 .unwrap();
260
261 crate::test_util::await_topic_count(
262 &tb,
263 "us-east.orders",
264 1,
265 std::time::Duration::from_secs(15),
266 )
267 .await;
268
269 worker.shutdown().await;
270 assert!(crate::test_util::topic_record_count(&tb, "us-east.orders").await >= 1);
271 }
272}