Skip to main content

crabka_replicator/
worker.rs

1//! Per-flow replication worker: drives one directional connect pipeline
2//! (`SourceConsumer` -> `TargetSink`) plus the heartbeat and checkpoint
3//! background tasks, with build-retry resilience and clean shutdown.
4
5use 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
20/// Maximum wall-clock time to keep retrying the pipeline build before giving up.
21const MAX_BUILD_ELAPSED: Duration = Duration::from_secs(30);
22/// Initial backoff between build attempts; doubles each retry up to [`MAX_BACKOFF`].
23const INITIAL_BACKOFF: Duration = Duration::from_millis(250);
24/// Cap on the per-retry backoff interval.
25const MAX_BACKOFF: Duration = Duration::from_secs(8);
26
27/// Parameters to start a [`FlowWorker`]. The supervisor resolves selectors into
28/// a concrete topic list and passes the per-flow cluster addresses, aliases,
29/// naming policy, residency policies, and security here.
30pub struct FlowWorkerParams {
31    /// Unique flow name (e.g. `"us-east__eu-west"`); seeds the consumer group id
32    /// and the checkpoint-store key.
33    pub flow_name: String,
34    /// Bootstrap address of the source cluster.
35    pub source_bootstrap: String,
36    /// Bootstrap address of the target cluster.
37    pub target_bootstrap: String,
38    /// Alias of the source cluster (stamped as provenance / used for MM2 names).
39    pub source_alias: String,
40    /// Alias of the target cluster (written into heartbeat records).
41    pub target_alias: String,
42    /// How source topics are renamed on the target.
43    pub naming: NamingPolicy,
44    /// Already-resolved source topic list (supervisor resolves selectors).
45    pub topics: Vec<String>,
46    /// Compliance zones of the target cluster (used for residency checks).
47    pub target_zones: Vec<String>,
48    /// Residency policies to enforce on the sink.
49    pub policies: Vec<PolicyConfig>,
50    /// Selector for which consumer groups the checkpoint task translates.
51    pub group_selector: Selector,
52    /// Optional TLS/SASL security for the source cluster.
53    pub security_source: Option<crabka_client_core::security::ClientSecurity>,
54    /// Optional TLS/SASL security for the target cluster.
55    pub security_target: Option<crabka_client_core::security::ClientSecurity>,
56}
57
58/// One directional replication flow: the running connect runtime plus the
59/// heartbeat and checkpoint background tasks.
60pub struct FlowWorker {
61    runtime: crabka_connect::ConnectorHandle,
62    heartbeat: HeartbeatTask,
63    checkpoint: CheckpointTask,
64}
65
66/// Current wall-clock time in milliseconds since the Unix epoch.
67fn 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
77/// Compute the next build-retry backoff: double the current interval, capped at
78/// [`MAX_BACKOFF`].
79fn next_backoff(current: Duration) -> Duration {
80    (current * 2).min(MAX_BACKOFF)
81}
82
83impl FlowWorker {
84    /// Build and start the pipeline, retrying transient build failures with
85    /// bounded exponential backoff (a cluster may be briefly unreachable).
86    ///
87    /// Backoff starts at ~250ms and doubles (capped at ~8s) until the cumulative
88    /// elapsed time exceeds ~30s, after which the last build error is returned.
89    ///
90    /// # Errors
91    ///
92    /// Returns the last build error if the pipeline cannot be constructed within
93    /// the retry budget.
94    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    /// A single build attempt: stand up the source, sink, checkpoint store,
120    /// connect runtime, and both background tasks.
121    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    /// Current connect-runtime state, for the supervisor to detect `Failed`
188    /// workers and restart them.
189    #[must_use]
190    pub fn state(&self) -> RuntimeState {
191        self.runtime.state()
192    }
193
194    /// Graceful stop: drain and shut down the connect runtime, then stop both
195    /// background tasks.
196    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        // Epoch millis for 2023-11-14T22:13:20Z; any sane current time exceeds
212        // it. The `-> -1 / 0 / 1` mutants all fall below this floor.
213        assert!(super::now_ms() > 1_700_000_000_000);
214    }
215
216    #[test]
217    fn next_backoff_doubles_and_caps() {
218        // Doubling: 250ms -> 500ms (the `*2`→`/2` mutant gives 125ms).
219        assert!(super::next_backoff(Duration::from_millis(250)) == Duration::from_millis(500));
220        // Cap: never exceeds MAX_BACKOFF even when already at the cap.
221        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}