reorder load_account and program_cache replenish

Author: HaoranYi

runtime/src/bank.rs

@@ -6737,7 +6737,6 @@ impl Bank {
             effective_epoch,
             self.epoch_schedule(),
             reload,
-            &HashMap::default(),
         )
     }
 

svm/src/account_loader.rs

@@ -32,6 +32,7 @@ use {
         transaction_context::{IndexOfAccount, TransactionAccount},
     },
     solana_system_program::{get_system_account_kind, SystemAccountKind},
+    std::collections::hash_map::Entry,
     std::{collections::HashMap, num::NonZeroUsize},
 };
 
@@ -108,15 +109,16 @@ pub fn validate_fee_payer(
 /// batch. Each tuple contains struct of information about accounts as
 /// its first element and an optional transaction nonce info as its
 /// second element.
+/// This function also populate program_accounts map.
 pub(crate) fn load_accounts<CB: TransactionProcessingCallback>(
     callbacks: &CB,
     txs: &[SanitizedTransaction],
     check_results: &[TransactionCheckResult],
     error_counters: &mut TransactionErrorMetrics,
     fee_structure: &FeeStructure,
     account_overrides: Option<&AccountOverrides>,
-    loaded_programs: &ProgramCacheForTxBatch,
-    accounts_map: &HashMap<Pubkey, Option<(AccountSharedData, Slot)>>,
+    program_owners: &[Pubkey],
+    program_accounts: &mut HashMap<Pubkey, u64>,
 ) -> Vec<TransactionLoadResult> {
     let feature_set = callbacks.get_feature_set();
     txs.iter()
@@ -146,8 +148,8 @@ pub(crate) fn load_accounts<CB: TransactionProcessingCallback>(
                     fee,
                     error_counters,
                     account_overrides,
-                    loaded_programs,
-                    accounts_map,
+                    program_owners,
+                    program_accounts,
                 ) {
                     Ok(loaded_transaction) => loaded_transaction,
                     Err(e) => return (Err(e), None),
@@ -183,20 +185,11 @@ fn load_transaction_accounts<CB: TransactionProcessingCallback>(
     fee: u64,
     error_counters: &mut TransactionErrorMetrics,
     account_overrides: Option<&AccountOverrides>,
-    loaded_programs: &ProgramCacheForTxBatch,
-    accounts_map: &HashMap<Pubkey, Option<(AccountSharedData, Slot)>>,
+    program_owners: &[Pubkey],
+    program_accounts: &mut HashMap<Pubkey, u64>,
 ) -> Result<LoadedTransaction> {
     let feature_set = callbacks.get_feature_set();
 
-    let load_account = |pubkey, should_cache| {
-        let maybe_found = if let Some(found) = accounts_map.get(pubkey) {
-            found.as_ref().cloned()
-        } else {
-            callbacks.load_account_with(pubkey, |_| should_cache)
-        };
-        maybe_found.map(|x| x.0)
-    };
-
     // There is no way to predict what program will execute without an error
     // If a fee can pay for execution then the program will be scheduled
     let mut validated_fee_payer = false;
@@ -210,13 +203,6 @@ fn load_transaction_accounts<CB: TransactionProcessingCallback>(
         get_requested_loaded_accounts_data_size_limit(message)?;
     let mut accumulated_accounts_data_size: usize = 0;
 
-    let instruction_accounts = message
-        .instructions()
-        .iter()
-        .flat_map(|instruction| &instruction.accounts)
-        .unique()
-        .collect::<Vec<&u8>>();
-
     let mut accounts = account_keys
         .iter()
         .enumerate()
@@ -226,26 +212,30 @@ fn load_transaction_accounts<CB: TransactionProcessingCallback>(
             let account = if solana_sdk::sysvar::instructions::check_id(key) {
                 construct_instructions_account(message)
             } else {
-                let instruction_account = u8::try_from(i)
-                    .map(|i| instruction_accounts.contains(&&i))
-                    .unwrap_or(false);
                 let (account_size, mut account, rent) = if let Some(account_override) =
                     account_overrides.and_then(|overrides| overrides.get(key))
                 {
                     (account_override.data().len(), account_override.clone(), 0)
-                } else if let Some(program) = (!instruction_account && !message.is_writable(i))
-                    .then_some(())
-                    .and_then(|_| loaded_programs.find(key))
-                {
-                    load_account(key, false).ok_or(TransactionError::AccountNotFound)?;
-                    // Optimization to skip loading of accounts which are only used as
-                    // programs in top-level instructions and not passed as instruction accounts.
-                    let program_account = account_shared_data_from_program(&program);
-                    (program.account_size, program_account, 0)
                 } else {
-                    load_account(key, !message.is_writable(i))
-                        .map(|mut account| {
-                            if message.is_writable(i) {
+                    let is_account_writable = message.is_writable(i);
+                    callbacks
+                        .load_account_with(key, |account| {
+                            // only cache account that is not writable and not having program owners.
+                            !is_account_writable && !program_owners.contains(account.owner())
+                        })
+                        .map(|(mut account, _slot)| {
+                            if program_owners.contains(account.owner()) {
+                                match program_accounts.entry(*key) {
+                                    Entry::Vacant(entry) => {
+                                        entry.insert(1);
+                                    }
+                                    Entry::Occupied(mut entry) => {
+                                        let count = entry.get_mut();
+                                        saturating_add_assign!(*count, 1);
+                                    }
+                                }
+                            }
+                            if is_account_writable {
                                 if !feature_set
                                     .is_active(&feature_set::disable_rent_fees_collection::id())
                                 {
@@ -563,15 +553,16 @@ mod tests {
             rent_collector: rent_collector.clone(),
             feature_set: Arc::new(feature_set.clone()),
         };
+        let mut program_accounts = HashMap::default();
         load_accounts(
             &callbacks,
             &[sanitized_tx],
             &[(Ok(()), None, Some(lamports_per_signature))],
             error_counters,
             fee_structure,
             None,
-            &ProgramCacheForTxBatch::default(),
-            &HashMap::default(),
+            &[],
+            &mut program_accounts,
         )
     }
 
@@ -1052,15 +1043,17 @@ mod tests {
             rent_collector: RentCollector::default(),
             feature_set: Arc::new(FeatureSet::all_enabled()),
         };
+
+        let mut program_accounts = HashMap::default();
         load_accounts(
             &callbacks,
             &[tx],
             &[(Ok(()), None, Some(10))],
             &mut error_counters,
             &FeeStructure::default(),
             account_overrides,
-            &ProgramCacheForTxBatch::default(),
-            &HashMap::default(),
+            &[],
+            &mut program_accounts,
         )
     }
 
@@ -1451,21 +1444,21 @@ mod tests {
         let sanitized_message = new_unchecked_sanitized_message(message);
         let mock_bank = TestCallbacks::default();
         let mut error_counter = TransactionErrorMetrics::default();
-        let loaded_programs = ProgramCacheForTxBatch::default();
 
         let sanitized_transaction = SanitizedTransaction::new_for_tests(
             sanitized_message,
             vec![Signature::new_unique()],
             false,
         );
+
         let result = load_transaction_accounts(
             &mock_bank,
             sanitized_transaction.message(),
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
         assert_eq!(result.err(), Some(TransactionError::AccountNotFound));
@@ -1508,8 +1501,8 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
         mock_bank
             .accounts_map
@@ -1560,8 +1553,8 @@ mod tests {
         mock_bank.accounts_map.insert(key1.pubkey(), account_data);
 
         let mut error_counter = TransactionErrorMetrics::default();
-        let mut loaded_programs = ProgramCacheForTxBatch::default();
-        loaded_programs.replenish(key2.pubkey(), Arc::new(ProgramCacheEntry::default()));
+        // let mut loaded_programs = ProgramCacheForTxBatch::default();
+        // loaded_programs.replenish(key2.pubkey(), Arc::new(ProgramCacheEntry::default()));
 
         let sanitized_transaction = SanitizedTransaction::new_for_tests(
             sanitized_message,
@@ -1574,11 +1567,17 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
-        assert_eq!(result.err(), Some(TransactionError::AccountNotFound));
+        // The original code throws "AccountNotFound" error. With this PR, it
+        // throws "ProgramAccountNotFound" error. But I don't think this is a
+        // concern for breaking consensus. As the original "AccountNotFound"
+        // error should not be triggered in real world. When we can find "key"
+        // in program_cache, which is after replenish, then it must be in
+        // accounts-db.
+        assert_eq!(result.err(), Some(TransactionError::ProgramAccountNotFound));
     }
 
     #[test]
@@ -1604,7 +1603,6 @@ mod tests {
         mock_bank.accounts_map.insert(key1.pubkey(), account_data);
 
         let mut error_counter = TransactionErrorMetrics::default();
-        let loaded_programs = ProgramCacheForTxBatch::default();
 
         let sanitized_transaction = SanitizedTransaction::new_for_tests(
             sanitized_message,
@@ -1617,8 +1615,8 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
         assert_eq!(result.err(), Some(TransactionError::ProgramAccountNotFound));
@@ -1647,7 +1645,6 @@ mod tests {
         mock_bank.accounts_map.insert(key1.pubkey(), account_data);
 
         let mut error_counter = TransactionErrorMetrics::default();
-        let loaded_programs = ProgramCacheForTxBatch::default();
 
         let sanitized_transaction = SanitizedTransaction::new_for_tests(
             sanitized_message,
@@ -1660,8 +1657,8 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
         assert_eq!(
@@ -1697,7 +1694,6 @@ mod tests {
         account_data.set_lamports(200);
         mock_bank.accounts_map.insert(key2.pubkey(), account_data);
         let mut error_counter = TransactionErrorMetrics::default();
-        let loaded_programs = ProgramCacheForTxBatch::default();
 
         let sanitized_transaction = SanitizedTransaction::new_for_tests(
             sanitized_message,
@@ -1710,8 +1706,8 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
         mock_bank
             .accounts_map
@@ -1765,7 +1761,6 @@ mod tests {
         account_data.set_lamports(200);
         mock_bank.accounts_map.insert(key2.pubkey(), account_data);
         let mut error_counter = TransactionErrorMetrics::default();
-        let loaded_programs = ProgramCacheForTxBatch::default();
 
         let sanitized_transaction = SanitizedTransaction::new_for_tests(
             sanitized_message,
@@ -1778,8 +1773,8 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
         mock_bank
             .accounts_map
@@ -1835,8 +1830,8 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
         mock_bank
             .accounts_map
@@ -1897,8 +1892,8 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
         mock_bank
             .accounts_map
@@ -1985,8 +1980,8 @@ mod tests {
             32,
             &mut error_counter,
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
         mock_bank
             .accounts_map
@@ -2057,8 +2052,8 @@ mod tests {
             &mut error_counters,
             &FeeStructure::default(),
             None,
-            &ProgramCacheForTxBatch::default(),
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
         let compute_budget = ComputeBudget::new(u64::from(
@@ -2141,8 +2136,8 @@ mod tests {
             &mut error_counter,
             &FeeStructure::default(),
             None,
-            &loaded_programs,
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
         let mut account_data = AccountSharedData::default();
@@ -2215,8 +2210,8 @@ mod tests {
             &mut TransactionErrorMetrics::default(),
             &fee_structure,
             None,
-            &ProgramCacheForTxBatch::default(),
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
         assert_eq!(
@@ -2234,8 +2229,8 @@ mod tests {
             &mut TransactionErrorMetrics::default(),
             &fee_structure,
             None,
-            &ProgramCacheForTxBatch::default(),
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
         assert_eq!(result, vec![(Err(TransactionError::AccountNotFound), None)]);
@@ -2253,8 +2248,8 @@ mod tests {
             &mut TransactionErrorMetrics::default(),
             &fee_structure,
             None,
-            &ProgramCacheForTxBatch::default(),
-            &HashMap::default(),
+            &[],
+            &mut HashMap::default(),
         );
 
         assert_eq!(