Parallel download for snapshots

andll · andll · commit aef9fbb06ab5 · 2026-03-22T22:00:43.000-07:00
diff --git a/crates/sui-core/src/authority/authority_store.rs b/crates/sui-core/src/authority/authority_store.rs
@@ -636,17 +636,51 @@ impl AuthorityStore {
         Ok(())
     }
 
-    pub fn bulk_insert_live_objects(
-        perpetual_db: &AuthorityPerpetualTables,
-        live_objects: impl Iterator<Item = LiveObject>,
+    pub async fn bulk_insert_live_objects(
+        perpetual_db: Arc<AuthorityPerpetualTables>,
+        objects: Vec<LiveObject>,
         expected_sha3_digest: &[u8; 32],
     ) -> SuiResult<()> {
+        // Verify SHA3 over the full object set before inserting.
         let mut hasher = Sha3_256::default();
+        for object in &objects {
+            hasher.update(object.object_reference().2.inner());
+        }
+        let sha3_digest = hasher.finalize().digest;
+        if *expected_sha3_digest != sha3_digest {
+            error!(
+                "Sha does not match! expected: {:?}, actual: {:?}",
+                expected_sha3_digest, sha3_digest
+            );
+            return Err(SuiError::from("Sha does not match"));
+        }
+
+        const NUM_PARALLEL_CHUNKS: usize = 8;
+        let chunk_size = ((objects.len() + NUM_PARALLEL_CHUNKS - 1) / NUM_PARALLEL_CHUNKS).max(1);
+        let mut remaining = objects;
+        let mut handles = Vec::new();
+        while !remaining.is_empty() {
+            let take = chunk_size.min(remaining.len());
+            let chunk: Vec<LiveObject> = remaining.drain(..take).collect();
+            let db = perpetual_db.clone();
+            handles.push(tokio::task::spawn_blocking(move || {
+                Self::insert_objects_chunk(db, chunk)
+            }));
+        }
+        for handle in handles {
+            handle.await.expect("insert task panicked")?;
+        }
+        Ok(())
+    }
+
+    fn insert_objects_chunk(
+        perpetual_db: Arc<AuthorityPerpetualTables>,
+        objects: Vec<LiveObject>,
+    ) -> SuiResult<()> {
         let mut batch = perpetual_db.objects.batch();
         let mut written = 0usize;
         const MAX_BATCH_SIZE: usize = 100_000;
-        for object in live_objects {
-            hasher.update(object.object_reference().2.inner());
+        for object in objects {
             match object {
                 LiveObject::Normal(object) => {
                     let store_object_wrapper = get_store_object(object.clone());
@@ -683,14 +717,6 @@ impl AuthorityStore {
                 written = 0;
             }
         }
-        let sha3_digest = hasher.finalize().digest;
-        if *expected_sha3_digest != sha3_digest {
-            error!(
-                "Sha does not match! expected: {:?}, actual: {:?}",
-                expected_sha3_digest, sha3_digest
-            );
-            return Err(SuiError::from("Sha does not match"));
-        }
         batch.write()?;
         Ok(())
     }
diff --git a/crates/sui-snapshot/src/reader.rs b/crates/sui-snapshot/src/reader.rs
@@ -294,7 +294,7 @@ impl StateSnapshotReaderV1 {
 
     pub async fn read(
         &mut self,
-        perpetual_db: &AuthorityPerpetualTables,
+        perpetual_db: Arc<AuthorityPerpetualTables>,
         abort_registration: AbortRegistration,
         sender: Option<tokio::sync::mpsc::Sender<(GlobalStateHash, u64)>>,
     ) -> Result<()> {
@@ -306,7 +306,7 @@ impl StateSnapshotReaderV1 {
         let (sha3_digests, num_part_files) = self.compute_checksum().await?;
         let accum_handle =
             sender.map(|sender| self.spawn_accumulation_tasks(sender, num_part_files));
-        self.sync_live_objects(perpetual_db, abort_registration, sha3_digests)
+        self.sync_live_objects(perpetual_db.clone(), abort_registration, sha3_digests)
             .await?;
         if let Some(handle) = accum_handle {
             handle.await?;
@@ -479,7 +479,7 @@ impl StateSnapshotReaderV1 {
 
     async fn sync_live_objects(
         &self,
-        perpetual_db: &AuthorityPerpetualTables,
+        perpetual_db: Arc<AuthorityPerpetualTables>,
         abort_registration: AbortRegistration,
         sha3_digests: Arc<Mutex<DigestByBucketAndPartition>>,
     ) -> Result<(), anyhow::Error> {
@@ -507,7 +507,7 @@ impl StateSnapshotReaderV1 {
         );
         let obj_progress_bar_clone = obj_progress_bar.clone();
         let instant = Instant::now();
-        let downloaded_bytes = AtomicUsize::new(0);
+        let downloaded_bytes = Arc::new(AtomicUsize::new(0));
 
         let ret = Abortable::new(
             async move {
@@ -538,25 +538,29 @@ impl StateSnapshotReaderV1 {
                     .boxed()
                     .buffer_unordered(concurrency)
                     .try_for_each(|(bytes, file_metadata, sha3_digest)| {
-                        let bytes_len = bytes.len();
-                        let result: Result<(), anyhow::Error> =
-                            LiveObjectIter::new(&file_metadata, bytes).map(|obj_iter| {
-                                AuthorityStore::bulk_insert_live_objects(
-                                    perpetual_db,
-                                    obj_iter,
-                                    &sha3_digest,
-                                )
-                                .expect("Failed to insert live objects");
-                            });
-                        downloaded_bytes.fetch_add(bytes_len, Ordering::Relaxed);
-                        obj_progress_bar_clone.inc(1);
-                        obj_progress_bar_clone.set_message(format!(
-                            "Download speed: {} MiB/s",
-                            downloaded_bytes.load(Ordering::Relaxed) as f64
-                                / (1024 * 1024) as f64
-                                / instant.elapsed().as_secs_f64(),
-                        ));
-                        futures::future::ready(result)
+                        let perpetual_db = perpetual_db.clone();
+                        let obj_progress_bar_clone = obj_progress_bar_clone.clone();
+                        let downloaded_bytes = downloaded_bytes.clone();
+                        async move {
+                            let bytes_len = bytes.len();
+                            let objects: Vec<LiveObject> =
+                                LiveObjectIter::new(&file_metadata, bytes)?.collect();
+                            AuthorityStore::bulk_insert_live_objects(
+                                perpetual_db,
+                                objects,
+                                &sha3_digest,
+                            )
+                            .await?;
+                            downloaded_bytes.fetch_add(bytes_len, Ordering::Relaxed);
+                            obj_progress_bar_clone.inc(1);
+                            obj_progress_bar_clone.set_message(format!(
+                                "Download speed: {} MiB/s",
+                                downloaded_bytes.load(Ordering::Relaxed) as f64
+                                    / (1024 * 1024) as f64
+                                    / instant.elapsed().as_secs_f64(),
+                            ));
+                            Ok(())
+                        }
                     })
                     .await
             },
diff --git a/crates/sui-snapshot/src/tests.rs b/crates/sui-snapshot/src/tests.rs
@@ -120,10 +120,14 @@ async fn test_snapshot_basic() -> Result<(), anyhow::Error> {
         3,     // max_retries
     )
     .await?;
-    let restored_perpetual_db = AuthorityPerpetualTables::open(&restored_db_path, None, None);
+    let restored_perpetual_db = Arc::new(AuthorityPerpetualTables::open(
+        &restored_db_path,
+        None,
+        None,
+    ));
     let (_abort_handle, abort_registration) = AbortHandle::new_pair();
     snapshot_reader
-        .read(&restored_perpetual_db, abort_registration, None)
+        .read(restored_perpetual_db.clone(), abort_registration, None)
         .await?;
     compare_live_objects(&perpetual_db, &restored_perpetual_db, true)?;
     Ok(())
@@ -176,10 +180,14 @@ async fn test_snapshot_empty_db() -> Result<(), anyhow::Error> {
         3,     // max_retries
     )
     .await?;
-    let restored_perpetual_db = AuthorityPerpetualTables::open(&restored_db_path, None, None);
+    let restored_perpetual_db = Arc::new(AuthorityPerpetualTables::open(
+        &restored_db_path,
+        None,
+        None,
+    ));
     let (_abort_handle, abort_registration) = AbortHandle::new_pair();
     snapshot_reader
-        .read(&restored_perpetual_db, abort_registration, None)
+        .read(restored_perpetual_db.clone(), abort_registration, None)
         .await?;
     compare_live_objects(
         &perpetual_db,
@@ -303,10 +311,14 @@ async fn test_snapshot_restore_from_archive() -> Result<(), anyhow::Error> {
         3,     // max_retries
     )
     .await?;
-    let restored_perpetual_db = AuthorityPerpetualTables::open(&restored_db_path, None, None);
+    let restored_perpetual_db = Arc::new(AuthorityPerpetualTables::open(
+        &restored_db_path,
+        None,
+        None,
+    ));
     let (_abort_handle, abort_registration) = AbortHandle::new_pair();
     snapshot_reader
-        .read(&restored_perpetual_db, abort_registration, None)
+        .read(restored_perpetual_db.clone(), abort_registration, None)
         .await?;
     compare_live_objects(&perpetual_db, &restored_perpetual_db, true)?;
     Ok(())
diff --git a/crates/sui-tool/src/lib.rs b/crates/sui-tool/src/lib.rs
@@ -925,7 +925,7 @@ pub async fn download_formal_snapshot(
         .await
         .unwrap_or_else(|err| panic!("Failed to create reader: {}", err));
         reader
-            .read(&perpetual_db_clone, abort_registration, Some(sender))
+            .read(perpetual_db_clone.clone(), abort_registration, Some(sender))
             .await
             .unwrap_or_else(|err| panic!("Failed during read: {}", err));
         info!("Snapshot download complete");
