simplify bank drop calls (#24142)

* simplify bank drop calls

* clippy: import

* Update ledger/src/blockstore_processor.rs

Co-authored-by: Brooks Prumo <brooks@prumo.org>

* Update runtime/src/accounts_background_service.rs

Co-authored-by: Brooks Prumo <brooks@prumo.org>

* Update runtime/src/bank.rs

Co-authored-by: Brooks Prumo <brooks@prumo.org>

* cleanup

* format

* use msb of bank_id to indicates that we are dropping

* clippy

* restore bank id

* clippy

* revert is-serialized_with_abs flag

* assert

* clippy

* whitespace

* fix bank drop callback check

* more fix

* remove msb dropping implementation

* fix

Co-authored-by: Brooks Prumo <brooks@prumo.org>

ledger/src/blockstore_processor.rs

@@ -1296,12 +1296,11 @@ fn load_frozen_forks(
                 if last_free.elapsed() > Duration::from_secs(10) {
                     // Purge account state for all dropped banks
                     for (pruned_slot, pruned_bank_id) in pruned_banks_receiver.try_iter() {
-                        // Simulate this purge being from the AccountsBackgroundService
+                        // Simulate this purge is from AccountBackgroundService
-                        let is_from_abs = true;
+                        new_root_bank.rc.accounts.accounts_db.purge_slot(
-                        new_root_bank.rc.accounts.purge_slot(
                             pruned_slot,
                             pruned_bank_id,
-                            is_from_abs,
+                            true,
                         );
                     }
 

runtime/src/accounts.rs

@@ -1156,13 +1156,6 @@ impl Accounts {
         self.accounts_db.store_cached(slot, &accounts_to_store);
     }
 
-    /// Purge a slot if it is not a root
-    /// Root slots cannot be purged
-    /// `is_from_abs` is true if the caller is the AccountsBackgroundService
-    pub fn purge_slot(&self, slot: Slot, bank_id: BankId, is_from_abs: bool) {
-        self.accounts_db.purge_slot(slot, bank_id, is_from_abs);
-    }
-
     /// Add a slot to root.  Root slots cannot be purged
     pub fn add_root(&self, slot: Slot) -> AccountsAddRootTiming {
         self.accounts_db.add_root(slot)

runtime/src/accounts_background_service.rs

@@ -341,14 +341,15 @@ impl AbsRequestHandler {
             })
     }
 
-    /// `is_from_abs` is true if the caller is the AccountsBackgroundService
+    pub fn handle_pruned_banks(&self, bank: &Bank, is_serialized_with_abs: bool) -> usize {
-    pub fn handle_pruned_banks(&self, bank: &Bank, is_from_abs: bool) -> usize {
         let mut count = 0;
         for (pruned_slot, pruned_bank_id) in self.pruned_banks_receiver.try_iter() {
             count += 1;
-            bank.rc
+            bank.rc.accounts.accounts_db.purge_slot(
-                .accounts
+                pruned_slot,
-                .purge_slot(pruned_slot, pruned_bank_id, is_from_abs);
+                pruned_bank_id,
+                is_serialized_with_abs,
+            );
         }
 
         count

runtime/src/accounts_db.rs

@@ -1090,6 +1090,7 @@ pub struct AccountsDb {
     #[cfg(test)]
     load_limit: AtomicU64,
 
+    /// true if drop_callback is attached to the bank.
     is_bank_drop_callback_enabled: AtomicBool,
 
     /// Set of slots currently being flushed by `flush_slot_cache()` or removed
@@ -4043,11 +4044,17 @@ impl AccountsDb {
 
     /// This should only be called after the `Bank::drop()` runs in bank.rs, See BANK_DROP_SAFETY
     /// comment below for more explanation.
-    /// `is_from_abs` is true if the caller is the AccountsBackgroundService
+    ///   * `is_serialized_with_abs` - indicates whehter this call runs sequentially with all other
-    pub fn purge_slot(&self, slot: Slot, bank_id: BankId, is_from_abs: bool) {
+    ///        accounts_db relevant calls, such as shrinking, purging etc., in account background
-        if self.is_bank_drop_callback_enabled.load(Ordering::Acquire) && !is_from_abs {
+    ///        service.
-            panic!("bad drop callpath detected; Bank::drop() must run serially with other logic in ABS like clean_accounts()")
+    pub fn purge_slot(&self, slot: Slot, bank_id: BankId, is_serialized_with_abs: bool) {
+        if self.is_bank_drop_callback_enabled.load(Ordering::Acquire) && !is_serialized_with_abs {
+            panic!(
+                "bad drop callpath detected; Bank::drop() must run serially with other logic in
+                ABS like clean_accounts()"
+            )
         }
+
         // BANK_DROP_SAFETY: Because this function only runs once the bank is dropped,
         // we know that there are no longer any ongoing scans on this bank, because scans require
         // and hold a reference to the bank at the tip of the fork they're scanning. Hence it's
@@ -7582,6 +7589,7 @@ pub mod tests {
         std::{
             iter::FromIterator,
             str::FromStr,
+            sync::atomic::AtomicBool,
             thread::{self, Builder, JoinHandle},
         },
     };
@@ -13599,8 +13607,7 @@ pub mod tests {
             .is_some());
 
         // Simulate purge_slot() all from AccountsBackgroundService
-        let is_from_abs = true;
+        accounts.purge_slot(slot0, 0, true);
-        accounts.purge_slot(slot0, 0, is_from_abs);
 
         // Now clean should clean up the remaining key
         accounts.clean_accounts(None, false, None);

runtime/src/bank.rs

@@ -1233,6 +1233,7 @@ pub struct Bank {
 
     pub feature_set: Arc<FeatureSet>,
 
+    /// callback function only to be called when dropping and should only be called once
     pub drop_callback: RwLock<OptionalDropCallback>,
 
     pub freeze_started: AtomicBool,
@@ -1384,7 +1385,7 @@ impl Bank {
             ),
             transaction_log_collector: Arc::<RwLock<TransactionLogCollector>>::default(),
             feature_set: Arc::<FeatureSet>::default(),
-            drop_callback: RwLock::<OptionalDropCallback>::default(),
+            drop_callback: RwLock::new(OptionalDropCallback(None)),
             freeze_started: AtomicBool::default(),
             vote_only_bank: false,
             cost_tracker: RwLock::<CostTracker>::default(),
@@ -6807,6 +6808,7 @@ impl Drop for Bank {
             // Default case for tests
             self.rc
                 .accounts
+                .accounts_db
                 .purge_slot(self.slot(), self.bank_id(), false);
         }
     }
@@ -14777,9 +14779,7 @@ pub(crate) mod tests {
                         current_major_fork_bank.clean_accounts(false, false, None);
                         // Move purge here so that Bank::drop()->purge_slots() doesn't race
                         // with clean. Simulates the call from AccountsBackgroundService
-                        let is_abs_service = true;
+                        abs_request_handler.handle_pruned_banks(&current_major_fork_bank, true);
-                        abs_request_handler
-                            .handle_pruned_banks(&current_major_fork_bank, is_abs_service);
                     }
                 },
                 Some(Box::new(SendDroppedBankCallback::new(