From 93ae1775038553c56478f390324c5bd3ad88a791 Mon Sep 17 00:00:00 2001
From: Michael Vines <mvines@gmail.com>
Date: Sun, 20 Dec 2020 22:37:33 -0800
Subject: [PATCH] Do not delete ALL other snapshots before downloading a new
 snapshot

---
 download-utils/src/lib.rs     | 80 +++++++++++++++--------------------
 runtime/src/snapshot_utils.rs | 19 +++++----
 2 files changed, 45 insertions(+), 54 deletions(-)

diff --git a/download-utils/src/lib.rs b/download-utils/src/lib.rs
index 40d903e4c6..441c6be9a0 100644
--- a/download-utils/src/lib.rs
+++ b/download-utils/src/lib.rs
@@ -168,52 +168,40 @@ pub fn download_snapshot(
     desired_snapshot_hash: (Slot, Hash),
     use_progress_bar: bool,
 ) -> Result<(), String> {
-    // Remove all snapshot not matching the desired hash
-    let snapshot_packages = snapshot_utils::get_snapshot_archives(ledger_path);
-    let mut found_package = false;
-    for (snapshot_package, (snapshot_slot, snapshot_hash, _compression)) in snapshot_packages.iter()
-    {
-        if (*snapshot_slot, *snapshot_hash) != desired_snapshot_hash {
-            info!("Removing old snapshot: {:?}", snapshot_package);
-            fs::remove_file(snapshot_package)
-                .unwrap_or_else(|err| info!("Failed to remove old snapshot: {:}", err));
-        } else {
-            found_package = true;
+    snapshot_utils::purge_old_snapshot_archives(ledger_path);
+
+    for compression in &[
+        CompressionType::Zstd,
+        CompressionType::Gzip,
+        CompressionType::Bzip2,
+    ] {
+        let desired_snapshot_package = snapshot_utils::get_snapshot_archive_path(
+            ledger_path,
+            &desired_snapshot_hash,
+            compression,
+        );
+
+        if desired_snapshot_package.is_file() {
+            return Ok(());
+        }
+
+        if download_file(
+            &format!(
+                "http://{}/{}",
+                rpc_addr,
+                desired_snapshot_package
+                    .file_name()
+                    .unwrap()
+                    .to_str()
+                    .unwrap()
+            ),
+            &desired_snapshot_package,
+            use_progress_bar,
+        )
+        .is_ok()
+        {
+            return Ok(());
         }
     }
-
-    if found_package {
-        Ok(())
-    } else {
-        for compression in &[
-            CompressionType::Zstd,
-            CompressionType::Gzip,
-            CompressionType::Bzip2,
-        ] {
-            let desired_snapshot_package = snapshot_utils::get_snapshot_archive_path(
-                ledger_path,
-                &desired_snapshot_hash,
-                compression,
-            );
-
-            if download_file(
-                &format!(
-                    "http://{}/{}",
-                    rpc_addr,
-                    desired_snapshot_package
-                        .file_name()
-                        .unwrap()
-                        .to_str()
-                        .unwrap()
-                ),
-                &desired_snapshot_package,
-                use_progress_bar,
-            )
-            .is_ok()
-            {
-                return Ok(());
-            }
-        }
-        Err("Snapshot couldn't be downloaded".to_string())
-    }
+    Err("Snapshot couldn't be downloaded".to_string())
 }
diff --git a/runtime/src/snapshot_utils.rs b/runtime/src/snapshot_utils.rs
index 898d481d24..b89a9acef6 100644
--- a/runtime/src/snapshot_utils.rs
+++ b/runtime/src/snapshot_utils.rs
@@ -363,14 +363,7 @@ pub fn archive_snapshot_package(snapshot_package: &AccountsPackage) -> Result<()
     let metadata = fs::metadata(&archive_path)?;
     fs::rename(&archive_path, &snapshot_package.tar_output_file)?;
 
-    // Keep around at most three snapshot archives
-    let mut archives = get_snapshot_archives(snapshot_package.tar_output_file.parent().unwrap());
-    // Keep the oldest snapshot so we can always play the ledger from it.
-    archives.pop();
-    for old_archive in archives.into_iter().skip(2) {
-        fs::remove_file(old_archive.0)
-            .unwrap_or_else(|err| info!("Failed to remove old snapshot: {:}", err));
-    }
+    purge_old_snapshot_archives(snapshot_package.tar_output_file.parent().unwrap());
 
     timer.stop();
     info!(
@@ -746,6 +739,16 @@ pub fn get_highest_snapshot_archive_path<P: AsRef<Path>>(
     archives.into_iter().next()
 }
 
+pub fn purge_old_snapshot_archives<P: AsRef<Path>>(snapshot_output_dir: P) {
+    let mut archives = get_snapshot_archives(snapshot_output_dir);
+    // Keep the oldest snapshot so we can always play the ledger from it.
+    archives.pop();
+    for old_archive in archives.into_iter().skip(2) {
+        fs::remove_file(old_archive.0)
+            .unwrap_or_else(|err| info!("Failed to remove old snapshot: {:}", err));
+    }
+}
+
 pub fn untar_snapshot_in<P: AsRef<Path>, Q: AsRef<Path>>(
     snapshot_tar: P,
     unpack_dir: Q,