simplify bank drop calls

ledger/src/blockstore_processor.rs

@@ -1292,12 +1292,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
-                        let is_from_abs = true;
-                        new_root_bank.rc.accounts.purge_slot(
+                        // Simulate this purge is from AccountBackgroundService
+                        new_root_bank.rc.accounts.accounts_db.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

@@ -56,7 +56,12 @@ pub struct SendDroppedBankCallback {
 
 impl DropCallback for SendDroppedBankCallback {
     fn callback(&self, bank: &Bank) {
-        if let Err(e) = self.sender.send((bank.slot(), bank.bank_id())) {
+        if bank.bank_id() >> 63 == 1 {
+            panic!("bad drop callpath detected; Bank::drop() must run serially with other logic in ABS like clean_accounts()")
+        }
+
+        let bank_id = bank.bank_id() & 0xefff_ffff_ffff_ffff;
+        if let Err(e) = self.sender.send((bank.slot(), bank_id)) {
             warn!("Error sending dropped banks: {:?}", e);
         }
     }
@@ -341,14 +346,15 @@ impl AbsRequestHandler {
             })
     }
 
-    /// `is_from_abs` is true if the caller is the AccountsBackgroundService
-    pub fn handle_pruned_banks(&self, bank: &Bank, is_from_abs: bool) -> usize {
+    pub fn handle_pruned_banks(&self, bank: &Bank, is_serialized_with_abs: bool) -> usize {
         let mut count = 0;
         for (pruned_slot, pruned_bank_id) in self.pruned_banks_receiver.try_iter() {
             count += 1;
-            bank.rc
-                .accounts
-                .purge_slot(pruned_slot, pruned_bank_id, is_from_abs);
+            bank.rc.accounts.accounts_db.purge_slot(
+                pruned_slot,
+                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,12 @@ 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
-    pub fn purge_slot(&self, slot: Slot, bank_id: BankId, is_from_abs: bool) {
-        if self.is_bank_drop_callback_enabled.load(Ordering::SeqCst) && !is_from_abs {
+    /// * `is_serialized_with_abs` - indicates whehter this call is run sequentially with other calls in account background service
+    pub fn purge_slot(&self, slot: Slot, bank_id: BankId, is_serialized_with_abs: bool) {
+        if self.is_bank_drop_callback_enabled.load(Ordering::SeqCst) && !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
@@ -7583,6 +7585,7 @@ pub mod tests {
         std::{
             iter::FromIterator,
             str::FromStr,
+            sync::atomic::AtomicBool,
             thread::{self, Builder, JoinHandle},
         },
     };
@@ -13600,8 +13603,7 @@ pub mod tests {
             .is_some());
 
         // Simulate purge_slot() all from AccountsBackgroundService
-        let is_from_abs = true;
-        accounts.purge_slot(slot0, 0, is_from_abs);
+        accounts.purge_slot(slot0, 0, true);
 
         // Now clean should clean up the remaining key
         accounts.clean_accounts(None, false, None);

runtime/src/bank.rs

@@ -1382,7 +1382,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(),
@@ -6800,12 +6800,19 @@ pub struct TotalAccountsStats {
 
 impl Drop for Bank {
     fn drop(&mut self) {
+        // Flip the MSB of bank_id to inidcates that we are dropping the bank.
+        // This indicates that there are no other reference to bank, so that
+        // drop_callback can assert it. The downside is that we can only support
+        // bank_id range of 63bits.
+        assert!(self.bank_id & 0x1 << 63 == 0);
+        self.bank_id |= 0x1 << 63;
         if let Some(drop_callback) = self.drop_callback.read().unwrap().0.as_ref() {
             drop_callback.callback(self);
         } else {
             // Default case for tests
             self.rc
                 .accounts
+                .accounts_db
                 .purge_slot(self.slot(), self.bank_id(), false);
         }
     }
@@ -14774,9 +14781,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, is_abs_service);
+                        abs_request_handler.handle_pruned_banks(&current_major_fork_bank, true);
                     }
                 },
                 Some(Box::new(SendDroppedBankCallback::new(