feat(resharding) - delete snapshot immediately after resharding is finished

chain/chain/src/chain.rs

@@ -3486,33 +3486,53 @@ impl Chain {
         )))
     }
 
-    /// Function to create a new snapshot if needed
+    /// Function to create or delete a snapshot if necessary.
     fn process_snapshot(&mut self) -> Result<(), Error> {
         let (make_snapshot, delete_snapshot) = self.should_make_or_delete_snapshot()?;
         if !make_snapshot && !delete_snapshot {
             return Ok(());
         }
-        if let Some(snapshot_callbacks) = &self.snapshot_callbacks {
-            if make_snapshot {
-                let head = self.head()?;
-                let epoch_height =
-                    self.epoch_manager.get_epoch_height_from_prev_block(&head.prev_block_hash)?;
-                let shard_uids = self
-                    .epoch_manager
-                    .get_shard_layout_from_prev_block(&head.prev_block_hash)?
-                    .shard_uids()
-                    .collect();
-                let last_block = self.get_block(&head.last_block_hash)?;
-                let make_snapshot_callback = &snapshot_callbacks.make_snapshot_callback;
-                make_snapshot_callback(head.prev_block_hash, epoch_height, shard_uids, last_block);
-            } else if delete_snapshot {
-                let delete_snapshot_callback = &snapshot_callbacks.delete_snapshot_callback;
-                delete_snapshot_callback();
-            }
+        let Some(snapshot_callbacks) = &self.snapshot_callbacks else { return Ok(()) };
+        if make_snapshot {
+            let head = self.head()?;
+            let prev_hash = head.prev_block_hash;
+            let epoch_height = self.epoch_manager.get_epoch_height_from_prev_block(&prev_hash)?;
+            let shard_layout = &self.epoch_manager.get_shard_layout_from_prev_block(&prev_hash)?;
+            let shard_uids = shard_layout.shard_uids().collect();
+            let last_block = self.get_block(&head.last_block_hash)?;
+            let make_snapshot_callback = &snapshot_callbacks.make_snapshot_callback;
+            make_snapshot_callback(prev_hash, epoch_height, shard_uids, last_block);
+        } else if delete_snapshot {
+            let delete_snapshot_callback = &snapshot_callbacks.delete_snapshot_callback;
+            delete_snapshot_callback();
         }
         Ok(())
     }
 
+    // Similar to `process_snapshot` but only called after resharding is done.
+    // This is to speed up the snapshot removal once resharding is finished in
+    // order to minimize the storage overhead.
+    pub fn process_snapshot_after_resharding(&mut self) -> Result<(), Error> {
+        let Some(snapshot_callbacks) = &self.snapshot_callbacks else { return Ok(()) };
+
+        let tries = self.runtime_adapter.get_tries();
+        let snapshot_config = tries.state_snapshot_config();
+        let delete_snapshot = match snapshot_config.state_snapshot_type {
+            // Do not delete snapshot if the node is configured to snapshot every epoch.
+            StateSnapshotType::EveryEpoch => false,
+            // Delete the snapshot if it was created only for resharding.
+            StateSnapshotType::ForReshardingOnly => true,
+        };
+
+        if delete_snapshot {
+            tracing::debug!(target: "resharding", "deleting snapshot after resharding");
+            let delete_snapshot_callback = &snapshot_callbacks.delete_snapshot_callback;
+            delete_snapshot_callback();
+        }
+
+        Ok(())
+    }
+
     /// Function to check whether we need to create a new snapshot while processing the current block
     /// Note that this functions is called as a part of block preprocesing, so the head is not updated to current block
     fn should_make_or_delete_snapshot(&mut self) -> Result<(bool, bool), Error> {

chain/client/src/sync/state.rs

@@ -702,6 +702,7 @@ impl StateSync {
         )?;
 
         if all_done {
+            chain.process_snapshot_after_resharding()?;
             Ok(StateSyncResult::Completed)
         } else {
             Ok(StateSyncResult::InProgress)