load_accounts_index_for_shrink uses scan (#26260)

runtime/src/accounts_db.rs

@@ -29,9 +29,9 @@ use {
         },
         accounts_index::{
             AccountIndexGetResult, AccountSecondaryIndexes, AccountsIndex, AccountsIndexConfig,
-            AccountsIndexRootsStats, IndexKey, IndexValue, IsCached, RefCount, ScanConfig,
-            ScanResult, SlotList, SlotSlice, ZeroLamport, ACCOUNTS_INDEX_CONFIG_FOR_BENCHMARKS,
-            ACCOUNTS_INDEX_CONFIG_FOR_TESTING,
+            AccountsIndexRootsStats, AccountsIndexScanResult, IndexKey, IndexValue, IsCached,
+            RefCount, ScanConfig, ScanResult, SlotList, SlotSlice, ZeroLamport,
+            ACCOUNTS_INDEX_CONFIG_FOR_BENCHMARKS, ACCOUNTS_INDEX_CONFIG_FOR_TESTING,
         },
         accounts_index_storage::Startup,
         accounts_update_notifier_interface::AccountsUpdateNotifier,
@@ -2592,7 +2592,11 @@ impl AccountsDb {
                                 if !useless {
                                     useful += 1;
                                 }
-                                !useless
+                                if useless {
+                                    AccountsIndexScanResult::None
+                                } else {
+                                    AccountsIndexScanResult::KeepInMemory
+                                }
                             },
                         );
                         found_not_zero_accum.fetch_add(found_not_zero, Ordering::Relaxed);
@@ -3014,15 +3018,18 @@ impl AccountsDb {
 
         let mut alive = 0;
         let mut dead = 0;
-
+        let mut index = 0;
+        self.accounts_index.scan(
             accounts[..std::cmp::min(accounts.len(), count)]
                 .iter()
-            .for_each(|pair| {
-                let pubkey = &pair.0;
+                .map(|(key, _)| key),
+            // return true if we want this item to remain in the cache
+            |exists, slot_list, pubkey, _ref_count| {
+                let mut result = AccountsIndexScanResult::None;
+                if exists {
+                    let pair = &accounts[index];
                     let stored_account = &pair.1;
-                let lookup = self.accounts_index.get_account_read_entry(pubkey);
-                if let Some(locked_entry) = lookup {
-                    let is_alive = locked_entry.slot_list().iter().any(|(_slot, acct_info)| {
+                    let is_alive = slot_list.iter().any(|(_slot, acct_info)| {
                         acct_info.matches_storage_location(
                             stored_account.store_id,
                             stored_account.account.offset,
@@ -3036,7 +3043,7 @@ impl AccountsDb {
                         if let Some(unrefed_pubkeys) = &mut unrefed_pubkeys {
                             unrefed_pubkeys.push(pubkey);
                         }
-                        locked_entry.unref();
+                        result = AccountsIndexScanResult::Unref;
                         dead += 1;
                     } else {
                         alive_accounts.push(pair);
@@ -3044,7 +3051,11 @@ impl AccountsDb {
                         alive += 1;
                     }
                 }
-            });
+                index += 1;
+                result
+            },
+        );
+        assert_eq!(index, std::cmp::min(accounts.len(), count));
         self.shrink_stats
             .alive_accounts
             .fetch_add(alive, Ordering::Relaxed);

runtime/src/accounts_index.rs

@@ -639,6 +639,15 @@ impl ScanSlotTracker {
     }
 }
 
+pub enum AccountsIndexScanResult {
+    /// if the entry is not in the in-memory index, do not add it, make no modifications to it
+    None,
+    /// keep the entry in the in-memory index
+    KeepInMemory,
+    /// reduce refcount by 1
+    Unref,
+}
+
 #[derive(Debug)]
 pub struct AccountsIndex<T: IndexValue> {
     pub account_maps: LockMapType<T>,
@@ -1295,13 +1304,12 @@ impl<T: IndexValue> AccountsIndex<T> {
     ///   call `callback`
     pub(crate) fn scan<'a, F, I>(&'a self, pubkeys: I, mut callback: F)
     where
-        // return true if accounts index entry should be put in in_mem cache
         // params:
         //  exists: false if not in index at all
         //  index in slot list where best slot was found or None if nothing found by root criteria
         //  pubkey looked up
         //  refcount of entry in index
-        F: FnMut(bool, &SlotList<T>, &Pubkey, RefCount) -> bool,
+        F: FnMut(bool, &SlotList<T>, &'a Pubkey, RefCount) -> AccountsIndexScanResult,
         I: IntoIterator<Item = &'a Pubkey>,
     {
         let empty_slot_list = vec![];
@@ -1315,13 +1323,24 @@ impl<T: IndexValue> AccountsIndex<T> {
                 last_bin = bin;
             }
             lock.as_ref().unwrap().get_internal(pubkey, |entry| {
-                let cache = match entry {
+                let mut cache = false;
+                match entry {
                     Some(locked_entry) => {
                         let slot_list = &locked_entry.slot_list.read().unwrap();
-                        callback(true, slot_list, pubkey, locked_entry.ref_count())
+                        let result = callback(true, slot_list, pubkey, locked_entry.ref_count());
+                        cache = match result {
+                            AccountsIndexScanResult::Unref => {
+                                locked_entry.add_un_ref(false);
+                                true
                             }
-                    None => callback(false, &empty_slot_list, pubkey, RefCount::MAX),
+                            AccountsIndexScanResult::KeepInMemory => true,
+                            AccountsIndexScanResult::None => false,
                         };
+                    }
+                    None => {
+                        callback(false, &empty_slot_list, pubkey, RefCount::MAX);
+                    }
+                }
                 (cache, ())
             });
         });