From 51fbc1a73edb7e5d9380c11ee8fca975e797fe51 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Mon, 11 Jan 2021 11:38:55 -0800
Subject: [PATCH] Use standard tmp-snapshot- file prefix for the "new_state"
 archive for better cleanup/consistency

---
 runtime/src/snapshot_utils.rs | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs
index a6c0262170..ab985d8109 100644
--- a/runtime/src/snapshot_utils.rs
+++ b/runtime/src/snapshot_utils.rs
@@ -39,7 +39,7 @@ pub const MAX_SNAPSHOTS: usize = 8; // Save some snapshots but not too many
 const MAX_SNAPSHOT_DATA_FILE_SIZE: u64 = 32 * 1024 * 1024 * 1024; // 32 GiB
 const VERSION_STRING_V1_2_0: &str = "1.2.0";
 const DEFAULT_SNAPSHOT_VERSION: SnapshotVersion = SnapshotVersion::V1_2_0;
-const TMP_SNAPSHOT_DIR_PREFIX: &str = "tmp-snapshot-";
+const TMP_SNAPSHOT_PREFIX: &str = "tmp-snapshot-";
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
 pub enum SnapshotVersion {
@@ -151,7 +151,7 @@ pub fn package_snapshot<P: AsRef<Path>, Q: AsRef<Path>>(
 ) -> Result<AccountsPackage> {
     // Hard link all the snapshots we need for this package
     let snapshot_tmpdir = tempfile::Builder::new()
-        .prefix(&format!("{}{}-", TMP_SNAPSHOT_DIR_PREFIX, bank.slot()))
+        .prefix(&format!("{}{}-", TMP_SNAPSHOT_PREFIX, bank.slot()))
         .tempdir_in(snapshot_path)?;
 
     // Create a snapshot package
@@ -212,9 +212,14 @@ pub fn remove_tmp_snapshot_archives(snapshot_path: &Path) {
                 .file_name()
                 .into_string()
                 .unwrap_or_else(|_| String::new())
-                .starts_with(TMP_SNAPSHOT_DIR_PREFIX)
+                .starts_with(TMP_SNAPSHOT_PREFIX)
             {
-                fs::remove_dir_all(entry.path()).unwrap_or_else(|err| {
+                if entry.path().is_file() {
+                    fs::remove_file(entry.path())
+                } else {
+                    fs::remove_dir_all(entry.path())
+                }
+                .unwrap_or_else(|err| {
                     warn!("Failed to remove {}: {}", entry.path().display(), err)
                 });
             }
@@ -249,7 +254,7 @@ pub fn archive_snapshot_package(snapshot_package: &AccountsPackage) -> Result<()
     let staging_dir = tempfile::Builder::new()
         .prefix(&format!(
             "{}{}-",
-            TMP_SNAPSHOT_DIR_PREFIX, snapshot_package.slot
+            TMP_SNAPSHOT_PREFIX, snapshot_package.slot
         ))
         .tempdir_in(tar_dir)?;
 
@@ -295,7 +300,10 @@ pub fn archive_snapshot_package(snapshot_package: &AccountsPackage) -> Result<()
     // Tar the staging directory into the archive at `archive_path`
     //
     // system `tar` program is used for -S (sparse file support)
-    let archive_path = tar_dir.join(format!("new_state{}", file_ext));
+    let archive_path = tar_dir.join(format!(
+        "{}{}{}",
+        TMP_SNAPSHOT_PREFIX, snapshot_package.slot, file_ext
+    ));
 
     let mut tar = process::Command::new("tar")
         .args(&[
@@ -590,7 +598,7 @@ pub fn bank_from_archive<P: AsRef<Path>>(
 ) -> Result<Bank> {
     // Untar the snapshot into a temporary directory
     let unpack_dir = tempfile::Builder::new()
-        .prefix(TMP_SNAPSHOT_DIR_PREFIX)
+        .prefix(TMP_SNAPSHOT_PREFIX)
         .tempdir_in(snapshot_path)?;
     untar_snapshot_in(&snapshot_tar, &unpack_dir, archive_format)?;