From 2f55547d37b7dd5389035119c7bf2c900e318f92 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)"
 <75863576+jeffwashington@users.noreply.github.com>
Date: Fri, 21 May 2021 09:26:46 -0500
Subject: [PATCH] remove flatten in storage scan (#17354)

* remove flatten in storage scan

* fix tests
---
 runtime/src/accounts_db.rs   | 51 +++++++++++-------------------------
 runtime/src/accounts_hash.rs |  9 +------
 2 files changed, 17 insertions(+), 43 deletions(-)

diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index 2aee0280db..8be4ddf9bf 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -4088,35 +4088,31 @@ impl AccountsDb {
     /// Scan through all the account storage in parallel
     fn scan_account_storage_no_bank<F, B>(
         snapshot_storages: &[SnapshotStorage],
-        stats: &mut crate::accounts_hash::HashStats,
         scan_func: F,
     ) -> Vec<B>
     where
         F: Fn(LoadedAccount, &mut B, Slot) + Send + Sync,
         B: Send + Default,
     {
-        let mut time = Measure::start("flatten");
-        let items: Vec<_> = snapshot_storages.iter().flatten().collect();
-        time.stop();
-        stats.pre_scan_flatten_time_total_us += time.as_us();
-
         // Without chunks, we end up with 1 output vec for each outer snapshot storage.
         // This results in too many vectors to be efficient.
         const MAX_ITEMS_PER_CHUNK: usize = 5_000;
-        items
+        snapshot_storages
             .par_chunks(MAX_ITEMS_PER_CHUNK)
-            .map(|storages: &[&Arc<AccountStorageEntry>]| {
+            .map(|storages: &[Vec<Arc<AccountStorageEntry>>]| {
                 let mut retval = B::default();
 
-                for storage in storages {
-                    let accounts = storage.accounts.accounts(0);
-                    accounts.into_iter().for_each(|stored_account| {
-                        scan_func(
-                            LoadedAccount::Stored(stored_account),
-                            &mut retval,
-                            storage.slot(),
-                        )
-                    });
+                for sub_storages in storages {
+                    for storage in sub_storages {
+                        let accounts = storage.accounts.accounts(0);
+                        accounts.into_iter().for_each(|stored_account| {
+                            scan_func(
+                                LoadedAccount::Stored(stored_account),
+                                &mut retval,
+                                storage.slot(),
+                            )
+                        });
+                    }
                 }
                 retval
             })
@@ -4181,7 +4177,6 @@ impl AccountsDb {
         stats.num_snapshot_storage = storage.len();
         let result: Vec<Vec<Vec<CalculateHashIntermediate>>> = Self::scan_account_storage_no_bank(
             &storage,
-            &mut stats,
             |loaded_account: LoadedAccount,
              accum: &mut Vec<Vec<CalculateHashIntermediate>>,
              slot: Slot| {
@@ -5763,6 +5758,7 @@ pub mod tests {
 
     #[test]
     fn test_accountsdb_scan_snapshot_stores() {
+        solana_logger::setup();
         let (mut storages, raw_expected) = sample_storages_and_accounts();
 
         let bins = 1;
@@ -5832,15 +5828,8 @@ pub mod tests {
 
         // enough stores to get to 2nd chunk
         let bins = 1;
-        let (_temp_dirs, paths) = get_temp_accounts_paths(1).unwrap();
-        let slot_expected: Slot = 0;
-        let size: usize = 123;
-        let data = AccountStorageEntry::new(&paths[0], slot_expected, 0, size as u64);
-
-        let arc = Arc::new(data);
-
         const MAX_ITEMS_PER_CHUNK: usize = 5_000;
-        storages[0].splice(0..0, vec![arc; MAX_ITEMS_PER_CHUNK]);
+        storages.splice(0..0, vec![vec![]; MAX_ITEMS_PER_CHUNK]);
 
         let mut stats = HashStats::default();
         let result = AccountsDb::scan_snapshot_stores(
@@ -5935,15 +5924,8 @@ pub mod tests {
         // enough stores to get to 2nd chunk
         // range is for only 1 bin out of 256.
         let bins = 256;
-        let (_temp_dirs, paths) = get_temp_accounts_paths(1).unwrap();
-        let slot_expected: Slot = 0;
-        let size: usize = 123;
-        let data = AccountStorageEntry::new(&paths[0], slot_expected, 0, size as u64);
-
-        let arc = Arc::new(data);
-
         const MAX_ITEMS_PER_CHUNK: usize = 5_000;
-        storages[0].splice(0..0, vec![arc; MAX_ITEMS_PER_CHUNK]);
+        storages.splice(0..0, vec![vec![]; MAX_ITEMS_PER_CHUNK]);
 
         let mut stats = HashStats::default();
         let result = AccountsDb::scan_snapshot_stores(
@@ -6030,7 +6012,6 @@ pub mod tests {
         let calls = AtomicU64::new(0);
         let result = AccountsDb::scan_account_storage_no_bank(
             &storages,
-            &mut HashStats::default(),
             |loaded_account: LoadedAccount, accum: &mut Vec<u64>, slot: Slot| {
                 calls.fetch_add(1, Ordering::Relaxed);
                 assert_eq!(loaded_account.pubkey(), &pubkey);
diff --git a/runtime/src/accounts_hash.rs b/runtime/src/accounts_hash.rs
index c340463785..86d46b39cc 100644
--- a/runtime/src/accounts_hash.rs
+++ b/runtime/src/accounts_hash.rs
@@ -25,7 +25,6 @@ pub struct HashStats {
     pub hash_time_total_us: u64,
     pub sort_time_total_us: u64,
     pub flatten_time_total_us: u64,
-    pub pre_scan_flatten_time_total_us: u64,
     pub hash_total: usize,
     pub unreduced_entries: usize,
     pub num_snapshot_storage: usize,
@@ -36,8 +35,7 @@ impl HashStats {
             + self.zeros_time_total_us
             + self.hash_time_total_us
             + self.sort_time_total_us
-            + self.flatten_time_total_us
-            + self.pre_scan_flatten_time_total_us;
+            + self.flatten_time_total_us;
         datapoint_info!(
             "calculate_accounts_hash_without_index",
             ("accounts_scan", self.scan_time_total_us, i64),
@@ -52,11 +50,6 @@ impl HashStats {
                 self.num_snapshot_storage as i64,
                 i64
             ),
-            (
-                "pre_scan_flatten",
-                self.pre_scan_flatten_time_total_us as i64,
-                i64
-            ),
             ("total", total_time_us as i64, i64),
         );
     }