Repair alternate versions of dead slots (#9805)

Co-authored-by: Carl <carl@solana.com>

runtime/src/accounts_db.rs

@@ -478,7 +478,7 @@ pub struct AccountsDB {
 
     pub bank_hashes: RwLock<HashMap<Slot, BankHashInfo>>,
 
-    pub dead_slots: RwLock<HashSet<Slot>>,
+    dead_slots: RwLock<HashSet<Slot>>,
 
     stats: AccountsStats,
 }
@@ -1266,6 +1266,39 @@ impl AccountsDB {
         }
     }
 
+    pub fn remove_unrooted_slot(&self, remove_slot: Slot) {
+        if self.accounts_index.read().unwrap().is_root(remove_slot) {
+            panic!("Trying to remove accounts for rooted slot {}", remove_slot);
+        }
+
+        let pubkey_sets: Vec<HashSet<Pubkey>> = self.scan_account_storage(
+            remove_slot,
+            |stored_account: &StoredAccount, _, accum: &mut HashSet<Pubkey>| {
+                accum.insert(stored_account.meta.pubkey);
+            },
+        );
+
+        // Purge this slot from the accounts index
+        let mut reclaims = vec![];
+        {
+            let pubkeys = pubkey_sets.iter().flatten();
+            let accounts_index = self.accounts_index.read().unwrap();
+
+            for pubkey in pubkeys {
+                accounts_index.clean_unrooted_entries_by_slot(remove_slot, pubkey, &mut reclaims);
+            }
+        }
+
+        self.handle_reclaims(&reclaims);
+
+        // 1) Remove old bank hash from self.bank_hashes
+        // 2) Purge this slot's storage entries from self.storage
+        self.process_dead_slots();
+
+        // Sanity check storage entries are removed from the index
+        assert!(self.storage.read().unwrap().0.get(&remove_slot).is_none());
+    }
+
     pub fn hash_stored_account(slot: Slot, account: &StoredAccount) -> Hash {
         Self::hash_account_data(
             slot,
@@ -2199,6 +2232,80 @@ pub mod tests {
         assert_eq!(db0.load_slow(&ancestors, &key), Some((account0, 0)));
     }
 
+    #[test]
+    fn test_remove_unrooted_slot() {
+        let unrooted_slot = 9;
+        let db = AccountsDB::new(Vec::new());
+        let key = Pubkey::default();
+        let account0 = Account::new(1, 0, &key);
+        let ancestors: HashMap<_, _> = vec![(unrooted_slot, 1)].into_iter().collect();
+        db.store(unrooted_slot, &[(&key, &account0)]);
+        db.bank_hashes
+            .write()
+            .unwrap()
+            .insert(unrooted_slot, BankHashInfo::default());
+        assert!(db
+            .accounts_index
+            .read()
+            .unwrap()
+            .get(&key, &ancestors)
+            .is_some());
+        assert_load_account(&db, unrooted_slot, key, 1);
+
+        // Purge the slot
+        db.remove_unrooted_slot(unrooted_slot);
+        assert!(db.load_slow(&ancestors, &key).is_none());
+        assert!(db.bank_hashes.read().unwrap().get(&unrooted_slot).is_none());
+        assert!(db.storage.read().unwrap().0.get(&unrooted_slot).is_none());
+        assert!(db
+            .accounts_index
+            .read()
+            .unwrap()
+            .account_maps
+            .get(&key)
+            .map(|pubkey_entry| pubkey_entry.1.read().unwrap().is_empty())
+            .unwrap_or(true));
+        assert!(db
+            .accounts_index
+            .read()
+            .unwrap()
+            .get(&key, &ancestors)
+            .is_none());
+
+        // Test we can store for the same slot again and get the right information
+        let account0 = Account::new(2, 0, &key);
+        db.store(unrooted_slot, &[(&key, &account0)]);
+        assert_load_account(&db, unrooted_slot, key, 2);
+    }
+
+    #[test]
+    fn test_remove_unrooted_slot_snapshot() {
+        let unrooted_slot = 9;
+        let db = AccountsDB::new(Vec::new());
+        let key = Pubkey::new_rand();
+        let account0 = Account::new(1, 0, &key);
+        db.store(unrooted_slot, &[(&key, &account0)]);
+
+        // Purge the slot
+        db.remove_unrooted_slot(unrooted_slot);
+
+        // Add a new root
+        let key2 = Pubkey::new_rand();
+        let new_root = unrooted_slot + 1;
+        db.store(new_root, &[(&key2, &account0)]);
+        db.add_root(new_root);
+
+        // Simulate reconstruction from snapshot
+        let db = reconstruct_accounts_db_via_serialization(&db, new_root);
+
+        // Check root account exists
+        assert_load_account(&db, new_root, key2, 1);
+
+        // Check purged account stays gone
+        let unrooted_slot_ancestors: HashMap<_, _> = vec![(unrooted_slot, 1)].into_iter().collect();
+        assert!(db.load_slow(&unrooted_slot_ancestors, &key).is_none());
+    }
+
     fn create_account(
         accounts: &AccountsDB,
         pubkeys: &mut Vec<Pubkey>,

runtime/src/accounts_index.rs

@@ -172,6 +172,23 @@ impl<T: Clone> AccountsIndex<T> {
         }
     }
 
+    pub fn clean_unrooted_entries_by_slot(
+        &self,
+        purge_slot: Slot,
+        pubkey: &Pubkey,
+        reclaims: &mut SlotList<T>,
+    ) {
+        if let Some(entry) = self.account_maps.get(pubkey) {
+            let mut list = entry.1.write().unwrap();
+            list.retain(|(slot, entry)| {
+                if *slot == purge_slot {
+                    reclaims.push((*slot, entry.clone()));
+                }
+                *slot != purge_slot
+            });
+        }
+    }
+
     pub fn add_index(&mut self, slot: Slot, pubkey: &Pubkey, account_info: T) {
         let entry = self
             .account_maps

runtime/src/bank.rs

@@ -355,6 +355,9 @@ pub struct Bank {
     /// Rewards that were paid out immediately after this bank was created
     #[serde(skip)]
     pub rewards: Option<Vec<(Pubkey, i64)>>,
+
+    #[serde(skip)]
+    pub skip_drop: AtomicBool,
 }
 
 impl Default for BlockhashQueue {
@@ -466,6 +469,7 @@ impl Bank {
             hard_forks: parent.hard_forks.clone(),
             last_vote_sync: AtomicU64::new(parent.last_vote_sync.load(Ordering::Relaxed)),
             rewards: None,
+            skip_drop: AtomicBool::new(false),
         };
 
         datapoint_info!(
@@ -966,6 +970,14 @@ impl Bank {
         self.src.status_cache.write().unwrap().clear_signatures();
     }
 
+    pub fn clear_slot_signatures(&self, slot: Slot) {
+        self.src
+            .status_cache
+            .write()
+            .unwrap()
+            .clear_slot_signatures(slot);
+    }
+
     pub fn can_commit(result: &Result<()>) -> bool {
         match result {
             Ok(_) => true,
@@ -1056,6 +1068,10 @@ impl Bank {
         }
     }
 
+    pub fn remove_unrooted_slot(&self, slot: Slot) {
+        self.rc.accounts.accounts_db.remove_unrooted_slot(slot)
+    }
+
     fn load_accounts(
         &self,
         txs: &[Transaction],
@@ -2259,7 +2275,9 @@ impl Bank {
 impl Drop for Bank {
     fn drop(&mut self) {
         // For root slots this is a noop
-        self.rc.accounts.purge_slot(self.slot());
+        if !self.skip_drop.load(Ordering::Relaxed) {
+            self.rc.accounts.purge_slot(self.slot());
+        }
     }
 }
 

runtime/src/genesis_utils.rs

@@ -28,6 +28,14 @@ impl ValidatorVoteKeypairs {
             stake_keypair,
         }
     }
+
+    pub fn new_rand() -> Self {
+        Self {
+            node_keypair: Keypair::new(),
+            vote_keypair: Keypair::new(),
+            stake_keypair: Keypair::new(),
+        }
+    }
 }
 
 pub struct GenesisConfigInfo {

runtime/src/status_cache.rs

@@ -9,7 +9,7 @@ use solana_sdk::{
     signature::Signature,
 };
 use std::{
-    collections::{HashMap, HashSet},
+    collections::{hash_map::Entry, HashMap, HashSet},
     sync::{Arc, Mutex},
 };
 
@@ -82,6 +82,46 @@ impl<T: Serialize + Clone + PartialEq> PartialEq for StatusCache<T> {
 }
 
 impl<T: Serialize + Clone> StatusCache<T> {
+    pub fn clear_slot_signatures(&mut self, slot: Slot) {
+        let slot_deltas = self.slot_deltas.remove(&slot);
+        if let Some(slot_deltas) = slot_deltas {
+            let slot_deltas = slot_deltas.lock().unwrap();
+            for (blockhash, (_, signature_list)) in slot_deltas.iter() {
+                // Any blockhash that exists in self.slot_deltas must also exist
+                // in self.cache, because in self.purge_roots(), when an entry
+                // (b, (max_slot, _, _)) is removed from self.cache, this implies
+                // all entries in self.slot_deltas < max_slot are also removed
+                if let Entry::Occupied(mut o_blockhash_entries) = self.cache.entry(*blockhash) {
+                    let (_, _, all_sig_maps) = o_blockhash_entries.get_mut();
+
+                    for (sig_slice, _) in signature_list {
+                        if let Entry::Occupied(mut o_sig_list) = all_sig_maps.entry(*sig_slice) {
+                            let sig_list = o_sig_list.get_mut();
+                            sig_list.retain(|(updated_slot, _)| *updated_slot != slot);
+                            if sig_list.is_empty() {
+                                o_sig_list.remove_entry();
+                            }
+                        } else {
+                            panic!(
+                                "Map for signature must exist if siganture exists in self.slot_deltas, slot: {}",
+                                slot
+                            )
+                        }
+                    }
+
+                    if all_sig_maps.is_empty() {
+                        o_blockhash_entries.remove_entry();
+                    }
+                } else {
+                    panic!(
+                        "Blockhash must exist if it exists in self.slot_deltas, slot: {}",
+                        slot
+                    )
+                }
+            }
+        }
+    }
+
     /// Check if the signature from a transaction is in any of the forks in the ancestors set.
     pub fn get_signature_status(
         &self,
@@ -408,6 +448,8 @@ mod tests {
             status_cache.add_root(i as u64);
         }
         let slots: Vec<_> = (0_u64..MAX_CACHE_ENTRIES as u64 + 1).collect();
+        assert_eq!(status_cache.slot_deltas.len(), 1);
+        assert!(status_cache.slot_deltas.get(&1).is_some());
         let slot_deltas = status_cache.slot_deltas(&slots);
         let cache = StatusCache::from_slot_deltas(&slot_deltas);
         assert_eq!(cache, status_cache);
@@ -417,4 +459,51 @@ mod tests {
     fn test_age_sanity() {
         assert!(MAX_CACHE_ENTRIES <= MAX_RECENT_BLOCKHASHES);
     }
+
+    #[test]
+    fn test_clear_slot_signatures() {
+        let sig = Signature::default();
+        let mut status_cache = BankStatusCache::default();
+        let blockhash = hash(Hash::default().as_ref());
+        let blockhash2 = hash(blockhash.as_ref());
+        status_cache.insert(&blockhash, &sig, 0, ());
+        status_cache.insert(&blockhash, &sig, 1, ());
+        status_cache.insert(&blockhash2, &sig, 1, ());
+
+        let mut ancestors0 = HashMap::new();
+        ancestors0.insert(0, 0);
+        let mut ancestors1 = HashMap::new();
+        ancestors1.insert(1, 0);
+
+        // Clear slot 0 related data
+        assert!(status_cache
+            .get_signature_status(&sig, &blockhash, &ancestors0)
+            .is_some());
+        status_cache.clear_slot_signatures(0);
+        assert!(status_cache
+            .get_signature_status(&sig, &blockhash, &ancestors0)
+            .is_none());
+        assert!(status_cache
+            .get_signature_status(&sig, &blockhash, &ancestors1)
+            .is_some());
+        assert!(status_cache
+            .get_signature_status(&sig, &blockhash2, &ancestors1)
+            .is_some());
+
+        // Check that the slot delta for slot 0 is gone, but slot 1 still
+        // exists
+        assert!(status_cache.slot_deltas.get(&0).is_none());
+        assert!(status_cache.slot_deltas.get(&1).is_some());
+
+        // Clear slot 1 related data
+        status_cache.clear_slot_signatures(1);
+        assert!(status_cache.slot_deltas.is_empty());
+        assert!(status_cache
+            .get_signature_status(&sig, &blockhash, &ancestors1)
+            .is_none());
+        assert!(status_cache
+            .get_signature_status(&sig, &blockhash2, &ancestors1)
+            .is_none());
+        assert!(status_cache.cache.is_empty());
+    }
 }