snapshot restore: add option to reuse downloaded files

crates/sui-indexer/src/restorer/formal_snapshot.rs

@@ -2,7 +2,6 @@
 // SPDX-License-Identifier: Apache-2.0
 
 use std::collections::BTreeMap;
-use std::fs;
 use std::num::NonZeroUsize;
 use std::path::PathBuf;
 use std::sync::Arc;
@@ -55,17 +54,6 @@ impl IndexerFormalSnapshotRestorer {
 
         let base_path = PathBuf::from(restore_config.snapshot_download_dir.clone());
         let snapshot_dir = base_path.join("snapshot");
-        if snapshot_dir.exists() {
-            fs::remove_dir_all(snapshot_dir.clone()).unwrap();
-            info!(
-                "Deleted all files from snapshot directory: {:?}",
-                snapshot_dir
-            );
-        } else {
-            fs::create_dir(snapshot_dir.clone()).unwrap();
-            info!("Created snapshot directory: {:?}", snapshot_dir);
-        }
-
         let local_store_config = ObjectStoreConfig {
             object_store: Some(ObjectStoreType::File),
             directory: Some(snapshot_dir.clone().to_path_buf()),
@@ -80,6 +68,7 @@ impl IndexerFormalSnapshotRestorer {
             usize::MAX,
             NonZeroUsize::new(restore_config.object_store_concurrent_limit).unwrap(),
             m.clone(),
+            true, // skip_reset_local_store
         )
         .await
         .unwrap_or_else(|err| panic!("Failed to create reader: {}", err));

crates/sui-mvr-indexer/src/restorer/formal_snapshot.rs

@@ -80,6 +80,7 @@ impl IndexerFormalSnapshotRestorer {
             usize::MAX,
             NonZeroUsize::new(restore_config.object_store_concurrent_limit).unwrap(),
             m.clone(),
+            true, // skip_reset_local_store
         )
         .await
         .unwrap_or_else(|err| panic!("Failed to create reader: {}", err));

crates/sui-snapshot/src/reader.rs

@@ -29,7 +29,7 @@ use sui_core::authority::AuthorityStore;
 use sui_storage::blob::{Blob, BlobEncoding};
 use sui_storage::object_store::http::HttpDownloaderBuilder;
 use sui_storage::object_store::util::{copy_file, copy_files, path_to_filesystem};
-use sui_storage::object_store::{ObjectStoreGetExt, ObjectStorePutExt};
+use sui_storage::object_store::{ObjectStoreGetExt, ObjectStoreListExt, ObjectStorePutExt};
 use sui_types::accumulator::Accumulator;
 use sui_types::base_types::{ObjectDigest, ObjectID, ObjectRef, SequenceNumber};
 use tokio::sync::Mutex;
@@ -61,6 +61,7 @@ impl StateSnapshotReaderV1 {
         indirect_objects_threshold: usize,
         download_concurrency: NonZeroUsize,
         m: MultiProgress,
+        skip_reset_local_store: bool,
     ) -> Result<Self> {
         let epoch_dir = format!("epoch_{}", epoch);
         let remote_object_store = if remote_store_config.no_sign_request {
@@ -70,16 +71,20 @@ impl StateSnapshotReaderV1 {
         };
         let local_object_store: Arc<dyn ObjectStorePutExt> =
             local_store_config.make().map(Arc::new)?;
+        let local_object_store_list: Arc<dyn ObjectStoreListExt> =
+            local_store_config.make().map(Arc::new)?;
         let local_staging_dir_root = local_store_config
             .directory
             .as_ref()
             .context("No directory specified")?
             .clone();
-        let local_epoch_dir_path = local_staging_dir_root.join(&epoch_dir);
-        if local_epoch_dir_path.exists() {
-            fs::remove_dir_all(&local_epoch_dir_path)?;
+        if !skip_reset_local_store {
+            let local_epoch_dir_path = local_staging_dir_root.join(&epoch_dir);
+            if local_epoch_dir_path.exists() {
+                fs::remove_dir_all(&local_epoch_dir_path)?;
+            }
+            fs::create_dir_all(&local_epoch_dir_path)?;
         }
-        fs::create_dir_all(&local_epoch_dir_path)?;
         // Download MANIFEST first
         let manifest_file_path = Path::from(epoch_dir.clone()).child("MANIFEST");
         copy_file(
@@ -136,24 +141,42 @@ impl StateSnapshotReaderV1 {
             })
             .collect();
 
+        let files_to_download = if skip_reset_local_store {
+            let mut list_stream = local_object_store_list
+                .list_objects(Some(&epoch_dir_path))
+                .await;
+            let mut existing_files = std::collections::HashSet::new();
+            while let Some(Ok(meta)) = list_stream.next().await {
+                existing_files.insert(meta.location);
+            }
+            let mut missing_files = Vec::new();
+            for file in &files {
+                if !existing_files.contains(file) {
+                    missing_files.push(file.clone());
+                }
+            }
+            missing_files
+        } else {
+            files
+        };
         let progress_bar = m.add(
-            ProgressBar::new(files.len() as u64).with_style(
+            ProgressBar::new(files_to_download.len() as u64).with_style(
                 ProgressStyle::with_template(
-                    "[{elapsed_precise}] {wide_bar} {pos} out of {len} .ref files done ({msg})",
+                    "[{elapsed_precise}] {wide_bar} {pos} out of {len} missing .ref files done ({msg})",
                 )
                 .unwrap(),
             ),
         );
         copy_files(
-            &files,
-            &files,
+            &files_to_download,
+            &files_to_download,
             &remote_object_store,
             &local_object_store,
             download_concurrency,
             Some(progress_bar.clone()),
         )
         .await?;
-        progress_bar.finish_with_message("ref files download complete");
+        progress_bar.finish_with_message("Missing ref files download complete");
         Ok(StateSnapshotReaderV1 {
             epoch,
             local_staging_dir_root,

crates/sui-snapshot/src/tests.rs

@@ -112,6 +112,7 @@ async fn test_snapshot_basic() -> Result<(), anyhow::Error> {
         usize::MAX,
         NonZeroUsize::new(1).unwrap(),
         MultiProgress::new(),
+        false, // skip_reset_local_store
     )
     .await?;
     let restored_perpetual_db = AuthorityPerpetualTables::open(&restored_db_path, None);
@@ -167,6 +168,7 @@ async fn test_snapshot_empty_db() -> Result<(), anyhow::Error> {
         usize::MAX,
         NonZeroUsize::new(1).unwrap(),
         MultiProgress::new(),
+        false, // skip_reset_local_store
     )
     .await?;
     let restored_perpetual_db = AuthorityPerpetualTables::open(&restored_db_path, None);

crates/sui-tool/src/lib.rs

@@ -890,6 +890,7 @@ pub async fn download_formal_snapshot(
             usize::MAX,
             NonZeroUsize::new(num_parallel_downloads).unwrap(),
             m_clone,
+            false, // skip_reset_local_store
         )
         .await
         .unwrap_or_else(|err| panic!("Failed to create reader: {}", err));