1#![allow(clippy::disallowed_types)] use std::collections::HashMap;
6use std::fmt;
7use std::num::NonZeroU64;
8use std::sync::{Arc, OnceLock};
9
10use crate::error::ConnectorError;
11
12enum OffsetTree {
13 Leaf(Arc<str>),
14 Branch { older: Arc<Self>, newer: Arc<Self> },
15}
16
17impl OffsetTree {
18 fn write_to(&self, output: &mut String, separator: &str, first: &mut bool) {
19 match self {
20 Self::Leaf(value) => {
21 if *first {
22 *first = false;
23 } else {
24 output.push_str(separator);
25 }
26 output.push_str(value);
27 }
28 Self::Branch { older, newer } => {
29 older.write_to(output, separator, first);
30 newer.write_to(output, separator, first);
31 }
32 }
33 }
34}
35
36#[derive(Clone)]
41pub struct PersistentOffset {
42 trees: Arc<Vec<Option<Arc<OffsetTree>>>>,
45 prefix: Arc<str>,
46 separator: Arc<str>,
47 suffix: Arc<str>,
48 fragment_count: usize,
49 fragment_bytes: usize,
50}
51
52impl PersistentOffset {
53 #[must_use]
58 pub fn new(
59 prefix: impl Into<Arc<str>>,
60 separator: impl Into<Arc<str>>,
61 suffix: impl Into<Arc<str>>,
62 ) -> Self {
63 Self {
64 trees: Arc::new(Vec::new()),
65 prefix: prefix.into(),
66 separator: separator.into(),
67 suffix: suffix.into(),
68 fragment_count: 0,
69 fragment_bytes: 0,
70 }
71 }
72
73 pub fn push_fragment(&mut self, fragment: impl Into<Arc<str>>) {
76 let value = fragment.into();
77 self.fragment_bytes = self.fragment_bytes.saturating_add(value.len());
81 self.fragment_count = self.fragment_count.saturating_add(1);
82
83 let trees = Arc::make_mut(&mut self.trees);
84 let mut carry = Arc::new(OffsetTree::Leaf(value));
85 let mut level = 0;
86 loop {
87 if level == trees.len() {
88 trees.push(Some(carry));
89 break;
90 }
91 if let Some(older) = trees[level].take() {
92 carry = Arc::new(OffsetTree::Branch {
93 older,
94 newer: carry,
95 });
96 level += 1;
97 } else {
98 trees[level] = Some(carry);
99 break;
100 }
101 }
102 }
103
104 #[must_use]
106 pub const fn fragment_count(&self) -> usize {
107 self.fragment_count
108 }
109
110 fn materialize(&self) -> String {
111 let separator_bytes = self
112 .separator
113 .len()
114 .saturating_mul(self.fragment_count.saturating_sub(1));
115 let capacity = self
116 .prefix
117 .len()
118 .saturating_add(self.fragment_bytes)
119 .saturating_add(separator_bytes)
120 .saturating_add(self.suffix.len());
121 let mut output = String::with_capacity(capacity);
122 output.push_str(&self.prefix);
123 let mut first = true;
124 for tree in self.trees.iter().rev().flatten() {
125 tree.write_to(&mut output, &self.separator, &mut first);
126 }
127 output.push_str(&self.suffix);
128 output
129 }
130}
131
132impl fmt::Debug for PersistentOffset {
133 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
134 f.debug_struct("PersistentOffset")
135 .field("fragment_count", &self.fragment_count)
136 .field("fragment_bytes", &self.fragment_bytes)
137 .finish_non_exhaustive()
138 }
139}
140
141pub struct SourceCheckpoint {
154 offsets: HashMap<String, String>,
156
157 persistent_offsets: HashMap<String, PersistentOffset>,
159
160 materialized_offsets: Arc<OnceLock<HashMap<String, String>>>,
162
163 metadata: HashMap<String, String>,
165
166 input_channels: Option<Arc<[Vec<u8>]>>,
168
169 assignment_version: Option<NonZeroU64>,
171}
172
173#[derive(Clone, Debug, PartialEq, Eq)]
179pub struct SourceCheckpointDelta {
180 assignment_version: NonZeroU64,
181 input_channels: Arc<[Vec<u8>]>,
182 offset_changes: HashMap<String, Option<String>>,
183}
184
185impl SourceCheckpointDelta {
186 pub fn new(
191 assignment_version: NonZeroU64,
192 input_channels: Arc<[Vec<u8>]>,
193 offset_changes: HashMap<String, Option<String>>,
194 ) -> Result<Self, ConnectorError> {
195 if offset_changes.is_empty() {
196 return Err(ConnectorError::InvalidState {
197 expected: "at least one changed source offset".into(),
198 actual: "an empty incremental source cursor".into(),
199 });
200 }
201 Ok(Self {
202 assignment_version,
203 input_channels,
204 offset_changes,
205 })
206 }
207
208 #[must_use]
210 pub const fn assignment_version(&self) -> NonZeroU64 {
211 self.assignment_version
212 }
213
214 #[must_use]
216 pub const fn input_channels_arc(&self) -> &Arc<[Vec<u8>]> {
217 &self.input_channels
218 }
219
220 #[cfg(all(test, feature = "kafka"))]
221 pub(crate) fn changes(&self) -> &HashMap<String, Option<String>> {
222 &self.offset_changes
223 }
224
225 pub fn validate_base(&self, base: &SourceCheckpoint) -> Result<(), ConnectorError> {
232 let Some(base_version) = base.assignment_version else {
233 return Err(ConnectorError::InvalidState {
234 expected: format!(
235 "complete source checkpoint for assignment {}",
236 self.assignment_version
237 ),
238 actual: "an unbound source checkpoint".into(),
239 });
240 };
241 if base_version != self.assignment_version {
242 return Err(ConnectorError::InvalidState {
243 expected: format!("source assignment {}", self.assignment_version),
244 actual: format!("source assignment {base_version}"),
245 });
246 }
247 let Some(base_channels) = base.input_channels.as_ref() else {
248 return Err(ConnectorError::InvalidState {
249 expected: "a complete source checkpoint with input channels".into(),
250 actual: "a source checkpoint without input channels".into(),
251 });
252 };
253 if !Arc::ptr_eq(base_channels, &self.input_channels)
254 && base_channels.as_ref() != self.input_channels.as_ref()
255 {
256 return Err(ConnectorError::InvalidState {
257 expected: "the input-channel roster of the complete source checkpoint".into(),
258 actual: "a different incremental source roster".into(),
259 });
260 }
261 Ok(())
262 }
263
264 pub fn merge(&mut self, newer: Self) -> Result<(), ConnectorError> {
270 if self.assignment_version != newer.assignment_version
271 || (!Arc::ptr_eq(&self.input_channels, &newer.input_channels)
272 && self.input_channels.as_ref() != newer.input_channels.as_ref())
273 {
274 return Err(ConnectorError::InvalidState {
275 expected: format!(
276 "source assignment {} and its input channels",
277 self.assignment_version
278 ),
279 actual: format!(
280 "source assignment {} or a different input-channel roster",
281 newer.assignment_version
282 ),
283 });
284 }
285 self.offset_changes.extend(newer.offset_changes);
286 Ok(())
287 }
288}
289
290impl Clone for SourceCheckpoint {
291 fn clone(&self) -> Self {
292 Self {
293 offsets: self.offsets.clone(),
294 persistent_offsets: self.persistent_offsets.clone(),
295 materialized_offsets: Arc::clone(&self.materialized_offsets),
296 metadata: self.metadata.clone(),
297 input_channels: self.input_channels.clone(),
298 assignment_version: self.assignment_version,
299 }
300 }
301}
302
303impl Default for SourceCheckpoint {
304 fn default() -> Self {
305 Self {
306 offsets: HashMap::new(),
307 persistent_offsets: HashMap::new(),
308 materialized_offsets: Arc::new(OnceLock::new()),
309 metadata: HashMap::new(),
310 input_channels: None,
311 assignment_version: None,
312 }
313 }
314}
315
316impl fmt::Debug for SourceCheckpoint {
317 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
318 let input_channel_count = self.input_channels.as_ref().map(|channels| channels.len());
319 f.debug_struct("SourceCheckpoint")
320 .field("offsets", &self.offsets)
321 .field(
322 "persistent_offset_keys",
323 &self.persistent_offsets.keys().collect::<Vec<_>>(),
324 )
325 .field("metadata", &self.metadata)
326 .field("input_channel_count", &input_channel_count)
327 .field("assignment_version", &self.assignment_version)
328 .finish_non_exhaustive()
329 }
330}
331
332impl PartialEq for SourceCheckpoint {
333 fn eq(&self, other: &Self) -> bool {
334 self.offsets() == other.offsets()
335 && self.metadata == other.metadata
336 && self.input_channels == other.input_channels
337 && self.assignment_version == other.assignment_version
338 }
339}
340
341impl Eq for SourceCheckpoint {}
342
343impl SourceCheckpoint {
344 #[must_use]
346 pub fn new() -> Self {
347 Self::default()
348 }
349
350 #[must_use]
352 pub fn with_offsets(offsets: HashMap<String, String>) -> Self {
353 Self {
354 offsets,
355 persistent_offsets: HashMap::new(),
356 materialized_offsets: Arc::new(OnceLock::new()),
357 metadata: HashMap::new(),
358 input_channels: None,
359 assignment_version: None,
360 }
361 }
362
363 pub fn set_offset(&mut self, key: impl Into<String>, value: impl Into<String>) {
365 let key = key.into();
366 self.persistent_offsets.remove(&key);
367 self.offsets.insert(key, value.into());
368 self.invalidate_materialized_offsets();
369 }
370
371 pub fn set_persistent_offset(&mut self, key: impl Into<String>, value: PersistentOffset) {
377 let key = key.into();
378 self.offsets.remove(&key);
379 self.persistent_offsets.insert(key, value);
380 self.invalidate_materialized_offsets();
381 }
382
383 #[must_use]
385 pub fn get_offset(&self, key: &str) -> Option<&str> {
386 if let Some(value) = self.offsets.get(key) {
387 return Some(value);
388 }
389 if !self.persistent_offsets.contains_key(key) {
390 return None;
391 }
392 self.offsets().get(key).map(String::as_str)
393 }
394
395 #[must_use]
397 pub fn offsets(&self) -> &HashMap<String, String> {
398 if self.persistent_offsets.is_empty() {
399 return &self.offsets;
400 }
401 self.materialized_offsets.get_or_init(|| {
402 let mut offsets = self.offsets.clone();
403 offsets.extend(
404 self.persistent_offsets
405 .iter()
406 .map(|(key, value)| (key.clone(), value.materialize())),
407 );
408 offsets
409 })
410 }
411
412 #[must_use]
419 pub fn durable_offsets(&self) -> HashMap<String, String> {
420 if let Some(offsets) = self.materialized_offsets.get() {
421 return offsets.clone();
422 }
423 let mut offsets = self.offsets.clone();
424 offsets.extend(
425 self.persistent_offsets
426 .iter()
427 .map(|(key, value)| (key.clone(), value.materialize())),
428 );
429 offsets
430 }
431
432 pub fn set_metadata(&mut self, key: impl Into<String>, value: impl Into<String>) {
434 self.metadata.insert(key.into(), value.into());
435 }
436
437 #[must_use]
439 pub fn get_metadata(&self, key: &str) -> Option<&str> {
440 self.metadata.get(key).map(String::as_str)
441 }
442
443 #[must_use]
445 pub fn metadata(&self) -> &HashMap<String, String> {
446 &self.metadata
447 }
448
449 pub fn set_input_channels(
455 &mut self,
456 channels: impl Into<Arc<[Vec<u8>]>>,
457 ) -> Result<(), ConnectorError> {
458 let channels = channels.into();
459 if channels.iter().any(Vec::is_empty) {
460 return Err(ConnectorError::InvalidState {
461 expected: "non-empty opaque input-channel identities".into(),
462 actual: "an empty input-channel identity".into(),
463 });
464 }
465 if channels
466 .iter()
467 .any(|channel| channel == laminar_core::checkpoint::SINGLETON_WATERMARK_CHANNEL)
468 {
469 return Err(ConnectorError::InvalidState {
470 expected: "physical input-channel identities".into(),
471 actual: "the reserved logical watermark channel".into(),
472 });
473 }
474 if !channels.windows(2).all(|pair| pair[0] < pair[1]) {
475 return Err(ConnectorError::InvalidState {
476 expected: "strictly ordered unique input-channel identities".into(),
477 actual: "noncanonical input-channel inventory".into(),
478 });
479 }
480 self.input_channels = Some(channels);
481 Ok(())
482 }
483
484 #[must_use]
486 pub fn input_channels(&self) -> Option<&[Vec<u8>]> {
487 self.input_channels.as_deref()
488 }
489
490 #[must_use]
492 pub const fn input_channels_arc(&self) -> Option<&Arc<[Vec<u8>]>> {
493 self.input_channels.as_ref()
494 }
495
496 pub fn bind_assignment_version(&mut self, assignment_version: NonZeroU64) {
498 self.assignment_version = Some(assignment_version);
499 }
500
501 #[must_use]
503 pub const fn assignment_version(&self) -> Option<NonZeroU64> {
504 self.assignment_version
505 }
506
507 pub fn apply_delta(&mut self, delta: SourceCheckpointDelta) -> Result<(), ConnectorError> {
514 delta.validate_base(self)?;
515 for (key, value) in delta.offset_changes {
516 self.persistent_offsets.remove(&key);
517 if let Some(value) = value {
518 self.offsets.insert(key, value);
519 } else {
520 self.offsets.remove(&key);
521 }
522 }
523 self.invalidate_materialized_offsets();
524 Ok(())
525 }
526
527 #[must_use]
529 pub fn is_empty(&self) -> bool {
530 self.offsets.is_empty() && self.persistent_offsets.is_empty()
531 }
532
533 fn invalidate_materialized_offsets(&mut self) {
534 self.materialized_offsets = Arc::new(OnceLock::new());
535 }
536
537 #[cfg(test)]
538 fn offsets_are_materialized(&self) -> bool {
539 self.materialized_offsets.get().is_some()
540 }
541}
542
543#[cfg(test)]
544mod tests {
545 use super::*;
546
547 #[test]
548 fn test_source_checkpoint_basic() {
549 let mut cp = SourceCheckpoint::new();
550 cp.set_offset("events:0", "1234");
551 cp.set_offset("events:1", "5678");
552
553 assert_eq!(cp.get_offset("events:0"), Some("1234"));
554 assert_eq!(cp.get_offset("events:1"), Some("5678"));
555 assert_eq!(cp.get_offset("events:2"), None);
556 assert!(!cp.is_empty());
557 }
558
559 #[test]
560 fn test_source_checkpoint_with_offsets() {
561 let mut offsets = HashMap::new();
562 offsets.insert("lsn".to_string(), "0/1234ABCD".to_string());
563
564 let cp = SourceCheckpoint::with_offsets(offsets);
565 assert_eq!(cp.get_offset("lsn"), Some("0/1234ABCD"));
566 }
567
568 #[test]
569 fn test_source_checkpoint_metadata() {
570 let mut cp = SourceCheckpoint::new();
571 cp.set_metadata("connector", "kafka");
572 cp.set_metadata("topic", "events");
573
574 assert_eq!(cp.get_metadata("connector"), Some("kafka"));
575 assert_eq!(cp.get_metadata("topic"), Some("events"));
576 }
577
578 #[test]
579 fn input_channels_are_canonical_and_part_of_checkpoint_identity() {
580 let mut checkpoint = SourceCheckpoint::new();
581 checkpoint
582 .set_input_channels(vec![b"partition-1".to_vec(), b"partition-2".to_vec()])
583 .unwrap();
584
585 assert_eq!(
586 checkpoint.input_channels(),
587 Some([b"partition-1".to_vec(), b"partition-2".to_vec()].as_slice())
588 );
589 assert_ne!(checkpoint, SourceCheckpoint::new());
590 let cloned = checkpoint.clone();
591 assert_eq!(cloned, checkpoint);
592 assert!(Arc::ptr_eq(
593 cloned.input_channels.as_ref().unwrap(),
594 checkpoint.input_channels.as_ref().unwrap()
595 ));
596
597 assert!(SourceCheckpoint::new()
598 .set_input_channels(vec![b"partition-2".to_vec(), b"partition-1".to_vec()])
599 .is_err());
600 assert!(SourceCheckpoint::new()
601 .set_input_channels(vec![Vec::new()])
602 .is_err());
603 assert!(SourceCheckpoint::new()
604 .set_input_channels(vec![
605 laminar_core::checkpoint::SINGLETON_WATERMARK_CHANNEL.to_vec(),
606 ])
607 .is_err());
608
609 let mut owned_empty = SourceCheckpoint::new();
610 owned_empty.set_input_channels(Vec::new()).unwrap();
611 assert_eq!(owned_empty.input_channels(), Some([].as_slice()));
612 assert_ne!(owned_empty, SourceCheckpoint::new());
613 }
614
615 #[test]
616 fn source_checkpoint_assignment_version_is_typed_and_part_of_identity() {
617 let mut checkpoint = SourceCheckpoint::new();
618 assert_eq!(checkpoint.assignment_version(), None);
619
620 let version = NonZeroU64::new(7).unwrap();
621 checkpoint.bind_assignment_version(version);
622 assert_eq!(checkpoint.assignment_version(), Some(version));
623 assert_eq!(checkpoint.clone(), checkpoint);
624 assert!(format!("{checkpoint:?}").contains("assignment_version: Some(7)"));
625
626 let unbound = SourceCheckpoint::new();
627 assert_ne!(checkpoint, unbound);
628 assert_eq!(SourceCheckpoint::default().assignment_version(), None);
629 }
630
631 #[test]
632 fn test_empty_checkpoint() {
633 let cp = SourceCheckpoint::new();
634 assert!(cp.is_empty());
635 }
636
637 #[test]
638 fn persistent_offset_is_lazy_and_durable_conversion_does_not_cache_a_copy() {
639 let mut value = PersistentOffset::new("[", ",", "]");
640 value.push_fragment(r#""first""#);
641 value.push_fragment(r#""second""#);
642 let mut cp = SourceCheckpoint::new();
643 cp.set_offset("small", "7");
644 cp.set_persistent_offset("large", value);
645
646 assert!(!cp.offsets_are_materialized());
647 assert_eq!(cp.get_offset("small"), Some("7"));
648 assert!(!cp.offsets_are_materialized());
649 let durable = cp.durable_offsets();
650 assert_eq!(
651 durable.get("large").map(String::as_str),
652 Some(r#"["first","second"]"#)
653 );
654 assert!(!cp.offsets_are_materialized());
655
656 assert_eq!(cp.get_offset("large"), Some(r#"["first","second"]"#));
657 assert!(cp.offsets_are_materialized());
658 }
659
660 #[test]
661 fn persistent_offset_snapshots_are_immutable() {
662 let mut value = PersistentOffset::new("[", ",", "]");
663 value.push_fragment("1");
664 let snapshot = value.clone();
665 value.push_fragment("2");
666
667 let mut old = SourceCheckpoint::new();
668 old.set_persistent_offset("values", snapshot);
669 let mut new = SourceCheckpoint::new();
670 new.set_persistent_offset("values", value);
671 assert_eq!(old.get_offset("values"), Some("[1]"));
672 assert_eq!(new.get_offset("values"), Some("[1,2]"));
673 }
674
675 #[test]
676 fn persistent_offset_preserves_order_across_tree_carries() {
677 let mut value = PersistentOffset::new("[", ",", "]");
678 for fragment in 0..19 {
679 value.push_fragment(fragment.to_string());
680 }
681 let mut checkpoint = SourceCheckpoint::new();
682 checkpoint.set_persistent_offset("values", value);
683 assert_eq!(
684 checkpoint.get_offset("values"),
685 Some("[0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,16,17,18]")
686 );
687 }
688
689 #[test]
690 fn eager_offsets_do_not_allocate_a_materialized_copy() {
691 let mut checkpoint = SourceCheckpoint::new();
692 checkpoint.set_offset("partition", "42");
693 assert_eq!(
694 checkpoint.offsets().get("partition").map(String::as_str),
695 Some("42")
696 );
697 assert!(!checkpoint.offsets_are_materialized());
698 }
699
700 #[test]
701 fn mutating_a_checkpoint_does_not_change_its_clone() {
702 let mut original = SourceCheckpoint::new();
703 original.set_offset("position", "1");
704 let snapshot = original.clone();
705 original.set_offset("position", "2");
706
707 assert_eq!(snapshot.get_offset("position"), Some("1"));
708 assert_eq!(original.get_offset("position"), Some("2"));
709 }
710}