Adds get_and_then() & family to AccountsIndex (#35218)

accounts-db/src/accounts_db.rs

@@ -9924,11 +9924,11 @@ pub mod tests {
                 .store(data.len(), Ordering::Relaxed);
 
             let genesis_config = GenesisConfig::default();
-            assert!(db.accounts_index.get_account_read_entry(&pubkey).is_none());
+            assert!(!db.accounts_index.contains(&pubkey));
             let result = db.generate_index(None, false, &genesis_config);
             // index entry should only contain a single entry for the pubkey since index cannot hold more than 1 entry per slot
-            let entry = db.accounts_index.get_account_read_entry(&pubkey).unwrap();
-            assert_eq!(entry.slot_list().len(), 1);
+            let entry = db.accounts_index.get_cloned(&pubkey).unwrap();
+            assert_eq!(entry.slot_list.read().unwrap().len(), 1);
             assert_eq!(append_vec.alive_bytes(), expected_alive_bytes);
             // total # accounts in append vec
             assert_eq!(append_vec.approx_stored_count(), 2);
@@ -11409,16 +11409,14 @@ pub mod tests {
         let key = Pubkey::default();
         let account0 = AccountSharedData::new(1, 0, &key);
         let ancestors = vec![(unrooted_slot, 1)].into_iter().collect();
+        assert!(!db.accounts_index.contains(&key));
         if is_cached {
             db.store_cached((unrooted_slot, &[(&key, &account0)][..]), None);
         } else {
             db.store_for_tests(unrooted_slot, &[(&key, &account0)]);
         }
         assert!(db.get_bank_hash_stats(unrooted_slot).is_some());
-        assert!(db
-            .accounts_index
-            .get(&key, Some(&ancestors), None)
-            .is_some());
+        assert!(db.accounts_index.contains(&key));
         db.assert_load_account(unrooted_slot, key, 1);
 
         // Purge the slot
@@ -11427,11 +11425,7 @@ pub mod tests {
         assert!(db.get_bank_hash_stats(unrooted_slot).is_none());
         assert!(db.accounts_cache.slot_cache(unrooted_slot).is_none());
         assert!(db.storage.get_slot_storage_entry(unrooted_slot).is_none());
-        assert!(db.accounts_index.get_account_read_entry(&key).is_none());
-        assert!(db
-            .accounts_index
-            .get(&key, Some(&ancestors), None)
-            .is_none());
+        assert!(!db.accounts_index.contains(&key));
 
         // Test we can store for the same slot again and get the right information
         let account0 = AccountSharedData::new(2, 0, &key);
@@ -12189,11 +12183,8 @@ pub mod tests {
 
         // The earlier entry for pubkey in the account index is purged,
         let (slot_list_len, index_slot) = {
-            let account_entry = accounts
-                .accounts_index
-                .get_account_read_entry(&pubkey)
-                .unwrap();
-            let slot_list = account_entry.slot_list();
+            let account_entry = accounts.accounts_index.get_cloned(&pubkey).unwrap();
+            let slot_list = account_entry.slot_list.read().unwrap();
             (slot_list.len(), slot_list[0].0)
         };
         assert_eq!(slot_list_len, 1);
@@ -12258,10 +12249,7 @@ pub mod tests {
         accounts.print_accounts_stats("post_purge");
 
         // Make sure the index is for pubkey cleared
-        assert!(accounts
-            .accounts_index
-            .get_account_read_entry(&pubkey)
-            .is_none());
+        assert!(!accounts.accounts_index.contains(&pubkey));
 
         // slot 1 & 2 should not have any stores
         assert_no_stores(&accounts, 1);
@@ -13833,10 +13821,7 @@ pub mod tests {
         assert!(db
             .load_without_fixed_root(&ancestors, &unrooted_key)
             .is_some());
-        assert!(db
-            .accounts_index
-            .get_account_read_entry(&unrooted_key)
-            .is_some());
+        assert!(db.accounts_index.contains(&unrooted_key));
         assert_eq!(db.accounts_cache.num_slots(), 1);
         assert!(db.accounts_cache.slot_cache(unrooted_slot).is_some());
         assert_eq!(
@@ -14368,12 +14353,13 @@ pub mod tests {
             let before_size = storage0.alive_bytes.load(Ordering::Acquire);
             let account_info = accounts_db
                 .accounts_index
-                .get_account_read_entry(account.pubkey())
-                .map(|locked_entry| {
-                    // Should only be one entry per key, since every key was only stored to slot 0
-                    locked_entry.slot_list()[0]
-                })
-                .unwrap();
+                .get_cloned(account.pubkey())
+                .unwrap()
+                .slot_list
+                .read()
+                .unwrap()
+                // Should only be one entry per key, since every key was only stored to slot 0
+                [0];
             assert_eq!(account_info.0, slot);
             let reclaims = [account_info];
             accounts_db.remove_dead_accounts(reclaims.iter(), None, None, true);
@@ -15335,8 +15321,7 @@ pub mod tests {
                 assert_no_storages_at_slot(&db, *slot);
                 assert!(db.accounts_cache.slot_cache(*slot).is_none());
                 let account_in_slot = slot_to_pubkey_map[slot];
-                let item = db.accounts_index.get_account_read_entry(&account_in_slot);
-                assert!(item.is_none(), "item: {item:?}");
+                assert!(!db.accounts_index.contains(&account_in_slot));
             }
 
             // Wait for flush to finish before starting next trial
@@ -15722,20 +15707,14 @@ pub mod tests {
         // The later rooted zero-lamport update to `shared_key` cannot be cleaned
         // because it is kept alive by the unrooted slot.
         accounts.clean_accounts_for_tests();
-        assert!(accounts
-            .accounts_index
-            .get_account_read_entry(&shared_key)
-            .is_some());
+        assert!(accounts.accounts_index.contains(&shared_key));
 
         // Simulate purge_slot() all from AccountsBackgroundService
         accounts.purge_slot(slot0, 0, true);
 
         // Now clean should clean up the remaining key
         accounts.clean_accounts_for_tests();
-        assert!(accounts
-            .accounts_index
-            .get_account_read_entry(&shared_key)
-            .is_none());
+        assert!(!accounts.accounts_index.contains(&shared_key));
         assert_no_storages_at_slot(&accounts, slot0);
     }
 

accounts-db/src/accounts_index.rs

@@ -1138,6 +1138,30 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
             .map(ReadAccountMapEntry::from_account_map_entry)
     }
 
+    /// Gets the index's entry for `pubkey` and applies `callback` to it
+    ///
+    /// If `callback`'s boolean return value is true, add this entry to the in-mem cache.
+    pub fn get_and_then<R>(
+        &self,
+        pubkey: &Pubkey,
+        callback: impl FnOnce(Option<&AccountMapEntry<T>>) -> (bool, R),
+    ) -> R {
+        self.get_bin(pubkey).get_internal(pubkey, callback)
+    }
+
+    /// Gets the index's entry for `pubkey` and clones it
+    ///
+    /// Prefer `get_and_then()` whenever possible.
+    /// NOTE: The entry is *not* added to the in-mem cache.
+    pub fn get_cloned(&self, pubkey: &Pubkey) -> Option<AccountMapEntry<T>> {
+        self.get_and_then(pubkey, |entry| (false, entry.cloned()))
+    }
+
+    /// Is `pubkey` in the index?
+    pub fn contains(&self, pubkey: &Pubkey) -> bool {
+        self.get_and_then(pubkey, |entry| (false, entry.is_some()))
+    }
+
     fn slot_list_mut<RT>(
         &self,
         pubkey: &Pubkey,