Avoid account index entry Arc clone in shrinking (#35010)

* avoid account index entry Arc clone in shrink

* use scan to addref

* update code comments for scan fn

* expect

* warn

* update log message

---------

Co-authored-by: HaoranYi <haoran.yi@solana.com>

accounts-db/src/accounts_db.rs

@@ -3983,16 +3983,39 @@ impl AccountsDb {
             shrink_collect.alive_total_bytes as u64,
             shrink_collect.capacity,
         ) {
+            warn!(
+                "Unexpected shrink for slot {} alive {} capacity {}, \
+                likely caused by a bug for calculating alive bytes.",
+                slot, shrink_collect.alive_total_bytes, shrink_collect.capacity
+            );
+
             self.shrink_stats
                 .skipped_shrink
                 .fetch_add(1, Ordering::Relaxed);
-            for pubkey in shrink_collect.unrefed_pubkeys {
+
-                if let Some(locked_entry) = self.accounts_index.get_account_read_entry(pubkey) {
+            self.accounts_index.scan(
+                shrink_collect.unrefed_pubkeys.into_iter(),
+                |pubkey, _slot_refs, entry| {
                     // pubkeys in `unrefed_pubkeys` were unref'd in `shrink_collect` above under the assumption that we would shrink everything.
                     // Since shrink is not occurring, we need to addref the pubkeys to get the system back to the prior state since the account still exists at this slot.
-                    locked_entry.addref();
+                    if let Some(entry) = entry {
-                }
+                        entry.addref();
-            }
+                    } else {
+                        // We also expect that the accounts index must contain an
+                        // entry for `pubkey`. Log a warning for now. In future,
+                        // we will panic when this happens.
+                        warn!("pubkey {pubkey} in slot {slot} was NOT found in accounts index during shrink");
+                        datapoint_warn!(
+                            "accounts_db-shink_pubkey_missing_from_index",
+                            ("store_slot", slot, i64),
+                            ("pubkey", pubkey.to_string(), String),
+                        )
+                    }
+                    AccountsIndexScanResult::OnlyKeepInMemoryIfDirty
+                },
+                None,
+                true,
+            );
             return;
         }
 

accounts-db/src/accounts_index.rs

@@ -1358,10 +1358,24 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
         self.storage.get_startup_remaining_items_to_flush_estimate()
     }
 
-    /// For each pubkey, find the slot list in the accounts index
+    /// Scan AccountsIndex for a given iterator of Pubkeys.
-    ///   apply 'avoid_callback_result' if specified.
+    ///
-    ///   otherwise, call `callback`
+    /// This fn takes 4 arguments.
-    /// if 'provide_entry_in_callback' is true, populate callback with the Arc of the entry itself.
+    ///  - an iterator of pubkeys to scan
+    ///  - callback fn to run for each pubkey in the accounts index
+    ///  - avoid_callback_result. If it is Some(default), then callback is ignored and
+    ///    default is returned instead.
+    ///  - provide_entry_in_callback. If true, populate the ref of the Arc of the
+    ///    index entry to `callback` fn. Otherwise, provide None.
+    ///
+    /// The `callback` fn must return `AccountsIndexScanResult`, which is
+    /// used to indicates whether the AccountIndex Entry should be added to
+    /// in-memory cache. The `callback` fn takes in 3 arguments:
+    ///   - the first an immutable ref of the pubkey,
+    ///   - the second an option of the SlotList and RefCount
+    ///   - the third an option of the AccountMapEntry, which is only populated
+    ///     when `provide_entry_in_callback` is true. Otherwise, it will be
+    ///     None.
     pub(crate) fn scan<'a, F, I>(
         &self,
         pubkeys: I,
@@ -1369,15 +1383,6 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
         avoid_callback_result: Option<AccountsIndexScanResult>,
         provide_entry_in_callback: bool,
     ) where
-        // params:
-        //  pubkey looked up
-        //  slots_refs is Option<(slot_list, ref_count)>
-        //    None if 'pubkey' is not in accounts index.
-        //   slot_list: comes from accounts index for 'pubkey'
-        //   ref_count: refcount of entry in index
-        //   entry, if 'provide_entry_in_callback' is true
-        // if 'avoid_callback_result' is Some(_), then callback is NOT called
-        //  and _ is returned as if callback were called.
         F: FnMut(
             &'a Pubkey,
             Option<(&SlotList<T>, RefCount)>,