Further polish implementation and fix/add tests

core/src/accounts_background_service.rs

@@ -23,7 +23,7 @@ impl AccountsBackgroundService {
     pub fn new(bank_forks: Arc<RwLock<BankForks>>, exit: &Arc<AtomicBool>) -> Self {
         info!("AccountsBackgroundService active");
         let exit = exit.clone();
-        let mut shrinking_budget = 0;
+        let mut consumed_budget = 0;
         let t_background = Builder::new()
             .name("solana-accounts-background".to_string())
             .spawn(move || loop {
@@ -34,17 +34,8 @@ impl AccountsBackgroundService {
 
                 bank.process_dead_slots();
 
-                if shrinking_budget == 0 {
-                    let shrunken_account_count = bank.process_stale_slot();
-                    if shrunken_account_count > 0 {
-                        datapoint_info!(
-                            "stale_slot_shrink",
-                            ("accounts", shrunken_account_count, i64)
-                        );
-                        shrinking_budget += shrunken_account_count;
-                    }
-                }
-                shrinking_budget = shrinking_budget.saturating_sub(SHRUNKEN_ACCOUNT_PER_INTERVAL);
+                consumed_budget = bank
+                    .process_stale_slot_with_budget(consumed_budget, SHRUNKEN_ACCOUNT_PER_INTERVAL);
 
                 sleep(Duration::from_millis(INTERVAL_MS));
             })

runtime/src/accounts_db.rs

@@ -747,10 +747,9 @@ impl AccountsDB {
             let mut accounts_index = self.accounts_index.write().unwrap();
             accounts_index.reset_uncleaned_roots()
         };
-        {
-            let mut candidates = self.shrink_candidate_slots.lock().unwrap();
-            candidates.extend(previous_roots);
-        }
+
+        let mut candidates = self.shrink_candidate_slots.lock().unwrap();
+        candidates.extend(previous_roots);
     }
 
     fn inc_store_counts(
@@ -1111,7 +1110,8 @@ impl AccountsDB {
 
     // Infinitely returns rooted roots in cyclic order
     fn next_shrink_slot(&self) -> Option<Slot> {
-        // hold a lock to keep reset_uncleaned_roots() from updating candidates
+        // hold a lock to keep reset_uncleaned_roots() from updating candidates;
+        // we might update in this fn it if it's empty
         let mut candidates = self.shrink_candidate_slots.lock().unwrap();
         let next = candidates.pop();
 
@@ -4084,6 +4084,10 @@ pub mod tests {
         accounts.add_root(1);
         accounts.add_root(2);
 
+        accounts.reset_uncleaned_roots();
+        let actual_slots = accounts.shrink_candidate_slots.lock().unwrap().clone();
+        assert_eq!(actual_slots, vec![] as Vec<Slot>);
+
         accounts.reset_uncleaned_roots();
         let mut actual_slots = accounts.shrink_candidate_slots.lock().unwrap().clone();
         actual_slots.sort();
@@ -4202,12 +4206,15 @@ pub mod tests {
             pubkey_count,
             accounts.all_account_count_in_append_vec(shrink_slot)
         );
+
+        // Only, try to shrink stale slots.
         accounts.shrink_all_stale_slots();
         assert_eq!(
             pubkey_count,
             accounts.all_account_count_in_append_vec(shrink_slot)
         );
 
+        // Now, do full-shrink.
         accounts.shrink_all_slots();
         assert_eq!(
             pubkey_count_after_shrink,

runtime/src/accounts_index.rs

@@ -227,7 +227,6 @@ impl<'a, T: 'a + Clone> AccountsIndex<T> {
     pub fn add_root(&mut self, slot: Slot) {
         self.roots.insert(slot);
         self.uncleaned_roots.insert(slot);
-        self.previous_uncleaned_roots.insert(slot);
     }
     /// Remove the slot when the storage for the slot is freed
     /// Accounts no longer reference this slot.
@@ -240,6 +239,7 @@ impl<'a, T: 'a + Clone> AccountsIndex<T> {
     pub fn reset_uncleaned_roots(&mut self) -> Vec<Slot> {
         let empty = HashSet::new();
         let new_previous = std::mem::replace(&mut self.uncleaned_roots, empty);
+        dbg!("{:?}", &new_previous);
         std::mem::replace(&mut self.previous_uncleaned_roots, new_previous)
             .into_iter()
             .collect()
@@ -366,10 +366,33 @@ mod tests {
     fn test_clean_and_unclean_slot() {
         let mut index = AccountsIndex::<bool>::default();
         assert_eq!(0, index.uncleaned_roots.len());
+        index.add_root(0);
         index.add_root(1);
-        assert_eq!(1, index.uncleaned_roots.len());
-        index.clean_dead_slot(1);
+        assert_eq!(2, index.uncleaned_roots.len());
+
+        assert_eq!(0, index.previous_uncleaned_roots.len());
+        index.reset_uncleaned_roots();
+        assert_eq!(2, index.roots.len());
         assert_eq!(0, index.uncleaned_roots.len());
+        assert_eq!(2, index.previous_uncleaned_roots.len());
+
+        index.add_root(2);
+        index.add_root(3);
+        assert_eq!(4, index.roots.len());
+        assert_eq!(2, index.uncleaned_roots.len());
+        assert_eq!(2, index.previous_uncleaned_roots.len());
+
+        index.clean_dead_slot(1);
+        assert_eq!(3, index.roots.len());
+        assert_eq!(2, index.uncleaned_roots.len());
+        //eprintln!("{:?}", &index.previous_uncleaned_roots);
+        assert_eq!(1, index.previous_uncleaned_roots.len());
+
+        index.clean_dead_slot(2);
+        assert_eq!(2, index.roots.len());
+        assert_eq!(1, index.uncleaned_roots.len());
+        //eprintln!("{:?}", &index.previous_uncleaned_roots);
+        assert_eq!(1, index.previous_uncleaned_roots.len());
     }
 
     #[test]

runtime/src/bank.rs

@@ -2570,10 +2570,28 @@ impl Bank {
         self.rc.accounts.accounts_db.process_dead_slots();
     }
 
-    pub fn process_stale_slot(&self) -> usize {
+    fn process_stale_slot(&self) -> usize {
         self.rc.accounts.accounts_db.process_stale_slot()
     }
 
+    pub fn process_stale_slot_with_budget(
+        &self,
+        mut consumed_budget: usize,
+        budget_recovery_delta: usize,
+    ) -> usize {
+        if consumed_budget == 0 {
+            let shrunken_account_count = self.process_stale_slot();
+            if shrunken_account_count > 0 {
+                datapoint_info!(
+                    "stale_slot_shrink",
+                    ("accounts", shrunken_account_count, i64)
+                );
+                consumed_budget += shrunken_account_count;
+            }
+        }
+        consumed_budget.saturating_sub(budget_recovery_delta)
+    }
+
     pub fn shrink_all_slots(&self) {
         self.rc.accounts.accounts_db.shrink_all_slots();
     }
@@ -7005,4 +7023,45 @@ mod tests {
         }
         info!("results: {:?}", results);
     }
+
+    #[test]
+    fn test_process_stale_slot_with_budget() {
+        solana_logger::setup();
+
+        let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000_000);
+        let pubkey1 = Pubkey::new_rand();
+        let pubkey2 = Pubkey::new_rand();
+
+        let bank = Arc::new(Bank::new(&genesis_config));
+        bank.lazy_rent_collection.store(true, Ordering::Relaxed);
+        assert_eq!(bank.process_stale_slot_with_budget(0, 0), 0);
+        assert_eq!(bank.process_stale_slot_with_budget(133, 0), 133);
+
+        assert_eq!(bank.process_stale_slot_with_budget(0, 100), 0);
+        assert_eq!(bank.process_stale_slot_with_budget(33, 100), 0);
+        assert_eq!(bank.process_stale_slot_with_budget(133, 100), 33);
+
+        bank.squash();
+
+        let some_lamports = 123;
+        let bank = Arc::new(new_from_parent(&bank));
+        bank.deposit(&pubkey1, some_lamports);
+        bank.deposit(&pubkey2, some_lamports);
+
+        let bank = Arc::new(new_from_parent(&bank));
+        bank.deposit(&pubkey1, some_lamports);
+        bank.squash();
+        bank.clean_accounts();
+        let force_to_return_alive_account = 0;
+        assert_eq!(
+            bank.process_stale_slot_with_budget(22, force_to_return_alive_account),
+            22
+        );
+
+        let mut consumed_budgets = (0..3)
+            .map(|_| bank.process_stale_slot_with_budget(0, force_to_return_alive_account))
+            .collect::<Vec<_>>();
+        consumed_budgets.sort();
+        assert_eq!(consumed_budgets, vec![0, 1, 8]);
+    }
 }