From edc2ab3e48dde069d6f737b3c739b4491d33182a Mon Sep 17 00:00:00 2001
From: steviez <steven@solana.com>
Date: Thu, 6 May 2021 14:13:44 -0500
Subject: [PATCH] Test account index and store alignment (#17038)

* Use ReclaimResult::Default() instead of building subtypes

* Add test to ensure account_db store and index are aligned
---
 runtime/src/accounts_db.rs | 49 ++++++++++++++++++++++++++++++++++++--
 1 file changed, 47 insertions(+), 2 deletions(-)

diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
index 0649ed450f..20a91ec6ec 100644
--- a/runtime/src/accounts_db.rs
+++ b/runtime/src/accounts_db.rs
@@ -1314,7 +1314,7 @@ impl AccountsDb {
         max_clean_root: Option<Slot>,
     ) -> ReclaimResult {
         if purges.is_empty() {
-            return (HashMap::new(), HashMap::new());
+            return ReclaimResult::default();
         }
         // This number isn't carefully chosen; just guessed randomly such that
         // the hot loop will be the order of ~Xms.
@@ -1345,7 +1345,7 @@ impl AccountsDb {
         // and those stores may be used for background hashing.
         let reset_accounts = false;
 
-        let mut reclaim_result = (HashMap::new(), HashMap::new());
+        let mut reclaim_result = ReclaimResult::default();
         self.handle_reclaims(
             &reclaims,
             None,
@@ -6694,6 +6694,51 @@ pub mod tests {
         assert_eq!(accounts.alive_account_count_in_slot(1), 0);
     }
 
+    #[test]
+    fn test_clean_multiple_zero_lamport_decrements_index_ref_count() {
+        solana_logger::setup();
+
+        let accounts = AccountsDb::new(Vec::new(), &ClusterType::Development);
+        let pubkey1 = solana_sdk::pubkey::new_rand();
+        let pubkey2 = solana_sdk::pubkey::new_rand();
+        let zero_lamport_account =
+            AccountSharedData::new(0, 0, AccountSharedData::default().owner());
+
+        // Store 2 accounts in slot 0, then update account 1 in two more slots
+        accounts.store_uncached(0, &[(&pubkey1, &zero_lamport_account)]);
+        accounts.store_uncached(0, &[(&pubkey2, &zero_lamport_account)]);
+        accounts.store_uncached(1, &[(&pubkey1, &zero_lamport_account)]);
+        accounts.store_uncached(2, &[(&pubkey1, &zero_lamport_account)]);
+        // Root all slots
+        accounts.add_root(0);
+        accounts.add_root(1);
+        accounts.add_root(2);
+
+        // Account ref counts should match how many slots they were stored in
+        // Account 1 = 3 slots; account 2 = 1 slot
+        assert_eq!(accounts.accounts_index.ref_count_from_storage(&pubkey1), 3);
+        assert_eq!(accounts.accounts_index.ref_count_from_storage(&pubkey2), 1);
+
+        accounts.clean_accounts(None);
+        // Slots 0 and 1 should each have been cleaned because all of their
+        // accounts are zero lamports
+        assert!(accounts.storage.get_slot_stores(0).is_none());
+        assert!(accounts.storage.get_slot_stores(1).is_none());
+        // Slot 2 only has a zero lamport account as well. But, calc_delete_dependencies()
+        // should exclude slot 2 from the clean due to changes in other slots
+        assert!(accounts.storage.get_slot_stores(2).is_some());
+        // Index ref counts should be consistent with the slot stores. Account 1 ref count
+        // should be 1 since slot 2 is the only alive slot; account 2 should have a ref
+        // count of 0 due to slot 0 being dead
+        assert_eq!(accounts.accounts_index.ref_count_from_storage(&pubkey1), 1);
+        assert_eq!(accounts.accounts_index.ref_count_from_storage(&pubkey2), 0);
+
+        accounts.clean_accounts(None);
+        // Slot 2 will now be cleaned, which will leave account 1 with a ref count of 0
+        assert!(accounts.storage.get_slot_stores(2).is_none());
+        assert_eq!(accounts.accounts_index.ref_count_from_storage(&pubkey1), 0);
+    }
+
     #[test]
     fn test_clean_zero_lamport_and_old_roots() {
         solana_logger::setup();