From 254790f8c815728558a8061058f8d23b404cd77b Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Mon, 30 Nov 2020 17:37:38 -0800
Subject: [PATCH] Fully clean up temporary snapshot staging directories

---
 runtime/src/snapshot_utils.rs | 38 +++++++++++++++++++++++++++++------
 validator/src/main.rs         |  1 +
 2 files changed, 33 insertions(+), 6 deletions(-)

diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs
index 56c1fa88df..9c8f451974 100644
--- a/runtime/src/snapshot_utils.rs
+++ b/runtime/src/snapshot_utils.rs
@@ -44,6 +44,7 @@ pub const TAR_VERSION_FILE: &str = "version";
 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-";
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]
 pub enum SnapshotVersion {
@@ -170,7 +171,9 @@ pub fn package_snapshot<P: AsRef<Path>, Q: AsRef<Path>>(
     snapshot_version: SnapshotVersion,
 ) -> Result<AccountsPackage> {
     // Hard link all the snapshots we need for this package
-    let snapshot_hard_links_dir = tempfile::tempdir_in(snapshot_path)?;
+    let snapshot_hard_links_dir = tempfile::Builder::new()
+        .prefix(TMP_SNAPSHOT_DIR_PREFIX)
+        .tempdir_in(snapshot_path)?;
 
     // Create a snapshot package
     info!(
@@ -213,6 +216,25 @@ fn get_compression_ext(compression: &CompressionType) -> &'static str {
     }
 }
 
+// If the validator is halted in the middle of `archive_snapshot_package` the temporary staging directory
+// won't be cleaned up.  Call this function to clean them up
+pub fn remove_tmp_snapshot_archives(snapshot_path: &Path) {
+    if let Ok(entries) = fs::read_dir(&snapshot_path) {
+        for entry in entries.filter_map(|entry| entry.ok()) {
+            if entry
+                .file_name()
+                .into_string()
+                .unwrap_or_else(|_| String::new())
+                .starts_with(TMP_SNAPSHOT_DIR_PREFIX)
+            {
+                fs::remove_dir_all(entry.path()).unwrap_or_else(|err| {
+                    warn!("Failed to remove {}: {}", entry.path().display(), err)
+                });
+            }
+        }
+    }
+}
+
 pub fn archive_snapshot_package(snapshot_package: &AccountsPackage) -> Result<()> {
     info!(
         "Generating snapshot archive for slot {}",
@@ -234,7 +256,10 @@ pub fn archive_snapshot_package(snapshot_package: &AccountsPackage) -> Result<()
     fs::create_dir_all(tar_dir)?;
 
     // Create the staging directories
-    let staging_dir = tempfile::tempdir_in(tar_dir)?;
+    let staging_dir = tempfile::Builder::new()
+        .prefix(TMP_SNAPSHOT_DIR_PREFIX)
+        .tempdir_in(tar_dir)?;
+
     let staging_accounts_dir = staging_dir.path().join(TAR_ACCOUNTS_DIR);
     let staging_snapshots_dir = staging_dir.path().join(TAR_SNAPSHOTS_DIR);
     let staging_version_file = staging_dir.path().join(TAR_VERSION_FILE);
@@ -580,8 +605,10 @@ pub fn bank_from_archive<P: AsRef<Path>>(
     debug_keys: Option<Arc<HashSet<Pubkey>>>,
     additional_builtins: Option<&Builtins>,
 ) -> Result<Bank> {
-    // Untar the snapshot into a temp directory under `snapshot_config.snapshot_path()`
-    let unpack_dir = tempfile::tempdir_in(snapshot_path)?;
+    // Untar the snapshot into a temporary directory
+    let unpack_dir = tempfile::Builder::new()
+        .prefix(TMP_SNAPSHOT_DIR_PREFIX)
+        .tempdir_in(snapshot_path)?;
     untar_snapshot_in(&snapshot_tar, &unpack_dir, compression)?;
 
     let mut measure = Measure::start("bank rebuild from snapshot");
@@ -885,8 +912,7 @@ pub fn snapshot_bank(
     let latest_slot_snapshot_paths = slot_snapshot_paths
         .last()
         .expect("no snapshots found in config snapshot_path");
-    // We only care about the last bank's snapshot.
-    // We'll ask the bank for MAX_CACHE_ENTRIES (on the rooted path) worth of statuses
+
     let package = package_snapshot(
         &root_bank,
         latest_slot_snapshot_paths,
diff --git a/validator/src/main.rs b/validator/src/main.rs
index 6bfa33edfc..0981b422a2 100644
--- a/validator/src/main.rs
+++ b/validator/src/main.rs
@@ -828,6 +828,7 @@ fn create_validator(
         enable_recycler_warming();
     }
     solana_ledger::entry::init_poh();
+    solana_runtime::snapshot_utils::remove_tmp_snapshot_archives(ledger_path);
 
     if let Some(ref cluster_entrypoint) = cluster_entrypoint {
         rpc_bootstrap(