Refactor function (#13294) (#13520)

Co-authored-by: Carl Lin <carl@solana.com>
(cherry picked from commit 2c2432fddc)

Co-authored-by: carllin <wumu727@gmail.com>

runtime/src/accounts_db.rs

@@ -40,6 +40,7 @@ use solana_sdk::{
 };
 use std::convert::TryFrom;
 use std::{
+    boxed::Box,
     collections::{HashMap, HashSet},
     convert::TryInto,
     io::{Error as IOError, Result as IOResult},
@@ -104,6 +105,7 @@ pub(crate) type SlotStores = Arc<RwLock<HashMap<usize, Arc<AccountStorageEntry>>
 type AccountSlots = HashMap<Pubkey, HashSet<Slot>>;
 type AppendVecOffsets = HashMap<AppendVecId, HashSet<usize>>;
 type ReclaimResult = (AccountSlots, AppendVecOffsets);
+type StorageFinder<'a> = Box<dyn Fn(Slot) -> Arc<AccountStorageEntry> + 'a>;
 
 trait Versioned {
     fn version(&self) -> u64;
@@ -397,6 +399,12 @@ struct FrozenAccountInfo {
     pub lamports: u64, // Account balance cannot be lower than this amount
 }
 
+#[derive(Default)]
+pub struct StoreAccountsTiming {
+    store_accounts_elapsed: u64,
+    update_index_elapsed: u64,
+    handle_reclaims_elapsed: u64,
+}
 // This structure handles the load/store of the accounts
 #[derive(Debug)]
 pub struct AccountsDB {
@@ -1060,10 +1068,8 @@ impl AccountsDB {
         let mut dead_storages = vec![];
         let mut find_alive_elapsed = 0;
         let mut create_and_insert_store_elapsed = 0;
-        let mut store_accounts_elapsed = 0;
-        let mut update_index_elapsed = 0;
-        let mut handle_reclaims_elapsed = 0;
         let mut write_storage_elapsed = 0;
+        let mut store_accounts_timing = StoreAccountsTiming::default();
         if aligned_total > 0 {
             let mut start = Measure::start("find_alive_elapsed");
             let mut accounts = Vec::with_capacity(alive_accounts.len());
@@ -1087,26 +1093,13 @@ impl AccountsDB {
             // here, we're writing back alive_accounts. That should be an atomic operation
             // without use of rather wide locks in this whole function, because we're
             // mutating rooted slots; There should be no writers to them.
-            let mut start = Measure::start("store_accounts_elapsed");
+            store_accounts_timing = self.store_accounts_custom(
-            let infos = self.store_accounts_to(
                 slot,
                 &accounts,
                 &hashes,
-                |_| shrunken_store.clone(),
+                Some(Box::new(move |_| shrunken_store.clone())),
-                write_versions.into_iter(),
+                Some(Box::new(write_versions.into_iter())),
             );
-            start.stop();
-            store_accounts_elapsed = start.as_us();
-
-            let mut start = Measure::start("update_index_elapsed");
-            let reclaims = self.update_index(slot, infos, &accounts);
-            start.stop();
-            update_index_elapsed = start.as_us();
-
-            let mut start = Measure::start("handle_reclaims_elapsed");
-            self.handle_reclaims(&reclaims, Some(slot), true, None);
-            start.stop();
-            handle_reclaims_elapsed = start.as_us();
 
             let mut start = Measure::start("write_storage_elapsed");
             if let Some(slot_stores) = self.storage.get_slot_stores(slot) {
@@ -1136,9 +1129,21 @@ impl AccountsDB {
                 create_and_insert_store_elapsed,
                 i64
             ),
-            ("store_accounts_elapsed", store_accounts_elapsed, i64),
+            (
-            ("update_index_elapsed", update_index_elapsed, i64),
+                "store_accounts_elapsed",
-            ("handle_reclaims_elapsed", handle_reclaims_elapsed, i64),
+                store_accounts_timing.store_accounts_elapsed,
+                i64
+            ),
+            (
+                "update_index_elapsed",
+                store_accounts_timing.update_index_elapsed,
+                i64
+            ),
+            (
+                "handle_reclaims_elapsed",
+                store_accounts_timing.handle_reclaims_elapsed,
+                i64
+            ),
             ("write_storage_elapsed", write_storage_elapsed, i64),
             ("rewrite_elapsed", rewrite_elapsed.as_us(), i64),
             (
@@ -1714,29 +1719,6 @@ impl AccountsDB {
             .fetch_add(count as u64, Ordering::Relaxed)
     }
 
-    fn store_accounts(
-        &self,
-        slot: Slot,
-        accounts: &[(&Pubkey, &Account)],
-        hashes: &[Hash],
-    ) -> Vec<AccountInfo> {
-        let mut current_version = self.bulk_assign_write_version(accounts.len());
-        let write_version_producer = std::iter::from_fn(move || {
-            let ret = current_version;
-            current_version += 1;
-            Some(ret)
-        });
-
-        let storage_finder = |slot| self.find_storage_candidate(slot);
-        self.store_accounts_to(
-            slot,
-            accounts,
-            hashes,
-            storage_finder,
-            write_version_producer,
-        )
-    }
-
     fn store_accounts_to<F: FnMut(Slot) -> Arc<AccountStorageEntry>, P: Iterator<Item = u64>>(
         &self,
         slot: Slot,
@@ -2323,12 +2305,58 @@ impl AccountsDB {
                 .cluster_type
                 .expect("Cluster type must be set at initialization"),
         );
-        self.store_with_hashes(slot, accounts, &hashes);
+        self.store_accounts_default(slot, accounts, &hashes);
     }
 
-    fn store_with_hashes(&self, slot: Slot, accounts: &[(&Pubkey, &Account)], hashes: &[Hash]) {
+    fn store_accounts_default<'a>(
-        let infos = self.store_accounts(slot, accounts, hashes);
+        &'a self,
+        slot: Slot,
+        accounts: &[(&Pubkey, &Account)],
+        hashes: &[Hash],
+    ) {
+        self.store_accounts_custom(
+            slot,
+            accounts,
+            hashes,
+            None::<StorageFinder>,
+            None::<Box<dyn Iterator<Item = u64>>>,
+        );
+    }
+
+    fn store_accounts_custom<'a>(
+        &'a self,
+        slot: Slot,
+        accounts: &[(&Pubkey, &Account)],
+        hashes: &[Hash],
+        storage_finder: Option<StorageFinder<'a>>,
+        write_version_producer: Option<Box<dyn Iterator<Item = u64>>>,
+    ) -> StoreAccountsTiming {
+        let storage_finder: StorageFinder<'a> = storage_finder
+            .unwrap_or_else(|| Box::new(move |slot| self.find_storage_candidate(slot)));
+
+        let write_version_producer: Box<dyn Iterator<Item = u64>> = write_version_producer
+            .unwrap_or_else(|| {
+                let mut current_version = self.bulk_assign_write_version(accounts.len());
+                Box::new(std::iter::from_fn(move || {
+                    let ret = current_version;
+                    current_version += 1;
+                    Some(ret)
+                }))
+            });
+
+        let mut store_accounts_time = Measure::start("store_accounts");
+        let infos = self.store_accounts_to(
+            slot,
+            accounts,
+            hashes,
+            storage_finder,
+            write_version_producer,
+        );
+        store_accounts_time.stop();
+
+        let mut update_index_time = Measure::start("update_index");
         let reclaims = self.update_index(slot, infos, accounts);
+        update_index_time.stop();
 
         // A store for a single slot should:
         // 1) Only make "reclaims" for the same slot
@@ -2338,7 +2366,15 @@ impl AccountsDB {
         //    b)From 1) we know no other slots are included in the "reclaims"
         //
         // From 1) and 2) we guarantee passing Some(slot), true is safe
+        let mut handle_reclaims_time = Measure::start("handle_reclaims");
         self.handle_reclaims(&reclaims, Some(slot), true, None);
+        handle_reclaims_time.stop();
+
+        StoreAccountsTiming {
+            store_accounts_elapsed: store_accounts_time.as_us(),
+            update_index_elapsed: update_index_time.as_us(),
+            handle_reclaims_elapsed: handle_reclaims_time.as_us(),
+        }
     }
 
     pub fn add_root(&self, slot: Slot) {
@@ -4320,7 +4356,7 @@ pub mod tests {
         db.hash_accounts(some_slot, accounts, &ClusterType::Development);
         // provide bogus account hashes
         let some_hash = Hash::new(&[0xca; HASH_BYTES]);
-        db.store_with_hashes(some_slot, accounts, &[some_hash]);
+        db.store_accounts_default(some_slot, accounts, &[some_hash]);
         db.add_root(some_slot);
         assert_matches!(
             db.verify_bank_hash_and_lamports(some_slot, &ancestors, 1),