Reduce account index lookups during clean (#16689)

* reduce account index lookups in clean

* rename

* rename enum

* hold locks during removal from zero pubkey list

* merge with zero lamport fix

* tests

runtime/src/accounts_db.rs

@@ -22,8 +22,8 @@ use crate::{
     accounts_cache::{AccountsCache, CachedAccount, SlotCache},
     accounts_hash::{AccountsHash, CalculateHashIntermediate, HashStats, PreviousPass},
     accounts_index::{
-        AccountIndex, AccountsIndex, AccountsIndexRootsStats, Ancestors, IndexKey, IsCached,
+        AccountIndex, AccountIndexGetResult, AccountsIndex, AccountsIndexRootsStats, Ancestors,
-        SlotList, SlotSlice, ZeroLamport,
+        IndexKey, IsCached, SlotList, SlotSlice, ZeroLamport,
     },
     append_vec::{AppendVec, StoredAccountMeta, StoredMeta},
     contains::Contains,
@@ -1559,9 +1559,8 @@ impl AccountsDb {
                         let mut purges_in_root = Vec::new();
                         let mut purges = HashMap::new();
                         for pubkey in pubkeys {
-                            if let Some((locked_entry, index)) =
+                            match self.accounts_index.get(pubkey, None, max_clean_root) {
-                                self.accounts_index.get(pubkey, None, max_clean_root)
+                                AccountIndexGetResult::Found(locked_entry, index) => {
-                            {
                                     let slot_list = locked_entry.slot_list();
                                     let (slot, account_info) = &slot_list[index];
                                     if account_info.lamports == 0 {
@@ -1574,14 +1573,13 @@ impl AccountsDb {
                                         // prune zero_lamport_pubkey set which should contain all 0-lamport
                                         // keys whether rooted or not. A 0-lamport update may become rooted
                                         // in the future.
-                                    let has_zero_lamport_accounts = slot_list
+                                        if !slot_list
                                             .iter()
-                                        .any(|(_slot, account_info)| account_info.lamports == 0);
+                                            .any(|(_slot, account_info)| account_info.lamports == 0)
-                                    if !has_zero_lamport_accounts {
+                                        {
                                             self.accounts_index.remove_zero_lamport_key(pubkey);
                                         }
                                     }
-
                                     // Release the lock
                                     let slot = *slot;
                                     drop(locked_entry);
@@ -1594,13 +1592,16 @@ impl AccountsDb {
                                         }
                                         purges_in_root.push(*pubkey);
                                     }
-                            } else {
+                                }
-                                let r_accounts_index =
+                                AccountIndexGetResult::NotFoundOnFork => {
-                                    self.accounts_index.account_maps.read().unwrap();
+                                    // do nothing - pubkey is in index, but not found in a root slot
-                                if !r_accounts_index.contains_key(pubkey) {
+                                }
+                                AccountIndexGetResult::Missing(lock) => {
+                                    // pubkey is missing from index, so remove from zero_lamports_list
                                     self.accounts_index.remove_zero_lamport_key(pubkey);
+                                    drop(lock);
                                 }
-                            }
+                            };
                         }
                         (purges, purges_in_root)
                     })
@@ -2334,8 +2335,16 @@ impl AccountsDb {
         max_root: Option<Slot>,
         clone_in_lock: bool,
     ) -> Option<(Slot, AppendVecId, usize, Option<LoadedAccountAccessor<'a>>)> {
-        let (lock, index) = self.accounts_index.get(pubkey, Some(ancestors), max_root)?;
+        let (lock, index) = match self.accounts_index.get(pubkey, Some(ancestors), max_root) {
-        // Notice the subtle `?` at previous line, we bail out pretty early for missing.
+            AccountIndexGetResult::Found(lock, index) => (lock, index),
+            // we bail out pretty early for missing.
+            AccountIndexGetResult::NotFoundOnFork => {
+                return None;
+            }
+            AccountIndexGetResult::Missing(_) => {
+                return None;
+            }
+        };
 
         let slot_list = lock.slot_list();
         let (
@@ -3883,7 +3892,7 @@ impl AccountsDb {
                         let result: Vec<Hash> = pubkeys
                             .iter()
                             .filter_map(|pubkey| {
-                                if let Some((lock, index)) =
+                                if let AccountIndexGetResult::Found(lock, index) =
                                     self.accounts_index.get(pubkey, Some(ancestors), Some(slot))
                                 {
                                     let (slot, account_info) = &lock.slot_list()[index];

runtime/src/accounts_index.rs

@@ -85,6 +85,12 @@ impl<T> AccountMapEntryInner<T> {
     }
 }
 
+pub enum AccountIndexGetResult<'a, T: 'static, U> {
+    Found(ReadAccountMapEntry<T>, usize),
+    NotFoundOnFork,
+    Missing(std::sync::RwLockReadGuard<'a, AccountMap<U, AccountMapEntry<T>>>),
+}
+
 #[self_referencing]
 pub struct ReadAccountMapEntry<T: 'static> {
     owned_entry: AccountMapEntry<T>,
@@ -694,7 +700,9 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         for pubkey in index.get(index_key) {
             // Maybe these reads from the AccountsIndex can be batched every time it
             // grabs the read lock as well...
-            if let Some((list_r, index)) = self.get(&pubkey, Some(ancestors), max_root) {
+            if let AccountIndexGetResult::Found(list_r, index) =
+                self.get(&pubkey, Some(ancestors), max_root)
+            {
                 func(
                     &pubkey,
                     (&list_r.slot_list()[index].1, list_r.slot_list()[index].0),
@@ -936,13 +944,25 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         pubkey: &Pubkey,
         ancestors: Option<&Ancestors>,
         max_root: Option<Slot>,
-    ) -> Option<(ReadAccountMapEntry<T>, usize)> {
+    ) -> AccountIndexGetResult<'_, T, Pubkey> {
-        self.get_account_read_entry(pubkey)
+        let read_lock = self.account_maps.read().unwrap();
-            .and_then(|locked_entry| {
+        let account = read_lock
-                let found_index =
+            .get(pubkey)
-                    self.latest_slot(ancestors, &locked_entry.slot_list(), max_root)?;
+            .cloned()
-                Some((locked_entry, found_index))
+            .map(ReadAccountMapEntry::from_account_map_entry);
-            })
+
+        match account {
+            Some(locked_entry) => {
+                drop(read_lock);
+                let slot_list = locked_entry.slot_list();
+                let found_index = self.latest_slot(ancestors, slot_list, max_root);
+                match found_index {
+                    Some(found_index) => AccountIndexGetResult::Found(locked_entry, found_index),
+                    None => AccountIndexGetResult::NotFoundOnFork,
+                }
+            }
+            None => AccountIndexGetResult::Missing(read_lock),
+        }
     }
 
     // Get the maximum root <= `max_allowed_root` from the given `slice`
@@ -1320,6 +1340,32 @@ pub mod tests {
         account_indexes
     }
 
+    impl<'a, T: 'static, U> AccountIndexGetResult<'a, T, U> {
+        pub fn unwrap(self) -> (ReadAccountMapEntry<T>, usize) {
+            match self {
+                AccountIndexGetResult::Found(lock, size) => (lock, size),
+                _ => {
+                    panic!("trying to unwrap AccountIndexGetResult with non-Success result");
+                }
+            }
+        }
+
+        pub fn is_none(&self) -> bool {
+            !self.is_some()
+        }
+
+        pub fn is_some(&self) -> bool {
+            matches!(self, AccountIndexGetResult::Found(_lock, _size))
+        }
+
+        pub fn map<V, F: FnOnce((ReadAccountMapEntry<T>, usize)) -> V>(self, f: F) -> Option<V> {
+            match self {
+                AccountIndexGetResult::Found(lock, size) => Some(f((lock, size))),
+                _ => None,
+            }
+        }
+    }
+
     fn create_dashmap_secondary_index_state() -> (usize, usize, HashSet<AccountIndex>) {
         {
             // Check that we're actually testing the correct variant