From 9ddb89226c6d6b6b9604c35583de167e5a8bf808 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Thu, 23 Feb 2023 15:20:33 -0600
Subject: [PATCH] remove_dead_accounts uses slot instead of AppendVecId
 (#30471)

---
 runtime/src/account_storage.rs | 17 +++++++++++------
 runtime/src/accounts_db.rs     | 12 +++++++-----
 2 files changed, 18 insertions(+), 11 deletions(-)

diff --git a/runtime/src/account_storage.rs b/runtime/src/account_storage.rs
index e40dbd3e38..4cf61d15dc 100644
--- a/runtime/src/account_storage.rs
+++ b/runtime/src/account_storage.rs
@@ -67,10 +67,15 @@ impl AccountStorage {
             .or_else(lookup_in_map)
     }
 
+    /// assert if shrink in progress is active
+    pub(crate) fn assert_no_shrink_in_progress(&self) {
+        assert!(self.shrink_in_progress_map.is_empty());
+    }
+
     /// return the append vec for 'slot' if it exists
     /// This is only ever called when shrink is not possibly running and there is a max of 1 append vec per slot.
     pub(crate) fn get_slot_storage_entry(&self, slot: Slot) -> Option<Arc<AccountStorageEntry>> {
-        assert!(self.shrink_in_progress_map.is_empty());
+        self.assert_no_shrink_in_progress();
         self.get_slot_storage_entry_shrinking_in_progress_ok(slot)
     }
 
@@ -83,21 +88,21 @@ impl AccountStorage {
     }
 
     pub(crate) fn all_slots(&self) -> Vec<Slot> {
-        assert!(self.shrink_in_progress_map.is_empty());
+        self.assert_no_shrink_in_progress();
         self.map.iter().map(|iter_item| *iter_item.key()).collect()
     }
 
     /// returns true if there is no entry for 'slot'
     #[cfg(test)]
     pub(crate) fn is_empty_entry(&self, slot: Slot) -> bool {
-        assert!(self.shrink_in_progress_map.is_empty());
+        self.assert_no_shrink_in_progress();
         self.map.get(&slot).is_none()
     }
 
     /// initialize the storage map to 'all_storages'
     pub(crate) fn initialize(&mut self, all_storages: AccountStorageMap) {
         assert!(self.map.is_empty());
-        assert!(self.shrink_in_progress_map.is_empty());
+        self.assert_no_shrink_in_progress();
         self.map.extend(all_storages.into_iter())
     }
 
@@ -114,12 +119,12 @@ impl AccountStorage {
 
     /// iterate through all (slot, append-vec)
     pub(crate) fn iter(&self) -> AccountStorageIter<'_> {
-        assert!(self.shrink_in_progress_map.is_empty());
+        self.assert_no_shrink_in_progress();
         AccountStorageIter::new(self)
     }
 
     pub(crate) fn insert(&self, slot: Slot, store: Arc<AccountStorageEntry>) {
-        assert!(self.shrink_in_progress_map.is_empty());
+        self.assert_no_shrink_in_progress();
         assert!(self
             .map
             .insert(
diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index 681bc07a7f..2f7f77c33d 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -740,8 +740,8 @@ pub type AtomicAppendVecId = AtomicU32;
 pub type AppendVecId = u32;
 
 type AccountSlots = HashMap<Pubkey, HashSet<Slot>>;
-type AppendVecOffsets = HashMap<AppendVecId, HashSet<usize>>;
-type ReclaimResult = (AccountSlots, AppendVecOffsets);
+type SlotOffsets = HashMap<Slot, HashSet<usize>>;
+type ReclaimResult = (AccountSlots, SlotOffsets);
 type PubkeysRemovedFromAccountsIndex = HashSet<Pubkey>;
 type ShrinkCandidates = HashMap<Slot, Arc<AccountStorageEntry>>;
 
@@ -3371,7 +3371,7 @@ impl AccountsDb {
                 // Check if this update in `slot` to the account with `key` was reclaimed earlier by
                 // `clean_accounts_older_than_root()`
                 let was_reclaimed = removed_accounts
-                    .get(&account_info.store_id())
+                    .get(slot)
                     .map(|store_removed| store_removed.contains(&account_info.offset()))
                     .unwrap_or(false);
                 if was_reclaimed {
@@ -7941,12 +7941,14 @@ impl AccountsDb {
         &'a self,
         reclaims: I,
         expected_slot: Option<Slot>,
-        mut reclaimed_offsets: Option<&mut AppendVecOffsets>,
+        mut reclaimed_offsets: Option<&mut SlotOffsets>,
         reset_accounts: bool,
     ) -> HashSet<Slot>
     where
         I: Iterator<Item = &'a (Slot, AccountInfo)>,
     {
+        self.storage.assert_no_shrink_in_progress();
+
         let mut dead_slots = HashSet::new();
         let mut new_shrink_candidates: ShrinkCandidates = HashMap::new();
         let mut measure = Measure::start("remove");
@@ -7955,7 +7957,7 @@ impl AccountsDb {
             assert!(!account_info.is_cached());
             if let Some(ref mut reclaimed_offsets) = reclaimed_offsets {
                 reclaimed_offsets
-                    .entry(account_info.store_id())
+                    .entry(*slot)
                     .or_default()
                     .insert(account_info.offset());
             }