From 4b314be5bd4ba343702437d88a38a1516bf9c544 Mon Sep 17 00:00:00 2001
From: "Jeff Washington (jwash)"
 <75863576+jeffwashington@users.noreply.github.com>
Date: Thu, 24 Jun 2021 14:52:11 -0500
Subject: [PATCH] rework dirty_pubkeys from
 insert_new_if_missing_into_primary_index (#18200)

---
 runtime/src/accounts_db.rs    | 23 ++++++-----------------
 runtime/src/accounts_index.rs | 16 ++++++++--------
 2 files changed, 14 insertions(+), 25 deletions(-)

diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index 1d96dbb710..49bcc856f5 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -5872,13 +5872,10 @@ impl AccountsDb {
 
     fn generate_index_for_slot<'a>(&self, accounts_map: GenerateIndexAccountsMap<'a>, slot: &Slot) {
         if !accounts_map.is_empty() {
-            let len = accounts_map.len();
-
-            let mut items = Vec::with_capacity(len);
-            let mut dirty_pubkeys = accounts_map
+            let items = accounts_map
                 .iter()
                 .map(|(pubkey, (_, store_id, stored_account))| {
-                    items.push((
+                    (
                         pubkey,
                         AccountInfo {
                             store_id: *store_id,
@@ -5886,25 +5883,17 @@ impl AccountsDb {
                             stored_size: stored_account.stored_size,
                             lamports: stored_account.account_meta.lamports,
                         },
-                    ));
-                    *pubkey
+                    )
                 })
                 .collect::<Vec<_>>();
 
-            let items_len = items.len();
-            let dirty_pubkey_mask = self
+            let dirty_pubkeys = self
                 .accounts_index
                 .insert_new_if_missing_into_primary_index(*slot, items);
 
-            assert_eq!(dirty_pubkey_mask.len(), items_len);
-
-            let mut dirty_pubkey_mask_iter = dirty_pubkey_mask.iter();
-
-            // dirty_pubkey_mask will return true if an item has multiple rooted entries for
+            // dirty_pubkeys will contain a pubkey if an item has multiple rooted entries for
             // a given pubkey. If there is just a single item, there is no cleaning to
-            // be done on that pubkey. Prune the touched pubkey set here for only those
-            // pubkeys with multiple updates.
-            dirty_pubkeys.retain(|_k| *dirty_pubkey_mask_iter.next().unwrap());
+            // be done on that pubkey. Use only those pubkeys with multiple updates.
             if !dirty_pubkeys.is_empty() {
                 self.uncleaned_pubkeys.insert(*slot, dirty_pubkeys);
             }
diff --git a/runtime/src/accounts_index.rs b/runtime/src/accounts_index.rs
index 334e5e570d..ca8df60611 100644
--- a/runtime/src/accounts_index.rs
+++ b/runtime/src/accounts_index.rs
@@ -1357,7 +1357,8 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         &self,
         slot: Slot,
         items: Vec<(&Pubkey, T)>,
-    ) -> Vec<bool> {
+    ) -> Vec<Pubkey> {
+        let item_len = items.len();
         let potentially_new_items = items
             .iter()
             .map(|(_pubkey, account_info)| {
@@ -1367,12 +1368,12 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
             .collect::<Vec<_>>(); // collect here so we have created all data prior to obtaining lock
 
         let mut _reclaims = SlotList::new();
-
+        let mut duplicate_keys = Vec::with_capacity(item_len / 100); // just an estimate
         let mut w_account_maps = self.get_account_maps_write_lock();
         items
             .into_iter()
             .zip(potentially_new_items.into_iter())
-            .map(|((pubkey, account_info), new_item)| {
+            .for_each(|((pubkey, account_info), new_item)| {
                 let account_entry = self.insert_new_entry_if_missing_with_lock(
                     pubkey,
                     &mut w_account_maps,
@@ -1380,12 +1381,11 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
                 );
                 if let Some(mut w_account_entry) = account_entry {
                     w_account_entry.update(slot, account_info, &mut _reclaims);
-                    true
-                } else {
-                    false
+                    duplicate_keys.push(*pubkey);
                 }
-            })
-            .collect()
+            });
+
+        duplicate_keys
     }
 
     // Updates the given pubkey at the given slot with the new account information.