From 5d9130a3c462e49bbdde2eee47cff6e3631e2d0c Mon Sep 17 00:00:00 2001
From: Ryo Onodera <ryoqun@gmail.com>
Date: Mon, 2 Mar 2020 22:49:09 +0900
Subject: [PATCH] Hack to skip cleanup_dead_slots upon snapshot load

---
 ledger/src/snapshot_utils.rs |  4 +++-
 runtime/src/accounts_db.rs   | 46 +++++++++++++++++++++++++++++++-----
 runtime/src/bank.rs          |  8 +++++++
 3 files changed, 51 insertions(+), 7 deletions(-)

diff --git a/ledger/src/snapshot_utils.rs b/ledger/src/snapshot_utils.rs
index e849b3eae2..c2b4574718 100644
--- a/ledger/src/snapshot_utils.rs
+++ b/ledger/src/snapshot_utils.rs
@@ -447,16 +447,18 @@ pub fn bank_from_archive<P: AsRef<Path>>(
     let mut snapshot_version = String::new();
     File::open(unpacked_version_file).and_then(|mut f| f.read_to_string(&mut snapshot_version))?;
 
-    let bank = rebuild_bank_from_snapshots(
+    let mut bank = rebuild_bank_from_snapshots(
         snapshot_version.trim(),
         account_paths,
         &unpacked_snapshots_dir,
         unpacked_accounts_dir,
     )?;
 
+    bank.work_around_dead_slots_cleaning_bug(true);
     if !bank.verify_snapshot_bank() {
         panic!("Snapshot bank for slot {} failed to verify", bank.slot());
     }
+    bank.work_around_dead_slots_cleaning_bug(false);
     measure.stop();
     info!("{}", measure);
 
diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index cdcdee7368..e2715356b8 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -452,6 +452,8 @@ pub struct AccountsDB {
     min_num_stores: usize,
 
     pub bank_hashes: RwLock<HashMap<Slot, BankHashInfo>>,
+
+    pub dont_cleanup_dead_slots: AtomicBool,
 }
 
 impl Default for AccountsDB {
@@ -474,6 +476,7 @@ impl Default for AccountsDB {
                 .unwrap(),
             min_num_stores: num_threads,
             bank_hashes: RwLock::new(bank_hashes),
+            dont_cleanup_dead_slots: AtomicBool::new(false),
         }
     }
 }
@@ -1296,6 +1299,10 @@ impl AccountsDB {
     }
 
     fn cleanup_dead_slots(&self, dead_slots: &mut HashSet<Slot>) {
+        if self.dont_cleanup_dead_slots.load(Ordering::Relaxed) {
+            return;
+        }
+
         if !dead_slots.is_empty() {
             {
                 let mut index = self.accounts_index.write().unwrap();
@@ -2290,10 +2297,10 @@ pub mod tests {
         assert_load_account(&accounts, current_slot, pubkey, zero_lamport);
     }
 
-    #[test]
-    fn test_accounts_purge_chained() {
-        solana_logger::setup();
-
+    fn with_chained_zero_lamport_accounts<F>(f: F)
+    where
+        F: Fn(AccountsDB, Slot) -> (AccountsDB, bool),
+    {
         let some_lamport = 223;
         let zero_lamport = 0;
         let dummy_lamport = 999;
@@ -2332,8 +2339,12 @@ pub mod tests {
         accounts.store(current_slot, &[(&dummy_pubkey, &dummy_account)]);
         accounts.add_root(current_slot);
 
-        purge_zero_lamport_accounts(&accounts, current_slot);
-        let accounts = reconstruct_accounts_db_via_serialization(&accounts, current_slot);
+        let (accounts, skip_account_assertion) = f(accounts, current_slot);
+
+        assert_eq!(4, accounts.accounts_index.read().unwrap().roots.len());
+        if skip_account_assertion {
+            return;
+        }
 
         assert_load_account(&accounts, current_slot, pubkey, some_lamport);
         assert_load_account(&accounts, current_slot, purged_pubkey1, 0);
@@ -2341,6 +2352,29 @@ pub mod tests {
         assert_load_account(&accounts, current_slot, dummy_pubkey, dummy_lamport);
     }
 
+    #[test]
+    fn test_accounts_purge_chained_purge_before_snapshot_restore() {
+        solana_logger::setup();
+        with_chained_zero_lamport_accounts(|accounts, current_slot| {
+            purge_zero_lamport_accounts(&accounts, current_slot);
+            let accounts = reconstruct_accounts_db_via_serialization(&accounts, current_slot);
+            (accounts, false)
+        });
+    }
+
+    #[test]
+    fn test_accounts_purge_chained_purge_after_snapshot_restore() {
+        solana_logger::setup();
+        with_chained_zero_lamport_accounts(|accounts, current_slot| {
+            let accounts = reconstruct_accounts_db_via_serialization(&accounts, current_slot);
+            accounts
+                .dont_cleanup_dead_slots
+                .store(true, Ordering::Relaxed);
+            purge_zero_lamport_accounts(&accounts, current_slot);
+            (accounts, true)
+        });
+    }
+
     #[test]
     #[ignore]
     fn test_store_account_stress() {
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index d9a3bc9262..4efae52294 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -1730,6 +1730,14 @@ impl Bank {
         self.rc.parent = RwLock::new(Some(parent.clone()));
     }
 
+    pub fn work_around_dead_slots_cleaning_bug(&mut self, flag: bool) {
+        self.rc
+            .accounts
+            .accounts_db
+            .dont_cleanup_dead_slots
+            .store(flag, Ordering::Relaxed);
+    }
+
     pub fn set_inflation(&self, inflation: Inflation) {
         *self.inflation.write().unwrap() = inflation;
     }