From 49402b7d82307bfc55392cc55fd99634749c2baa Mon Sep 17 00:00:00 2001
From: sakridge <sakridge@gmail.com>
Date: Thu, 27 May 2021 16:31:24 +0200
Subject: [PATCH] Only add keys with multiple index entries to the clean set
 (#17473)

---
 runtime/src/accounts_db.rs    | 21 ++++++++++++++++++---
 runtime/src/accounts_index.rs | 10 +++++++---
 runtime/src/bank.rs           |  4 ++++
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index a4593aa132..8db96999da 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -5205,7 +5205,7 @@ impl AccountsDb {
                         let len = accounts_map.len();
 
                         let mut items = Vec::with_capacity(len);
-                        let dirty_keys = accounts_map
+                        let mut dirty_pubkeys = accounts_map
                             .iter()
                             .map(|(pubkey, (_, store_id, stored_account))| {
                                 items.push((
@@ -5220,10 +5220,25 @@ impl AccountsDb {
                                 *pubkey
                             })
                             .collect::<Vec<_>>();
-                        self.uncleaned_pubkeys.insert(*slot, dirty_keys);
 
-                        self.accounts_index
+                        let items_len = items.len();
+                        let dirty_pubkey_mask = 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
+                        // 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());
+                        if !dirty_pubkeys.is_empty() {
+                            self.uncleaned_pubkeys.insert(*slot, dirty_pubkeys);
+                        }
+
                         if !self.account_indexes.is_empty() {
                             for (pubkey, (_, _store_id, stored_account)) in accounts_map.iter() {
                                 self.accounts_index.update_secondary_indexes(
diff --git a/runtime/src/accounts_index.rs b/runtime/src/accounts_index.rs
index 2e7de4f28c..b3920ade71 100644
--- a/runtime/src/accounts_index.rs
+++ b/runtime/src/accounts_index.rs
@@ -1205,7 +1205,7 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         &self,
         slot: Slot,
         items: Vec<(&Pubkey, T)>,
-    ) {
+    ) -> Vec<bool> {
         let potentially_new_items = items
             .iter()
             .map(|(_pubkey, account_info)| {
@@ -1220,7 +1220,7 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         items
             .into_iter()
             .zip(potentially_new_items.into_iter())
-            .for_each(|((pubkey, account_info), new_item)| {
+            .map(|((pubkey, account_info), new_item)| {
                 let account_entry = self.insert_new_entry_if_missing_with_lock(
                     pubkey,
                     &mut w_account_maps,
@@ -1231,8 +1231,12 @@ 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
                 }
-            });
+            })
+            .collect()
     }
 
     // Updates the given pubkey at the given slot with the new account information.
diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs
index 167c1d9308..74fd6d6001 100644
--- a/runtime/src/bank.rs
+++ b/runtime/src/bank.rs
@@ -4624,22 +4624,26 @@ impl Bank {
     /// A snapshot bank should be purged of 0 lamport accounts which are not part of the hash
     /// calculation and could shield other real accounts.
     pub fn verify_snapshot_bank(&self) -> bool {
+        info!("cleaning..");
         let mut clean_time = Measure::start("clean");
         if self.slot() > 0 {
             self.clean_accounts(true, true);
         }
         clean_time.stop();
 
+        info!("shrinking..");
         let mut shrink_all_slots_time = Measure::start("shrink_all_slots");
         if self.slot() > 0 {
             self.shrink_all_slots(true);
         }
         shrink_all_slots_time.stop();
 
+        info!("verify_bank_hash..");
         let mut verify_time = Measure::start("verify_bank_hash");
         let mut verify = self.verify_bank_hash();
         verify_time.stop();
 
+        info!("verify_hash..");
         let mut verify2_time = Measure::start("verify_hash");
         // Order and short-circuiting is significant; verify_hash requires a valid bank hash
         verify = verify && self.verify_hash();