laminar_connectors/
checkpoint.rs

1//! Source checkpoints: the minimal state needed to resume a connector
2//! where it left off after a restart (Kafka offsets, CDC LSN/GTID, etc.).
3#![allow(clippy::disallowed_types)] // cold path: connector checkpoint
4
5use std::collections::HashMap;
6use std::fmt;
7use std::num::NonZeroU64;
8use std::sync::{Arc, OnceLock};
9
10enum OffsetTree {
11    Leaf(Arc<str>),
12    Branch { older: Arc<Self>, newer: Arc<Self> },
13}
14
15impl OffsetTree {
16    fn write_to(&self, output: &mut String, separator: &str, first: &mut bool) {
17        match self {
18            Self::Leaf(value) => {
19                if *first {
20                    *first = false;
21                } else {
22                    output.push_str(separator);
23                }
24                output.push_str(value);
25            }
26            Self::Branch { older, newer } => {
27                older.write_to(output, separator, first);
28                newer.write_to(output, separator, first);
29            }
30        }
31    }
32}
33
34/// Immutable, append-only serialized value for a large source offset.
35///
36/// Cloning or taking a checkpoint snapshot is `O(1)`. Appending path-copies at
37/// most `O(log N)` tree roots, and materialization preserves insertion order.
38#[derive(Clone)]
39pub struct PersistentOffset {
40    /// Binary-counter forest. Slot `n` holds an immutable tree with `2^n`
41    /// fragments; roots from highest to lowest are chronological.
42    trees: Arc<Vec<Option<Arc<OffsetTree>>>>,
43    prefix: Arc<str>,
44    separator: Arc<str>,
45    suffix: Arc<str>,
46    fragment_count: usize,
47    fragment_bytes: usize,
48}
49
50impl PersistentOffset {
51    /// Creates an empty persistent serialized value.
52    ///
53    /// For example, a JSON array uses `"["`, `","`, and
54    /// `"]"`. Each pushed fragment must itself be a complete JSON value.
55    #[must_use]
56    pub fn new(
57        prefix: impl Into<Arc<str>>,
58        separator: impl Into<Arc<str>>,
59        suffix: impl Into<Arc<str>>,
60    ) -> Self {
61        Self {
62            trees: Arc::new(Vec::new()),
63            prefix: prefix.into(),
64            separator: separator.into(),
65            suffix: suffix.into(),
66            fragment_count: 0,
67            fragment_bytes: 0,
68        }
69    }
70
71    /// Appends one already-serialized fragment in `O(log N)` without changing
72    /// any previously cloned snapshot.
73    pub fn push_fragment(&mut self, fragment: impl Into<Arc<str>>) {
74        let value = fragment.into();
75        // These counters are allocation hints/diagnostics, never correctness
76        // state. Saturation keeps an impossible address-space overflow from
77        // turning connector input into a process panic.
78        self.fragment_bytes = self.fragment_bytes.saturating_add(value.len());
79        self.fragment_count = self.fragment_count.saturating_add(1);
80
81        let trees = Arc::make_mut(&mut self.trees);
82        let mut carry = Arc::new(OffsetTree::Leaf(value));
83        let mut level = 0;
84        loop {
85            if level == trees.len() {
86                trees.push(Some(carry));
87                break;
88            }
89            if let Some(older) = trees[level].take() {
90                carry = Arc::new(OffsetTree::Branch {
91                    older,
92                    newer: carry,
93                });
94                level += 1;
95            } else {
96                trees[level] = Some(carry);
97                break;
98            }
99        }
100    }
101
102    /// Number of serialized fragments in this value.
103    #[must_use]
104    pub const fn fragment_count(&self) -> usize {
105        self.fragment_count
106    }
107
108    fn materialize(&self) -> String {
109        let separator_bytes = self
110            .separator
111            .len()
112            .saturating_mul(self.fragment_count.saturating_sub(1));
113        let capacity = self
114            .prefix
115            .len()
116            .saturating_add(self.fragment_bytes)
117            .saturating_add(separator_bytes)
118            .saturating_add(self.suffix.len());
119        let mut output = String::with_capacity(capacity);
120        output.push_str(&self.prefix);
121        let mut first = true;
122        for tree in self.trees.iter().rev().flatten() {
123            tree.write_to(&mut output, &self.separator, &mut first);
124        }
125        output.push_str(&self.suffix);
126        output
127    }
128}
129
130impl fmt::Debug for PersistentOffset {
131    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
132        f.debug_struct("PersistentOffset")
133            .field("fragment_count", &self.fragment_count)
134            .field("fragment_bytes", &self.fragment_bytes)
135            .finish_non_exhaustive()
136    }
137}
138
139/// Checkpoint state for a source connector.
140///
141/// Captures the connector's position using string key-value pairs.
142/// This is flexible enough to represent:
143/// - Kafka: `{"events:0": "1234", "events:1": "5678"}`
144/// - `PostgreSQL` CDC: `{"lsn": "0/1234ABCD"}`
145/// - File: `{"manifest": "[...]", "file_progress": "..."}`
146///
147/// This payload deliberately has no checkpoint epoch or attempt identifier.
148/// The barrier protocol and top-level manifest bind it to an exact attempt.
149/// Ordinary offsets remain eager; connectors with large append-only positions
150/// can opt into [`PersistentOffset`] without changing the compatibility API.
151pub struct SourceCheckpoint {
152    /// Small connector-specific offset data.
153    offsets: HashMap<String, String>,
154
155    /// Large values represented as immutable serialized fragment logs.
156    persistent_offsets: HashMap<String, PersistentOffset>,
157
158    /// Compatibility view for callers that explicitly request all offsets.
159    materialized_offsets: Arc<OnceLock<HashMap<String, String>>>,
160
161    /// Optional metadata for the checkpoint.
162    metadata: HashMap<String, String>,
163
164    /// Provider-neutral assignment version that owns this source cut.
165    assignment_version: Option<NonZeroU64>,
166}
167
168impl Clone for SourceCheckpoint {
169    fn clone(&self) -> Self {
170        Self {
171            offsets: self.offsets.clone(),
172            persistent_offsets: self.persistent_offsets.clone(),
173            materialized_offsets: Arc::clone(&self.materialized_offsets),
174            metadata: self.metadata.clone(),
175            assignment_version: self.assignment_version,
176        }
177    }
178}
179
180impl Default for SourceCheckpoint {
181    fn default() -> Self {
182        Self {
183            offsets: HashMap::new(),
184            persistent_offsets: HashMap::new(),
185            materialized_offsets: Arc::new(OnceLock::new()),
186            metadata: HashMap::new(),
187            assignment_version: None,
188        }
189    }
190}
191
192impl fmt::Debug for SourceCheckpoint {
193    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
194        f.debug_struct("SourceCheckpoint")
195            .field("offsets", &self.offsets)
196            .field(
197                "persistent_offset_keys",
198                &self.persistent_offsets.keys().collect::<Vec<_>>(),
199            )
200            .field("metadata", &self.metadata)
201            .field("assignment_version", &self.assignment_version)
202            .finish_non_exhaustive()
203    }
204}
205
206impl PartialEq for SourceCheckpoint {
207    fn eq(&self, other: &Self) -> bool {
208        self.offsets() == other.offsets()
209            && self.metadata == other.metadata
210            && self.assignment_version == other.assignment_version
211    }
212}
213
214impl Eq for SourceCheckpoint {}
215
216impl SourceCheckpoint {
217    /// Creates an empty checkpoint.
218    #[must_use]
219    pub fn new() -> Self {
220        Self::default()
221    }
222
223    /// Creates a checkpoint with the given connector-specific offsets.
224    #[must_use]
225    pub fn with_offsets(offsets: HashMap<String, String>) -> Self {
226        Self {
227            offsets,
228            persistent_offsets: HashMap::new(),
229            materialized_offsets: Arc::new(OnceLock::new()),
230            metadata: HashMap::new(),
231            assignment_version: None,
232        }
233    }
234
235    /// Sets an offset value.
236    pub fn set_offset(&mut self, key: impl Into<String>, value: impl Into<String>) {
237        let key = key.into();
238        self.persistent_offsets.remove(&key);
239        self.offsets.insert(key, value.into());
240        self.invalidate_materialized_offsets();
241    }
242
243    /// Sets a large offset value backed by an immutable fragment log.
244    ///
245    /// This replaces any eager value under the same key. The value remains
246    /// lazy across checkpoint clones until [`Self::get_offset`],
247    /// [`Self::offsets`], or [`Self::durable_offsets`] explicitly requests it.
248    pub fn set_persistent_offset(&mut self, key: impl Into<String>, value: PersistentOffset) {
249        let key = key.into();
250        self.offsets.remove(&key);
251        self.persistent_offsets.insert(key, value);
252        self.invalidate_materialized_offsets();
253    }
254
255    /// Gets an offset value.
256    #[must_use]
257    pub fn get_offset(&self, key: &str) -> Option<&str> {
258        if let Some(value) = self.offsets.get(key) {
259            return Some(value);
260        }
261        if !self.persistent_offsets.contains_key(key) {
262            return None;
263        }
264        self.offsets().get(key).map(String::as_str)
265    }
266
267    /// Returns all offsets.
268    #[must_use]
269    pub fn offsets(&self) -> &HashMap<String, String> {
270        if self.persistent_offsets.is_empty() {
271            return &self.offsets;
272        }
273        self.materialized_offsets.get_or_init(|| {
274            let mut offsets = self.offsets.clone();
275            offsets.extend(
276                self.persistent_offsets
277                    .iter()
278                    .map(|(key, value)| (key.clone(), value.materialize())),
279            );
280            offsets
281        })
282    }
283
284    /// Materializes all offsets into an owned map for durable checkpoint
285    /// conversion.
286    ///
287    /// If the compatibility view has not been requested, persistent values are
288    /// written directly into the returned map and are not retained as a second
289    /// full copy inside this checkpoint.
290    #[must_use]
291    pub fn durable_offsets(&self) -> HashMap<String, String> {
292        if let Some(offsets) = self.materialized_offsets.get() {
293            return offsets.clone();
294        }
295        let mut offsets = self.offsets.clone();
296        offsets.extend(
297            self.persistent_offsets
298                .iter()
299                .map(|(key, value)| (key.clone(), value.materialize())),
300        );
301        offsets
302    }
303
304    /// Sets metadata on the checkpoint.
305    pub fn set_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
306        self.metadata.insert(key.into(), value.into());
307    }
308
309    /// Gets metadata from the checkpoint.
310    #[must_use]
311    pub fn get_metadata(&self, key: &str) -> Option<&str> {
312        self.metadata.get(key).map(String::as_str)
313    }
314
315    /// Returns all metadata.
316    #[must_use]
317    pub fn metadata(&self) -> &HashMap<String, String> {
318        &self.metadata
319    }
320
321    /// Binds this source cut to the exact partition-assignment publication that owns it.
322    pub fn bind_assignment_version(&mut self, assignment_version: NonZeroU64) {
323        self.assignment_version = Some(assignment_version);
324    }
325
326    /// Returns the partition-assignment publication that owns this source cut.
327    #[must_use]
328    pub const fn assignment_version(&self) -> Option<NonZeroU64> {
329        self.assignment_version
330    }
331
332    /// Returns `true` if the checkpoint has no offsets.
333    #[must_use]
334    pub fn is_empty(&self) -> bool {
335        self.offsets.is_empty() && self.persistent_offsets.is_empty()
336    }
337
338    fn invalidate_materialized_offsets(&mut self) {
339        self.materialized_offsets = Arc::new(OnceLock::new());
340    }
341
342    #[cfg(test)]
343    fn offsets_are_materialized(&self) -> bool {
344        self.materialized_offsets.get().is_some()
345    }
346}
347
348#[cfg(test)]
349mod tests {
350    use super::*;
351
352    #[test]
353    fn test_source_checkpoint_basic() {
354        let mut cp = SourceCheckpoint::new();
355        cp.set_offset("events:0", "1234");
356        cp.set_offset("events:1", "5678");
357
358        assert_eq!(cp.get_offset("events:0"), Some("1234"));
359        assert_eq!(cp.get_offset("events:1"), Some("5678"));
360        assert_eq!(cp.get_offset("events:2"), None);
361        assert!(!cp.is_empty());
362    }
363
364    #[test]
365    fn test_source_checkpoint_with_offsets() {
366        let mut offsets = HashMap::new();
367        offsets.insert("lsn".to_string(), "0/1234ABCD".to_string());
368
369        let cp = SourceCheckpoint::with_offsets(offsets);
370        assert_eq!(cp.get_offset("lsn"), Some("0/1234ABCD"));
371    }
372
373    #[test]
374    fn test_source_checkpoint_metadata() {
375        let mut cp = SourceCheckpoint::new();
376        cp.set_metadata("connector", "kafka");
377        cp.set_metadata("topic", "events");
378
379        assert_eq!(cp.get_metadata("connector"), Some("kafka"));
380        assert_eq!(cp.get_metadata("topic"), Some("events"));
381    }
382
383    #[test]
384    fn source_checkpoint_assignment_version_is_typed_and_part_of_identity() {
385        let mut checkpoint = SourceCheckpoint::new();
386        assert_eq!(checkpoint.assignment_version(), None);
387
388        let version = NonZeroU64::new(7).unwrap();
389        checkpoint.bind_assignment_version(version);
390        assert_eq!(checkpoint.assignment_version(), Some(version));
391        assert_eq!(checkpoint.clone(), checkpoint);
392        assert!(format!("{checkpoint:?}").contains("assignment_version: Some(7)"));
393
394        let unbound = SourceCheckpoint::new();
395        assert_ne!(checkpoint, unbound);
396        assert_eq!(SourceCheckpoint::default().assignment_version(), None);
397    }
398
399    #[test]
400    fn test_empty_checkpoint() {
401        let cp = SourceCheckpoint::new();
402        assert!(cp.is_empty());
403    }
404
405    #[test]
406    fn persistent_offset_is_lazy_and_durable_conversion_does_not_cache_a_copy() {
407        let mut value = PersistentOffset::new("[", ",", "]");
408        value.push_fragment(r#""first""#);
409        value.push_fragment(r#""second""#);
410        let mut cp = SourceCheckpoint::new();
411        cp.set_offset("small", "7");
412        cp.set_persistent_offset("large", value);
413
414        assert!(!cp.offsets_are_materialized());
415        assert_eq!(cp.get_offset("small"), Some("7"));
416        assert!(!cp.offsets_are_materialized());
417        let durable = cp.durable_offsets();
418        assert_eq!(
419            durable.get("large").map(String::as_str),
420            Some(r#"["first","second"]"#)
421        );
422        assert!(!cp.offsets_are_materialized());
423
424        assert_eq!(cp.get_offset("large"), Some(r#"["first","second"]"#));
425        assert!(cp.offsets_are_materialized());
426    }
427
428    #[test]
429    fn persistent_offset_snapshots_are_immutable() {
430        let mut value = PersistentOffset::new("[", ",", "]");
431        value.push_fragment("1");
432        let snapshot = value.clone();
433        value.push_fragment("2");
434
435        let mut old = SourceCheckpoint::new();
436        old.set_persistent_offset("values", snapshot);
437        let mut new = SourceCheckpoint::new();
438        new.set_persistent_offset("values", value);
439        assert_eq!(old.get_offset("values"), Some("[1]"));
440        assert_eq!(new.get_offset("values"), Some("[1,2]"));
441    }
442
443    #[test]
444    fn persistent_offset_preserves_order_across_tree_carries() {
445        let mut value = PersistentOffset::new("[", ",", "]");
446        for fragment in 0..19 {
447            value.push_fragment(fragment.to_string());
448        }
449        let mut checkpoint = SourceCheckpoint::new();
450        checkpoint.set_persistent_offset("values", value);
451        assert_eq!(
452            checkpoint.get_offset("values"),
453            Some("[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18]")
454        );
455    }
456
457    #[test]
458    fn eager_offsets_do_not_allocate_a_compatibility_copy() {
459        let mut checkpoint = SourceCheckpoint::new();
460        checkpoint.set_offset("partition", "42");
461        assert_eq!(
462            checkpoint.offsets().get("partition").map(String::as_str),
463            Some("42")
464        );
465        assert!(!checkpoint.offsets_are_materialized());
466    }
467
468    #[test]
469    fn mutating_a_checkpoint_does_not_change_its_clone() {
470        let mut original = SourceCheckpoint::new();
471        original.set_offset("position", "1");
472        let snapshot = original.clone();
473        original.set_offset("position", "2");
474
475        assert_eq!(snapshot.get_offset("position"), Some("1"));
476        assert_eq!(original.get_offset("position"), Some("2"));
477    }
478}
laminar_connectors/checkpoint.rs

laminar_connectors/
checkpoint.rs
