Clean unrooted unfrozen banks (#16580) (#17003)

Cleanup pubkeys when an unrooted, unfrozen bank is dropped.  This is a
continuation of PR #16911.

runtime/src/accounts_db.rs

@@ -4272,6 +4272,49 @@ impl AccountsDb {
         }
     }
 
+    /// Perform the scan for pubkeys that were written to in a slot
+    fn do_scan_slot_for_dirty_pubkeys(
+        &self,
+        slot: Slot,
+    ) -> ScanStorageResult<Pubkey, DashSet<Pubkey>> {
+        self.scan_account_storage(
+            slot,
+            |loaded_account: LoadedAccount| Some(*loaded_account.pubkey()),
+            |accum: &DashSet<Pubkey>, loaded_account: LoadedAccount| {
+                accum.insert(*loaded_account.pubkey());
+            },
+        )
+    }
+
+    /// Reduce the scan result of dirty pubkeys after calling `scan_account_storage()` into a
+    /// single vec of Pubkeys.
+    fn do_reduce_scan_slot_for_dirty_pubkeys(
+        scan_result: ScanStorageResult<Pubkey, DashSet<Pubkey>>,
+    ) -> Vec<Pubkey> {
+        match scan_result {
+            ScanStorageResult::Cached(cached_result) => cached_result,
+            ScanStorageResult::Stored(stored_result) => {
+                stored_result.into_iter().collect::<Vec<_>>()
+            }
+        }
+    }
+
+    /// Scan a slot for dirty pubkeys
+    fn scan_slot_for_dirty_pubkeys(&self, slot: Slot) -> Vec<Pubkey> {
+        let dirty_pubkeys = self.do_scan_slot_for_dirty_pubkeys(slot);
+        Self::do_reduce_scan_slot_for_dirty_pubkeys(dirty_pubkeys)
+    }
+
+    /// Scan a slot in the account storage for dirty pubkeys and insert them into the list of
+    /// uncleaned pubkeys
+    ///
+    /// This function is called in Bank::drop() when the bank is _not_ frozen, so that its pubkeys
+    /// are considered for cleanup.
+    pub fn scan_slot_and_insert_dirty_pubkeys_into_uncleaned_pubkeys(&self, slot: Slot) {
+        let dirty_pubkeys = self.scan_slot_for_dirty_pubkeys(slot);
+        self.uncleaned_pubkeys.insert(slot, dirty_pubkeys);
+    }
+
     pub fn get_accounts_delta_hash(&self, slot: Slot) -> Hash {
         let mut scan = Measure::start("scan");
 

runtime/src/bank.rs

@@ -5075,20 +5075,42 @@ impl Bank {
                 .is_active(&feature_set::consistent_recent_blockhashes_sysvar::id()),
         }
     }
+
+    /// Bank cleanup
+    ///
+    /// If the bank is unfrozen and then dropped, additional cleanup is needed.  In particular,
+    /// cleaning up the pubkeys that are only in this bank.  To do that, call into AccountsDb to
+    /// scan for dirty pubkeys and add them to the uncleaned pubkeys list so they will be cleaned
+    /// up in AccountsDb::clean_accounts().
+    fn cleanup(&self) {
+        if self.is_frozen() {
+            // nothing to do here
+            return;
+        }
+
+        self.rc
+            .accounts
+            .accounts_db
+            .scan_slot_and_insert_dirty_pubkeys_into_uncleaned_pubkeys(self.slot);
+    }
 }
 
 impl Drop for Bank {
     fn drop(&mut self) {
-        if !self.skip_drop.load(Relaxed) {
-            if let Some(drop_callback) = self.drop_callback.read().unwrap().0.as_ref() {
-                drop_callback.callback(self);
-            } else {
-                // Default case
-                // 1. Tests
-                // 2. At startup when replaying blockstore and there's no
-                // AccountsBackgroundService to perform cleanups yet.
-                self.rc.accounts.purge_slot(self.slot());
-            }
+        if self.skip_drop.load(Relaxed) {
+            return;
+        }
+
+        self.cleanup();
+
+        if let Some(drop_callback) = self.drop_callback.read().unwrap().0.as_ref() {
+            drop_callback.callback(self);
+        } else {
+            // Default case
+            // 1. Tests
+            // 2. At startup when replaying blockstore and there's no
+            // AccountsBackgroundService to perform cleanups yet.
+            self.rc.accounts.purge_slot(self.slot());
         }
     }
 }
@@ -12537,8 +12559,25 @@ pub(crate) mod tests {
     }
 
     #[test]
-    fn test_clean_unrooted_dropped_banks() {
-        //! Test that unrooted banks are cleaned up properly
+    fn test_clean_dropped_unrooted_frozen_banks() {
+        solana_logger::setup();
+        do_test_clean_dropped_unrooted_banks(FreezeBank1::Yes);
+    }
+
+    #[test]
+    fn test_clean_dropped_unrooted_unfrozen_banks() {
+        solana_logger::setup();
+        do_test_clean_dropped_unrooted_banks(FreezeBank1::No);
+    }
+
+    /// A simple enum to toggle freezing Bank1 or not.  Used in the clean_dropped_unrooted tests.
+    enum FreezeBank1 {
+        No,
+        Yes,
+    }
+
+    fn do_test_clean_dropped_unrooted_banks(freeze_bank1: FreezeBank1) {
+        //! Test that dropped unrooted banks are cleaned up properly
         //!
         //! slot 0:       bank0 (rooted)
         //!               /   \
@@ -12548,40 +12587,102 @@ pub(crate) mod tests {
         //!
         //! In the scenario above, when `clean_accounts()` is called on bank2, the keys that exist
         //! _only_ in bank1 should be cleaned up, since those keys are unreachable.
-        //
-        solana_logger::setup();
+        //!
+        //! The following scenarios are tested:
+        //!
+        //! 1. A key is written _only_ in an unrooted bank (key1)
+        //!     - In this case, key1 should be cleaned up
+        //! 2. A key is written in both an unrooted _and_ rooted bank (key3)
+        //!     - In this case, key3's ref-count should be decremented correctly
+        //! 3. A key with zero lamports is _only_ in an unrooted bank (key4)
+        //!     - In this case, key4 should be cleaned up
+        //! 4. A key with zero lamports is in both an unrooted _and_ rooted bank (key5)
+        //!     - In this case, key5's ref-count should be decremented correctly
 
         let (genesis_config, mint_keypair) = create_genesis_config(100);
         let bank0 = Arc::new(Bank::new(&genesis_config));
 
         let collector = Pubkey::new_unique();
-        let pubkey1 = Pubkey::new_unique();
-        let pubkey2 = Pubkey::new_unique();
+        let owner = Pubkey::new_unique();
 
-        bank0.transfer(2, &mint_keypair, &pubkey2).unwrap();
+        let key1 = Keypair::new(); // only touched in bank1
+        let key2 = Keypair::new(); // only touched in bank2
+        let key3 = Keypair::new(); // touched in both bank1 and bank2
+        let key4 = Keypair::new(); // in only bank1, and has zero lamports
+        let key5 = Keypair::new(); // in both bank1 and bank2, and has zero lamports
+
+        bank0.transfer(2, &mint_keypair, &key2.pubkey()).unwrap();
         bank0.freeze();
 
         let slot = 1;
         let bank1 = Bank::new_from_parent(&bank0, &collector, slot);
-        bank1.transfer(3, &mint_keypair, &pubkey1).unwrap();
-        bank1.freeze();
+        bank1.transfer(3, &mint_keypair, &key1.pubkey()).unwrap();
+        bank1.store_account(&key4.pubkey(), &AccountSharedData::new(0, 0, &owner));
+        bank1.store_account(&key5.pubkey(), &AccountSharedData::new(0, 0, &owner));
+
+        if let FreezeBank1::Yes = freeze_bank1 {
+            bank1.freeze();
+        }
 
         let slot = slot + 1;
         let bank2 = Bank::new_from_parent(&bank0, &collector, slot);
-        bank2.transfer(4, &mint_keypair, &pubkey2).unwrap();
+        bank2.transfer(4, &mint_keypair, &key2.pubkey()).unwrap();
+        bank2.transfer(6, &mint_keypair, &key3.pubkey()).unwrap();
+        bank2.store_account(&key5.pubkey(), &AccountSharedData::new(0, 0, &owner));
+
         bank2.freeze(); // the freeze here is not strictly necessary, but more for illustration
         bank2.squash();
 
         drop(bank1);
-
         bank2.clean_accounts(false);
+
+        let expected_ref_count_for_cleaned_up_keys = 0;
+        let expected_ref_count_for_keys_only_in_slot_2 = bank2
+            .rc
+            .accounts
+            .accounts_db
+            .accounts_index
+            .ref_count_from_storage(&key2.pubkey());
+
         assert_eq!(
             bank2
                 .rc
                 .accounts
                 .accounts_db
                 .accounts_index
-                .ref_count_from_storage(&pubkey1),
+                .ref_count_from_storage(&key1.pubkey()),
+            expected_ref_count_for_cleaned_up_keys
+        );
+        assert_ne!(
+            bank2
+                .rc
+                .accounts
+                .accounts_db
+                .accounts_index
+                .ref_count_from_storage(&key3.pubkey()),
+            expected_ref_count_for_cleaned_up_keys
+        );
+        assert_eq!(
+            bank2
+                .rc
+                .accounts
+                .accounts_db
+                .accounts_index
+                .ref_count_from_storage(&key4.pubkey()),
+            expected_ref_count_for_cleaned_up_keys
+        );
+        assert_eq!(
+            bank2
+                .rc
+                .accounts
+                .accounts_db
+                .accounts_index
+                .ref_count_from_storage(&key5.pubkey()),
+            expected_ref_count_for_keys_only_in_slot_2
+        );
+
+        assert_eq!(
+            bank2.rc.accounts.accounts_db.alive_account_count_in_slot(1),
             0
         );
     }