diff --git a/runtime/src/accounts_db.rs b/runtime/src/accounts_db.rs index 407ac792df..0649ed450f 100644 --- a/runtime/src/accounts_db.rs +++ b/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> { + self.scan_account_storage( + slot, + |loaded_account: LoadedAccount| Some(*loaded_account.pubkey()), + |accum: &DashSet, 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>, + ) -> Vec { + match scan_result { + ScanStorageResult::Cached(cached_result) => cached_result, + ScanStorageResult::Stored(stored_result) => { + stored_result.into_iter().collect::>() + } + } + } + + /// Scan a slot for dirty pubkeys + fn scan_slot_for_dirty_pubkeys(&self, slot: Slot) -> Vec { + 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"); diff --git a/runtime/src/bank.rs b/runtime/src/bank.rs index 271c86358f..d46036a455 100644 --- a/runtime/src/bank.rs +++ b/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 ); }