Account for possibility of cache flush in load() (#15454)

* Account for possibility of cache flush in load()

* More cleaning

* More cleaning

* Remove unused method and some comment cleaning

* Fix typo

* Make the detected impossible purge race panic()!

* Finally revert to original .expect()

* Fix typos...

* Add assertion for max_root for easier reasoning

* Reframe races with LoadHint as possible opt.

* Fix test

* Make race bug tests run longer for less flaky

* Delay the clone-in-lock slow path even for RPC

* Make get_account panic-free & add its onchain ver.

* Fix rebase conflicts...

* Clean up

* Clean up comment

* Revert fn name change

* Fix flaky test...

* fmt...

Co-authored-by: Ryo Onodera <ryoqun@gmail.com>

runtime/benches/accounts.rs

@@ -236,7 +236,7 @@ fn bench_concurrent_read_write(bencher: &mut Bencher) {
                 let i = rng.gen_range(0, pubkeys.len());
                 test::black_box(
                     accounts
-                        .load_slow(&Ancestors::default(), &pubkeys[i])
+                        .load_without_fixed_root(&Ancestors::new(), &pubkeys[i])
                         .unwrap(),
                 );
             }

runtime/src/accounts.rs

@@ -1,5 +1,7 @@
 use crate::{
-    accounts_db::{AccountsDb, BankHashInfo, ErrorCounters, LoadedAccount, ScanStorageResult},
+    accounts_db::{
+        AccountsDb, BankHashInfo, ErrorCounters, LoadHint, LoadedAccount, ScanStorageResult,
+    },
     accounts_index::{AccountIndex, Ancestors, IndexKey},
     bank::{
         NonceRollbackFull, NonceRollbackInfo, TransactionCheckResult, TransactionExecutionResult,
@@ -214,7 +216,7 @@ impl Accounts {
                     } else {
                         let (account, rent) = self
                             .accounts_db
-                            .load(ancestors, key)
+                            .load_with_fixed_root(ancestors, key)
                             .map(|(mut account, _)| {
                                 if message.is_writable(i, demote_sysvar_write_locks) {
                                     let rent_due = rent_collector
@@ -234,7 +236,7 @@ impl Accounts {
                             {
                                 if let Some(account) = self
                                     .accounts_db
-                                    .load(ancestors, &programdata_address)
+                                    .load_with_fixed_root(ancestors, &programdata_address)
                                     .map(|(account, _)| account)
                                 {
                                     account_deps.push((programdata_address, account));
@@ -341,7 +343,7 @@ impl Accounts {
 
             let program = match self
                 .accounts_db
-                .load(ancestors, &program_id)
+                .load_with_fixed_root(ancestors, &program_id)
                 .map(|(account, _)| account)
             {
                 Some(program) => program,
@@ -366,7 +368,7 @@ impl Accounts {
                 {
                     if let Some(program) = self
                         .accounts_db
-                        .load(ancestors, &programdata_address)
+                        .load_with_fixed_root(ancestors, &programdata_address)
                         .map(|(account, _)| account)
                     {
                         accounts.insert(0, (programdata_address, program));
@@ -450,14 +452,10 @@ impl Accounts {
             .collect()
     }
 
-    /// Slow because lock is held for 1 operation instead of many
+    fn filter_zero_lamport_account(
-    pub fn load_slow(
+        account: AccountSharedData,
-        &self,
+        slot: Slot,
-        ancestors: &Ancestors,
-        pubkey: &Pubkey,
     ) -> Option<(AccountSharedData, Slot)> {
-        let (account, slot) = self.accounts_db.load_slow(ancestors, pubkey)?;
-
         if account.lamports > 0 {
             Some((account, slot))
         } else {
@@ -465,6 +463,33 @@ impl Accounts {
         }
     }
 
+    /// Slow because lock is held for 1 operation instead of many
+    fn load_slow(
+        &self,
+        ancestors: &Ancestors,
+        pubkey: &Pubkey,
+        load_hint: LoadHint,
+    ) -> Option<(AccountSharedData, Slot)> {
+        let (account, slot) = self.accounts_db.load(ancestors, pubkey, load_hint)?;
+        Self::filter_zero_lamport_account(account, slot)
+    }
+
+    pub fn load_with_fixed_root(
+        &self,
+        ancestors: &Ancestors,
+        pubkey: &Pubkey,
+    ) -> Option<(AccountSharedData, Slot)> {
+        self.load_slow(ancestors, pubkey, LoadHint::FixedMaxRoot)
+    }
+
+    pub fn load_without_fixed_root(
+        &self,
+        ancestors: &Ancestors,
+        pubkey: &Pubkey,
+    ) -> Option<(AccountSharedData, Slot)> {
+        self.load_slow(ancestors, pubkey, LoadHint::Unspecified)
+    }
+
     /// scans underlying accounts_db for this delta (slot) with a map function
     ///   from LoadedAccount to B
     /// returns only the latest/current version of B for this slot
@@ -527,7 +552,7 @@ impl Accounts {
             };
 
             if hit {
-                Some((*stored_account.pubkey(), stored_account.account()))
+                Some((*stored_account.pubkey(), stored_account.take_account()))
             } else {
                 None
             }

runtime/src/accounts_background_service.rs

@@ -444,9 +444,11 @@ mod test {
         );
         assert!(bank0.get_account(&account_key).is_some());
         pruned_banks_sender.send(0).unwrap();
+
+        assert!(!bank0.rc.accounts.scan_slot(0, |_| Some(())).is_empty());
+
         AccountsBackgroundService::remove_dead_slots(&bank0, &request_handler, &mut 0, &mut 0);
 
-        // Slot should be removed
+        assert!(bank0.rc.accounts.scan_slot(0, |_| Some(())).is_empty());
-        assert!(bank0.get_account(&account_key).is_none());
     }
 }

runtime/src/accounts_index.rs

@@ -554,7 +554,8 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
         bank 5's parent reference keeps bank 4 alive, which will prevent the `Bank::drop()` from
         running and cleaning up bank 4. Furthermore, no cleans can happen past the saved max_root == 1,
         so a potential newer max root at 3 will not clean up any of the ancestors > 1, so slot 4
-        will not be cleaned in the middle of the scan either.
+        will not be cleaned in the middle of the scan either. (NOTE similar reasoning is employed for
+        assert!() justification in AccountsDb::retry_to_get_account_accessor)
         */
         match scan_type {
             ScanTypes::Unindexed(range) => {
@@ -854,21 +855,23 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
     where
         C: Contains<'a, Slot>,
     {
-        let res = {
+        let is_empty = {
             let mut write_account_map_entry = self.get_account_write_entry(pubkey).unwrap();
             write_account_map_entry.slot_list_mut(|slot_list| {
                 slot_list.retain(|(slot, item)| {
                     let should_purge = slots_to_purge.contains(&slot);
                     if should_purge {
                         reclaims.push((*slot, item.clone()));
+                        false
+                    } else {
+                        true
                     }
-                    !should_purge
                 });
                 slot_list.is_empty()
             })
         };
         self.purge_secondary_indexes_by_inner_key(pubkey, Some(slots_to_purge), account_indexes);
-        res
+        is_empty
     }
 
     pub fn min_ongoing_scan_root(&self) -> Option<Slot> {
@@ -1129,14 +1132,15 @@ impl<T: 'static + Clone + IsCached + ZeroLamport> AccountsIndex<T> {
             if should_purge {
                 reclaims.push((*slot, value.clone()));
                 purged_slots.insert(*slot);
+                false
+            } else {
+                true
             }
-            !should_purge
         });
 
         self.purge_secondary_indexes_by_inner_key(pubkey, Some(&purged_slots), account_indexes);
     }
 
-    // `is_cached` closure is needed to work around the generic (`T`) indexed type.
     pub fn clean_rooted_entries(
         &self,
         pubkey: &Pubkey,

runtime/src/bank.rs

@@ -1367,7 +1367,7 @@ impl Bank {
     where
         F: Fn(&Option<AccountSharedData>) -> AccountSharedData,
     {
-        let old_account = self.get_sysvar_account(pubkey);
+        let old_account = self.get_sysvar_account_with_fixed_root(pubkey);
         let new_account = updater(&old_account);
 
         self.store_account_and_update_capitalization(pubkey, &new_account);
@@ -1584,7 +1584,7 @@ impl Bank {
             .into_iter()
             .for_each(|(stake_pubkey, _delegation)| {
                 examined_count += 1;
-                if let Some(mut stake_account) = self.get_account(&stake_pubkey) {
+                if let Some(mut stake_account) = self.get_account_with_fixed_root(&stake_pubkey) {
                     if let Ok(result) =
                         stake_state::rewrite_stakes(&mut stake_account, &self.rent_collector.rent)
                     {
@@ -1769,8 +1769,8 @@ impl Bank {
             .iter()
             .for_each(|(stake_pubkey, delegation)| {
                 match (
-                    self.get_account(&stake_pubkey),
+                    self.get_account_with_fixed_root(&stake_pubkey),
-                    self.get_account(&delegation.voter_pubkey),
+                    self.get_account_with_fixed_root(&delegation.voter_pubkey),
                 ) {
                     (Some(stake_account), Some(vote_account)) => {
                         // call tracer to catch any illegal data if any
@@ -2198,27 +2198,28 @@ impl Bank {
 
     // NOTE: must hold idempotent for the same set of arguments
     pub fn add_native_program(&self, name: &str, program_id: &Pubkey, must_replace: bool) {
-        let existing_genuine_program = if let Some(mut account) = self.get_account(&program_id) {
+        let existing_genuine_program =
-            // it's very unlikely to be squatted at program_id as non-system account because of burden to
+            if let Some(mut account) = self.get_account_with_fixed_root(&program_id) {
-            // find victim's pubkey/hash. So, when account.owner is indeed native_loader's, it's
+                // it's very unlikely to be squatted at program_id as non-system account because of burden to
-            // safe to assume it's a genuine program.
+                // find victim's pubkey/hash. So, when account.owner is indeed native_loader's, it's
-            if native_loader::check_id(&account.owner) {
+                // safe to assume it's a genuine program.
-                Some(account)
+                if native_loader::check_id(&account.owner) {
+                    Some(account)
+                } else {
+                    // malicious account is pre-occupying at program_id
+                    // forcibly burn and purge it
+
+                    self.capitalization.fetch_sub(account.lamports, Relaxed);
+
+                    // Resetting account balance to 0 is needed to really purge from AccountsDb and
+                    // flush the Stakes cache
+                    account.lamports = 0;
+                    self.store_account(&program_id, &account);
+                    None
+                }
             } else {
-                // malicious account is pre-occupying at program_id
-                // forcibly burn and purge it
-
-                self.capitalization.fetch_sub(account.lamports, Relaxed);
-
-                // Resetting account balance to 0 is needed to really purge from AccountsDb and
-                // flush the Stakes cache
-                account.lamports = 0;
-                self.store_account(&program_id, &account);
                 None
-            }
+            };
-        } else {
-            None
-        };
 
         if must_replace {
             // updating native program
@@ -3361,7 +3362,9 @@ impl Bank {
                     rent_share
                 };
                 if !enforce_fix || rent_to_be_paid > 0 {
-                    let mut account = self.get_account(&pubkey).unwrap_or_default();
+                    let mut account = self
+                        .get_account_with_fixed_root(&pubkey)
+                        .unwrap_or_default();
                     account.lamports += rent_to_be_paid;
                     self.store_account(&pubkey, &account);
                     rewards.push((
@@ -3429,7 +3432,9 @@ impl Bank {
     }
 
     fn run_incinerator(&self) {
-        if let Some((account, _)) = self.get_account_modified_since_parent(&incinerator::id()) {
+        if let Some((account, _)) =
+            self.get_account_modified_since_parent_with_fixed_root(&incinerator::id())
+        {
             self.capitalization.fetch_sub(account.lamports, Relaxed);
             self.store_account(&incinerator::id(), &AccountSharedData::default());
         }
@@ -3986,7 +3991,7 @@ impl Bank {
         pubkey: &Pubkey,
         new_account: &AccountSharedData,
     ) {
-        if let Some(old_account) = self.get_account(&pubkey) {
+        if let Some(old_account) = self.get_account_with_fixed_root(&pubkey) {
             match new_account.lamports.cmp(&old_account.lamports) {
                 std::cmp::Ordering::Greater => {
                     self.capitalization
@@ -4006,7 +4011,7 @@ impl Bank {
     }
 
     pub fn withdraw(&self, pubkey: &Pubkey, lamports: u64) -> Result<()> {
-        match self.get_account(pubkey) {
+        match self.get_account_with_fixed_root(pubkey) {
             Some(mut account) => {
                 let min_balance = match get_system_account_kind(&account) {
                     Some(SystemAccountKind::Nonce) => self
@@ -4031,7 +4036,7 @@ impl Bank {
     pub fn deposit(&self, pubkey: &Pubkey, lamports: u64) -> u64 {
         // This doesn't collect rents intentionally.
         // Rents should only be applied to actual TXes
-        let mut account = self.get_account(pubkey).unwrap_or_default();
+        let mut account = self.get_account_with_fixed_root(pubkey).unwrap_or_default();
         account.lamports += lamports;
         self.store_account(pubkey, &account);
         account.lamports
@@ -4082,13 +4087,46 @@ impl Bank {
         self.hard_forks.clone()
     }
 
+    // Hi! leaky abstraction here....
+    // try to use get_account_with_fixed_root() if it's called ONLY from on-chain runtime account
+    // processing. That alternative fn provides more safety.
     pub fn get_account(&self, pubkey: &Pubkey) -> Option<AccountSharedData> {
         self.get_account_modified_slot(pubkey)
             .map(|(acc, _slot)| acc)
     }
 
+    // Hi! leaky abstraction here....
+    // use this over get_account() if it's called ONLY from on-chain runtime account
+    // processing (i.e. from in-band replay/banking stage; that ensures root is *fixed* while
+    // running).
+    // pro: safer assertion can be enabled inside AccountsDb
+    // con: panics!() if called from off-chain processing
+    pub fn get_account_with_fixed_root(&self, pubkey: &Pubkey) -> Option<AccountSharedData> {
+        self.load_slow_with_fixed_root(&self.ancestors, pubkey)
+            .map(|(acc, _slot)| acc)
+    }
+
     pub fn get_account_modified_slot(&self, pubkey: &Pubkey) -> Option<(AccountSharedData, Slot)> {
-        self.rc.accounts.load_slow(&self.ancestors, pubkey)
+        self.load_slow(&self.ancestors, pubkey)
+    }
+
+    fn load_slow(
+        &self,
+        ancestors: &Ancestors,
+        pubkey: &Pubkey,
+    ) -> Option<(AccountSharedData, Slot)> {
+        // get_account (= primary this fn caller) may be called from on-chain Bank code even if we
+        // try hard to use get_account_with_fixed_root for that purpose...
+        // so pass safer LoadHint:Unspecified here as a fallback
+        self.rc.accounts.load_without_fixed_root(ancestors, pubkey)
+    }
+
+    fn load_slow_with_fixed_root(
+        &self,
+        ancestors: &Ancestors,
+        pubkey: &Pubkey,
+    ) -> Option<(AccountSharedData, Slot)> {
+        self.rc.accounts.load_with_fixed_root(ancestors, pubkey)
     }
 
     // Exclude self to really fetch the parent Bank's account hash and data.
@@ -4098,12 +4136,12 @@ impl Bank {
     // multiple times with the same parent_slot in the case of forking.
     //
     // Generally, all of sysvar update granularity should be slot boundaries.
-    fn get_sysvar_account(&self, pubkey: &Pubkey) -> Option<AccountSharedData> {
+    fn get_sysvar_account_with_fixed_root(&self, pubkey: &Pubkey) -> Option<AccountSharedData> {
         let mut ancestors = self.ancestors.clone();
         ancestors.remove(&self.slot());
         self.rc
             .accounts
-            .load_slow(&ancestors, pubkey)
+            .load_with_fixed_root(&ancestors, pubkey)
             .map(|(acc, _slot)| acc)
     }
 
@@ -4170,12 +4208,13 @@ impl Bank {
         self.rc.accounts.load_by_program_slot(self.slot(), None)
     }
 
-    pub fn get_account_modified_since_parent(
+    // if you want get_account_modified_since_parent without fixed_root, please define so...
+    fn get_account_modified_since_parent_with_fixed_root(
         &self,
         pubkey: &Pubkey,
     ) -> Option<(AccountSharedData, Slot)> {
         let just_self: Ancestors = vec![(self.slot(), 0)].into_iter().collect();
-        if let Some((account, slot)) = self.rc.accounts.load_slow(&just_self, pubkey) {
+        if let Some((account, slot)) = self.load_slow_with_fixed_root(&just_self, pubkey) {
             if slot == self.slot() {
                 return Some((account, slot));
             }
@@ -4789,7 +4828,7 @@ impl Bank {
 
         for feature_id in &self.feature_set.inactive {
             let mut activated = None;
-            if let Some(mut account) = self.get_account(feature_id) {
+            if let Some(mut account) = self.get_account_with_fixed_root(feature_id) {
                 if let Some(mut feature) = feature::from_account(&account) {
                     match feature.activated_at {
                         None => {
@@ -4851,9 +4890,9 @@ impl Bank {
     }
 
     fn apply_spl_token_v2_self_transfer_fix(&mut self) {
-        if let Some(old_account) = self.get_account(&inline_spl_token_v2_0::id()) {
+        if let Some(old_account) = self.get_account_with_fixed_root(&inline_spl_token_v2_0::id()) {
             if let Some(new_account) =
-                self.get_account(&inline_spl_token_v2_0::new_token_program::id())
+                self.get_account_with_fixed_root(&inline_spl_token_v2_0::new_token_program::id())
             {
                 datapoint_info!(
                     "bank-apply_spl_token_v2_self_transfer_fix",
@@ -4897,7 +4936,7 @@ impl Bank {
             // https://github.com/solana-labs/solana-program-library/issues/374, ensure that the
             // spl-token 2 native mint account is owned by the spl-token 2 program.
             let store = if let Some(existing_native_mint_account) =
-                self.get_account(&inline_spl_token_v2_0::native_mint::id())
+                self.get_account_with_fixed_root(&inline_spl_token_v2_0::native_mint::id())
             {
                 if existing_native_mint_account.owner == solana_sdk::system_program::id() {
                     native_mint_account.lamports = existing_native_mint_account.lamports;
@@ -4933,7 +4972,7 @@ impl Bank {
 
         if purge_window_epoch {
             for reward_pubkey in self.rewards_pool_pubkeys.iter() {
-                if let Some(mut reward_account) = self.get_account(&reward_pubkey) {
+                if let Some(mut reward_account) = self.get_account_with_fixed_root(&reward_pubkey) {
                     if reward_account.lamports == u64::MAX {
                         reward_account.lamports = 0;
                         self.store_account(&reward_pubkey, &reward_account);
@@ -8273,27 +8312,29 @@ pub(crate) mod tests {
     }
 
     #[test]
-    fn test_bank_get_account_modified_since_parent() {
+    fn test_bank_get_account_modified_since_parent_with_fixed_root() {
         let pubkey = solana_sdk::pubkey::new_rand();
 
         let (genesis_config, mint_keypair) = create_genesis_config(500);
         let bank1 = Arc::new(Bank::new(&genesis_config));
         bank1.transfer(1, &mint_keypair, &pubkey).unwrap();
-        let result = bank1.get_account_modified_since_parent(&pubkey);
+        let result = bank1.get_account_modified_since_parent_with_fixed_root(&pubkey);
         assert!(result.is_some());
         let (account, slot) = result.unwrap();
         assert_eq!(account.lamports, 1);
         assert_eq!(slot, 0);
 
         let bank2 = Arc::new(Bank::new_from_parent(&bank1, &Pubkey::default(), 1));
-        assert!(bank2.get_account_modified_since_parent(&pubkey).is_none());
+        assert!(bank2
+            .get_account_modified_since_parent_with_fixed_root(&pubkey)
+            .is_none());
         bank2.transfer(100, &mint_keypair, &pubkey).unwrap();
-        let result = bank1.get_account_modified_since_parent(&pubkey);
+        let result = bank1.get_account_modified_since_parent_with_fixed_root(&pubkey);
         assert!(result.is_some());
         let (account, slot) = result.unwrap();
         assert_eq!(account.lamports, 1);
         assert_eq!(slot, 0);
-        let result = bank2.get_account_modified_since_parent(&pubkey);
+        let result = bank2.get_account_modified_since_parent_with_fixed_root(&pubkey);
         assert!(result.is_some());
         let (account, slot) = result.unwrap();
         assert_eq!(account.lamports, 101);
@@ -8302,7 +8343,10 @@ pub(crate) mod tests {
         bank1.squash();
 
         let bank3 = Bank::new_from_parent(&bank2, &Pubkey::default(), 3);
-        assert_eq!(None, bank3.get_account_modified_since_parent(&pubkey));
+        assert_eq!(
+            None,
+            bank3.get_account_modified_since_parent_with_fixed_root(&pubkey)
+        );
     }
 
     #[test]

runtime/src/message_processor.rs

@@ -434,7 +434,7 @@ impl<'a> InvokeContext for ThisInvokeContext<'a> {
             // Load it
             result = self
                 .account_db
-                .load_slow(self.ancestors, id)
+                .load_with_fixed_root(self.ancestors, id)
                 .map(|(account, _)| Rc::new(account.data().clone()));
             // Cache it
             self.sysvars.push((*id, result.clone()));

runtime/src/serde_snapshot/tests.rs

@@ -46,7 +46,7 @@ fn check_accounts(accounts: &Accounts, pubkeys: &[Pubkey], num: usize) {
     for _ in 1..num {
         let idx = thread_rng().gen_range(0, num - 1);
         let ancestors = vec![(0, 0)].into_iter().collect();
-        let account = accounts.load_slow(&ancestors, &pubkeys[idx]);
+        let account = accounts.load_without_fixed_root(&ancestors, &pubkeys[idx]);
         let account1 = Some((
             AccountSharedData::new((idx + 1) as u64, 0, &AccountSharedData::default().owner),
             0,

runtime/tests/accounts.rs

@@ -1,7 +1,10 @@
 use log::*;
 use rand::{thread_rng, Rng};
 use rayon::prelude::*;
-use solana_runtime::{accounts_db::AccountsDb, accounts_index::Ancestors};
+use solana_runtime::{
+    accounts_db::{AccountsDb, LoadHint},
+    accounts_index::Ancestors,
+};
 use solana_sdk::genesis_config::ClusterType;
 use solana_sdk::{account::AccountSharedData, clock::Slot, pubkey::Pubkey};
 use std::collections::HashSet;
@@ -112,7 +115,8 @@ fn test_bad_bank_hash() {
 
         for (key, account) in &account_refs {
             assert_eq!(
-                db.load_account_hash(&ancestors, &key),
+                db.load_account_hash(&ancestors, &key, None, LoadHint::Unspecified)
+                    .unwrap(),
                 AccountsDb::hash_account(some_slot, &account, &key)
             );
         }