Skip to main content

laminar_db/checkpoint_coordinator/
coordinator.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::Arc;
3use std::time::Duration;
4
5use laminar_connectors::checkpoint::SourceCheckpoint;
6#[cfg(feature = "cluster")]
7use laminar_core::checkpoint::CommittedCheckpointRef;
8use laminar_core::checkpoint::{
9    CheckpointAttempt, CheckpointManifest, CheckpointStore, ConnectorCheckpoint, PipelineIdentity,
10};
11use tracing::warn;
12
13use super::retention::run_gc_worker;
14use super::{
15    CheckpointConfig, CheckpointCoordinator, CheckpointPhase, DbError, EpochAllocator,
16    REFERENCED_CHUNK_REBASE_THRESHOLD,
17};
18
19impl CheckpointCoordinator {
20    pub fn new(config: CheckpointConfig, store: Box<dyn CheckpointStore>) -> Result<Self, DbError> {
21        laminar_core::checkpoint::checkpoint_store::validate_max_checkpoint_node_data_bytes(
22            config.max_node_data_bytes,
23        )
24        .map_err(|error| DbError::Config(format!("checkpoint.max_node_data_bytes: {error}")))?;
25        if store.max_node_data_bytes() != config.max_node_data_bytes {
26            return Err(DbError::Config(format!(
27                "checkpoint store node-data limit {} does not match checkpoint.max_node_data_bytes {}",
28                store.max_node_data_bytes(),
29                config.max_node_data_bytes
30            )));
31        }
32        if store.participant_id() == 0 {
33            return Err(DbError::Config(
34                "checkpoint participant ID must be nonzero".into(),
35            ));
36        }
37        let vnode_count = store.key_group_count().get();
38        #[cfg(feature = "cluster")]
39        let recovery_graph_payload_limit =
40            usize::try_from(config.max_node_data_bytes).unwrap_or(usize::MAX);
41        let store: Arc<dyn CheckpointStore> = Arc::from(store);
42        let (gc_requests, gc_receiver) = tokio::sync::watch::channel(None);
43        let read_pins = Arc::<super::committed_reads::ReadPins>::default();
44        let gc_task = tokio::spawn(run_gc_worker(
45            Arc::clone(&store),
46            gc_receiver,
47            Arc::clone(&read_pins),
48        ));
49        Ok(Self {
50            allocator: Arc::new(EpochAllocator::new()),
51            config,
52            store,
53            phase: CheckpointPhase::Idle,
54            sinks: Vec::new(),
55            assignment_version: 0,
56            assignment_scoped_sources: HashSet::new(),
57            owned_vnodes: (0..u32::from(vnode_count)).collect(),
58            pipeline_identity: None,
59            deployment_id: None,
60            decision_store: None,
61            active_sink_witness: None,
62            active_sink_artifact_intents: None,
63            prepared: HashMap::new(),
64            last_committed_manifest: None,
65            view_log_predecessor: None,
66            last_committed_ref: None,
67            last_committed_frontiers: None,
68            failure_requires_recovery: false,
69            local_watermark: laminar_core::checkpoint::CheckpointWatermark::Uninitialized,
70            #[cfg(feature = "cluster")]
71            cluster_controller: None,
72            #[cfg(feature = "cluster")]
73            recovery_graph_payload_limit,
74            checkpoints_completed: 0,
75            checkpoints_failed: 0,
76            last_checkpoint_duration: None,
77            duration_histogram: DurationHistogram::new(),
78            total_bytes_written: 0,
79            prom: None,
80            gc_requests,
81            gc_task,
82            committed_cut: tokio::sync::watch::Sender::new(None),
83            #[cfg(feature = "cluster")]
84            head_follower: None,
85            read_pins,
86        })
87    }
88
89    pub fn set_decision_store(
90        &mut self,
91        store: Arc<laminar_core::checkpoint_decision::CheckpointDecisionStore>,
92    ) -> Result<(), DbError> {
93        self.allocator.bind_decision_store(Arc::clone(&store))?;
94        self.decision_store = Some(store);
95        Ok(())
96    }
97
98    pub async fn bind_durable_decision_store(
99        &mut self,
100        store: Arc<laminar_core::checkpoint_decision::CheckpointDecisionStore>,
101    ) -> Result<(), DbError> {
102        let deployment_id = store
103            .load_or_create_deployment_id()
104            .await
105            .map_err(|error| {
106                DbError::Checkpoint(format!("load checkpoint deployment identity: {error}"))
107            })?;
108        self.set_decision_store(store)?;
109        self.bind_deployment_id(deployment_id)
110    }
111
112    pub(crate) fn bind_pipeline_identity(
113        &mut self,
114        identity: PipelineIdentity,
115    ) -> Result<(), DbError> {
116        match self.pipeline_identity.as_ref() {
117            None => {
118                self.pipeline_identity = Some(identity);
119                Ok(())
120            }
121            Some(current) if current == &identity => Ok(()),
122            Some(_) => Err(DbError::Checkpoint(
123                "pipeline identity cannot change while checkpointing is active".into(),
124            )),
125        }
126    }
127
128    pub(crate) fn bound_pipeline_identity(&self) -> Result<PipelineIdentity, DbError> {
129        self.pipeline_identity.clone().ok_or_else(|| {
130            DbError::Checkpoint("checkpoint coordinator has no pipeline identity".into())
131        })
132    }
133
134    pub(crate) fn bind_deployment_id(&mut self, deployment_id: String) -> Result<(), DbError> {
135        let canonical = uuid::Uuid::parse_str(&deployment_id)
136            .is_ok_and(|id| !id.is_nil() && id.to_string() == deployment_id);
137        if !canonical {
138            return Err(DbError::Checkpoint(
139                "checkpoint deployment identity must be a canonical non-nil UUID".into(),
140            ));
141        }
142        match self.deployment_id.as_ref() {
143            None => {
144                self.deployment_id = Some(deployment_id);
145                Ok(())
146            }
147            Some(current) if current == &deployment_id => Ok(()),
148            Some(_) => Err(DbError::Checkpoint(
149                "deployment identity cannot change while checkpointing is active".into(),
150            )),
151        }
152    }
153
154    pub(super) fn expected_pipeline_identity(&self) -> Result<PipelineIdentity, DbError> {
155        self.bound_pipeline_identity()
156    }
157
158    pub(super) fn expected_deployment_id(&self) -> Result<&str, DbError> {
159        self.deployment_id.as_deref().ok_or_else(|| {
160            DbError::Checkpoint("checkpoint deployment identity is not bound".into())
161        })
162    }
163
164    pub(crate) fn bound_deployment_id(&self) -> Result<&str, DbError> {
165        self.expected_deployment_id()
166    }
167
168    #[cfg(feature = "cluster")]
169    pub fn set_cluster_controller(
170        &mut self,
171        controller: Arc<laminar_core::cluster::control::ClusterController>,
172    ) {
173        self.cluster_controller = Some(controller);
174    }
175
176    #[must_use]
177    pub(crate) fn participant_id(&self) -> u64 {
178        self.store.participant_id()
179    }
180
181    pub fn set_assignment_version(&mut self, version: u64) {
182        self.assignment_version = version;
183    }
184
185    pub(crate) fn set_assignment_scoped_sources(
186        &mut self,
187        sources: impl IntoIterator<Item = String>,
188    ) {
189        self.assignment_scoped_sources = sources.into_iter().collect();
190    }
191
192    pub fn set_vnode_set(&mut self, mut vnodes: Vec<u32>) {
193        vnodes.sort_unstable();
194        vnodes.dedup();
195        self.owned_vnodes = vnodes;
196    }
197
198    pub fn set_local_watermark(
199        &mut self,
200        watermark: laminar_core::checkpoint::CheckpointWatermark,
201    ) {
202        self.local_watermark = watermark;
203    }
204
205    pub fn set_metrics(&mut self, prom: Arc<crate::engine_metrics::EngineMetrics>) {
206        self.prom = Some(prom);
207    }
208
209    fn emit_checkpoint_metrics(
210        &self,
211        success: bool,
212        attempt: CheckpointAttempt,
213        duration: Duration,
214        checkpoint_bytes: Option<u64>,
215    ) {
216        let Some(metrics) = self.prom.as_ref() else {
217            return;
218        };
219        if success {
220            metrics.checkpoints_completed.inc();
221            if let Some(checkpoint_bytes) = checkpoint_bytes {
222                metrics
223                    .checkpoint_size_bytes
224                    .set(i64::try_from(checkpoint_bytes).unwrap_or(i64::MAX));
225            }
226        } else {
227            metrics.checkpoints_failed.inc();
228            warn!(
229                checkpoint_id = attempt.checkpoint_id,
230                epoch = attempt.epoch,
231                "checkpoint failure metric recorded"
232            );
233        }
234        metrics
235            .checkpoint_epoch
236            .set(i64::try_from(attempt.epoch).unwrap_or(i64::MAX));
237        metrics.checkpoint_duration.observe(duration.as_secs_f64());
238    }
239
240    pub(super) fn record_checkpoint_outcome(
241        &mut self,
242        success: bool,
243        attempt: CheckpointAttempt,
244        duration: Duration,
245        checkpoint_bytes: Option<u64>,
246    ) {
247        if success {
248            self.checkpoints_completed = self.checkpoints_completed.saturating_add(1);
249        } else {
250            self.checkpoints_failed = self.checkpoints_failed.saturating_add(1);
251        }
252        self.last_checkpoint_duration = Some(duration);
253        self.duration_histogram.record(duration);
254        self.emit_checkpoint_metrics(success, attempt, duration, checkpoint_bytes);
255    }
256
257    #[must_use]
258    pub fn phase(&self) -> CheckpointPhase {
259        self.phase
260    }
261
262    #[must_use]
263    pub fn epoch(&self) -> u64 {
264        self.allocator.peek_epoch()
265    }
266
267    #[must_use]
268    pub fn config(&self) -> &CheckpointConfig {
269        &self.config
270    }
271
272    #[must_use]
273    pub(crate) fn epoch_allocator(&self) -> Arc<EpochAllocator> {
274        Arc::clone(&self.allocator)
275    }
276
277    #[must_use]
278    pub fn store(&self) -> &dyn CheckpointStore {
279        self.store.as_ref()
280    }
281}
282
283impl std::fmt::Debug for CheckpointCoordinator {
284    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
285        formatter
286            .debug_struct("CheckpointCoordinator")
287            .field("phase", &self.phase)
288            .field("participant", &self.store.participant_id())
289            .field("sinks", &self.sinks.len())
290            .field("completed", &self.checkpoints_completed)
291            .field("failed", &self.checkpoints_failed)
292            .finish_non_exhaustive()
293    }
294}
295
296impl Drop for CheckpointCoordinator {
297    fn drop(&mut self) {
298        self.gc_task.abort();
299        #[cfg(feature = "cluster")]
300        if let Some(follower) = &self.head_follower {
301            follower.abort();
302        }
303    }
304}
305
306pub(super) struct DurationHistogram {
307    samples: Box<[u64; Self::CAPACITY]>,
308    cursor: usize,
309    count: usize,
310}
311
312impl DurationHistogram {
313    const CAPACITY: usize = 100;
314
315    fn new() -> Self {
316        Self {
317            samples: Box::new([0; Self::CAPACITY]),
318            cursor: 0,
319            count: 0,
320        }
321    }
322
323    fn record(&mut self, duration: Duration) {
324        self.samples[self.cursor] = u64::try_from(duration.as_micros()).unwrap_or(u64::MAX);
325        self.cursor = (self.cursor + 1) % Self::CAPACITY;
326        self.count = self.count.saturating_add(1).min(Self::CAPACITY);
327    }
328
329    fn percentiles(&self) -> (u64, u64, u64) {
330        if self.count == 0 {
331            return (0, 0, 0);
332        }
333        let mut values = self.samples[..self.count].to_vec();
334        values.sort_unstable();
335        let at = |numerator: usize| {
336            let index = (self.count.saturating_sub(1) * numerator).div_ceil(100);
337            values[index]
338        };
339        (at(50), at(95), at(99))
340    }
341}
342
343#[derive(Debug, Clone, serde::Serialize)]
344pub struct CheckpointStats {
345    pub completed: u64,
346    pub failed: u64,
347    pub last_duration: Option<Duration>,
348    pub duration_p50_ms: u64,
349    pub duration_p95_ms: u64,
350    pub duration_p99_ms: u64,
351    pub total_bytes_written: u64,
352    pub current_phase: CheckpointPhase,
353    pub current_epoch: u64,
354}
355
356impl CheckpointCoordinator {
357    #[cfg(feature = "cluster")]
358    pub(crate) fn last_committed_ref(&self) -> Option<&CommittedCheckpointRef> {
359        self.last_committed_ref.as_ref()
360    }
361
362    #[cfg(all(test, feature = "cluster"))]
363    pub(crate) fn set_last_committed_ref_for_test(&mut self, reference: CommittedCheckpointRef) {
364        self.last_committed_frontiers = None;
365        self.last_committed_ref = Some(reference);
366    }
367
368    pub(crate) fn last_committed_manifest(&self) -> Option<&CheckpointManifest> {
369        self.last_committed_manifest.as_deref()
370    }
371
372    pub(crate) fn committed_manifest_needs_vnode_rebase(&self, attempt: CheckpointAttempt) -> bool {
373        self.last_committed_manifest
374            .as_ref()
375            .is_some_and(|manifest| {
376                manifest.checkpoint_id == attempt.checkpoint_id
377                    && manifest.epoch == attempt.epoch
378                    && (manifest.referenced_chunks.len() >= REFERENCED_CHUNK_REBASE_THRESHOLD
379                        // A migration retains its historical parent reference. The first target
380                        // checkpoint must capture full state under its own pipeline identity.
381                        || self
382                            .pipeline_identity
383                            .as_ref()
384                            .is_some_and(|bound| bound != &manifest.pipeline_identity))
385            })
386    }
387
388    #[must_use]
389    pub fn stats(&self) -> CheckpointStats {
390        let (p50, p95, p99) = self.duration_histogram.percentiles();
391        CheckpointStats {
392            completed: self.checkpoints_completed,
393            failed: self.checkpoints_failed,
394            last_duration: self.last_checkpoint_duration,
395            duration_p50_ms: p50 / 1_000,
396            duration_p95_ms: p95 / 1_000,
397            duration_p99_ms: p99 / 1_000,
398            total_bytes_written: self.total_bytes_written,
399            current_phase: self.phase,
400            current_epoch: self.allocator.peek_epoch(),
401        }
402    }
403}
404
405#[must_use]
406pub(crate) fn source_to_connector_checkpoint(cp: &SourceCheckpoint) -> ConnectorCheckpoint {
407    ConnectorCheckpoint {
408        offsets: cp.durable_offsets(),
409        metadata: cp.metadata().clone(),
410        input_channels: cp.input_channels().map(<[Vec<u8>]>::to_vec),
411        source_assignment_version: cp.assignment_version(),
412    }
413}
414
415#[must_use]
416pub(crate) fn connector_to_source_checkpoint(cp: &ConnectorCheckpoint) -> SourceCheckpoint {
417    let mut source = SourceCheckpoint::with_offsets(cp.offsets.clone());
418    for (key, value) in &cp.metadata {
419        source.set_metadata(key.clone(), value.clone());
420    }
421    if let Some(channels) = &cp.input_channels {
422        source
423            .set_input_channels(channels.clone())
424            .expect("validated connector checkpoint input-channel inventory");
425    }
426    if let Some(version) = cp.source_assignment_version {
427        source.bind_assignment_version(version);
428    }
429    source
430}