From 5bfd1d461346b2e9742c5ee93276309eb6a90eef Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Tue, 3 Jan 2023 14:10:34 -0600
Subject: [PATCH] add AccountStorage.get_slot_storage_entry (#29480)

---
 runtime/src/account_storage.rs | 11 +++++++++++
 runtime/src/accounts_db.rs     | 34 ++++++++++++++--------------------
 2 files changed, 25 insertions(+), 20 deletions(-)

diff --git a/runtime/src/account_storage.rs b/runtime/src/account_storage.rs
index 33a9ac252f1670..974224ea34aaa3 100644
--- a/runtime/src/account_storage.rs
+++ b/runtime/src/account_storage.rs
@@ -32,6 +32,17 @@ impl AccountStorage {
         self.map.get(&slot).map(|result| result.value().clone())
     }
 
+    /// return the append vec for 'slot' if it exists
+    /// This is only ever called when shrink is not possibly running and there is a max of 1 append vec per slot.
+    pub(crate) fn get_slot_storage_entry(&self, slot: Slot) -> Option<Arc<AccountStorageEntry>> {
+        self.get_slot_stores(slot).and_then(|res| {
+            let read = res.read().unwrap();
+            assert!(read.len() <= 1);
+            read.values().next().cloned()
+        })
+    }
+
+    /// return all append vecs for 'slot' if any exist
     pub(crate) fn get_slot_storage_entries(&self, slot: Slot) -> Option<SnapshotStorage> {
         self.get_slot_stores(slot)
             .map(|res| res.read().unwrap().values().cloned().collect())
diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index 52d985878a1485..cb9e5ae940d36c 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -2960,11 +2960,7 @@ impl AccountsDb {
             if slot > max_slot_inclusive {
                 return;
             }
-            for storage in self
-                .storage
-                .get_slot_storage_entries(slot)
-                .unwrap_or_default()
-            {
+            if let Some(storage) = self.storage.get_slot_storage_entry(slot) {
                 storage.all_accounts().iter().for_each(|account| {
                     let pk = account.pubkey();
                     match pubkey_refcount.entry(*pk) {
@@ -4833,7 +4829,7 @@ impl AccountsDb {
         Ok(used_index)
     }
 
-    /// Scan a specific slot through all the account storage in parallel
+    /// Scan a specific slot through all the account storage
     pub fn scan_account_storage<R, B>(
         &self,
         slot: Slot,
@@ -4875,17 +4871,12 @@ impl AccountsDb {
             // If the slot is not in the cache, then all the account information must have
             // been flushed. This is guaranteed because we only remove the rooted slot from
             // the cache *after* we've finished flushing in `flush_slot_cache`.
-            let storage_maps = self
-                .storage
-                .get_slot_storage_entries(slot)
-                .unwrap_or_default();
-            self.thread_pool.install(|| {
-                storage_maps.par_iter().for_each(|storage| {
-                    storage.accounts.account_iter().for_each(|account| {
-                        storage_scan_func(&retval, LoadedAccount::Stored(account))
-                    })
-                });
-            });
+            if let Some(storage) = self.storage.get_slot_storage_entry(slot) {
+                storage
+                    .accounts
+                    .account_iter()
+                    .for_each(|account| storage_scan_func(&retval, LoadedAccount::Stored(account)));
+            }
 
             ScanStorageResult::Stored(retval)
         }
@@ -8963,10 +8954,13 @@ impl AccountsDb {
                     for (index, slot) in slots.iter().enumerate() {
                         let mut scan_time = Measure::start("scan");
                         log_status.report(index as u64);
-                        let storage_maps = self.storage.get_slot_storage_entries(*slot);
-                        let accounts_map = storage_maps
+                        let storage = self
+                            .storage
+                            .get_slot_storage_entry(*slot)
+                            .map(|storage| vec![storage]);
+                        let accounts_map = storage
                             .as_ref()
-                            .map(|storage_maps| self.process_storage_slot(storage_maps))
+                            .map(|storage| self.process_storage_slot(storage))
                             .unwrap_or_default();
 
                         scan_time.stop();