unref accounts from destroyed forks

runtime/src/accounts_db.rs

@@ -300,7 +300,7 @@ impl AccountStorage {
             .and_then(|storage_map| storage_map.read().unwrap().get(&store_id).cloned())
     }
 
-    fn get_slot_stores(&self, slot: Slot) -> Option<SlotStores> {
+    pub fn get_slot_stores(&self, slot: Slot) -> Option<SlotStores> {
         self.0.get(&slot).map(|result| result.value().clone())
     }
 
@@ -376,6 +376,8 @@ pub struct AccountStorageEntry {
     approx_store_count: AtomicUsize,
 
     alive_bytes: AtomicUsize,
+
+    pub(crate) unref_done: AtomicBool,
 }
 
 impl AccountStorageEntry {
@@ -391,6 +393,7 @@ impl AccountStorageEntry {
             count_and_status: RwLock::new((0, AccountStorageStatus::Available)),
             approx_store_count: AtomicUsize::new(0),
             alive_bytes: AtomicUsize::new(0),
+            unref_done: AtomicBool::new(false),
         }
     }
 
@@ -407,6 +410,7 @@ impl AccountStorageEntry {
             count_and_status: RwLock::new((0, AccountStorageStatus::Available)),
             approx_store_count: AtomicUsize::new(num_accounts),
             alive_bytes: AtomicUsize::new(0),
+            unref_done: AtomicBool::new(false),
         }
     }
 
@@ -439,6 +443,7 @@ impl AccountStorageEntry {
         self.id.store(id, Ordering::Relaxed);
         self.approx_store_count.store(0, Ordering::Relaxed);
         self.alive_bytes.store(0, Ordering::Relaxed);
+        self.unref_done.store(false, Ordering::Relaxed);
     }
 
     pub fn status(&self) -> AccountStorageStatus {
@@ -2619,6 +2624,16 @@ impl AccountsDb {
         }
         remove_storages_elapsed.stop();
 
+        for s in &all_removed_slot_storages {
+            for (_store_id, store) in s.read().unwrap().iter() {
+                if !store.unref_done.swap(true, Ordering::Relaxed) {
+                    for a in store.accounts.accounts(0) {
+                        self.accounts_index.unref_from_storage(&a.meta.pubkey);
+                    }
+                }
+            }
+        }
+
         let num_stored_slots_removed = all_removed_slot_storages.len();
 
         let recycle_stores_write_elapsed =
@@ -3071,6 +3086,10 @@ impl AccountsDb {
             .fetch_add(recycle_stores_write_elapsed.as_us(), Ordering::Relaxed);
     }
 
+    pub fn flush_accounts_cache_slot(&self, slot: Slot) {
+        self.flush_slot_cache(slot, None::<&mut fn(&_, &_) -> bool>);
+    }
+
     // `force_flush` flushes all the cached roots `<= requested_flush_root`. It also then
     // flushes:
     // 1) Any remaining roots if there are > MAX_CACHE_SLOTS remaining slots in the cache,
@@ -4032,6 +4051,9 @@ impl AccountsDb {
         for slot in dead_slots.iter() {
             if let Some(slot_storage) = self.storage.get_slot_stores(*slot) {
                 for store in slot_storage.read().unwrap().values() {
+                    let already_unrefd = store.unref_done.swap(true, Ordering::Relaxed);
+                    assert!(!already_unrefd);
+
                     stores.push(store.clone());
                 }
             }
@@ -5848,7 +5870,7 @@ pub mod tests {
             }
         }
 
-        fn ref_count_for_pubkey(&self, pubkey: &Pubkey) -> RefCount {
+        pub fn ref_count_for_pubkey(&self, pubkey: &Pubkey) -> RefCount {
             self.accounts_index.ref_count_from_storage(&pubkey)
         }
     }

runtime/src/bank.rs

@@ -4051,6 +4051,14 @@ impl Bank {
             .flush_accounts_cache(false, Some(self.slot()))
     }
 
+    #[cfg(test)]
+    pub fn flush_accounts_cache_slot(&self) {
+        self.rc
+            .accounts
+            .accounts_db
+            .flush_accounts_cache_slot(self.slot())
+    }
+
     pub fn expire_old_recycle_stores(&self) {
         self.rc.accounts.accounts_db.expire_old_recycle_stores()
     }
@@ -10559,6 +10567,83 @@ pub(crate) mod tests {
         pubkey0_size as usize
     }
 
+    #[test]
+    fn test_clean_nonrooted() {
+        solana_logger::setup();
+
+        let (genesis_config, _mint_keypair) = create_genesis_config(1_000_000_000);
+        let pubkey0 = Pubkey::new(&[0; 32]);
+        let pubkey1 = Pubkey::new(&[1; 32]);
+
+        info!("pubkey0: {}", pubkey0);
+        info!("pubkey1: {}", pubkey1);
+
+        // Set root for bank 0, with caching enabled
+        let mut bank0 = Arc::new(Bank::new_with_config(
+            &genesis_config,
+            AccountSecondaryIndexes::default(),
+            true,
+        ));
+
+        let account_zero = AccountSharedData::new(0, 0, &Pubkey::new_unique());
+
+        goto_end_of_slot(Arc::<Bank>::get_mut(&mut bank0).unwrap());
+        bank0.freeze();
+        bank0.squash();
+        // Flush now so that accounts cache cleaning doesn't clean up bank 0 when later
+        // slots add updates to the cache
+        bank0.force_flush_accounts_cache();
+
+        // Store some lamports in bank 1
+        let some_lamports = 123;
+        let mut bank1 = Arc::new(Bank::new_from_parent(&bank0, &Pubkey::default(), 1));
+        bank1.deposit(&pubkey0, some_lamports);
+        goto_end_of_slot(Arc::<Bank>::get_mut(&mut bank1).unwrap());
+        bank1.freeze();
+        bank1.flush_accounts_cache_slot();
+
+        bank1.print_accounts_stats();
+
+        // Store some lamports for pubkey1 in bank 2, root bank 2
+        // bank2's parent is bank0
+        let mut bank2 = Arc::new(Bank::new_from_parent(&bank0, &Pubkey::default(), 2));
+        bank2.deposit(&pubkey1, some_lamports);
+        bank2.store_account(&pubkey0, &account_zero);
+        goto_end_of_slot(Arc::<Bank>::get_mut(&mut bank2).unwrap());
+        bank2.freeze();
+        bank2.squash();
+        bank2.force_flush_accounts_cache();
+
+        bank2.print_accounts_stats();
+        drop(bank1);
+
+        // Clean accounts, which should add earlier slots to the shrink
+        // candidate set
+        bank2.clean_accounts(false);
+
+        let mut bank3 = Arc::new(Bank::new_from_parent(&bank2, &Pubkey::default(), 3));
+        bank3.deposit(&pubkey1, some_lamports + 1);
+        goto_end_of_slot(Arc::<Bank>::get_mut(&mut bank3).unwrap());
+        bank3.freeze();
+        bank3.squash();
+        bank3.force_flush_accounts_cache();
+
+        bank3.clean_accounts(false);
+        assert_eq!(
+            bank3.rc.accounts.accounts_db.ref_count_for_pubkey(&pubkey0),
+            2
+        );
+        assert!(bank3
+            .rc
+            .accounts
+            .accounts_db
+            .storage
+            .get_slot_stores(1)
+            .is_none());
+
+        bank3.print_accounts_stats();
+    }
+
     #[test]
     fn test_shrink_candidate_slots_cached() {
         solana_logger::setup();