Fix race in remove_unrooted_slot (#10607) (#10617)

* Fix race

* clippy fixes

* Rename and add comment

Co-authored-by: Carl <carl@solana.com>
(cherry picked from commit 8bd62d78eb)

Co-authored-by: carllin <wumu727@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

runtime/src/accounts_db.rs

@@ -42,7 +42,7 @@ use std::{
     ops::RangeBounds,
     path::{Path, PathBuf},
     sync::atomic::{AtomicBool, AtomicU64, AtomicUsize, Ordering},
-    sync::{Arc, Mutex, RwLock},
+    sync::{Arc, Mutex, RwLock, RwLockWriteGuard},
     time::Instant,
 };
 use tempfile::TempDir;
@@ -535,7 +535,7 @@ impl AccountsDB {
         inc_new_counter_info!("clean-old-root-par-clean-ms", clean_rooted.as_ms() as usize);
 
         let mut measure = Measure::start("clean_old_root_reclaims");
-        self.handle_reclaims(&reclaims);
+        self.handle_reclaims_maybe_cleanup(&reclaims);
         measure.stop();
         debug!("{} {}", clean_rooted, measure);
         inc_new_counter_info!("clean-old-root-reclaim-ms", measure.as_ms() as usize);
@@ -732,7 +732,7 @@ impl AccountsDB {
             }
         }
 
-        self.handle_reclaims(&reclaims);
+        self.handle_reclaims_maybe_cleanup(&reclaims);
         reclaims_time.stop();
         debug!(
             "clean_accounts: {} {} {} {}",
@@ -740,7 +740,7 @@ impl AccountsDB {
         );
     }
 
-    fn handle_reclaims(&self, reclaims: SlotSlice<AccountInfo>) {
+    fn handle_reclaims_maybe_cleanup(&self, reclaims: SlotSlice<AccountInfo>) {
         let mut dead_accounts = Measure::start("reclaims::remove_dead_accounts");
         let dead_slots = self.remove_dead_accounts(reclaims);
         dead_accounts.stop();
@@ -750,13 +750,27 @@ impl AccountsDB {
             dead_slots_w.len()
         };
         if dead_slots_len > 5000 {
-            self.process_dead_slots();
+            self.process_dead_slots(None);
         }
     }
 
-    pub fn process_dead_slots(&self) {
+    // Atomicallly process reclaims and new dead_slots in this thread, gauranteeing
-        let empty = HashSet::new();
+    // complete data removal for slots in reclaims.
+    fn handle_reclaims_ensure_cleanup(&self, reclaims: SlotSlice<AccountInfo>) {
+        let mut dead_accounts = Measure::start("reclaims::remove_dead_accounts");
+        let dead_slots = self.remove_dead_accounts(reclaims);
+        dead_accounts.stop();
         let mut dead_slots_w = self.dead_slots.write().unwrap();
+        dead_slots_w.extend(dead_slots);
+        self.process_dead_slots(Some(dead_slots_w));
+    }
+
+    pub fn process_dead_slots<'a>(
+        &'a self,
+        dead_slots_w: Option<RwLockWriteGuard<'a, HashSet<Slot>>>,
+    ) {
+        let empty = HashSet::new();
+        let mut dead_slots_w = dead_slots_w.unwrap_or_else(|| self.dead_slots.write().unwrap());
         let dead_slots = std::mem::replace(&mut *dead_slots_w, empty);
         drop(dead_slots_w);
 
@@ -894,7 +908,7 @@ impl AccountsDB {
             );
             let reclaims = self.update_index(slot, infos, &accounts);
 
-            self.handle_reclaims(&reclaims);
+            self.handle_reclaims_maybe_cleanup(&reclaims);
 
             let mut storage = self.storage.write().unwrap();
             if let Some(slot_storage) = storage.0.get_mut(&slot) {
@@ -1188,13 +1202,9 @@ impl AccountsDB {
             }
         }
 
-        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();
+        self.handle_reclaims_ensure_cleanup(&reclaims);
-
-        // Sanity check storage entries are removed from the index
         assert!(self.storage.read().unwrap().0.get(&remove_slot).is_none());
     }
 
@@ -1828,7 +1838,7 @@ impl AccountsDB {
         update_index.stop();
         trace!("reclaim: {}", reclaims.len());
 
-        self.handle_reclaims(&reclaims);
+        self.handle_reclaims_maybe_cleanup(&reclaims);
     }
 
     pub fn add_root(&self, slot: Slot) {
@@ -2538,7 +2548,7 @@ pub mod tests {
         //slot is gone
         print_accounts("pre-clean", &accounts);
         accounts.clean_accounts();
-        accounts.process_dead_slots();
+        accounts.process_dead_slots(None);
         assert!(accounts.storage.read().unwrap().0.get(&0).is_none());
 
         //new value is there
@@ -2976,7 +2986,7 @@ pub mod tests {
         let hash = accounts.update_accounts_hash(current_slot, &ancestors);
 
         accounts.clean_accounts();
-        accounts.process_dead_slots();
+        accounts.process_dead_slots(None);
 
         assert_eq!(
             accounts.update_accounts_hash(current_slot, &ancestors),
@@ -3792,7 +3802,7 @@ pub mod tests {
         current_slot += 1;
         assert_eq!(4, accounts.ref_count_for_pubkey(&pubkey1));
         accounts.store(current_slot, &[(&pubkey1, &zero_lamport_account)]);
-        accounts.process_dead_slots();
+        accounts.process_dead_slots(None);
         assert_eq!(
             3, /* == 4 - 2 + 1 */
             accounts.ref_count_for_pubkey(&pubkey1)

runtime/src/bank.rs

@@ -2553,7 +2553,7 @@ impl Bank {
     }
 
     pub fn process_dead_slots(&self) {
-        self.rc.accounts.accounts_db.process_dead_slots();
+        self.rc.accounts.accounts_db.process_dead_slots(None);
     }
 
     pub fn shrink_all_slots(&self) {