add the account cache except for setup and usage

svm/Cargo.toml

@@ -10,6 +10,7 @@ license = { workspace = true }
 edition = { workspace = true }
 
 [dependencies]
+ahash = { workspace = true }
 itertools = { workspace = true }
 log = { workspace = true }
 percentage = { workspace = true }

svm/src/account_loader.rs

@@ -4,8 +4,10 @@ use {
         nonce_info::NonceInfo,
         rollback_accounts::RollbackAccounts,
         transaction_error_metrics::TransactionErrorMetrics,
+        transaction_execution_result::ExecutedTransaction,
         transaction_processing_callback::{AccountState, TransactionProcessingCallback},
     },
+    ahash::AHashMap,
     solana_compute_budget::compute_budget_limits::ComputeBudgetLimits,
     solana_feature_set::{self as feature_set, FeatureSet},
     solana_program_runtime::loaded_programs::{ProgramCacheEntry, ProgramCacheForTxBatch},
@@ -99,17 +101,20 @@ pub struct FeesOnlyTransaction {
 pub(crate) struct AccountLoader<'a, CB: TransactionProcessingCallback> {
     pub(crate) program_cache: ProgramCacheForTxBatch,
     account_overrides: Option<&'a AccountOverrides>,
+    account_cache: AHashMap<Pubkey, AccountCacheItem>,
     callbacks: &'a CB,
 }
 impl<'a, CB: TransactionProcessingCallback> AccountLoader<'a, CB> {
-    pub fn new(
+    pub fn new_with_account_cache_capacity(
         callbacks: &'a CB,
         account_overrides: Option<&'a AccountOverrides>,
         program_cache: ProgramCacheForTxBatch,
+        capacity: usize,
     ) -> AccountLoader<'a, CB> {
         Self {
             program_cache,
             account_overrides,
+            account_cache: AHashMap::with_capacity(capacity),
             callbacks,
         }
     }
@@ -161,6 +166,103 @@ impl<'a, CB: TransactionProcessingCallback> AccountLoader<'a, CB> {
             None
         }
     }
+
+    pub fn update_accounts_for_executed_tx(
+        &mut self,
+        message: &impl SVMMessage,
+        executed_transaction: &ExecutedTransaction,
+    ) {
+        if executed_transaction.was_successful() {
+            self.program_cache
+                .merge(&executed_transaction.programs_modified_by_tx);
+
+            self.update_accounts_for_successful_tx(
+                message,
+                &executed_transaction.loaded_transaction.accounts,
+            );
+        } else {
+            self.update_accounts_for_failed_tx(
+                message,
+                &executed_transaction.loaded_transaction.rollback_accounts,
+            );
+        }
+    }
+
+    pub fn update_accounts_for_failed_tx(
+        &mut self,
+        message: &impl SVMMessage,
+        rollback_accounts: &RollbackAccounts,
+    ) {
+        let fee_payer_address = message.fee_payer();
+        match rollback_accounts {
+            RollbackAccounts::FeePayerOnly { fee_payer_account } => {
+                self.update_account(fee_payer_address, fee_payer_account);
+            }
+            RollbackAccounts::SameNonceAndFeePayer { nonce } => {
+                self.update_account(nonce.address(), nonce.account());
+            }
+            RollbackAccounts::SeparateNonceAndFeePayer {
+                nonce,
+                fee_payer_account,
+            } => {
+                self.update_account(nonce.address(), nonce.account());
+                self.update_account(fee_payer_address, fee_payer_account);
+            }
+        }
+    }
+
+    fn update_accounts_for_successful_tx(
+        &mut self,
+        message: &impl SVMMessage,
+        transaction_accounts: &[TransactionAccount],
+    ) {
+        for (i, (address, account)) in (0..message.account_keys().len()).zip(transaction_accounts) {
+            if !message.is_writable(i) {
+                continue;
+            }
+
+            // Accounts that are invoked and also not passed as an instruction
+            // account to a program don't need to be stored because it's assumed
+            // to be impossible for a committable transaction to modify an
+            // invoked account if said account isn't passed to some program.
+            if message.is_invoked(i) && !message.is_instruction_account(i) {
+                continue;
+            }
+
+            self.update_account(address, account);
+        }
+    }
+
+    fn update_account(&mut self, account_key: &Pubkey, account: &AccountSharedData) {
+        // If an account has been dropped, we must hide the data from any future transactions.
+        // We NEVER remove items from cache, or else we will fetch stale data from accounts-db.
+        // Because `update_account()` is only called on writable accounts using the same rules
+        // as commit, this (correctly) does not shadow rent-paying zero-lamport read-only accounts.
+        let account = if account.lamports() == 0 {
+            &AccountSharedData::default()
+        } else {
+            account
+        };
+
+        // Accounts are inspected during loading, and we only update writable accounts.
+        // Ergo, `inspected_as_writable` is always true.
+        self.account_cache
+            .entry(*account_key)
+            .and_modify(|cache_item| {
+                debug_assert!(cache_item.inspected_as_writable);
+                cache_item.account = account.clone();
+            })
+            .or_insert_with(|| AccountCacheItem {
+                account: account.clone(),
+                inspected_as_writable: true,
+            });
+    }
+}
+
+#[derive(Debug, Clone)]
+struct AccountCacheItem {
+    account: AccountSharedData,
+    inspected_as_writable: bool,
 }
 
 #[derive(Debug, Copy, Clone, PartialEq, Eq)]

svm/src/transaction_processor.rs

@@ -297,10 +297,17 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
             program_cache_for_tx_batch
         });
 
-        let mut account_loader = AccountLoader::new(
+        // Determine a capacity for the internal account cache. This
+        // over-allocates but avoids ever reallocating, and spares us from
+        // deduplicating the account keys lists.
+        let account_keys_in_batch = sanitized_txs.iter().map(|tx| tx.account_keys().len()).sum();
+
+        // Create the account loader, which wraps all external account fetching.
+        let mut account_loader = AccountLoader::new_with_account_cache_capacity(
             callbacks,
             config.account_overrides,
             program_cache_for_tx_batch,
+            account_keys_in_batch,
         );
 
         let enable_transaction_loading_failure_fees = environment
@@ -349,6 +356,10 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
                 TransactionLoadResult::NotLoaded(err) => Err(err),
                 TransactionLoadResult::FeesOnly(fees_only_tx) => {
                     if enable_transaction_loading_failure_fees {
+                        // Update loaded accounts cache with nonce and fee-payer
+                        account_loader
+                            .update_accounts_for_failed_tx(tx, &fees_only_tx.rollback_accounts);
+
                         Ok(ProcessedTransaction::FeesOnly(Box::new(fees_only_tx)))
                     } else {
                         Err(fees_only_tx.load_error)
@@ -365,13 +376,10 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
                         config,
                     );
 
-                    // Update batch specific cache of the loaded programs with the modifications
-                    // made by the transaction, if it executed successfully.
-                    if executed_tx.was_successful() {
-                        account_loader
-                            .program_cache
-                            .merge(&executed_tx.programs_modified_by_tx);
-                    }
+                    // Update loaded accounts cache with account states which might have changed.
+                    // Also update local progrma cache with modifications made by the transaction,
+                    // if it executed successfully.
+                    account_loader.update_accounts_for_executed_tx(tx, &executed_tx);
 
                     Ok(ProcessedTransaction::Executed(Box::new(executed_tx)))
                 }