parallelize index_read in shrink (#19506)

runtime/src/accounts_db.rs

@@ -2266,11 +2266,21 @@ impl AccountsDb {
         }
 
         let mut index_read_elapsed = Measure::start("index_read_elapsed");
-        let mut alive_total = 0;
+        let alive_total_collect = AtomicUsize::new(0);
 
-        let mut alive_accounts: Vec<_> = Vec::with_capacity(stored_accounts.len());
+        let len = stored_accounts.len();
-        let mut unrefed_pubkeys = vec![];
+        let alive_accounts_collect = Mutex::new(Vec::with_capacity(len));
-        for (pubkey, stored_account) in &stored_accounts {
+        let unrefed_pubkeys_collect = Mutex::new(Vec::with_capacity(len));
+        self.thread_pool.install(|| {
+            let chunk_size = 50; // # accounts/thread
+            let chunks = len / chunk_size + 1;
+            (0..chunks).into_par_iter().for_each(|chunk| {
+                let mut alive_accounts = Vec::with_capacity(chunk_size);
+                let mut unrefed_pubkeys = Vec::with_capacity(chunk_size);
+                let mut alive_total = 0;
+                let skip = chunk * chunk_size;
+                stored_accounts.iter().skip(skip).take(chunk_size).for_each(
+                    |(pubkey, stored_account)| {
                         let lookup = self.accounts_index.get_account_read_entry(pubkey);
                         if let Some(locked_entry) = lookup {
                             let is_alive = locked_entry.slot_list().iter().any(|(_slot, i)| {
@@ -2295,7 +2305,25 @@ impl AccountsDb {
                                     .fetch_add(1, Ordering::Relaxed);
                             }
                         }
-        }
+                    },
+                );
+
+                // collect
+                alive_accounts_collect
+                    .lock()
+                    .unwrap()
+                    .append(&mut alive_accounts);
+                unrefed_pubkeys_collect
+                    .lock()
+                    .unwrap()
+                    .append(&mut unrefed_pubkeys);
+                alive_total_collect.fetch_add(alive_total, Ordering::Relaxed);
+            });
+        });
+
+        let alive_accounts = alive_accounts_collect.into_inner().unwrap();
+        let unrefed_pubkeys = unrefed_pubkeys_collect.into_inner().unwrap();
+        let alive_total = alive_total_collect.load(Ordering::Relaxed);
 
         index_read_elapsed.stop();
         let aligned_total: u64 = Self::page_align(alive_total as u64);