1use std::sync::Arc;
8use std::time::Duration;
9
10use arrow_array::RecordBatch;
11use laminar_connectors::checkpoint::SourceCheckpoint;
12use laminar_connectors::config::ConnectorConfig;
13use laminar_connectors::connector::{SourceConnector, SourceContract, SourcePosition};
14use laminar_core::checkpoint::{CheckpointBarrier, CheckpointBarrierInjector};
15use laminar_core::cluster::control::CheckpointAssignmentFence;
16use laminar_core::state::{CheckpointAttempt, CheckpointAttemptRelation};
17use rustc_hash::{FxHashMap, FxHashSet};
18
19#[cfg(feature = "cluster")]
20const CHECKPOINT_CONTROL_DIRECT_FALLBACK: Duration = Duration::from_millis(250);
21const CHECKPOINT_CONTROL_POLL_FALLBACK: Duration = Duration::from_millis(25);
22
23#[cfg(feature = "cluster")]
25#[doc(hidden)]
26pub struct CheckpointControlWake {
27 announcements: Option<
28 tokio::sync::watch::Receiver<Option<laminar_core::cluster::control::BarrierAnnouncement>>,
29 >,
30 fallback: Duration,
31}
32
33#[cfg(not(feature = "cluster"))]
34#[doc(hidden)]
35pub struct CheckpointControlWake {
36 _private: (),
37}
38
39#[cfg(feature = "cluster")]
40impl CheckpointControlWake {
41 #[must_use]
42 pub(crate) fn new(
43 announcements: Option<
44 tokio::sync::watch::Receiver<
45 Option<laminar_core::cluster::control::BarrierAnnouncement>,
46 >,
47 >,
48 ) -> Self {
49 let fallback = if announcements.is_some() {
50 CHECKPOINT_CONTROL_DIRECT_FALLBACK
51 } else {
52 CHECKPOINT_CONTROL_POLL_FALLBACK
53 };
54 Self {
55 announcements,
56 fallback,
57 }
58 }
59
60 pub(crate) async fn wait_until(&mut self, fallback_at: tokio::time::Instant) {
61 let Some(announcements) = self.announcements.as_mut() else {
62 tokio::time::sleep_until(fallback_at).await;
63 return;
64 };
65 tokio::select! {
66 biased;
67 changed = announcements.changed() => {
68 if changed.is_err() {
69 self.announcements = None;
70 self.fallback = CHECKPOINT_CONTROL_POLL_FALLBACK;
71 }
72 }
73 () = tokio::time::sleep_until(fallback_at) => {}
74 }
75 }
76
77 #[must_use]
78 pub(crate) const fn fallback(&self) -> Duration {
79 self.fallback
80 }
81
82 #[must_use]
83 pub(crate) const fn capacity_retry() -> Duration {
84 CHECKPOINT_CONTROL_POLL_FALLBACK
85 }
86}
87
88#[cfg(not(feature = "cluster"))]
89impl CheckpointControlWake {
90 pub(crate) async fn wait_until(&mut self, fallback_at: tokio::time::Instant) {
91 tokio::time::sleep_until(fallback_at).await;
92 }
93
94 #[must_use]
95 pub(crate) const fn fallback() -> Duration {
96 CHECKPOINT_CONTROL_POLL_FALLBACK
97 }
98
99 #[must_use]
100 pub(crate) const fn capacity_retry() -> Duration {
101 CHECKPOINT_CONTROL_POLL_FALLBACK
102 }
103}
104
105#[derive(Debug, Clone, Copy, PartialEq, Eq)]
110pub(crate) enum SourceBarrierSignal {
111 Release(CheckpointAttempt),
112 Stop,
113}
114
115#[derive(Clone)]
117#[doc(hidden)]
118pub struct SourceBarrierControl {
119 injector: CheckpointBarrierInjector,
120 release_tx: tokio::sync::watch::Sender<Option<SourceBarrierSignal>>,
121}
122
123impl SourceBarrierControl {
124 pub(crate) fn new(
125 injector: CheckpointBarrierInjector,
126 release_tx: tokio::sync::watch::Sender<Option<SourceBarrierSignal>>,
127 ) -> Self {
128 Self {
129 injector,
130 release_tx,
131 }
132 }
133
134 #[cfg(feature = "cluster")]
135 pub(crate) fn can_trigger(&self) -> bool {
136 self.injector.can_trigger()
137 }
138
139 #[cfg(feature = "cluster")]
140 pub(crate) fn trigger(&self, barrier: CheckpointBarrier) -> bool {
141 self.injector.trigger(barrier)
142 }
143
144 pub(crate) fn cancel_exact(&self, barrier: CheckpointBarrier) -> bool {
146 self.injector.cancel_exact(barrier)
147 }
148
149 pub(crate) fn release_exact(&self, attempt: CheckpointAttempt) {
151 if !attempt.is_canonical() {
152 return;
153 }
154 self.release_tx.send_if_modified(|signal| match *signal {
155 Some(SourceBarrierSignal::Stop) => false,
156 Some(SourceBarrierSignal::Release(released)) => match released.relation_to(attempt) {
157 CheckpointAttemptRelation::Older => {
158 *signal = Some(SourceBarrierSignal::Release(attempt));
159 true
160 }
161 CheckpointAttemptRelation::Exact
164 | CheckpointAttemptRelation::Newer
165 | CheckpointAttemptRelation::Conflict => false,
166 },
167 None => {
168 *signal = Some(SourceBarrierSignal::Release(attempt));
169 true
170 }
171 });
172 }
173
174 pub(crate) fn stop_hold(&self) {
175 let _ = self.release_tx.send(Some(SourceBarrierSignal::Stop));
176 }
177}
178
179#[cfg(test)]
180mod source_barrier_tests {
181 use super::*;
182
183 #[test]
184 fn noncanonical_release_cannot_publish_or_advance_source_release() {
185 let injector = CheckpointBarrierInjector::new();
186 let (release_tx, release_rx) = tokio::sync::watch::channel(None);
187 let control = SourceBarrierControl::new(injector, release_tx);
188 let invalid = CheckpointAttempt::new(6, 7);
189
190 control.release_exact(invalid);
191 assert_eq!(*release_rx.borrow(), None);
192
193 let canonical = CheckpointAttempt::canonical(5);
194 control.release_exact(canonical);
195 assert_eq!(
196 *release_rx.borrow(),
197 Some(SourceBarrierSignal::Release(canonical))
198 );
199
200 control.release_exact(invalid);
201 assert_eq!(
202 *release_rx.borrow(),
203 Some(SourceBarrierSignal::Release(canonical))
204 );
205 }
206}
207
208#[derive(Debug, Clone, Copy, PartialEq, Eq)]
211pub enum SkipReason {
212 PreservingReplayWindowAfterSinkTimeout,
214}
215
216impl std::fmt::Display for SkipReason {
217 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
218 f.write_str(match self {
219 SkipReason::PreservingReplayWindowAfterSinkTimeout => {
220 "preserving_replay_window_after_sink_timeout"
221 }
222 })
223 }
224}
225
226#[derive(Debug)]
228pub enum BarrierOutcome {
229 Committed(u64),
231 Async,
233 Skipped(SkipReason),
235 CancelledBeforeCapture,
237 Aborted,
239 Failed,
241}
242
243#[derive(Debug)]
250pub(crate) enum CheckpointCompletion {
251 Committed {
253 attempt: CheckpointAttempt,
255 result: crate::checkpoint_coordinator::CheckpointResult,
257 source_checkpoints: FxHashMap<String, SourceCheckpoint>,
259 },
260 Failed {
262 attempt: CheckpointAttempt,
264 error: String,
266 },
267}
268
269#[doc(hidden)]
274#[derive(Debug, Clone, PartialEq, Eq)]
275pub enum CheckpointControlOutcome {
276 Idle,
278 AdmissionFailed { error: String },
280 Started {
282 attempt: CheckpointAttempt,
283 captured: bool,
284 },
285 Aborted { attempt: CheckpointAttempt },
287 Cancelled { attempt: CheckpointAttempt },
289 Failed {
291 attempt: CheckpointAttempt,
292 error: String,
293 },
294}
295
296#[doc(hidden)]
298#[derive(Debug)]
299pub enum CheckpointAssignmentAdmission {
300 Ready(Option<CheckpointAssignmentFence>),
302 Deferred(String),
304 Fault(String),
306}
307
308impl CheckpointCompletion {
309 #[cfg(any(feature = "cluster", test))]
311 #[must_use]
312 pub(crate) fn new(
313 attempt: CheckpointAttempt,
314 source_checkpoints: FxHashMap<String, SourceCheckpoint>,
315 ) -> Self {
316 Self::Committed {
317 attempt,
318 result: crate::checkpoint_coordinator::CheckpointResult {
319 success: true,
320 checkpoint_id: attempt.checkpoint_id,
321 epoch: attempt.epoch,
322 duration: std::time::Duration::ZERO,
323 error: None,
324 failure_disposition: None,
325 },
326 source_checkpoints,
327 }
328 }
329
330 pub(crate) fn validated(
332 admitted: CheckpointAttempt,
333 result: crate::checkpoint_coordinator::CheckpointResult,
334 source_checkpoints: FxHashMap<String, SourceCheckpoint>,
335 ) -> Result<Self, String> {
336 let completed = CheckpointAttempt::new(result.epoch, result.checkpoint_id);
337 if completed != admitted {
338 return Err(format!(
339 "checkpoint completion identity mismatch: admitted epoch={} id={}, \
340 coordinator completed epoch={} id={}",
341 admitted.epoch, admitted.checkpoint_id, completed.epoch, completed.checkpoint_id,
342 ));
343 }
344 Ok(Self::Committed {
345 attempt: admitted,
346 result,
347 source_checkpoints,
348 })
349 }
350
351 pub(crate) fn failed(attempt: CheckpointAttempt, error: impl Into<String>) -> Self {
353 Self::Failed {
354 attempt,
355 error: error.into(),
356 }
357 }
358
359 #[must_use]
361 pub(crate) const fn attempt(&self) -> CheckpointAttempt {
362 match self {
363 Self::Committed { attempt, .. } | Self::Failed { attempt, .. } => *attempt,
364 }
365 }
366}
367
368#[derive(Debug, thiserror::Error)]
370pub enum CycleError {
371 #[error("{0}")]
373 Fatal(String),
374 #[error("{0}")]
376 Recovery(String),
377 #[error("{0}")]
379 Halt(String),
380}
381
382pub struct CycleOutcome {
385 pub results: FxHashMap<Arc<str>, Vec<RecordBatch>>,
387 pub any_failed: bool,
390 pub failed_sources: FxHashSet<Arc<str>>,
392 pub any_deferred: bool,
395 pub deferred_sources: FxHashSet<Arc<str>>,
398}
399
400impl CycleOutcome {
401 #[must_use]
403 pub fn clean(results: FxHashMap<Arc<str>, Vec<RecordBatch>>) -> Self {
404 Self {
405 results,
406 any_failed: false,
407 failed_sources: FxHashSet::default(),
408 any_deferred: false,
409 deferred_sources: FxHashSet::default(),
410 }
411 }
412}
413
414pub struct SourceRegistration {
416 pub name: String,
418 pub connector: Box<dyn SourceConnector>,
420 pub config: ConnectorConfig,
422 pub contract: SourceContract,
424 pub assignment_scoped: bool,
428 pub position: SourcePosition,
430}
431
432#[trait_variant::make(Send)]
435pub trait PipelineCallback: Send + 'static {
436 fn prepare_source_intake(&mut self) -> Result<(), String> {
444 Ok(())
445 }
446
447 async fn execute_cycle(
451 &mut self,
452 source_batches: &FxHashMap<Arc<str>, Vec<RecordBatch>>,
453 watermark: i64,
454 ) -> Result<CycleOutcome, CycleError>;
455
456 async fn drain_checkpoint_edges_until(
462 &mut self,
463 deadline: tokio::time::Instant,
464 ) -> Result<(), CycleError>;
465
466 fn push_to_streams(
472 &self,
473 results: &FxHashMap<Arc<str>, Vec<RecordBatch>>,
474 ) -> Result<(), CycleError> {
475 let _ = results;
476 Ok(())
477 }
478
479 fn update_mv_stores(
485 &self,
486 results: &FxHashMap<Arc<str>, Vec<RecordBatch>>,
487 ) -> Result<(), CycleError> {
488 let _ = results;
489 Ok(())
490 }
491
492 async fn write_to_sinks(
494 &mut self,
495 results: &FxHashMap<Arc<str>, Vec<RecordBatch>>,
496 deadline: Option<tokio::time::Instant>,
497 ) -> Result<(), CycleError>;
498
499 fn extract_watermark(&mut self, source_name: &str, batch: &RecordBatch);
501
502 fn filter_late_rows(&self, source_name: &str, batch: &RecordBatch) -> Option<RecordBatch>;
504
505 fn current_watermark(&self) -> i64;
507
508 fn is_leader(&self) -> bool {
510 true
511 }
512
513 fn is_recovering(&mut self) -> bool {
516 false
517 }
518
519 fn fault_on_cycle_error(&self) -> bool {
522 false
523 }
524
525 fn take_pipeline_fault(&mut self) -> Option<String> {
528 None
529 }
530
531 fn record_checkpoint_failure(&mut self, _checkpoint_id: u64, _reason: &str) {}
534
535 fn record_checkpoint_continuation_fault(&mut self, _attempt: CheckpointAttempt, _reason: &str) {
537 }
538
539 fn record_checkpoint_admission_failure(&mut self, _reason: &str) {}
541
542 fn settle_checkpoint_tail_tasks(
546 &mut self,
547 _abort: bool,
548 ) -> impl std::future::Future<Output = Result<(), String>> + Send {
549 std::future::ready(Ok(()))
550 }
551
552 fn reserve_checkpoint_attempt(
557 &mut self,
558 _attempt_started: std::time::Instant,
559 ) -> impl std::future::Future<Output = Result<CheckpointAttempt, String>> + Send {
560 std::future::ready(Err(
561 "checkpoint coordinator has no durable attempt allocator".into(),
562 ))
563 }
564
565 fn publish_checkpoint_prepare(
568 &mut self,
569 _attempt: CheckpointAttempt,
570 _attempt_started: std::time::Instant,
571 _assignment_fence: Option<CheckpointAssignmentFence>,
572 ) -> impl std::future::Future<Output = Result<(), String>> + Send {
573 std::future::ready(Ok(()))
574 }
575
576 fn abandon_checkpoint_attempt(
579 &mut self,
580 _attempt: CheckpointAttempt,
581 _reason: &str,
582 _assignment_fence: Option<CheckpointAssignmentFence>,
583 ) -> impl std::future::Future<Output = Result<(), String>> + Send;
584
585 fn cancel_source_barrier_attempt(
591 &mut self,
592 _attempt: CheckpointAttempt,
593 _reason: &str,
594 ) -> impl std::future::Future<Output = Result<(), String>> + Send;
595
596 fn resolve_authoritative_follower_abort(
604 &mut self,
605 _attempt: CheckpointAttempt,
606 ) -> Result<(), String> {
607 Err("authoritative follower Abort cleanup is not implemented by this callback".into())
608 }
609
610 fn checkpoint_assignment_for_admission(
612 &mut self,
613 ) -> impl std::future::Future<Output = CheckpointAssignmentAdmission> + Send {
614 std::future::ready(CheckpointAssignmentAdmission::Ready(None))
615 }
616
617 fn checkpoint_control_wake(&self) -> Option<CheckpointControlWake> {
620 None
621 }
622
623 fn tick_idle_watermark(&mut self) {}
625
626 async fn service_checkpoint_control(
631 &mut self,
632 source_offsets: FxHashMap<String, SourceCheckpoint>,
633 ) -> CheckpointControlOutcome;
634
635 async fn checkpoint_with_barrier(
637 &mut self,
638 source_checkpoints: FxHashMap<String, SourceCheckpoint>,
639 attempt: CheckpointAttempt,
640 attempt_started: std::time::Instant,
641 assignment_fence: Option<CheckpointAssignmentFence>,
642 ) -> BarrierOutcome;
643
644 fn record_cycle(&self, events_ingested: u64, batches: u64, elapsed_ns: u64);
646
647 fn note_cycle_error(&self) {}
649
650 fn apply_control(&mut self, msg: super::ControlMsg);
652
653 fn is_backpressured(&self) -> bool {
655 false
656 }
657
658 fn intake_paused(&self) -> bool {
662 false
663 }
664
665 fn has_deferred_input(&self) -> bool {
667 false
668 }
669
670 fn reserve_subscription_cut(&self, _attempt: CheckpointAttempt) -> Result<(), String> {
676 Ok(())
677 }
678
679 fn abort_subscription_cut(&self, _attempt: CheckpointAttempt) {}
681
682 fn publish_barrier(&self, _attempt: CheckpointAttempt) -> Result<(), String> {
688 Ok(())
689 }
690
691 fn invalidate_subscriptions(&self, _reason: &str) {}
693
694 fn settle_sink_epoch_for_shutdown(
697 &mut self,
698 ) -> impl std::future::Future<Output = Result<(), String>> + Send {
699 std::future::ready(Ok(()))
700 }
701
702 fn close_sinks(&mut self) -> impl std::future::Future<Output = Result<(), String>> + Send {
705 std::future::ready(Ok(()))
706 }
707
708 fn set_barrier_injectors(&mut self, injectors: Vec<SourceBarrierControl>) {
710 let _ = injectors;
711 }
712}
713
714#[cfg(all(test, feature = "cluster"))]
715mod tests {
716 use super::*;
717
718 async fn assert_quiet_wake_at(wake: &mut CheckpointControlWake, delay: Duration) {
719 let deadline = tokio::time::Instant::now() + delay;
720 assert!(tokio::time::timeout(
721 delay.saturating_sub(Duration::from_millis(1)),
722 wake.wait_until(deadline),
723 )
724 .await
725 .is_err());
726 tokio::time::timeout(Duration::from_millis(2), wake.wait_until(deadline))
727 .await
728 .expect("checkpoint control fallback exceeded its configured deadline");
729 }
730
731 #[tokio::test(start_paused = true)]
732 async fn checkpoint_control_wake_uses_bounded_direct_and_poll_fallbacks() {
733 let (_direct_tx, direct_rx) = tokio::sync::watch::channel(None);
734 let mut direct = CheckpointControlWake::new(Some(direct_rx));
735 assert_eq!(direct.fallback(), Duration::from_millis(250));
736 assert_quiet_wake_at(&mut direct, Duration::from_millis(250)).await;
737
738 let mut poll = CheckpointControlWake::new(None);
739 assert_eq!(poll.fallback(), Duration::from_millis(25));
740 assert_quiet_wake_at(&mut poll, Duration::from_millis(25)).await;
741 }
742
743 #[tokio::test(start_paused = true)]
744 async fn checkpoint_control_wake_degrades_when_direct_delivery_closes() {
745 let (direct_tx, direct_rx) = tokio::sync::watch::channel(None);
746 let mut wake = CheckpointControlWake::new(Some(direct_rx));
747 drop(direct_tx);
748
749 wake.wait_until(tokio::time::Instant::now() + Duration::from_millis(250))
750 .await;
751 assert_eq!(wake.fallback(), Duration::from_millis(25));
752 assert_quiet_wake_at(&mut wake, Duration::from_millis(25)).await;
753 }
754}