krabka_replicator/
checkpoint_store.rs

1//! Position recovery: persist the worker's [`SourceOffset`] to a compacted
2//! internal topic on the TARGET cluster, keyed by flow name.
3//!
4//! On restart, [`InternalTopicCheckpointStore::load`] reads the last value for
5//! the flow's key out of the compacted topic. It recovers the exact position
6//! that the worker had reached before it stopped.
7
8use async_trait::async_trait;
9use bytes::Bytes;
10use krabka_connect::{CheckpointStore, ConnectError, SourceOffset};
11
12use crate::config::{ClientResourcePolicy, ReplicatorRuntimePolicy};
13
14/// The internal compacted topic used to store replicator checkpoints.
15pub(crate) const STATE_TOPIC: &str = "krabka-replicator-offsets";
16
17/// A [`CheckpointStore`] backed by a compacted internal Kafka topic on the
18/// target cluster.
19///
20/// Each flow gets its own key, `flow_name`, within the shared compacted
21/// internal state topic. On restart, [`load`](Self::load) fetches the last
22/// value for that key and recovers the exact partition offsets that the worker
23/// had reached.
24pub struct InternalTopicCheckpointStore {
25    producer: krabka_client_producer::Producer,
26    target_bootstrap: String,
27    topic: String,
28    key: String,
29    security: Option<krabka_client_core::security::ClientSecurity>,
30    client_resource_policy: ClientResourcePolicy,
31    runtime_policy: ReplicatorRuntimePolicy,
32    /// Exactly-once sinks persist this checkpoint in the same transaction as
33    /// their target writes, so the runtime's subsequent save is only an
34    /// acknowledgement barrier.
35    sink_transaction: bool,
36}
37
38impl InternalTopicCheckpointStore {
39    /// Ensure the compacted offset topic exists on the target cluster, build a
40    /// producer, and return a store keyed by `flow_name`.
41    ///
42    /// # Errors
43    ///
44    /// Returns [`ConnectError::Offset`] if the topic cannot be created or the
45    /// producer cannot connect.
46    pub async fn start(
47        target_bootstrap: &str,
48        flow_name: &str,
49        security: Option<krabka_client_core::security::ClientSecurity>,
50    ) -> Result<Self, ConnectError> {
51        Self::start_with_policy(
52            target_bootstrap,
53            flow_name,
54            security,
55            ClientResourcePolicy::default(),
56        )
57        .await
58    }
59
60    /// Start with the deployment's client resource policy.
61    ///
62    /// # Errors
63    ///
64    /// Returns [`ConnectError::Offset`] if the topic cannot be created or the
65    /// producer cannot connect.
66    pub async fn start_with_policy(
67        target_bootstrap: &str,
68        flow_name: &str,
69        security: Option<krabka_client_core::security::ClientSecurity>,
70        client_resource_policy: ClientResourcePolicy,
71    ) -> Result<Self, ConnectError> {
72        Self::start_with_runtime_policy(
73            target_bootstrap,
74            flow_name,
75            security,
76            client_resource_policy,
77            ReplicatorRuntimePolicy::default(),
78        )
79        .await
80    }
81
82    pub(crate) async fn start_with_runtime_policy(
83        target_bootstrap: &str,
84        flow_name: &str,
85        security: Option<krabka_client_core::security::ClientSecurity>,
86        client_resource_policy: ClientResourcePolicy,
87        runtime_policy: ReplicatorRuntimePolicy,
88    ) -> Result<Self, ConnectError> {
89        crate::admin_util::ensure_compacted_topic_with_runtime_policy(
90            target_bootstrap,
91            STATE_TOPIC,
92            security.clone(),
93            client_resource_policy,
94            &runtime_policy,
95        )
96        .await
97        .map_err(ConnectError::Offset)?;
98
99        let builder = krabka_client_producer::Producer::builder()
100            .bootstrap(target_bootstrap)
101            .dispatch_queue_capacity(client_resource_policy.dispatch_queue_capacity.get())
102            .frame_max(client_resource_policy.frame_max.size())
103            .enable_idempotence(false)
104            .acks(krabka_client_producer::Acks::All);
105
106        let producer = match security.clone() {
107            Some(s) => builder.security(s).build().await,
108            None => builder.build().await,
109        }
110        .map_err(|e| ConnectError::Offset(e.to_string()))?;
111
112        Ok(Self {
113            producer,
114            target_bootstrap: target_bootstrap.to_string(),
115            topic: STATE_TOPIC.into(),
116            key: flow_name.into(),
117            security,
118            client_resource_policy,
119            runtime_policy,
120            sink_transaction: false,
121        })
122    }
123
124    /// Let the target sink own checkpoint persistence inside its transaction.
125    #[must_use]
126    pub(crate) fn with_sink_transaction(mut self, enabled: bool) -> Self {
127        self.sink_transaction = enabled;
128        self
129    }
130}
131
132#[async_trait]
133impl CheckpointStore for InternalTopicCheckpointStore {
134    #[tracing::instrument(
135        level = "debug",
136        skip_all,
137        fields(key = %self.key, positions = offset.position.len()),
138        err,
139    )]
140    async fn save(&self, offset: &SourceOffset) -> Result<(), ConnectError> {
141        if self.sink_transaction {
142            return Ok(());
143        }
144        let bytes = serde_json::to_vec(offset).map_err(|e| ConnectError::Offset(e.to_string()))?;
145
146        self.producer
147            .send(krabka_client_producer::ProducerRecord {
148                topic: self.topic.clone(),
149                partition: None,
150                key: Some(Bytes::copy_from_slice(self.key.as_bytes())),
151                value: Some(Bytes::from(bytes)),
152                headers: vec![],
153                timestamp_ms: None,
154            })
155            .await
156            .map_err(|e| ConnectError::Offset(e.to_string()))?;
157
158        self.producer
159            .flush()
160            .await
161            .map_err(|e| ConnectError::Offset(e.to_string()))?;
162
163        Ok(())
164    }
165
166    #[tracing::instrument(level = "info", skip_all, fields(key = %self.key), err)]
167    async fn load(&self) -> Result<Option<SourceOffset>, ConnectError> {
168        let latest = crate::admin_util::read_last_value_for_key_with_runtime_policy(
169            &self.target_bootstrap,
170            &self.topic,
171            self.key.as_bytes(),
172            self.security.clone(),
173            self.client_resource_policy,
174            &self.runtime_policy,
175        )
176        .await
177        .map_err(ConnectError::Offset)?;
178
179        match latest {
180            Some(bytes) => Ok(Some(
181                serde_json::from_slice(&bytes).map_err(|e| ConnectError::Offset(e.to_string()))?,
182            )),
183            None => Ok(None),
184        }
185    }
186}
187
188#[cfg(test)]
189mod tests {
190    use std::collections::BTreeMap;
191
192    use krabka_connect::{CheckpointStore, OffsetValue, SourceOffset};
193
194    use super::*;
195
196    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
197    async fn persists_and_reloads_position_from_target() {
198        let dir = tempfile::TempDir::new().unwrap();
199        let broker = krabka_broker::Broker::start(krabka_broker::BrokerConfig::for_tests(
200            dir.path().to_path_buf(),
201        ))
202        .await
203        .unwrap();
204        let target = broker.listen_addr().to_string();
205
206        let store = InternalTopicCheckpointStore::start(&target, "flow1", None)
207            .await
208            .unwrap();
209
210        let mut pos = BTreeMap::new();
211        pos.insert("orders-0".to_string(), OffsetValue::Long(42));
212        let off = SourceOffset::new(BTreeMap::new().into(), pos.into());
213
214        store.save(&off).await.unwrap();
215
216        let store2 = InternalTopicCheckpointStore::start(&target, "flow1", None)
217            .await
218            .unwrap();
219        let loaded = store2.load().await.unwrap().unwrap();
220        assert2::assert!(loaded == off);
221    }
222}
krabka_replicator/checkpoint_store.rs

krabka_replicator/
checkpoint_store.rs
