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

laminar_db/checkpoint_coordinator/
coordinator.rs
