From 4a66e3eddc3fea65f0e72c9433ba421f9f1fa90f 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 +++++++++++++++-------
 validator/src/bin/solana-test-validator.rs |  1 +
 2 files changed, 16 insertions(+), 7 deletions(-)

diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs
index a665bf2638..b83ce343e5 100644
--- a/runtime/src/snapshot_utils.rs
+++ b/runtime/src/snapshot_utils.rs
@@ -38,7 +38,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 {
@@ -147,7 +147,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
@@ -208,9 +208,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)
                 });
             }
@@ -245,7 +250,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)?;
 
@@ -291,7 +296,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(&[
@@ -586,7 +594,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)?;
 
diff --git a/validator/src/bin/solana-test-validator.rs b/validator/src/bin/solana-test-validator.rs
index 595e961a0d..cec95617c4 100644
--- a/validator/src/bin/solana-test-validator.rs
+++ b/validator/src/bin/solana-test-validator.rs
@@ -220,6 +220,7 @@ fn main() {
             exit(1);
         })
     }
+    solana_runtime::snapshot_utils::remove_tmp_snapshot_archives(&ledger_path);
 
     let validator_log_symlink = ledger_path.join("validator.log");
     let logfile = if output != Output::Log {