sui_core/
db_checkpoint_handler.rs

1// Copyright (c) Mysten Labs, Inc.
2// SPDX-License-Identifier: Apache-2.0
3
4use crate::authority::authority_store_pruner::{
5    AuthorityStorePruner, AuthorityStorePruningMetrics, EPOCH_DURATION_MS_FOR_TESTING,
6};
7use crate::authority::authority_store_tables::AuthorityPerpetualTables;
8use crate::checkpoints::CheckpointStore;
9use anyhow::Result;
10use bytes::Bytes;
11use futures::future::try_join_all;
12use object_store::path::Path;
13use object_store::{DynObjectStore, ObjectStoreExt};
14use prometheus::{IntGauge, Registry, register_int_gauge_with_registry};
15use std::fs;
16use std::num::NonZeroUsize;
17use std::path::PathBuf;
18use std::sync::Arc;
19use std::time::Duration;
20use sui_config::node::AuthorityStorePruningConfig;
21use sui_config::object_storage_config::{ObjectStoreConfig, ObjectStoreType};
22use sui_storage::object_store::util::{
23    copy_recursively, find_all_dirs_with_epoch_prefix, find_missing_epochs_dirs,
24    path_to_filesystem, put, run_manifest_update_loop, write_snapshot_manifest,
25};
26use tracing::{debug, error, info};
27
28pub const SUCCESS_MARKER: &str = "_SUCCESS";
29pub const TEST_MARKER: &str = "_TEST";
30pub const UPLOAD_COMPLETED_MARKER: &str = "_UPLOAD_COMPLETED";
31pub const STATE_SNAPSHOT_COMPLETED_MARKER: &str = "_STATE_SNAPSHOT_COMPLETED";
32
33pub struct DBCheckpointMetrics {
34    pub first_missing_db_checkpoint_epoch: IntGauge,
35    pub num_local_db_checkpoints: IntGauge,
36}
37
38impl DBCheckpointMetrics {
39    pub fn new(registry: &Registry) -> Arc<Self> {
40        let this = Self {
41            first_missing_db_checkpoint_epoch: register_int_gauge_with_registry!(
42                "first_missing_db_checkpoint_epoch",
43                "First epoch for which we have no db checkpoint in remote store",
44                registry
45            )
46            .unwrap(),
47            num_local_db_checkpoints: register_int_gauge_with_registry!(
48                "num_local_db_checkpoints",
49                "Number of RocksDB checkpoints currently residing on local disk (i.e. not yet garbage collected)",
50                registry
51            )
52            .unwrap(),
53        };
54        Arc::new(this)
55    }
56}
57
58pub struct DBCheckpointHandler {
59    /// Directory on local disk where db checkpoints are stored
60    input_object_store: Arc<DynObjectStore>,
61    /// DB checkpoint directory on local filesystem
62    input_root_path: PathBuf,
63    /// Bucket on cloud object store where db checkpoints will be copied
64    output_object_store: Option<Arc<DynObjectStore>>,
65    /// Time interval to check for presence of new db checkpoint
66    interval: Duration,
67    /// File markers which signal that local db checkpoint can be garbage collected
68    gc_markers: Vec<String>,
69    /// Boolean flag to enable/disable object pruning and manual compaction before upload
70    prune_and_compact_before_upload: bool,
71    /// If true, upload will block on state snapshot upload completed marker
72    state_snapshot_enabled: bool,
73    /// Pruning objects
74    pruning_config: AuthorityStorePruningConfig,
75    metrics: Arc<DBCheckpointMetrics>,
76}
77
78impl DBCheckpointHandler {
79    pub fn new(
80        input_path: &std::path::Path,
81        output_object_store_config: Option<&ObjectStoreConfig>,
82        interval_s: u64,
83        prune_and_compact_before_upload: bool,
84        pruning_config: AuthorityStorePruningConfig,
85        registry: &Registry,
86        state_snapshot_enabled: bool,
87    ) -> Result<Arc<Self>> {
88        let input_store_config = ObjectStoreConfig {
89            object_store: Some(ObjectStoreType::File),
90            directory: Some(input_path.to_path_buf()),
91            ..Default::default()
92        };
93        let mut gc_markers = vec![UPLOAD_COMPLETED_MARKER.to_string()];
94        if state_snapshot_enabled {
95            gc_markers.push(STATE_SNAPSHOT_COMPLETED_MARKER.to_string());
96        }
97        Ok(Arc::new(DBCheckpointHandler {
98            input_object_store: input_store_config.make()?,
99            input_root_path: input_path.to_path_buf(),
100            output_object_store: output_object_store_config
101                .map(|config| config.make().expect("Failed to make object store")),
102            interval: Duration::from_secs(interval_s),
103            gc_markers,
104            prune_and_compact_before_upload,
105            state_snapshot_enabled,
106            pruning_config,
107            metrics: DBCheckpointMetrics::new(registry),
108        }))
109    }
110    pub fn new_for_test(
111        input_object_store_config: &ObjectStoreConfig,
112        output_object_store_config: Option<&ObjectStoreConfig>,
113        interval_s: u64,
114        prune_and_compact_before_upload: bool,
115        state_snapshot_enabled: bool,
116    ) -> Result<Arc<Self>> {
117        Ok(Arc::new(DBCheckpointHandler {
118            input_object_store: input_object_store_config.make()?,
119            input_root_path: input_object_store_config
120                .directory
121                .as_ref()
122                .unwrap()
123                .clone(),
124            output_object_store: output_object_store_config
125                .map(|config| config.make().expect("Failed to make object store")),
126            interval: Duration::from_secs(interval_s),
127            gc_markers: vec![UPLOAD_COMPLETED_MARKER.to_string(), TEST_MARKER.to_string()],
128            prune_and_compact_before_upload,
129            state_snapshot_enabled,
130            pruning_config: AuthorityStorePruningConfig::default(),
131            metrics: DBCheckpointMetrics::new(&Registry::default()),
132        }))
133    }
134    pub fn start(self: Arc<Self>) -> tokio::sync::broadcast::Sender<()> {
135        let (kill_sender, _kill_receiver) = tokio::sync::broadcast::channel::<()>(1);
136        if let Some(output_object_store) = self.output_object_store.as_ref() {
137            let output_object_store = output_object_store.clone();
138            tokio::task::spawn(Self::run_db_checkpoint_upload_loop(
139                self.clone(),
140                kill_sender.subscribe(),
141            ));
142            tokio::task::spawn(run_manifest_update_loop(
143                output_object_store,
144                kill_sender.subscribe(),
145            ));
146        } else {
147            // if db checkpoint remote store is not specified, cleanup loop
148            // is run to immediately mark db checkpoint upload as successful
149            // so that they can be snapshotted and garbage collected
150            tokio::task::spawn(Self::run_db_checkpoint_cleanup_loop(
151                self.clone(),
152                kill_sender.subscribe(),
153            ));
154        }
155        tokio::task::spawn(Self::run_db_checkpoint_gc_loop(
156            self,
157            kill_sender.subscribe(),
158        ));
159        kill_sender
160    }
161    async fn run_db_checkpoint_upload_loop(
162        self: Arc<Self>,
163        mut recv: tokio::sync::broadcast::Receiver<()>,
164    ) -> Result<()> {
165        let mut interval = tokio::time::interval(self.interval);
166        info!("DB checkpoint upload loop started");
167        loop {
168            tokio::select! {
169                _now = interval.tick() => {
170                    let local_checkpoints_by_epoch =
171                        find_all_dirs_with_epoch_prefix(&self.input_object_store, None).await?;
172                    self.metrics.num_local_db_checkpoints.set(local_checkpoints_by_epoch.len() as i64);
173                    match find_missing_epochs_dirs(self.output_object_store.as_ref().unwrap(), SUCCESS_MARKER).await {
174                        Ok(epochs) => {
175                            self.metrics.first_missing_db_checkpoint_epoch.set(epochs.first().cloned().unwrap_or(0) as i64);
176                            if let Err(err) = self.upload_db_checkpoints_to_object_store(epochs).await {
177                                error!("Failed to upload db checkpoint to remote store with err: {:?}", err);
178                            }
179                        }
180                        Err(err) => {
181                            error!("Failed to find missing db checkpoints in remote store: {:?}", err);
182                        }
183                    }
184                },
185                 _ = recv.recv() => break,
186            }
187        }
188        Ok(())
189    }
190    async fn run_db_checkpoint_cleanup_loop(
191        self: Arc<Self>,
192        mut recv: tokio::sync::broadcast::Receiver<()>,
193    ) -> Result<()> {
194        let mut interval = tokio::time::interval(self.interval);
195        info!("DB checkpoint upload disabled. DB checkpoint cleanup loop started");
196        loop {
197            tokio::select! {
198                _now = interval.tick() => {
199                    let local_checkpoints_by_epoch =
200                        find_all_dirs_with_epoch_prefix(&self.input_object_store, None).await?;
201                    self.metrics.num_local_db_checkpoints.set(local_checkpoints_by_epoch.len() as i64);
202                    let mut dirs: Vec<_> = local_checkpoints_by_epoch.iter().collect();
203                    dirs.sort_by_key(|(epoch_num, _path)| *epoch_num);
204                    for (_, db_path) in dirs {
205                        // If db checkpoint marked as completed, skip
206                        let local_db_path = path_to_filesystem(self.input_root_path.clone(), db_path)?;
207                        let upload_completed_path = local_db_path.join(UPLOAD_COMPLETED_MARKER);
208                        if upload_completed_path.exists() {
209                            continue;
210                        }
211                        let bytes = Bytes::from_static(b"success");
212                        let upload_completed_marker = db_path.child(UPLOAD_COMPLETED_MARKER);
213                        put(&self.input_object_store,
214                            &upload_completed_marker,
215                            bytes.clone(),
216                        )
217                        .await?;
218                    }
219                },
220                 _ = recv.recv() => break,
221            }
222        }
223        Ok(())
224    }
225    async fn run_db_checkpoint_gc_loop(
226        self: Arc<Self>,
227        mut recv: tokio::sync::broadcast::Receiver<()>,
228    ) -> Result<()> {
229        let mut gc_interval = tokio::time::interval(Duration::from_secs(30));
230        info!("DB checkpoint garbage collection loop started");
231        loop {
232            tokio::select! {
233                _now = gc_interval.tick() => {
234                    if let Ok(deleted) = self.garbage_collect_old_db_checkpoints().await
235                        && !deleted.is_empty() {
236                            info!("Garbage collected local db checkpoints: {:?}", deleted);
237                        }
238                },
239                 _ = recv.recv() => break,
240            }
241        }
242        Ok(())
243    }
244
245    async fn prune_and_compact(
246        &self,
247        db_path: PathBuf,
248        epoch: u64,
249        epoch_duration_ms: u64,
250    ) -> Result<()> {
251        let perpetual_db = Arc::new(AuthorityPerpetualTables::open(
252            &db_path.join("store"),
253            None,
254            None,
255        ));
256        let checkpoint_store = Arc::new(CheckpointStore::new_for_db_checkpoint_handler(
257            &db_path.join("checkpoints"),
258        ));
259        let metrics = AuthorityStorePruningMetrics::new(&Registry::default());
260        info!(
261            "Pruning db checkpoint in {:?} for epoch: {epoch}",
262            db_path.display()
263        );
264        AuthorityStorePruner::prune_objects_for_eligible_epochs(
265            &perpetual_db,
266            &checkpoint_store,
267            None,
268            self.pruning_config.clone(),
269            metrics,
270            epoch_duration_ms,
271        )
272        .await?;
273        info!(
274            "Compacting db checkpoint in {:?} for epoch: {epoch}",
275            db_path.display()
276        );
277        AuthorityStorePruner::compact(&perpetual_db)?;
278        Ok(())
279    }
280    async fn upload_db_checkpoints_to_object_store(
281        &self,
282        missing_epochs: Vec<u64>,
283    ) -> Result<(), anyhow::Error> {
284        let last_missing_epoch = missing_epochs.last().cloned().unwrap_or(0);
285        let local_checkpoints_by_epoch =
286            find_all_dirs_with_epoch_prefix(&self.input_object_store, None).await?;
287        let mut dirs: Vec<_> = local_checkpoints_by_epoch.iter().collect();
288        dirs.sort_by_key(|(epoch_num, _path)| *epoch_num);
289        let object_store = self
290            .output_object_store
291            .as_ref()
292            .expect("Expected object store to exist")
293            .clone();
294        for (epoch, db_path) in dirs {
295            // Convert `db_path` to the local filesystem path to where db checkpoint is stored
296            let local_db_path = path_to_filesystem(self.input_root_path.clone(), db_path)?;
297            if missing_epochs.contains(epoch) || *epoch >= last_missing_epoch {
298                if self.state_snapshot_enabled {
299                    let snapshot_completed_marker =
300                        local_db_path.join(STATE_SNAPSHOT_COMPLETED_MARKER);
301                    if !snapshot_completed_marker.exists() {
302                        info!(
303                            "DB checkpoint upload for epoch {} to wait until state snasphot uploaded",
304                            *epoch
305                        );
306                        continue;
307                    }
308                }
309
310                if self.prune_and_compact_before_upload {
311                    // Invoke pruning and compaction on the db checkpoint
312                    self.prune_and_compact(local_db_path, *epoch, EPOCH_DURATION_MS_FOR_TESTING)
313                        .await?;
314                }
315
316                info!("Copying db checkpoint for epoch: {epoch} to remote storage");
317                copy_recursively(
318                    db_path,
319                    &self.input_object_store,
320                    &object_store,
321                    NonZeroUsize::new(20).unwrap(),
322                )
323                .await?;
324
325                // This writes a single "MANIFEST" file which contains a list of all files that make up a db snapshot
326                write_snapshot_manifest(db_path, &object_store, format!("epoch_{}/", epoch))
327                    .await?;
328                // Drop marker in the output directory that upload completed successfully
329                let bytes = Bytes::from_static(b"success");
330                let success_marker = db_path.child(SUCCESS_MARKER);
331                put(&object_store, &success_marker, bytes.clone()).await?;
332            }
333            let bytes = Bytes::from_static(b"success");
334            let upload_completed_marker = db_path.child(UPLOAD_COMPLETED_MARKER);
335            put(
336                &self.input_object_store,
337                &upload_completed_marker,
338                bytes.clone(),
339            )
340            .await?;
341        }
342        Ok(())
343    }
344
345    async fn garbage_collect_old_db_checkpoints(&self) -> Result<Vec<u64>> {
346        let local_checkpoints_by_epoch =
347            find_all_dirs_with_epoch_prefix(&self.input_object_store, None).await?;
348        let mut deleted = Vec::new();
349        for (epoch, path) in local_checkpoints_by_epoch.iter() {
350            let marker_paths: Vec<Path> = self
351                .gc_markers
352                .iter()
353                .map(|marker| path.child(marker.clone()))
354                .collect();
355            let all_markers_present = try_join_all(
356                marker_paths
357                    .iter()
358                    .map(|path| self.input_object_store.get(path)),
359            )
360            .await;
361            match all_markers_present {
362                // After state snapshots, gc will also need to wait for a state snapshot
363                // upload completed marker
364                Ok(_) => {
365                    info!("Deleting db checkpoint dir: {path} for epoch: {epoch}");
366                    deleted.push(*epoch);
367                    let local_fs_path = path_to_filesystem(self.input_root_path.clone(), path)?;
368                    fs::remove_dir_all(&local_fs_path)?;
369                }
370                Err(_) => {
371                    debug!("Not ready for deletion yet: {path}");
372                }
373            }
374        }
375        Ok(deleted)
376    }
377}
378
379#[cfg(test)]
380mod tests {
381    use crate::db_checkpoint_handler::{
382        DBCheckpointHandler, SUCCESS_MARKER, TEST_MARKER, UPLOAD_COMPLETED_MARKER,
383    };
384    use itertools::Itertools;
385    use std::fs;
386    use sui_config::object_storage_config::{ObjectStoreConfig, ObjectStoreType};
387    use sui_storage::object_store::util::{
388        find_all_dirs_with_epoch_prefix, find_missing_epochs_dirs, path_to_filesystem,
389    };
390    use tempfile::TempDir;
391
392    #[tokio::test]
393    async fn test_basic() -> anyhow::Result<()> {
394        let checkpoint_dir = TempDir::new()?;
395        let checkpoint_dir_path = checkpoint_dir.path();
396        let local_epoch0_checkpoint = checkpoint_dir_path.join("epoch_0");
397        fs::create_dir(&local_epoch0_checkpoint)?;
398        let file1 = local_epoch0_checkpoint.join("file1");
399        fs::write(file1, b"Lorem ipsum")?;
400        let file2 = local_epoch0_checkpoint.join("file2");
401        fs::write(file2, b"Lorem ipsum")?;
402        let nested_dir = local_epoch0_checkpoint.join("data");
403        fs::create_dir(&nested_dir)?;
404        let file3 = nested_dir.join("file3");
405        fs::write(file3, b"Lorem ipsum")?;
406
407        let remote_checkpoint_dir = TempDir::new()?;
408        let remote_checkpoint_dir_path = remote_checkpoint_dir.path();
409        let remote_epoch0_checkpoint = remote_checkpoint_dir_path.join("epoch_0");
410
411        let input_store_config = ObjectStoreConfig {
412            object_store: Some(ObjectStoreType::File),
413            directory: Some(checkpoint_dir_path.to_path_buf()),
414            ..Default::default()
415        };
416        let output_store_config = ObjectStoreConfig {
417            object_store: Some(ObjectStoreType::File),
418            directory: Some(remote_checkpoint_dir_path.to_path_buf()),
419            ..Default::default()
420        };
421        let db_checkpoint_handler = DBCheckpointHandler::new_for_test(
422            &input_store_config,
423            Some(&output_store_config),
424            10,
425            false,
426            false,
427        )?;
428        let local_checkpoints_by_epoch =
429            find_all_dirs_with_epoch_prefix(&db_checkpoint_handler.input_object_store, None)
430                .await?;
431        assert!(!local_checkpoints_by_epoch.is_empty());
432        assert_eq!(*local_checkpoints_by_epoch.first_key_value().unwrap().0, 0);
433        assert_eq!(
434            path_to_filesystem(
435                db_checkpoint_handler.input_root_path.clone(),
436                local_checkpoints_by_epoch.first_key_value().unwrap().1
437            )
438            .unwrap(),
439            std::fs::canonicalize(local_epoch0_checkpoint.clone()).unwrap()
440        );
441        let missing_epochs = find_missing_epochs_dirs(
442            db_checkpoint_handler.output_object_store.as_ref().unwrap(),
443            SUCCESS_MARKER,
444        )
445        .await?;
446        db_checkpoint_handler
447            .upload_db_checkpoints_to_object_store(missing_epochs)
448            .await?;
449
450        assert!(remote_epoch0_checkpoint.join("file1").exists());
451        assert!(remote_epoch0_checkpoint.join("file2").exists());
452        assert!(remote_epoch0_checkpoint.join("data").join("file3").exists());
453        assert!(remote_epoch0_checkpoint.join(SUCCESS_MARKER).exists());
454        assert!(
455            local_epoch0_checkpoint
456                .join(UPLOAD_COMPLETED_MARKER)
457                .exists()
458        );
459
460        // Drop an extra gc marker meant only for gc to trigger
461        let test_marker = local_epoch0_checkpoint.join(TEST_MARKER);
462        fs::write(test_marker, b"Lorem ipsum")?;
463        db_checkpoint_handler
464            .garbage_collect_old_db_checkpoints()
465            .await?;
466
467        assert!(!local_epoch0_checkpoint.join("file1").exists());
468        assert!(!local_epoch0_checkpoint.join("file1").exists());
469        assert!(!local_epoch0_checkpoint.join("file2").exists());
470        assert!(!local_epoch0_checkpoint.join("data").join("file3").exists());
471        Ok(())
472    }
473
474    #[tokio::test]
475    async fn test_upload_resumes() -> anyhow::Result<()> {
476        let checkpoint_dir = TempDir::new()?;
477        let checkpoint_dir_path = checkpoint_dir.path();
478        let local_epoch0_checkpoint = checkpoint_dir_path.join("epoch_0");
479
480        let remote_checkpoint_dir = TempDir::new()?;
481        let remote_checkpoint_dir_path = remote_checkpoint_dir.path();
482        let remote_epoch0_checkpoint = remote_checkpoint_dir_path.join("epoch_0");
483
484        let input_store_config = ObjectStoreConfig {
485            object_store: Some(ObjectStoreType::File),
486            directory: Some(checkpoint_dir_path.to_path_buf()),
487            ..Default::default()
488        };
489        let output_store_config = ObjectStoreConfig {
490            object_store: Some(ObjectStoreType::File),
491            directory: Some(remote_checkpoint_dir_path.to_path_buf()),
492            ..Default::default()
493        };
494        let db_checkpoint_handler = DBCheckpointHandler::new_for_test(
495            &input_store_config,
496            Some(&output_store_config),
497            10,
498            false,
499            false,
500        )?;
501
502        fs::create_dir(&local_epoch0_checkpoint)?;
503        let file1 = local_epoch0_checkpoint.join("file1");
504        fs::write(file1, b"Lorem ipsum")?;
505        let file2 = local_epoch0_checkpoint.join("file2");
506        fs::write(file2, b"Lorem ipsum")?;
507        let nested_dir = local_epoch0_checkpoint.join("data");
508        fs::create_dir(&nested_dir)?;
509        let file3 = nested_dir.join("file3");
510        fs::write(file3, b"Lorem ipsum")?;
511
512        let missing_epochs = find_missing_epochs_dirs(
513            db_checkpoint_handler.output_object_store.as_ref().unwrap(),
514            SUCCESS_MARKER,
515        )
516        .await?;
517        db_checkpoint_handler
518            .upload_db_checkpoints_to_object_store(missing_epochs)
519            .await?;
520        assert!(remote_epoch0_checkpoint.join("file1").exists());
521        assert!(remote_epoch0_checkpoint.join("file2").exists());
522        assert!(remote_epoch0_checkpoint.join("data").join("file3").exists());
523        assert!(remote_epoch0_checkpoint.join(SUCCESS_MARKER).exists());
524        assert!(
525            local_epoch0_checkpoint
526                .join(UPLOAD_COMPLETED_MARKER)
527                .exists()
528        );
529
530        // Add a new db checkpoint to the local checkpoint directory
531        let local_epoch1_checkpoint = checkpoint_dir_path.join("epoch_1");
532        fs::create_dir(&local_epoch1_checkpoint)?;
533        let file1 = local_epoch1_checkpoint.join("file1");
534        fs::write(file1, b"Lorem ipsum")?;
535        let file2 = local_epoch1_checkpoint.join("file2");
536        fs::write(file2, b"Lorem ipsum")?;
537        let nested_dir = local_epoch1_checkpoint.join("data");
538        fs::create_dir(&nested_dir)?;
539        let file3 = nested_dir.join("file3");
540        fs::write(file3, b"Lorem ipsum")?;
541
542        // Now delete the success marker from remote checkpointed directory
543        // This is the scenario where uploads stops mid way because system stopped
544        fs::remove_file(remote_epoch0_checkpoint.join(SUCCESS_MARKER))?;
545
546        // Checkpoint handler should copy checkpoint for epoch_0 first before copying
547        // epoch_1
548        let missing_epochs = find_missing_epochs_dirs(
549            db_checkpoint_handler.output_object_store.as_ref().unwrap(),
550            SUCCESS_MARKER,
551        )
552        .await?;
553        db_checkpoint_handler
554            .upload_db_checkpoints_to_object_store(missing_epochs)
555            .await?;
556        assert!(remote_epoch0_checkpoint.join("file1").exists());
557        assert!(remote_epoch0_checkpoint.join("file2").exists());
558        assert!(remote_epoch0_checkpoint.join("data").join("file3").exists());
559        assert!(remote_epoch0_checkpoint.join(SUCCESS_MARKER).exists());
560        assert!(
561            local_epoch0_checkpoint
562                .join(UPLOAD_COMPLETED_MARKER)
563                .exists()
564        );
565
566        let remote_epoch1_checkpoint = remote_checkpoint_dir_path.join("epoch_1");
567        assert!(remote_epoch1_checkpoint.join("file1").exists());
568        assert!(remote_epoch1_checkpoint.join("file2").exists());
569        assert!(remote_epoch1_checkpoint.join("data").join("file3").exists());
570        assert!(remote_epoch1_checkpoint.join(SUCCESS_MARKER).exists());
571        assert!(
572            local_epoch1_checkpoint
573                .join(UPLOAD_COMPLETED_MARKER)
574                .exists()
575        );
576
577        // Drop an extra gc marker meant only for gc to trigger
578        let test_marker = local_epoch0_checkpoint.join(TEST_MARKER);
579        fs::write(test_marker, b"Lorem ipsum")?;
580        let test_marker = local_epoch1_checkpoint.join(TEST_MARKER);
581        fs::write(test_marker, b"Lorem ipsum")?;
582
583        db_checkpoint_handler
584            .garbage_collect_old_db_checkpoints()
585            .await?;
586        assert!(!local_epoch0_checkpoint.join("file1").exists());
587        assert!(!local_epoch0_checkpoint.join("file1").exists());
588        assert!(!local_epoch0_checkpoint.join("file2").exists());
589        assert!(!local_epoch0_checkpoint.join("data").join("file3").exists());
590        assert!(!local_epoch1_checkpoint.join("file1").exists());
591        assert!(!local_epoch1_checkpoint.join("file1").exists());
592        assert!(!local_epoch1_checkpoint.join("file2").exists());
593        assert!(!local_epoch1_checkpoint.join("data").join("file3").exists());
594        Ok(())
595    }
596
597    #[tokio::test]
598    async fn test_missing_epochs() -> anyhow::Result<()> {
599        let checkpoint_dir = TempDir::new()?;
600        let checkpoint_dir_path = checkpoint_dir.path();
601        let local_epoch0_checkpoint = checkpoint_dir_path.join("epoch_0");
602        fs::create_dir(&local_epoch0_checkpoint)?;
603        let local_epoch1_checkpoint = checkpoint_dir_path.join("epoch_1");
604        fs::create_dir(&local_epoch1_checkpoint)?;
605        // Missing epoch 2
606        let local_epoch3_checkpoint = checkpoint_dir_path.join("epoch_3");
607        fs::create_dir(&local_epoch3_checkpoint)?;
608        let remote_checkpoint_dir = TempDir::new()?;
609        let remote_checkpoint_dir_path = remote_checkpoint_dir.path();
610
611        let input_store_config = ObjectStoreConfig {
612            object_store: Some(ObjectStoreType::File),
613            directory: Some(checkpoint_dir_path.to_path_buf()),
614            ..Default::default()
615        };
616
617        let output_store_config = ObjectStoreConfig {
618            object_store: Some(ObjectStoreType::File),
619            directory: Some(remote_checkpoint_dir_path.to_path_buf()),
620            ..Default::default()
621        };
622        let db_checkpoint_handler = DBCheckpointHandler::new_for_test(
623            &input_store_config,
624            Some(&output_store_config),
625            10,
626            false,
627            false,
628        )?;
629
630        let missing_epochs = find_missing_epochs_dirs(
631            db_checkpoint_handler.output_object_store.as_ref().unwrap(),
632            SUCCESS_MARKER,
633        )
634        .await?;
635        db_checkpoint_handler
636            .upload_db_checkpoints_to_object_store(missing_epochs)
637            .await?;
638
639        let first_missing_epoch = find_missing_epochs_dirs(
640            db_checkpoint_handler.output_object_store.as_ref().unwrap(),
641            SUCCESS_MARKER,
642        )
643        .await?
644        .first()
645        .cloned()
646        .unwrap();
647        assert_eq!(first_missing_epoch, 2);
648
649        let remote_epoch0_checkpoint = remote_checkpoint_dir_path.join("epoch_0");
650        fs::remove_file(remote_epoch0_checkpoint.join(SUCCESS_MARKER))?;
651
652        let first_missing_epoch = find_missing_epochs_dirs(
653            db_checkpoint_handler.output_object_store.as_ref().unwrap(),
654            SUCCESS_MARKER,
655        )
656        .await?
657        .first()
658        .cloned()
659        .unwrap();
660        assert_eq!(first_missing_epoch, 0);
661
662        Ok(())
663    }
664
665    #[tokio::test]
666    async fn test_range_missing_epochs() -> anyhow::Result<()> {
667        let checkpoint_dir = TempDir::new()?;
668        let checkpoint_dir_path = checkpoint_dir.path();
669        let local_epoch100_checkpoint = checkpoint_dir_path.join("epoch_100");
670        fs::create_dir(&local_epoch100_checkpoint)?;
671        let local_epoch200_checkpoint = checkpoint_dir_path.join("epoch_200");
672        fs::create_dir(&local_epoch200_checkpoint)?;
673        let remote_checkpoint_dir = TempDir::new()?;
674        let remote_checkpoint_dir_path = remote_checkpoint_dir.path();
675
676        let input_store_config = ObjectStoreConfig {
677            object_store: Some(ObjectStoreType::File),
678            directory: Some(checkpoint_dir_path.to_path_buf()),
679            ..Default::default()
680        };
681
682        let output_store_config = ObjectStoreConfig {
683            object_store: Some(ObjectStoreType::File),
684            directory: Some(remote_checkpoint_dir_path.to_path_buf()),
685            ..Default::default()
686        };
687        let db_checkpoint_handler = DBCheckpointHandler::new_for_test(
688            &input_store_config,
689            Some(&output_store_config),
690            10,
691            false,
692            false,
693        )?;
694
695        let missing_epochs = find_missing_epochs_dirs(
696            db_checkpoint_handler.output_object_store.as_ref().unwrap(),
697            SUCCESS_MARKER,
698        )
699        .await?;
700        assert_eq!(missing_epochs, vec![0]);
701        db_checkpoint_handler
702            .upload_db_checkpoints_to_object_store(missing_epochs)
703            .await?;
704
705        let missing_epochs = find_missing_epochs_dirs(
706            db_checkpoint_handler.output_object_store.as_ref().unwrap(),
707            SUCCESS_MARKER,
708        )
709        .await?;
710        let mut expected_missing_epochs: Vec<u64> = (0..100).collect();
711        expected_missing_epochs.extend((101..200).collect_vec().iter());
712        expected_missing_epochs.push(201);
713        assert_eq!(missing_epochs, expected_missing_epochs);
714        Ok(())
715    }
716}
sui_core/db_checkpoint_handler.rs

sui_core/
db_checkpoint_handler.rs
