From 0aa4dc904e1b53e74a1eabc2258c2aa5eb1009c6 Mon Sep 17 00:00:00 2001
From: sakridge <stephen@solana.com>
Date: Wed, 11 Dec 2019 11:11:31 -0800
Subject: [PATCH] Accounts index updates, remove pubkey when dead account
 (#7408)

---
 runtime/src/accounts_db.rs    | 34 ++++++++++++++++++++--------------
 runtime/src/accounts_index.rs | 28 +++++++++++++++++++++++++---
 2 files changed, 45 insertions(+), 17 deletions(-)

diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index 5c5a79c6e6..c29f55cf2b 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -603,8 +603,13 @@ impl AccountsDB {
 
         // Recalculate reclaims with new purge set
         let mut reclaims = Vec::new();
+        let mut dead_keys = Vec::new();
         for pubkey in purges.keys() {
-            reclaims.extend(accounts_index.purge(&pubkey));
+            let (new_reclaims, is_empty) = accounts_index.purge(&pubkey);
+            if is_empty {
+                dead_keys.push(*pubkey);
+            }
+            reclaims.extend(new_reclaims);
         }
 
         let last_root = accounts_index.last_root;
@@ -612,6 +617,13 @@ impl AccountsDB {
         drop(accounts_index);
         drop(storage);
 
+        if !dead_keys.is_empty() {
+            let mut accounts_index = self.accounts_index.write().unwrap();
+            for key in &dead_keys {
+                accounts_index.account_maps.remove(key);
+            }
+        }
+
         self.handle_reclaims(&reclaims, last_root);
     }
 
@@ -1916,19 +1928,13 @@ pub mod tests {
         print_accounts("post_purge", &accounts);
 
         // Make sure the index is for pubkey cleared
-        assert_eq!(
-            accounts
-                .accounts_index
-                .read()
-                .unwrap()
-                .account_maps
-                .get(&pubkey)
-                .unwrap()
-                .read()
-                .unwrap()
-                .len(),
-            0
-        );
+        assert!(accounts
+            .accounts_index
+            .read()
+            .unwrap()
+            .account_maps
+            .get(&pubkey)
+            .is_none());
 
         // slot 1 & 2 should not have any stores
         assert_no_stores(&accounts, 1);
diff --git a/runtime/src/accounts_index.rs b/runtime/src/accounts_index.rs
index 0254fe10fc..62ff863777 100644
--- a/runtime/src/accounts_index.rs
+++ b/runtime/src/accounts_index.rs
@@ -41,11 +41,13 @@ impl<T: Clone> AccountsIndex<T> {
         self.get_rooted_entries(&list)
     }
 
-    pub fn purge(&self, pubkey: &Pubkey) -> Vec<(Slot, T)> {
+    // filter any rooted entries and return them along with a bool that indicates
+    // if this account has no more entries.
+    pub fn purge(&self, pubkey: &Pubkey) -> (Vec<(Slot, T)>, bool) {
         let mut list = self.account_maps.get(&pubkey).unwrap().write().unwrap();
         let reclaims = self.get_rooted_entries(&list);
         list.retain(|(slot, _)| !self.is_root(*slot));
-        reclaims
+        (reclaims, list.is_empty())
     }
 
     // find the latest slot and T in a list for a given ancestor
@@ -133,7 +135,6 @@ impl<T: Clone> AccountsIndex<T> {
                     .cloned(),
             );
             slot_vec.retain(|(slot, _)| !Self::can_purge(max_root, *slot));
-
             None
         } else {
             Some(account_info)
@@ -382,4 +383,25 @@ mod tests {
         assert_eq!(num, 1);
         assert!(found_key);
     }
+
+    #[test]
+    fn test_purge() {
+        let key = Keypair::new();
+        let mut index = AccountsIndex::<u64>::default();
+        let mut gc = Vec::new();
+        assert_eq!(Some(12), index.update(1, &key.pubkey(), 12, &mut gc));
+
+        index.insert(1, &key.pubkey(), 12, &mut gc);
+
+        assert_eq!(None, index.update(1, &key.pubkey(), 10, &mut gc));
+
+        let purges = index.purge(&key.pubkey());
+        assert_eq!(purges, (vec![], false));
+        index.add_root(1);
+
+        let purges = index.purge(&key.pubkey());
+        assert_eq!(purges, (vec![(1, 10)], true));
+
+        assert_eq!(None, index.update(1, &key.pubkey(), 9, &mut gc));
+    }
 }