Hack to skip cleanup_dead_slots upon snapshot load

(cherry picked from commit 5d9130a)

ledger/src/snapshot_utils.rs

@@ -453,16 +453,18 @@ pub fn bank_from_archive<P: AsRef<Path>>(
         String::from("0.22.3")
     };
 
-    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);
 

runtime/src/accounts_db.rs

@@ -430,6 +430,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 {
@@ -450,6 +452,7 @@ impl Default for AccountsDB {
                 .unwrap(),
             min_num_stores: num_threads,
             bank_hashes: RwLock::new(HashMap::default()),
+            dont_cleanup_dead_slots: AtomicBool::new(false),
         }
     }
 }
@@ -1163,6 +1166,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();
@@ -2137,10 +2144,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;
@@ -2179,15 +2186,42 @@ 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);
         assert_load_account(&accounts, current_slot, purged_pubkey2, 0);
         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() {

runtime/src/bank.rs

@@ -1729,6 +1729,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;
     }