From 64021989021b833a4ca3f7e91bb714542ab2330e Mon Sep 17 00:00:00 2001
From: Brooks <brooks@solana.com>
Date: Wed, 28 Feb 2024 14:38:19 -0500
Subject: [PATCH] Replaces ReadAccountMapEntry in
 calculate_accounts_hash_from_index() (#35349)

---
 accounts-db/src/accounts_db.rs    | 47 +++++++++++--------------------
 accounts-db/src/accounts_index.rs |  2 +-
 2 files changed, 18 insertions(+), 31 deletions(-)

diff --git a/accounts-db/src/accounts_db.rs b/accounts-db/src/accounts_db.rs
index 5153d8585..3402c42f4 100644
--- a/accounts-db/src/accounts_db.rs
+++ b/accounts-db/src/accounts_db.rs
@@ -41,11 +41,10 @@ use {
             ZeroLamportAccounts,
         },
         accounts_index::{
-            AccountIndexGetResult, AccountMapEntry, AccountSecondaryIndexes, AccountsIndex,
-            AccountsIndexConfig, AccountsIndexRootsStats, AccountsIndexScanResult, DiskIndexValue,
-            IndexKey, IndexValue, IsCached, RefCount, ScanConfig, ScanResult, SlotList,
-            UpsertReclaim, ZeroLamport, ACCOUNTS_INDEX_CONFIG_FOR_BENCHMARKS,
-            ACCOUNTS_INDEX_CONFIG_FOR_TESTING,
+            AccountMapEntry, AccountSecondaryIndexes, AccountsIndex, AccountsIndexConfig,
+            AccountsIndexRootsStats, AccountsIndexScanResult, DiskIndexValue, IndexKey, IndexValue,
+            IsCached, RefCount, ScanConfig, ScanResult, SlotList, UpsertReclaim, ZeroLamport,
+            ACCOUNTS_INDEX_CONFIG_FOR_BENCHMARKS, ACCOUNTS_INDEX_CONFIG_FOR_TESTING,
         },
         accounts_index_storage::Startup,
         accounts_partition::RentPayingAccountsByPartition,
@@ -6864,28 +6863,20 @@ impl AccountsDb {
                     let result: Vec<Hash> = pubkeys
                         .iter()
                         .filter_map(|pubkey| {
-                            if let AccountIndexGetResult::Found(lock, index) =
-                                self.accounts_index.get(pubkey, config.ancestors, Some(max_slot))
-                            {
-                                let (slot, account_info) = &lock.slot_list()[index];
-                                if !account_info.is_zero_lamport() {
-                                    // Because we're keeping the `lock' here, there is no need
-                                    // to use retry_to_get_account_accessor()
-                                    // In other words, flusher/shrinker/cleaner is blocked to
-                                    // cause any Accessor(None) situation.
-                                    // Anyway this race condition concern is currently a moot
-                                    // point because calculate_accounts_hash() should not
-                                    // currently race with clean/shrink because the full hash
-                                    // is synchronous with clean/shrink in
-                                    // AccountsBackgroundService
+                            let index_entry = self.accounts_index.get_cloned(pubkey)?;
+                            self.accounts_index.get_account_info_with_and_then(
+                                &index_entry,
+                                config.ancestors,
+                                Some(max_slot),
+                                |(slot, account_info)| {
+                                    if account_info.is_zero_lamport() { return None; }
                                     self.get_account_accessor(
-                                        *slot,
+                                        slot,
                                         pubkey,
                                         &account_info.storage_location(),
                                     )
                                     .get_loaded_account()
-                                    .and_then(
-                                        |loaded_account| {
+                                    .and_then(|loaded_account| {
                                             let mut loaded_hash = loaded_account.loaded_hash();
                                             let balance = loaded_account.lamports();
                                             let hash_is_missing = loaded_hash == AccountHash(Hash::default());
@@ -6905,14 +6896,10 @@ impl AccountsDb {
 
                                             sum += balance as u128;
                                             Some(loaded_hash.0)
-                                        },
-                                    )
-                                } else {
-                                    None
-                                }
-                            } else {
-                                None
-                            }
+                                    })
+                                },
+                            )
+                            .flatten()
                         })
                         .collect();
                     let mut total = total_lamports.lock().unwrap();
diff --git a/accounts-db/src/accounts_index.rs b/accounts-db/src/accounts_index.rs
index 2f3ba4b58..5221ac434 100644
--- a/accounts-db/src/accounts_index.rs
+++ b/accounts-db/src/accounts_index.rs
@@ -1155,7 +1155,7 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
 
     /// Gets the account info (and slot) in `entry`, with `ancestors` and `max_root`,
     /// and applies `callback` to it
-    fn get_account_info_with_and_then<R>(
+    pub(crate) fn get_account_info_with_and_then<R>(
         &self,
         entry: &AccountMapEntryInner<T>,
         ancestors: Option<&Ancestors>,