Fix roots_and_ref_count (solana-labs#13827)

Co-authored-by: Carl Lin <[email protected]>
takeuchi5512 · Nov 26, 2020 · 955b99c · 955b99c
1 parent 761499d
commit 955b99c
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 8 deletions.
diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs
@@ -779,7 +779,8 @@ impl AccountsDB {
                         if account_info.lamports == 0 {
                             purges.insert(
                                 *pubkey,
-                                self.accounts_index.roots_and_ref_count(&locked_entry),
+                                self.accounts_index
+                                    .roots_and_ref_count(&locked_entry, max_clean_root),
                             );
                         }
 
@@ -5380,11 +5381,11 @@ pub mod tests {
         accounts_index.add_root(3);
         let mut purges = HashMap::new();
         let (key0_entry, _) = accounts_index.get(&key0, None, None).unwrap();
-        purges.insert(key0, accounts_index.roots_and_ref_count(&key0_entry));
+        purges.insert(key0, accounts_index.roots_and_ref_count(&key0_entry, None));
         let (key1_entry, _) = accounts_index.get(&key1, None, None).unwrap();
-        purges.insert(key1, accounts_index.roots_and_ref_count(&key1_entry));
+        purges.insert(key1, accounts_index.roots_and_ref_count(&key1_entry, None));
         let (key2_entry, _) = accounts_index.get(&key2, None, None).unwrap();
-        purges.insert(key2, accounts_index.roots_and_ref_count(&key2_entry));
+        purges.insert(key2, accounts_index.roots_and_ref_count(&key2_entry, None));
         for (key, (list, ref_count)) in &purges {
             info!(" purge {} ref_count {} =>", key, ref_count);
             for x in list {
@@ -5610,4 +5611,26 @@ pub mod tests {
             assert_eq!(accounts.load_slow(&ancestors, key).unwrap().0, account_ref);
         }
     }
+
+    #[test]
+    fn test_zero_lamport_new_root_not_cleaned() {
+        let db = AccountsDB::new(Vec::new(), &ClusterType::Development);
+        let account_key = Pubkey::new_unique();
+        let zero_lamport_account = Account::new(0, 0, &Account::default().owner);
+
+        // Store zero lamport account into slots 0 and 1, root both slots
+        db.store(0, &[(&account_key, &zero_lamport_account)]);
+        db.store(1, &[(&account_key, &zero_lamport_account)]);
+        db.add_root(0);
+        db.add_root(1);
+
+        // Only clean zero lamport accounts up to slot 0
+        db.clean_accounts(Some(0));
+
+        // Should still be able to find zero lamport account in slot 1
+        assert_eq!(
+            db.load_slow(&HashMap::new(), &account_key),
+            Some((zero_lamport_account, 1))
+        );
+    }
 }
diff --git a/runtime/src/accounts_index.rs b/runtime/src/accounts_index.rs
@@ -464,10 +464,10 @@ impl<T: 'static + Clone> AccountsIndex<T> {
         self.do_unchecked_scan_accounts(ancestors, func, Some(range));
     }
 
-    pub fn get_rooted_entries(&self, slice: SlotSlice<T>) -> SlotList<T> {
+    pub fn get_rooted_entries(&self, slice: SlotSlice<T>, max: Option<Slot>) -> SlotList<T> {
         slice
             .iter()
-            .filter(|(slot, _)| self.is_root(*slot))
+            .filter(|(slot, _)| self.is_root(*slot) && max.map_or(true, |max| *slot <= max))
             .cloned()
             .collect()
     }
@@ -476,9 +476,10 @@ impl<T: 'static + Clone> AccountsIndex<T> {
     pub fn roots_and_ref_count(
         &self,
         locked_account_entry: &ReadAccountMapEntry<T>,
+        max: Option<Slot>,
     ) -> (SlotList<T>, RefCount) {
         (
-            self.get_rooted_entries(&locked_account_entry.slot_list()),
+            self.get_rooted_entries(&locked_account_entry.slot_list(), max),
             locked_account_entry.ref_count().load(Ordering::Relaxed),
         )
     }
@@ -488,7 +489,7 @@ impl<T: 'static + Clone> AccountsIndex<T> {
     pub fn purge(&self, pubkey: &Pubkey) -> (SlotList<T>, bool) {
         let mut write_account_map_entry = self.get_account_write_entry(pubkey).unwrap();
         write_account_map_entry.slot_list_mut(|slot_list| {
-            let reclaims = self.get_rooted_entries(slot_list);
+            let reclaims = self.get_rooted_entries(slot_list, None);
             slot_list.retain(|(slot, _)| !self.is_root(*slot));
             (reclaims, slot_list.is_empty())
         })