From 35bd2df0a686a6bacc8944aab08313911526b977 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)" <jeff.washington@solana.com>
Date: Mon, 28 Aug 2023 10:20:29 -0700
Subject: [PATCH] speedup populate_and_retrieve_duplicate_keys_from_startup
 (#33013)

---
 accounts-db/src/accounts_db.rs    | 29 ++++++++++++++---------------
 accounts-db/src/accounts_index.rs |  9 +++++----
 2 files changed, 19 insertions(+), 19 deletions(-)

diff --git a/accounts-db/src/accounts_db.rs b/accounts-db/src/accounts_db.rs
index dda861499..56f7199da 100644
--- a/accounts-db/src/accounts_db.rs
+++ b/accounts-db/src/accounts_db.rs
@@ -9261,7 +9261,7 @@ impl AccountsDb {
                 .sum();
 
             let mut index_flush_us = 0;
-            let mut total_duplicate_slot_keys = 0;
+            let total_duplicate_slot_keys = AtomicU64::default();
             let mut populate_duplicate_keys_us = 0;
             if pass == 0 {
                 // tell accounts index we are done adding the initial accounts at startup
@@ -9273,20 +9273,19 @@ impl AccountsDb {
                 populate_duplicate_keys_us = measure_us!({
                     // this has to happen before visit_duplicate_pubkeys_during_startup below
                     // get duplicate keys from acct idx. We have to wait until we've finished flushing.
-                    for (slot, key) in self
-                        .accounts_index
-                        .populate_and_retrieve_duplicate_keys_from_startup()
-                        .into_iter()
-                        .flatten()
-                    {
-                        total_duplicate_slot_keys += 1;
-                        match self.uncleaned_pubkeys.entry(slot) {
-                            Occupied(mut occupied) => occupied.get_mut().push(key),
-                            Vacant(vacant) => {
-                                vacant.insert(vec![key]);
+                    self.accounts_index
+                        .populate_and_retrieve_duplicate_keys_from_startup(|slot_keys| {
+                            total_duplicate_slot_keys
+                                .fetch_add(slot_keys.len() as u64, Ordering::Relaxed);
+                            for (slot, key) in slot_keys {
+                                match self.uncleaned_pubkeys.entry(slot) {
+                                    Occupied(mut occupied) => occupied.get_mut().push(key),
+                                    Vacant(vacant) => {
+                                        vacant.insert(vec![key]);
+                                    }
+                                }
                             }
-                        }
-                    }
+                        });
                 })
                 .1;
             }
@@ -9302,7 +9301,7 @@ impl AccountsDb {
                 total_items,
                 rent_paying,
                 amount_to_top_off_rent,
-                total_duplicate_slot_keys,
+                total_duplicate_slot_keys: total_duplicate_slot_keys.load(Ordering::Relaxed),
                 populate_duplicate_keys_us,
                 total_including_duplicates: total_including_duplicates.load(Ordering::Relaxed),
                 storage_size_accounts_map_us: storage_info_timings.storage_size_accounts_map_us,
diff --git a/accounts-db/src/accounts_index.rs b/accounts-db/src/accounts_index.rs
index 76390761b..4a02ce5fa 100644
--- a/accounts-db/src/accounts_index.rs
+++ b/accounts-db/src/accounts_index.rs
@@ -1679,17 +1679,18 @@ impl<T: IndexValue, U: DiskIndexValue + From<T> + Into<T>> AccountsIndex<T, U> {
         (dirty_pubkeys, insertion_time.load(Ordering::Relaxed))
     }
 
-    /// return Vec<Vec<>> because the internal vecs are already allocated per bin
+    /// use Vec<> because the internal vecs are already allocated per bin
     pub(crate) fn populate_and_retrieve_duplicate_keys_from_startup(
         &self,
-    ) -> Vec<Vec<(Slot, Pubkey)>> {
+        f: impl Fn(Vec<(Slot, Pubkey)>) + Sync + Send,
+    ) {
         (0..self.bins())
             .into_par_iter()
             .map(|pubkey_bin| {
                 let r_account_maps = &self.account_maps[pubkey_bin];
                 r_account_maps.populate_and_retrieve_duplicate_keys_from_startup()
             })
-            .collect()
+            .for_each(f);
     }
 
     /// Updates the given pubkey at the given slot with the new account information.
@@ -2497,7 +2498,7 @@ pub mod tests {
             index.set_startup(Startup::Normal);
         }
         assert!(gc.is_empty());
-        index.populate_and_retrieve_duplicate_keys_from_startup();
+        index.populate_and_retrieve_duplicate_keys_from_startup(|_slot_keys| {});
 
         for lock in &[false, true] {
             let read_lock = if *lock {