Kaynağa Gözat

svm: collect transaction balances (#5588)

hana 6 ay önce
ebeveyn
işleme
cf6946d4f6

+ 3 - 2
Cargo.lock

@@ -8545,9 +8545,8 @@ dependencies = [
  "solana-transaction-status",
  "solana-vote",
  "solana-vote-program",
+ "spl-generic-token",
  "spl-pod",
- "spl-token",
- "spl-token-2022",
  "static_assertions",
  "strum",
  "strum_macros",
@@ -10447,6 +10446,8 @@ dependencies = [
  "solana-transaction-context 2.3.0",
  "solana-transaction-error",
  "solana-type-overrides",
+ "spl-generic-token",
+ "spl-token",
  "test-case",
  "thiserror 2.0.12",
 ]

+ 23 - 30
core/src/banking_stage/committer.rs

@@ -3,26 +3,25 @@ use {
     itertools::Itertools,
     solana_cost_model::cost_model::CostModel,
     solana_ledger::{
-        blockstore_processor::TransactionStatusSender, token_balances::collect_token_balances,
+        blockstore_processor::TransactionStatusSender,
+        transaction_balances::compile_collected_balances,
     },
     solana_measure::measure_us,
     solana_runtime::{
-        bank::{Bank, ProcessedTransactionCounts, TransactionBalancesSet},
+        bank::{Bank, ProcessedTransactionCounts},
         bank_utils,
         prioritization_fee_cache::PrioritizationFeeCache,
         transaction_batch::TransactionBatch,
         vote_sender_types::ReplayVoteSender,
     },
     solana_runtime_transaction::transaction_with_meta::TransactionWithMeta,
-    solana_sdk::{pubkey::Pubkey, saturating_add_assign},
+    solana_sdk::saturating_add_assign,
     solana_svm::{
+        transaction_balances::BalanceCollector,
         transaction_commit_result::{TransactionCommitResult, TransactionCommitResultExtensions},
         transaction_processing_result::TransactionProcessingResult,
     },
-    solana_transaction_status::{
-        token_balances::TransactionTokenBalancesSet, TransactionTokenBalance,
-    },
-    std::{collections::HashMap, sync::Arc},
+    std::sync::Arc,
 };
 
 #[derive(Clone, Debug, PartialEq, Eq)]
@@ -34,13 +33,6 @@ pub enum CommitTransactionDetails {
     NotCommitted,
 }
 
-#[derive(Default)]
-pub(super) struct PreBalanceInfo {
-    pub native: Vec<Vec<u64>>,
-    pub token: Vec<Vec<TransactionTokenBalance>>,
-    pub mint_decimals: HashMap<Pubkey, u8>,
-}
-
 #[derive(Clone)]
 pub struct Committer {
     transaction_status_sender: Option<TransactionStatusSender>,
@@ -71,7 +63,7 @@ impl Committer {
         processing_results: Vec<TransactionProcessingResult>,
         starting_transaction_index: Option<usize>,
         bank: &Arc<Bank>,
-        pre_balance_info: &mut PreBalanceInfo,
+        balance_collector: Option<BalanceCollector>,
         execute_and_commit_timings: &mut LeaderExecuteAndCommitTimings,
         processed_counts: &ProcessedTransactionCounts,
     ) -> (u64, Vec<CommitTransactionDetails>) {
@@ -117,7 +109,7 @@ impl Committer {
                 commit_results,
                 bank,
                 batch,
-                pre_balance_info,
+                balance_collector,
                 starting_transaction_index,
             );
         });
@@ -130,7 +122,7 @@ impl Committer {
         commit_results: Vec<TransactionCommitResult>,
         bank: &Arc<Bank>,
         batch: &TransactionBatch<impl TransactionWithMeta>,
-        pre_balance_info: &mut PreBalanceInfo,
+        balance_collector: Option<BalanceCollector>,
         starting_transaction_index: Option<usize>,
     ) {
         if let Some(transaction_status_sender) = &self.transaction_status_sender {
@@ -142,11 +134,6 @@ impl Committer {
                 .iter()
                 .map(|tx| tx.as_sanitized_transaction().into_owned())
                 .collect_vec();
-
-            let post_balances = bank.collect_balances(batch);
-            let post_token_balances =
-                collect_token_balances(bank, batch, &mut pre_balance_info.mint_decimals);
-
             let mut transaction_index = starting_transaction_index.unwrap_or_default();
             let (batch_transaction_indexes, tx_costs): (Vec<_>, Vec<_>) = commit_results
                 .iter()
@@ -173,18 +160,24 @@ impl Committer {
                 })
                 .unzip();
 
+            // There are two cases where balance_collector could be None:
+            // * Balance recording is disabled. If that were the case, there would
+            //   be no TransactionStatusSender, and we would not be in this branch.
+            // * The batch was aborted in its entirety in SVM. In that case, there
+            //   would be zero processed transactions, and commit_transactions()
+            //   would not have been called at all.
+            // Therefore this should always be true.
+            debug_assert!(balance_collector.is_some());
+
+            let (balances, token_balances) =
+                compile_collected_balances(balance_collector.unwrap_or_default());
+
             transaction_status_sender.send_transaction_status_batch(
                 bank.slot(),
                 txs,
                 commit_results,
-                TransactionBalancesSet::new(
-                    std::mem::take(&mut pre_balance_info.native),
-                    post_balances,
-                ),
-                TransactionTokenBalancesSet::new(
-                    std::mem::take(&mut pre_balance_info.token),
-                    post_token_balances,
-                ),
+                balances,
+                token_balances,
                 tx_costs,
                 batch_transaction_indexes,
             );

+ 0 - 10
core/src/banking_stage/consume_worker.rs

@@ -285,7 +285,6 @@ impl ConsumeWorkerMetrics {
     fn update_on_execute_and_commit_timings(
         &self,
         LeaderExecuteAndCommitTimings {
-            collect_balances_us,
             load_execute_us,
             freeze_lock_us,
             record_us,
@@ -294,9 +293,6 @@ impl ConsumeWorkerMetrics {
             ..
         }: &LeaderExecuteAndCommitTimings,
     ) {
-        self.timing_metrics
-            .collect_balances_us
-            .fetch_add(*collect_balances_us, Ordering::Relaxed);
         self.timing_metrics
             .load_execute_us_min
             .fetch_min(*load_execute_us, Ordering::Relaxed);
@@ -512,7 +508,6 @@ impl ConsumeWorkerCountMetrics {
 #[derive(Default)]
 struct ConsumeWorkerTimingMetrics {
     cost_model_us: AtomicU64,
-    collect_balances_us: AtomicU64,
     load_execute_us: AtomicU64,
     load_execute_us_min: AtomicU64,
     load_execute_us_max: AtomicU64,
@@ -535,11 +530,6 @@ impl ConsumeWorkerTimingMetrics {
                 self.cost_model_us.swap(0, Ordering::Relaxed),
                 i64
             ),
-            (
-                "collect_balances_us",
-                self.collect_balances_us.swap(0, Ordering::Relaxed),
-                i64
-            ),
             (
                 "load_execute_us",
                 self.load_execute_us.swap(0, Ordering::Relaxed),

+ 3 - 15
core/src/banking_stage/consumer.rs

@@ -1,13 +1,12 @@
 use {
     super::{
-        committer::{CommitTransactionDetails, Committer, PreBalanceInfo},
+        committer::{CommitTransactionDetails, Committer},
         leader_slot_timing_metrics::LeaderExecuteAndCommitTimings,
         qos_service::QosService,
         scheduler_messages::MaxAge,
     },
     itertools::Itertools,
     solana_fee::FeeFeatures,
-    solana_ledger::token_balances::collect_token_balances,
     solana_measure::measure_us,
     solana_poh::{
         poh_recorder::PohRecorderError,
@@ -232,18 +231,6 @@ impl Consumer {
         let transaction_status_sender_enabled = self.committer.transaction_status_sender_enabled();
         let mut execute_and_commit_timings = LeaderExecuteAndCommitTimings::default();
 
-        let mut pre_balance_info = PreBalanceInfo::default();
-        let (_, collect_balances_us) = measure_us!({
-            // If the extra meta-data services are enabled for RPC, collect the
-            // pre-balances for native and token programs.
-            if transaction_status_sender_enabled {
-                pre_balance_info.native = bank.collect_balances(batch);
-                pre_balance_info.token =
-                    collect_token_balances(bank, batch, &mut pre_balance_info.mint_decimals)
-            }
-        });
-        execute_and_commit_timings.collect_balances_us = collect_balances_us;
-
         let min_max = batch
             .sanitized_transactions()
             .iter()
@@ -316,6 +303,7 @@ impl Consumer {
         let LoadAndExecuteTransactionsOutput {
             processing_results,
             processed_counts,
+            balance_collector,
         } = load_and_execute_transactions_output;
 
         let actual_execute_time = execute_and_commit_timings
@@ -404,7 +392,7 @@ impl Consumer {
                     processing_results,
                     starting_transaction_index,
                     bank,
-                    &mut pre_balance_info,
+                    balance_collector,
                     &mut execute_and_commit_timings,
                     &processed_counts,
                 )

+ 0 - 3
core/src/banking_stage/leader_slot_timing_metrics.rs

@@ -5,7 +5,6 @@ use {
 
 #[derive(Default, Debug)]
 pub struct LeaderExecuteAndCommitTimings {
-    pub collect_balances_us: u64,
     pub load_execute_us: u64,
     pub freeze_lock_us: u64,
     pub record_us: u64,
@@ -17,7 +16,6 @@ pub struct LeaderExecuteAndCommitTimings {
 
 impl LeaderExecuteAndCommitTimings {
     pub fn accumulate(&mut self, other: &LeaderExecuteAndCommitTimings) {
-        self.collect_balances_us += other.collect_balances_us;
         self.load_execute_us += other.load_execute_us;
         self.freeze_lock_us += other.freeze_lock_us;
         self.record_us += other.record_us;
@@ -32,7 +30,6 @@ impl LeaderExecuteAndCommitTimings {
         datapoint_info!(
             "banking_stage-leader_slot_vote_execute_and_commit_timings",
             ("slot", slot as i64, i64),
-            ("collect_balances_us", self.collect_balances_us as i64, i64),
             ("load_execute_us", self.load_execute_us as i64, i64),
             ("freeze_lock_us", self.freeze_lock_us as i64, i64),
             ("record_us", self.record_us as i64, i64),

+ 0 - 1
core/tests/scheduler_cost_adjustment.rs

@@ -109,7 +109,6 @@ impl TestSetup {
             .load_execute_and_commit_transactions(
                 &batch,
                 MAX_PROCESSING_AGE,
-                false,
                 ExecutionRecordingConfig::new_single_setting(false),
                 &mut ExecuteTimings::default(),
                 None,

+ 1 - 2
ledger/Cargo.toml

@@ -75,8 +75,6 @@ solana-transaction-context = { workspace = true }
 solana-transaction-status = { workspace = true }
 solana-vote = { workspace = true }
 solana-vote-program = { workspace = true }
-spl-token = { workspace = true, features = ["no-entrypoint"] }
-spl-token-2022 = { workspace = true, features = ["no-entrypoint"] }
 static_assertions = { workspace = true }
 strum = { workspace = true, features = ["derive"] }
 strum_macros = { workspace = true }
@@ -107,6 +105,7 @@ solana-program-option = { workspace = true }
 solana-program-pack = { workspace = true }
 solana-runtime = { workspace = true, features = ["dev-context-only-utils"] }
 solana-vote = { workspace = true, features = ["dev-context-only-utils"] }
+spl-generic-token = { workspace = true }
 spl-pod = { workspace = true }
 test-case = { workspace = true }
 

+ 12 - 19
ledger/src/blockstore_processor.rs

@@ -5,7 +5,7 @@ use {
         blockstore_meta::SlotMeta,
         entry_notifier_service::{EntryNotification, EntryNotifierSender},
         leader_schedule_cache::LeaderScheduleCache,
-        token_balances::collect_token_balances,
+        transaction_balances::compile_collected_balances,
         use_snapshot_archives_at_startup::UseSnapshotArchivesAtStartup,
     },
     chrono_humanize::{Accuracy, HumanTime, Tense},
@@ -177,15 +177,7 @@ pub fn execute_batch<'a>(
     //   None    => block verification path(s)
     let block_verification = extra_pre_commit_callback.is_none();
     let record_transaction_meta = transaction_status_sender.is_some();
-
     let mut transaction_indexes = Cow::from(transaction_indexes);
-    let mut mint_decimals: HashMap<Pubkey, u8> = HashMap::new();
-
-    let pre_token_balances = if record_transaction_meta {
-        collect_token_balances(bank, batch, &mut mint_decimals)
-    } else {
-        vec![]
-    };
 
     let pre_commit_callback = |_timings: &mut _, processing_results: &_| -> PreCommitResult {
         match extra_pre_commit_callback {
@@ -235,12 +227,11 @@ pub fn execute_batch<'a>(
         }
     };
 
-    let (commit_results, balances) = batch
+    let (commit_results, balance_collector) = batch
         .bank()
         .load_execute_and_commit_transactions_with_pre_commit_callback(
             batch,
             MAX_PROCESSING_AGE,
-            transaction_status_sender.is_some(),
             ExecutionRecordingConfig::new_single_setting(transaction_status_sender.is_some()),
             timings,
             log_messages_bytes_limit,
@@ -291,14 +282,17 @@ pub fn execute_batch<'a>(
             .iter()
             .map(|tx| tx.as_sanitized_transaction().into_owned())
             .collect();
-        let post_token_balances = if record_transaction_meta {
-            collect_token_balances(bank, batch, &mut mint_decimals)
-        } else {
-            vec![]
-        };
 
-        let token_balances =
-            TransactionTokenBalancesSet::new(pre_token_balances, post_token_balances);
+        // There are two cases where balance_collector could be None:
+        // * Balance recording is disabled. If that were the case, there would
+        //   be no TransactionStatusSender, and we would not be in this branch.
+        // * The batch was aborted in its entirety in SVM. In that case, nothing
+        //   would have been committed.
+        // Therefore this should always be true.
+        debug_assert!(balance_collector.is_some());
+
+        let (balances, token_balances) =
+            compile_collected_balances(balance_collector.unwrap_or_default());
 
         // The length of costs vector needs to be consistent with all other
         // vectors that are sent over (such as `transactions`). So, replace the
@@ -4429,7 +4423,6 @@ pub mod tests {
         let (commit_results, _) = batch.bank().load_execute_and_commit_transactions(
             &batch,
             MAX_PROCESSING_AGE,
-            false,
             ExecutionRecordingConfig::new_single_setting(false),
             &mut ExecuteTimings::default(),
             None,

+ 1 - 1
ledger/src/lib.rs

@@ -31,8 +31,8 @@ mod shredder;
 pub mod sigverify_shreds;
 pub mod slot_stats;
 mod staking_utils;
-pub mod token_balances;
 mod transaction_address_lookup_table_scanner;
+pub mod transaction_balances;
 pub mod use_snapshot_archives_at_startup;
 
 #[macro_use]

+ 0 - 486
ledger/src/token_balances.rs

@@ -1,486 +0,0 @@
-use {
-    solana_account_decoder::{
-        parse_account_data::SplTokenAdditionalDataV2,
-        parse_token::{is_known_spl_token_id, token_amount_to_ui_amount_v3, UiTokenAmount},
-    },
-    solana_measure::measure::Measure,
-    solana_metrics::datapoint_debug,
-    solana_runtime::{bank::Bank, transaction_batch::TransactionBatch},
-    solana_sdk::{account::ReadableAccount, pubkey::Pubkey},
-    solana_svm_transaction::svm_message::SVMMessage,
-    solana_transaction_status::{
-        token_balances::TransactionTokenBalances, TransactionTokenBalance,
-    },
-    spl_token_2022::{
-        extension::StateWithExtensions,
-        state::{Account as TokenAccount, Mint},
-    },
-    std::collections::HashMap,
-};
-
-fn get_mint_decimals(bank: &Bank, mint: &Pubkey) -> Option<u8> {
-    if mint == &spl_token::native_mint::id() {
-        Some(spl_token::native_mint::DECIMALS)
-    } else {
-        let mint_account = bank.get_account(mint)?;
-
-        if !is_known_spl_token_id(mint_account.owner()) {
-            return None;
-        }
-
-        let decimals = StateWithExtensions::<Mint>::unpack(mint_account.data())
-            .map(|mint| mint.base.decimals)
-            .ok()?;
-
-        Some(decimals)
-    }
-}
-
-pub fn collect_token_balances(
-    bank: &Bank,
-    batch: &TransactionBatch<impl SVMMessage>,
-    mint_decimals: &mut HashMap<Pubkey, u8>,
-) -> TransactionTokenBalances {
-    let mut balances: TransactionTokenBalances = vec![];
-    let mut collect_time = Measure::start("collect_token_balances");
-
-    for transaction in batch.sanitized_transactions() {
-        let account_keys = transaction.account_keys();
-        let has_token_program = account_keys.iter().any(is_known_spl_token_id);
-
-        let mut transaction_balances: Vec<TransactionTokenBalance> = vec![];
-        if has_token_program {
-            for (index, account_id) in account_keys.iter().enumerate() {
-                if transaction.is_invoked(index) || is_known_spl_token_id(account_id) {
-                    continue;
-                }
-
-                if let Some(TokenBalanceData {
-                    mint,
-                    ui_token_amount,
-                    owner,
-                    program_id,
-                }) = collect_token_balance_from_account(bank, account_id, mint_decimals)
-                {
-                    transaction_balances.push(TransactionTokenBalance {
-                        account_index: index as u8,
-                        mint,
-                        ui_token_amount,
-                        owner,
-                        program_id,
-                    });
-                }
-            }
-        }
-        balances.push(transaction_balances);
-    }
-    collect_time.stop();
-    datapoint_debug!(
-        "collect_token_balances",
-        ("collect_time_us", collect_time.as_us(), i64),
-    );
-    balances
-}
-
-#[derive(Debug, PartialEq)]
-struct TokenBalanceData {
-    mint: String,
-    owner: String,
-    ui_token_amount: UiTokenAmount,
-    program_id: String,
-}
-
-fn collect_token_balance_from_account(
-    bank: &Bank,
-    account_id: &Pubkey,
-    mint_decimals: &mut HashMap<Pubkey, u8>,
-) -> Option<TokenBalanceData> {
-    let account = bank.get_account(account_id)?;
-
-    if !is_known_spl_token_id(account.owner()) {
-        return None;
-    }
-
-    let token_account = StateWithExtensions::<TokenAccount>::unpack(account.data()).ok()?;
-    let mint = token_account.base.mint;
-
-    let decimals = mint_decimals.get(&mint).cloned().or_else(|| {
-        let decimals = get_mint_decimals(bank, &mint)?;
-        mint_decimals.insert(mint, decimals);
-        Some(decimals)
-    })?;
-
-    Some(TokenBalanceData {
-        mint: token_account.base.mint.to_string(),
-        owner: token_account.base.owner.to_string(),
-        ui_token_amount: token_amount_to_ui_amount_v3(
-            token_account.base.amount,
-            // NOTE: Same as parsed instruction data, ledger data always uses
-            // the raw token amount, and does not calculate the UI amount with
-            // any consideration for interest.
-            &SplTokenAdditionalDataV2::with_decimals(decimals),
-        ),
-        program_id: account.owner().to_string(),
-    })
-}
-
-#[cfg(test)]
-mod test {
-    use {
-        super::*,
-        solana_program_option::COption,
-        solana_program_pack::Pack,
-        solana_sdk::{account::Account, genesis_config::create_genesis_config},
-        spl_pod::optional_keys::OptionalNonZeroPubkey,
-        spl_token_2022::extension::{
-            immutable_owner::ImmutableOwner, memo_transfer::MemoTransfer,
-            mint_close_authority::MintCloseAuthority, BaseStateWithExtensionsMut, ExtensionType,
-            StateWithExtensionsMut,
-        },
-        std::collections::BTreeMap,
-    };
-
-    #[test]
-    fn test_collect_token_balance_from_account() {
-        let (mut genesis_config, _mint_keypair) = create_genesis_config(500);
-
-        // Add a variety of accounts, token and not
-        let account = Account::new(42, 55, &Pubkey::new_unique());
-
-        let mint_data = Mint {
-            mint_authority: COption::None,
-            supply: 4242,
-            decimals: 2,
-            is_initialized: true,
-            freeze_authority: COption::None,
-        };
-        let mut data = [0; Mint::LEN];
-        Mint::pack(mint_data, &mut data).unwrap();
-        let mint_pubkey = Pubkey::new_unique();
-        let mint = Account {
-            lamports: 100,
-            data: data.to_vec(),
-            owner: spl_token::id(),
-            executable: false,
-            rent_epoch: 0,
-        };
-        let other_mint_pubkey = Pubkey::new_unique();
-        let other_mint = Account {
-            lamports: 100,
-            data: data.to_vec(),
-            owner: Pubkey::new_unique(), // !is_known_spl_token_id
-            executable: false,
-            rent_epoch: 0,
-        };
-
-        let token_owner = Pubkey::new_unique();
-        let token_data = TokenAccount {
-            mint: mint_pubkey,
-            owner: token_owner,
-            amount: 42,
-            delegate: COption::None,
-            state: spl_token_2022::state::AccountState::Initialized,
-            is_native: COption::Some(100),
-            delegated_amount: 0,
-            close_authority: COption::None,
-        };
-        let mut data = [0; TokenAccount::LEN];
-        TokenAccount::pack(token_data, &mut data).unwrap();
-
-        let spl_token_account = Account {
-            lamports: 100,
-            data: data.to_vec(),
-            owner: spl_token::id(),
-            executable: false,
-            rent_epoch: 0,
-        };
-        let other_account = Account {
-            lamports: 100,
-            data: data.to_vec(),
-            owner: Pubkey::new_unique(), // !is_known_spl_token_id
-            executable: false,
-            rent_epoch: 0,
-        };
-
-        let other_mint_data = TokenAccount {
-            mint: other_mint_pubkey,
-            owner: token_owner,
-            amount: 42,
-            delegate: COption::None,
-            state: spl_token_2022::state::AccountState::Initialized,
-            is_native: COption::Some(100),
-            delegated_amount: 0,
-            close_authority: COption::None,
-        };
-        let mut data = [0; TokenAccount::LEN];
-        TokenAccount::pack(other_mint_data, &mut data).unwrap();
-
-        let other_mint_token_account = Account {
-            lamports: 100,
-            data: data.to_vec(),
-            owner: spl_token::id(),
-            executable: false,
-            rent_epoch: 0,
-        };
-
-        let mut accounts = BTreeMap::new();
-
-        let account_pubkey = Pubkey::new_unique();
-        accounts.insert(account_pubkey, account);
-        accounts.insert(mint_pubkey, mint);
-        accounts.insert(other_mint_pubkey, other_mint);
-        let spl_token_account_pubkey = Pubkey::new_unique();
-        accounts.insert(spl_token_account_pubkey, spl_token_account);
-        let other_account_pubkey = Pubkey::new_unique();
-        accounts.insert(other_account_pubkey, other_account);
-        let other_mint_account_pubkey = Pubkey::new_unique();
-        accounts.insert(other_mint_account_pubkey, other_mint_token_account);
-
-        genesis_config.accounts = accounts;
-
-        let bank = Bank::new_for_tests(&genesis_config);
-        let mut mint_decimals = HashMap::new();
-
-        // Account is not owned by spl_token (nor does it have TokenAccount state)
-        assert_eq!(
-            collect_token_balance_from_account(&bank, &account_pubkey, &mut mint_decimals),
-            None
-        );
-
-        // Mint does not have TokenAccount state
-        assert_eq!(
-            collect_token_balance_from_account(&bank, &mint_pubkey, &mut mint_decimals),
-            None
-        );
-
-        // TokenAccount owned by spl_token::id() works
-        assert_eq!(
-            collect_token_balance_from_account(
-                &bank,
-                &spl_token_account_pubkey,
-                &mut mint_decimals
-            ),
-            Some(TokenBalanceData {
-                mint: mint_pubkey.to_string(),
-                owner: token_owner.to_string(),
-                ui_token_amount: UiTokenAmount {
-                    ui_amount: Some(0.42),
-                    decimals: 2,
-                    amount: "42".to_string(),
-                    ui_amount_string: "0.42".to_string(),
-                },
-                program_id: spl_token::id().to_string(),
-            })
-        );
-
-        // TokenAccount is not owned by known spl-token program_id
-        assert_eq!(
-            collect_token_balance_from_account(&bank, &other_account_pubkey, &mut mint_decimals),
-            None
-        );
-
-        // TokenAccount's mint is not owned by known spl-token program_id
-        assert_eq!(
-            collect_token_balance_from_account(
-                &bank,
-                &other_mint_account_pubkey,
-                &mut mint_decimals
-            ),
-            None
-        );
-    }
-
-    #[test]
-    fn test_collect_token_balance_from_spl_token_2022_account() {
-        let (mut genesis_config, _mint_keypair) = create_genesis_config(500);
-
-        // Add a variety of accounts, token and not
-        let account = Account::new(42, 55, &Pubkey::new_unique());
-
-        let mint_authority = Pubkey::new_unique();
-        let mint_size =
-            ExtensionType::try_calculate_account_len::<Mint>(&[ExtensionType::MintCloseAuthority])
-                .unwrap();
-        let mint_base = Mint {
-            mint_authority: COption::None,
-            supply: 4242,
-            decimals: 2,
-            is_initialized: true,
-            freeze_authority: COption::None,
-        };
-        let mut mint_data = vec![0; mint_size];
-        let mut mint_state =
-            StateWithExtensionsMut::<Mint>::unpack_uninitialized(&mut mint_data).unwrap();
-        mint_state.base = mint_base;
-        mint_state.pack_base();
-        mint_state.init_account_type().unwrap();
-        let mint_close_authority = mint_state
-            .init_extension::<MintCloseAuthority>(true)
-            .unwrap();
-        mint_close_authority.close_authority =
-            OptionalNonZeroPubkey::try_from(Some(mint_authority)).unwrap();
-
-        let mint_pubkey = Pubkey::new_unique();
-        let mint = Account {
-            lamports: 100,
-            data: mint_data.to_vec(),
-            owner: spl_token_2022::id(),
-            executable: false,
-            rent_epoch: 0,
-        };
-        let other_mint_pubkey = Pubkey::new_unique();
-        let other_mint = Account {
-            lamports: 100,
-            data: mint_data.to_vec(),
-            owner: Pubkey::new_unique(),
-            executable: false,
-            rent_epoch: 0,
-        };
-
-        let token_owner = Pubkey::new_unique();
-        let token_base = TokenAccount {
-            mint: mint_pubkey,
-            owner: token_owner,
-            amount: 42,
-            delegate: COption::None,
-            state: spl_token_2022::state::AccountState::Initialized,
-            is_native: COption::Some(100),
-            delegated_amount: 0,
-            close_authority: COption::None,
-        };
-        let account_size = ExtensionType::try_calculate_account_len::<TokenAccount>(&[
-            ExtensionType::ImmutableOwner,
-            ExtensionType::MemoTransfer,
-        ])
-        .unwrap();
-        let mut account_data = vec![0; account_size];
-        let mut account_state =
-            StateWithExtensionsMut::<TokenAccount>::unpack_uninitialized(&mut account_data)
-                .unwrap();
-        account_state.base = token_base;
-        account_state.pack_base();
-        account_state.init_account_type().unwrap();
-        account_state
-            .init_extension::<ImmutableOwner>(true)
-            .unwrap();
-        let memo_transfer = account_state.init_extension::<MemoTransfer>(true).unwrap();
-        memo_transfer.require_incoming_transfer_memos = true.into();
-
-        let spl_token_account = Account {
-            lamports: 100,
-            data: account_data.to_vec(),
-            owner: spl_token_2022::id(),
-            executable: false,
-            rent_epoch: 0,
-        };
-        let other_account = Account {
-            lamports: 100,
-            data: account_data.to_vec(),
-            owner: Pubkey::new_unique(),
-            executable: false,
-            rent_epoch: 0,
-        };
-
-        let other_mint_token_base = TokenAccount {
-            mint: other_mint_pubkey,
-            owner: token_owner,
-            amount: 42,
-            delegate: COption::None,
-            state: spl_token_2022::state::AccountState::Initialized,
-            is_native: COption::Some(100),
-            delegated_amount: 0,
-            close_authority: COption::None,
-        };
-        let account_size = ExtensionType::try_calculate_account_len::<TokenAccount>(&[
-            ExtensionType::ImmutableOwner,
-            ExtensionType::MemoTransfer,
-        ])
-        .unwrap();
-        let mut account_data = vec![0; account_size];
-        let mut account_state =
-            StateWithExtensionsMut::<TokenAccount>::unpack_uninitialized(&mut account_data)
-                .unwrap();
-        account_state.base = other_mint_token_base;
-        account_state.pack_base();
-        account_state.init_account_type().unwrap();
-        account_state
-            .init_extension::<ImmutableOwner>(true)
-            .unwrap();
-        let memo_transfer = account_state.init_extension::<MemoTransfer>(true).unwrap();
-        memo_transfer.require_incoming_transfer_memos = true.into();
-
-        let other_mint_token_account = Account {
-            lamports: 100,
-            data: account_data.to_vec(),
-            owner: spl_token_2022::id(),
-            executable: false,
-            rent_epoch: 0,
-        };
-
-        let mut accounts = BTreeMap::new();
-
-        let account_pubkey = Pubkey::new_unique();
-        accounts.insert(account_pubkey, account);
-        accounts.insert(mint_pubkey, mint);
-        accounts.insert(other_mint_pubkey, other_mint);
-        let spl_token_account_pubkey = Pubkey::new_unique();
-        accounts.insert(spl_token_account_pubkey, spl_token_account);
-        let other_account_pubkey = Pubkey::new_unique();
-        accounts.insert(other_account_pubkey, other_account);
-        let other_mint_account_pubkey = Pubkey::new_unique();
-        accounts.insert(other_mint_account_pubkey, other_mint_token_account);
-
-        genesis_config.accounts = accounts;
-
-        let bank = Bank::new_for_tests(&genesis_config);
-        let mut mint_decimals = HashMap::new();
-
-        // Account is not owned by spl_token (nor does it have TokenAccount state)
-        assert_eq!(
-            collect_token_balance_from_account(&bank, &account_pubkey, &mut mint_decimals),
-            None
-        );
-
-        // Mint does not have TokenAccount state
-        assert_eq!(
-            collect_token_balance_from_account(&bank, &mint_pubkey, &mut mint_decimals),
-            None
-        );
-
-        // TokenAccount owned by spl_token_2022::id() works
-        assert_eq!(
-            collect_token_balance_from_account(
-                &bank,
-                &spl_token_account_pubkey,
-                &mut mint_decimals
-            ),
-            Some(TokenBalanceData {
-                mint: mint_pubkey.to_string(),
-                owner: token_owner.to_string(),
-                ui_token_amount: UiTokenAmount {
-                    ui_amount: Some(0.42),
-                    decimals: 2,
-                    amount: "42".to_string(),
-                    ui_amount_string: "0.42".to_string(),
-                },
-                program_id: spl_token_2022::id().to_string(),
-            })
-        );
-
-        // TokenAccount is not owned by known spl-token program_id
-        assert_eq!(
-            collect_token_balance_from_account(&bank, &other_account_pubkey, &mut mint_decimals),
-            None
-        );
-
-        // TokenAccount's mint is not owned by known spl-token program_id
-        assert_eq!(
-            collect_token_balance_from_account(
-                &bank,
-                &other_mint_account_pubkey,
-                &mut mint_decimals
-            ),
-            None
-        );
-    }
-}

+ 161 - 0
ledger/src/transaction_balances.rs

@@ -0,0 +1,161 @@
+use {
+    solana_account_decoder::{
+        parse_account_data::SplTokenAdditionalDataV2, parse_token::token_amount_to_ui_amount_v3,
+    },
+    solana_runtime::bank::TransactionBalancesSet,
+    solana_svm::transaction_balances::{BalanceCollector, SvmTokenInfo},
+    solana_transaction_status::{
+        token_balances::TransactionTokenBalancesSet, TransactionTokenBalance,
+    },
+};
+
+// decompose the contents of BalanceCollector into the two structs required by TransactionStatusSender
+pub fn compile_collected_balances(
+    balance_collector: BalanceCollector,
+) -> (TransactionBalancesSet, TransactionTokenBalancesSet) {
+    let (native_pre, native_post, token_pre, token_post) = balance_collector.into_vecs();
+
+    let native_balances = TransactionBalancesSet::new(native_pre, native_post);
+    let token_balances = TransactionTokenBalancesSet::new(
+        collected_token_infos_to_token_balances(token_pre),
+        collected_token_infos_to_token_balances(token_post),
+    );
+
+    (native_balances, token_balances)
+}
+
+fn collected_token_infos_to_token_balances(
+    svm_infos: Vec<Vec<SvmTokenInfo>>,
+) -> Vec<Vec<TransactionTokenBalance>> {
+    svm_infos
+        .into_iter()
+        .map(|infos| {
+            infos
+                .into_iter()
+                .map(svm_token_info_to_token_balance)
+                .collect()
+        })
+        .collect()
+}
+
+fn svm_token_info_to_token_balance(svm_info: SvmTokenInfo) -> TransactionTokenBalance {
+    let SvmTokenInfo {
+        account_index,
+        mint,
+        amount,
+        owner,
+        program_id,
+        decimals,
+    } = svm_info;
+    TransactionTokenBalance {
+        account_index,
+        mint: mint.to_string(),
+        ui_token_amount: token_amount_to_ui_amount_v3(
+            amount,
+            // NOTE: Same as parsed instruction data, ledger data always uses
+            // the raw token amount, and does not calculate the UI amount with
+            // any consideration for interest.
+            &SplTokenAdditionalDataV2::with_decimals(decimals),
+        ),
+        owner: owner.to_string(),
+        program_id: program_id.to_string(),
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use {
+        super::*,
+        solana_account_decoder::parse_token::UiTokenAmount,
+        solana_pubkey::Pubkey,
+        spl_generic_token::{token, token_2022},
+    };
+
+    #[test]
+    fn test_compile_collected_balances() {
+        let native_pre = vec![vec![1, 2, 3], vec![4, 5, 6]];
+        let native_post = vec![vec![7, 8, 9], vec![10, 11, 0]];
+
+        let account_index = 1;
+        let mint1 = Pubkey::new_unique();
+        let mint2 = Pubkey::new_unique();
+        let owner1 = Pubkey::new_unique();
+        let owner2 = Pubkey::new_unique();
+        let amount1 = 10;
+        let amount2 = 200;
+        let decimals1 = 1;
+        let decimals2 = 2;
+
+        let token_info_before = SvmTokenInfo {
+            account_index,
+            mint: mint1,
+            amount: amount1,
+            owner: owner1,
+            program_id: token::id(),
+            decimals: decimals1,
+        };
+        let token_info_after = SvmTokenInfo {
+            account_index,
+            mint: mint2,
+            amount: amount2,
+            owner: owner2,
+            program_id: token_2022::id(),
+            decimals: decimals2,
+        };
+
+        let token_pre = vec![vec![token_info_before], vec![]];
+        let token_post = vec![vec![token_info_after], vec![]];
+
+        let token_balance_before = TransactionTokenBalance {
+            account_index,
+            mint: mint1.to_string(),
+            ui_token_amount: UiTokenAmount {
+                ui_amount: Some(1.0),
+                decimals: decimals1,
+                amount: amount1.to_string(),
+                ui_amount_string: "1".to_string(),
+            },
+            owner: owner1.to_string(),
+            program_id: token::id().to_string(),
+        };
+        let token_balance_after = TransactionTokenBalance {
+            account_index,
+            mint: mint2.to_string(),
+            ui_token_amount: UiTokenAmount {
+                ui_amount: Some(2.0),
+                decimals: decimals2,
+                amount: amount2.to_string(),
+                ui_amount_string: "2".to_string(),
+            },
+            owner: owner2.to_string(),
+            program_id: token_2022::id().to_string(),
+        };
+
+        let expected_native = TransactionBalancesSet::new(native_pre.clone(), native_post.clone());
+        let expected_token = TransactionTokenBalancesSet::new(
+            vec![vec![token_balance_before], vec![]],
+            vec![vec![token_balance_after], vec![]],
+        );
+
+        let balance_collector = BalanceCollector {
+            native_pre,
+            native_post,
+            token_pre,
+            token_post,
+        };
+
+        let (actual_native, actual_token) = compile_collected_balances(balance_collector);
+
+        assert_eq!(expected_native.pre_balances, actual_native.pre_balances);
+        assert_eq!(expected_native.post_balances, actual_native.post_balances);
+
+        assert_eq!(
+            expected_token.pre_token_balances,
+            actual_token.pre_token_balances
+        );
+        assert_eq!(
+            expected_token.post_token_balances,
+            actual_token.post_token_balances
+        );
+    }
+}

+ 1 - 2
programs/sbf/Cargo.lock

@@ -6638,8 +6638,6 @@ dependencies = [
  "solana-transaction-status",
  "solana-vote",
  "solana-vote-program",
- "spl-token",
- "spl-token-2022",
  "static_assertions",
  "strum",
  "strum_macros",
@@ -8681,6 +8679,7 @@ dependencies = [
  "solana-transaction-context 2.3.0",
  "solana-transaction-error",
  "solana-type-overrides",
+ "spl-generic-token",
  "thiserror 2.0.12",
 ]
 

+ 2 - 2
programs/sbf/tests/programs.rs

@@ -97,11 +97,11 @@ fn load_execute_and_commit_transaction(bank: &Bank, tx: Transaction) -> Transact
         .load_execute_and_commit_transactions(
             &tx_batch,
             MAX_PROCESSING_AGE,
-            false,
             ExecutionRecordingConfig {
                 enable_cpi_recording: true,
                 enable_log_recording: true,
                 enable_return_data_recording: false,
+                enable_transaction_balance_recording: false,
             },
             &mut ExecuteTimings::default(),
             None,
@@ -5067,11 +5067,11 @@ fn test_function_call_args() {
         .load_execute_and_commit_transactions(
             &tx_batch,
             MAX_PROCESSING_AGE,
-            false,
             ExecutionRecordingConfig {
                 enable_cpi_recording: false,
                 enable_log_recording: false,
                 enable_return_data_recording: true,
+                enable_transaction_balance_recording: false,
             },
             &mut ExecuteTimings::default(),
             None,

+ 12 - 25
runtime/src/bank.rs

@@ -155,6 +155,7 @@ use {
         account_loader::{collect_rent_from_account, LoadedTransaction},
         account_overrides::AccountOverrides,
         program_loader::load_program_with_pubkey,
+        transaction_balances::BalanceCollector,
         transaction_commit_result::{CommittedTransaction, TransactionCommitResult},
         transaction_error_metrics::TransactionErrorMetrics,
         transaction_execution_result::{
@@ -327,6 +328,9 @@ pub struct LoadAndExecuteTransactionsOutput {
     // Processed transaction counts used to update bank transaction counts and
     // for metrics reporting.
     pub processed_counts: ProcessedTransactionCounts,
+    // Balances accumulated for TransactionStatusSender when transaction
+    // balance recording is enabled.
+    pub balance_collector: Option<BalanceCollector>,
 }
 
 #[derive(Debug, PartialEq)]
@@ -3290,6 +3294,7 @@ impl Bank {
                     enable_cpi_recording,
                     enable_log_recording: true,
                     enable_return_data_recording: true,
+                    enable_transaction_balance_recording: false,
                 },
             },
         );
@@ -3498,6 +3503,7 @@ impl Bank {
         LoadAndExecuteTransactionsOutput {
             processing_results: sanitized_output.processing_results,
             processed_counts,
+            balance_collector: sanitized_output.balance_collector,
         }
     }
 
@@ -4553,15 +4559,13 @@ impl Bank {
         &self,
         batch: &TransactionBatch<impl TransactionWithMeta>,
         max_age: usize,
-        collect_balances: bool,
         recording_config: ExecutionRecordingConfig,
         timings: &mut ExecuteTimings,
         log_messages_bytes_limit: Option<usize>,
-    ) -> (Vec<TransactionCommitResult>, TransactionBalancesSet) {
+    ) -> (Vec<TransactionCommitResult>, Option<BalanceCollector>) {
         self.do_load_execute_and_commit_transactions_with_pre_commit_callback(
             batch,
             max_age,
-            collect_balances,
             recording_config,
             timings,
             log_messages_bytes_limit,
@@ -4574,7 +4578,6 @@ impl Bank {
         &'a self,
         batch: &TransactionBatch<impl TransactionWithMeta>,
         max_age: usize,
-        collect_balances: bool,
         recording_config: ExecutionRecordingConfig,
         timings: &mut ExecuteTimings,
         log_messages_bytes_limit: Option<usize>,
@@ -4582,11 +4585,10 @@ impl Bank {
             &mut ExecuteTimings,
             &[TransactionProcessingResult],
         ) -> PreCommitResult<'a>,
-    ) -> Result<(Vec<TransactionCommitResult>, TransactionBalancesSet)> {
+    ) -> Result<(Vec<TransactionCommitResult>, Option<BalanceCollector>)> {
         self.do_load_execute_and_commit_transactions_with_pre_commit_callback(
             batch,
             max_age,
-            collect_balances,
             recording_config,
             timings,
             log_messages_bytes_limit,
@@ -4598,23 +4600,17 @@ impl Bank {
         &'a self,
         batch: &TransactionBatch<impl TransactionWithMeta>,
         max_age: usize,
-        collect_balances: bool,
         recording_config: ExecutionRecordingConfig,
         timings: &mut ExecuteTimings,
         log_messages_bytes_limit: Option<usize>,
         pre_commit_callback: Option<
             impl FnOnce(&mut ExecuteTimings, &[TransactionProcessingResult]) -> PreCommitResult<'a>,
         >,
-    ) -> Result<(Vec<TransactionCommitResult>, TransactionBalancesSet)> {
-        let pre_balances = if collect_balances {
-            self.collect_balances(batch)
-        } else {
-            vec![]
-        };
-
+    ) -> Result<(Vec<TransactionCommitResult>, Option<BalanceCollector>)> {
         let LoadAndExecuteTransactionsOutput {
             processing_results,
             processed_counts,
+            balance_collector,
         } = self.load_and_execute_transactions(
             batch,
             max_age,
@@ -4645,15 +4641,7 @@ impl Bank {
             timings,
         );
         drop(freeze_lock);
-        let post_balances = if collect_balances {
-            self.collect_balances(batch)
-        } else {
-            vec![]
-        };
-        Ok((
-            commit_results,
-            TransactionBalancesSet::new(pre_balances, post_balances),
-        ))
+        Ok((commit_results, balance_collector))
     }
 
     /// Process a Transaction. This is used for unit tests and simply calls the vector
@@ -4677,11 +4665,11 @@ impl Bank {
         let (mut commit_results, ..) = self.load_execute_and_commit_transactions(
             &batch,
             MAX_PROCESSING_AGE,
-            false, // collect_balances
             ExecutionRecordingConfig {
                 enable_cpi_recording: false,
                 enable_log_recording: true,
                 enable_return_data_recording: true,
+                enable_transaction_balance_recording: false,
             },
             &mut ExecuteTimings::default(),
             Some(1000 * 1000),
@@ -4720,7 +4708,6 @@ impl Bank {
         self.load_execute_and_commit_transactions(
             batch,
             MAX_PROCESSING_AGE,
-            false,
             ExecutionRecordingConfig::new_single_setting(false),
             &mut ExecuteTimings::default(),
             None,

+ 12 - 9
runtime/src/bank/tests.rs

@@ -3164,7 +3164,6 @@ fn test_interleaving_locks() {
         .load_execute_and_commit_transactions(
             &lock_result,
             MAX_PROCESSING_AGE,
-            false,
             ExecutionRecordingConfig::new_single_setting(false),
             &mut ExecuteTimings::default(),
             None,
@@ -3250,7 +3249,6 @@ fn test_load_and_execute_commit_transactions_fees_only() {
         .load_execute_and_commit_transactions(
             &batch,
             MAX_PROCESSING_AGE,
-            true,
             ExecutionRecordingConfig::new_single_setting(true),
             &mut ExecuteTimings::default(),
             None,
@@ -5977,15 +5975,22 @@ fn test_pre_post_transaction_balances() {
     let txs = vec![tx0, tx1, tx2];
 
     let lock_result = bank0.prepare_batch_for_tests(txs);
-    let (commit_results, transaction_balances_set) = bank0.load_execute_and_commit_transactions(
+    let (commit_results, balance_collector) = bank0.load_execute_and_commit_transactions(
         &lock_result,
         MAX_PROCESSING_AGE,
-        true,
-        ExecutionRecordingConfig::new_single_setting(false),
+        ExecutionRecordingConfig {
+            enable_cpi_recording: false,
+            enable_log_recording: false,
+            enable_return_data_recording: false,
+            enable_transaction_balance_recording: true,
+        },
         &mut ExecuteTimings::default(),
         None,
     );
 
+    let (native_pre, native_post, _, _) = balance_collector.unwrap().into_vecs();
+    let transaction_balances_set = TransactionBalancesSet::new(native_pre, native_post);
+
     assert_eq!(transaction_balances_set.pre_balances.len(), 3);
     assert_eq!(transaction_balances_set.post_balances.len(), 3);
 
@@ -9456,11 +9461,11 @@ fn test_tx_log_order() {
         .load_execute_and_commit_transactions(
             &batch,
             MAX_PROCESSING_AGE,
-            false,
             ExecutionRecordingConfig {
                 enable_cpi_recording: false,
                 enable_log_recording: true,
                 enable_return_data_recording: false,
+                enable_transaction_balance_recording: false,
             },
             &mut ExecuteTimings::default(),
             None,
@@ -9564,11 +9569,11 @@ fn test_tx_return_data() {
             .load_execute_and_commit_transactions(
                 &batch,
                 MAX_PROCESSING_AGE,
-                false,
                 ExecutionRecordingConfig {
                     enable_cpi_recording: false,
                     enable_log_recording: false,
                     enable_return_data_recording: true,
+                    enable_transaction_balance_recording: false,
                 },
                 &mut ExecuteTimings::default(),
                 None,
@@ -9616,7 +9621,6 @@ fn test_load_and_execute_commit_transactions_rent_debits() {
             .load_execute_and_commit_transactions(
                 &batch,
                 MAX_PROCESSING_AGE,
-                false,
                 ExecutionRecordingConfig::new_single_setting(false),
                 &mut ExecuteTimings::default(),
                 None,
@@ -9640,7 +9644,6 @@ fn test_load_and_execute_commit_transactions_rent_debits() {
             .load_execute_and_commit_transactions(
                 &batch,
                 MAX_PROCESSING_AGE,
-                false,
                 ExecutionRecordingConfig::new_single_setting(false),
                 &mut ExecuteTimings::default(),
                 None,

+ 2 - 0
svm/Cargo.toml

@@ -51,6 +51,7 @@ solana-timings = { workspace = true }
 solana-transaction-context = { workspace = true }
 solana-transaction-error = { workspace = true }
 solana-type-overrides = { workspace = true }
+spl-generic-token = { workspace = true }
 thiserror = { workspace = true }
 
 [lib]
@@ -98,6 +99,7 @@ solana-system-transaction = { workspace = true }
 solana-sysvar = { workspace = true }
 solana-transaction = { workspace = true }
 solana-transaction-context = { workspace = true, features = ["dev-context-only-utils"] }
+spl-token = { workspace = true }
 test-case = { workspace = true }
 
 [package.metadata.docs.rs]

+ 1 - 2
svm/examples/Cargo.lock

@@ -6441,8 +6441,6 @@ dependencies = [
  "solana-transaction-status",
  "solana-vote",
  "solana-vote-program",
- "spl-token",
- "spl-token-2022",
  "static_assertions",
  "strum",
  "strum_macros",
@@ -7979,6 +7977,7 @@ dependencies = [
  "solana-transaction-context 2.3.0",
  "solana-transaction-error",
  "solana-type-overrides",
+ "spl-generic-token",
  "thiserror 2.0.12",
 ]
 

+ 1 - 0
svm/examples/json-rpc/server/src/rpc_process.rs

@@ -332,6 +332,7 @@ impl JsonRpcRequestProcessor {
                     enable_cpi_recording,
                     enable_log_recording: true,
                     enable_return_data_recording: true,
+                    enable_transaction_balance_recording: true,
                 },
             },
         );

+ 1 - 0
svm/src/lib.rs

@@ -8,6 +8,7 @@ pub mod nonce_info;
 pub mod program_loader;
 pub mod rollback_accounts;
 pub mod transaction_account_state_info;
+pub mod transaction_balances;
 pub mod transaction_commit_result;
 pub mod transaction_error_metrics;
 pub mod transaction_execution_result;

+ 211 - 0
svm/src/transaction_balances.rs

@@ -0,0 +1,211 @@
+#[cfg(feature = "dev-context-only-utils")]
+use qualifier_attr::field_qualifiers;
+use {
+    crate::{
+        account_loader::AccountLoader,
+        transaction_processing_callback::TransactionProcessingCallback,
+    },
+    solana_account::{AccountSharedData, ReadableAccount},
+    solana_pubkey::Pubkey,
+    solana_svm_transaction::svm_transaction::SVMTransaction,
+    spl_generic_token::{generic_token, is_known_spl_token_id},
+};
+
+// we use internal aliases for clarity, the external type aliases are often confusing
+type TxNativeBalances = Vec<u64>;
+type TxTokenBalances = Vec<SvmTokenInfo>;
+type BatchNativeBalances = Vec<TxNativeBalances>;
+type BatchTokenBalances = Vec<TxTokenBalances>;
+
+// to operate cleanly over Option<BalanceCollector> we use a trait impled on the outer and inner type
+pub(crate) trait BalanceCollectionRoutines {
+    fn collect_pre_balances<CB: TransactionProcessingCallback>(
+        &mut self,
+        account_loader: &mut AccountLoader<CB>,
+        transaction: &impl SVMTransaction,
+    );
+
+    fn collect_post_balances<CB: TransactionProcessingCallback>(
+        &mut self,
+        account_loader: &mut AccountLoader<CB>,
+        transaction: &impl SVMTransaction,
+    );
+}
+
+#[derive(Debug, Default)]
+#[cfg_attr(
+    feature = "dev-context-only-utils",
+    field_qualifiers(native_pre(pub), native_post(pub), token_pre(pub), token_post(pub),)
+)]
+pub struct BalanceCollector {
+    native_pre: BatchNativeBalances,
+    native_post: BatchNativeBalances,
+    token_pre: BatchTokenBalances,
+    token_post: BatchTokenBalances,
+}
+
+impl BalanceCollector {
+    // we always provide one vec for every transaction, even if the vecs are empty
+    pub(crate) fn new_with_transaction_count(transaction_count: usize) -> Self {
+        Self {
+            native_pre: Vec::with_capacity(transaction_count),
+            native_post: Vec::with_capacity(transaction_count),
+            token_pre: Vec::with_capacity(transaction_count),
+            token_post: Vec::with_capacity(transaction_count),
+        }
+    }
+
+    // we use this pattern to prevent anything outside svm mutating BalanceCollector internals
+    // with no public constructor, and only private fields, non-svm code can only disassemble the struct
+    pub fn into_vecs(
+        self,
+    ) -> (
+        BatchNativeBalances,
+        BatchNativeBalances,
+        BatchTokenBalances,
+        BatchTokenBalances,
+    ) {
+        (
+            self.native_pre,
+            self.native_post,
+            self.token_pre,
+            self.token_post,
+        )
+    }
+
+    // gather native lamport balances for all accounts
+    // and token balances for valid, initialized token accounts with valid, initialized mints
+    fn collect_balances<CB: TransactionProcessingCallback>(
+        &mut self,
+        account_loader: &mut AccountLoader<CB>,
+        transaction: &impl SVMTransaction,
+    ) -> (TxNativeBalances, TxTokenBalances) {
+        let mut native_balances = Vec::with_capacity(transaction.account_keys().len());
+        let mut token_balances = vec![];
+
+        let has_token_program = transaction.account_keys().iter().any(is_known_spl_token_id);
+
+        for (index, key) in transaction.account_keys().iter().enumerate() {
+            // we load as read-only to avoid triggering a bad account inspection
+            let Some(account) = account_loader
+                .load_account(key, false)
+                .map(|loaded| loaded.account)
+            else {
+                native_balances.push(0);
+                continue;
+            };
+
+            native_balances.push(account.lamports());
+
+            if has_token_program
+                && !transaction.is_invoked(index)
+                && !is_known_spl_token_id(key)
+                && is_known_spl_token_id(account.owner())
+            {
+                if let Some(token_info) =
+                    SvmTokenInfo::unpack_token_account(account_loader, &account, index)
+                {
+                    token_balances.push(token_info);
+                }
+            }
+        }
+
+        (native_balances, token_balances)
+    }
+
+    pub(crate) fn lengths_match_expected(&self, expected_len: usize) -> bool {
+        self.native_pre.len() == expected_len
+            && self.native_post.len() == expected_len
+            && self.token_pre.len() == expected_len
+            && self.token_post.len() == expected_len
+    }
+}
+
+impl BalanceCollectionRoutines for BalanceCollector {
+    fn collect_pre_balances<CB: TransactionProcessingCallback>(
+        &mut self,
+        account_loader: &mut AccountLoader<CB>,
+        transaction: &impl SVMTransaction,
+    ) {
+        let (native_balances, token_balances) = self.collect_balances(account_loader, transaction);
+        self.native_pre.push(native_balances);
+        self.token_pre.push(token_balances);
+    }
+
+    fn collect_post_balances<CB: TransactionProcessingCallback>(
+        &mut self,
+        account_loader: &mut AccountLoader<CB>,
+        transaction: &impl SVMTransaction,
+    ) {
+        let (native_balances, token_balances) = self.collect_balances(account_loader, transaction);
+        self.native_post.push(native_balances);
+        self.token_post.push(token_balances);
+    }
+}
+
+impl BalanceCollectionRoutines for Option<BalanceCollector> {
+    fn collect_pre_balances<CB: TransactionProcessingCallback>(
+        &mut self,
+        account_loader: &mut AccountLoader<CB>,
+        transaction: &impl SVMTransaction,
+    ) {
+        if let Some(inner) = self {
+            inner.collect_pre_balances(account_loader, transaction)
+        }
+    }
+
+    fn collect_post_balances<CB: TransactionProcessingCallback>(
+        &mut self,
+        account_loader: &mut AccountLoader<CB>,
+        transaction: &impl SVMTransaction,
+    ) {
+        if let Some(inner) = self {
+            inner.collect_post_balances(account_loader, transaction)
+        }
+    }
+}
+
+// this contains all the information we can provide to construct TransactionTokenBalance
+// that type, in ledger, depends on UiTokenAmount from account-decoder, so we cannot build it here
+#[derive(Debug, Clone)]
+pub struct SvmTokenInfo {
+    pub account_index: u8,
+    pub mint: Pubkey,
+    pub amount: u64,
+    pub owner: Pubkey,
+    pub program_id: Pubkey,
+    pub decimals: u8,
+}
+
+impl SvmTokenInfo {
+    fn unpack_token_account<CB: TransactionProcessingCallback>(
+        account_loader: &mut AccountLoader<CB>,
+        account: &AccountSharedData,
+        index: usize,
+    ) -> Option<Self> {
+        let program_id = *account.owner();
+        let generic_token::Account {
+            mint,
+            owner,
+            amount,
+        } = generic_token::Account::unpack(account.data(), &program_id)?;
+
+        // we load as read-only to avoid triggering a bad account inspection
+        let mint_account = account_loader.load_account(&mint, false)?.account;
+        if *mint_account.owner() != program_id {
+            return None;
+        }
+
+        let generic_token::Mint { decimals, .. } =
+            generic_token::Mint::unpack(mint_account.data(), &program_id)?;
+
+        Some(Self {
+            account_index: index.try_into().ok()?,
+            mint,
+            amount,
+            owner,
+            program_id,
+            decimals,
+        })
+    }
+}

+ 38 - 10
svm/src/transaction_processor.rs

@@ -13,6 +13,7 @@ use {
         program_loader::{get_program_modification_slot, load_program_with_pubkey},
         rollback_accounts::RollbackAccounts,
         transaction_account_state_info::TransactionAccountStateInfo,
+        transaction_balances::{BalanceCollectionRoutines, BalanceCollector},
         transaction_error_metrics::TransactionErrorMetrics,
         transaction_execution_result::{ExecutedTransaction, TransactionExecutionDetails},
         transaction_processing_result::{ProcessedTransaction, TransactionProcessingResult},
@@ -76,6 +77,9 @@ pub struct LoadAndExecuteSanitizedTransactionsOutput {
     /// could not be processed. Note processed transactions can still have a
     /// failure result meaning that the transaction will be rolled back.
     pub processing_results: Vec<TransactionProcessingResult>,
+    /// Balances accumulated for TransactionStatusSender when
+    /// transaction balance recording is enabled.
+    pub balance_collector: Option<BalanceCollector>,
 }
 
 /// Configuration of the recording capabilities for transaction execution
@@ -84,6 +88,7 @@ pub struct ExecutionRecordingConfig {
     pub enable_cpi_recording: bool,
     pub enable_log_recording: bool,
     pub enable_return_data_recording: bool,
+    pub enable_transaction_balance_recording: bool,
 }
 
 impl ExecutionRecordingConfig {
@@ -92,6 +97,7 @@ impl ExecutionRecordingConfig {
             enable_return_data_recording: option,
             enable_log_recording: option,
             enable_cpi_recording: option,
+            enable_transaction_balance_recording: option,
         }
     }
 }
@@ -341,6 +347,8 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
             }
             program_accounts_map
         });
+        execute_timings
+            .saturating_add_in_place(ExecuteTimingType::FilterExecutableUs, filter_executable_us);
 
         let (mut program_cache_for_tx_batch, program_cache_us) = measure_us!({
             let program_cache_for_tx_batch = self.replenish_program_cache(
@@ -358,11 +366,16 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
                     processing_results: (0..sanitized_txs.len())
                         .map(|_| Err(TransactionError::ProgramCacheHitMaxLimit))
                         .collect(),
+                    // If we abort the batch and balance recording is enabled, no balances should be
+                    // collected. If this is a leader thread, no batch will be committed.
+                    balance_collector: None,
                 };
             }
 
             program_cache_for_tx_batch
         });
+        execute_timings
+            .saturating_add_in_place(ExecuteTimingType::ProgramCacheUs, program_cache_us);
 
         // Determine a capacity for the internal account cache. This
         // over-allocates but avoids ever reallocating, and spares us from
@@ -377,14 +390,20 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
             account_keys_in_batch,
         );
 
-        let (mut validate_fees_us, mut load_us, mut execution_us): (u64, u64, u64) = (0, 0, 0);
+        // Create the transaction balance collector if recording is enabled.
+        let mut balance_collector = config
+            .recording_config
+            .enable_transaction_balance_recording
+            .then(|| BalanceCollector::new_with_transaction_count(sanitized_txs.len()));
+
+        let (mut load_us, mut execution_us): (u64, u64) = (0, 0);
 
         // Validate, execute, and collect results from each transaction in order.
         // With SIMD83, transactions must be executed in order, because transactions
         // in the same batch may modify the same accounts. Transaction order is
         // preserved within entries written to the ledger.
         for (tx, check_result) in sanitized_txs.iter().zip(check_results) {
-            let (validate_result, single_validate_fees_us) =
+            let (validate_result, validate_fees_us) =
                 measure_us!(check_result.and_then(|tx_details| {
                     Self::validate_transaction_nonce_and_fee_payer(
                         &mut account_loader,
@@ -397,7 +416,8 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
                         &mut error_metrics,
                     )
                 }));
-            validate_fees_us = validate_fees_us.saturating_add(single_validate_fees_us);
+            execute_timings
+                .saturating_add_in_place(ExecuteTimingType::ValidateFeesUs, validate_fees_us);
 
             let (load_result, single_load_us) = measure_us!(load_transaction(
                 &mut account_loader,
@@ -410,6 +430,11 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
             ));
             load_us = load_us.saturating_add(single_load_us);
 
+            let ((), collect_balances_us) =
+                measure_us!(balance_collector.collect_pre_balances(&mut account_loader, tx));
+            execute_timings
+                .saturating_add_in_place(ExecuteTimingType::CollectBalancesUs, collect_balances_us);
+
             let (processing_result, single_execution_us) = measure_us!(match load_result {
                 TransactionLoadResult::NotLoaded(err) => Err(err),
                 TransactionLoadResult::FeesOnly(fees_only_tx) => {
@@ -444,6 +469,11 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
             });
             execution_us = execution_us.saturating_add(single_execution_us);
 
+            let ((), collect_balances_us) =
+                measure_us!(balance_collector.collect_post_balances(&mut account_loader, tx));
+            execute_timings
+                .saturating_add_in_place(ExecuteTimingType::CollectBalancesUs, collect_balances_us);
+
             processing_results.push(processing_result);
         }
 
@@ -468,20 +498,18 @@ impl<FG: ForkGraph> TransactionBatchProcessor<FG> {
             execution_us,
             sanitized_txs.len(),
         );
-
-        execute_timings
-            .saturating_add_in_place(ExecuteTimingType::ValidateFeesUs, validate_fees_us);
-        execute_timings
-            .saturating_add_in_place(ExecuteTimingType::FilterExecutableUs, filter_executable_us);
-        execute_timings
-            .saturating_add_in_place(ExecuteTimingType::ProgramCacheUs, program_cache_us);
         execute_timings.saturating_add_in_place(ExecuteTimingType::LoadUs, load_us);
         execute_timings.saturating_add_in_place(ExecuteTimingType::ExecuteUs, execution_us);
 
+        if let Some(ref balance_collector) = balance_collector {
+            debug_assert!(balance_collector.lengths_match_expected(sanitized_txs.len()));
+        }
+
         LoadAndExecuteSanitizedTransactionsOutput {
             error_metrics,
             execute_timings,
             processing_results,
+            balance_collector,
         }
     }
 

+ 1 - 0
svm/tests/concurrent_tests.rs

@@ -257,6 +257,7 @@ fn svm_concurrent() {
                     enable_log_recording: true,
                     enable_return_data_recording: false,
                     enable_cpi_recording: false,
+                    enable_transaction_balance_recording: false,
                 },
                 ..Default::default()
             };

+ 378 - 2
svm/tests/integration_test.rs

@@ -37,7 +37,8 @@ use {
         transaction_execution_result::TransactionExecutionDetails,
         transaction_processing_result::{ProcessedTransaction, TransactionProcessingResult},
         transaction_processor::{
-            ExecutionRecordingConfig, TransactionBatchProcessor, TransactionProcessingConfig,
+            ExecutionRecordingConfig, LoadAndExecuteSanitizedTransactionsOutput,
+            TransactionBatchProcessor, TransactionProcessingConfig,
             TransactionProcessingEnvironment,
         },
     },
@@ -104,6 +105,7 @@ impl SvmTestEnvironment<'_> {
                 enable_log_recording: true,
                 enable_return_data_recording: true,
                 enable_cpi_recording: false,
+                enable_transaction_balance_recording: false,
             },
             ..Default::default()
         };
@@ -126,7 +128,7 @@ impl SvmTestEnvironment<'_> {
         }
     }
 
-    pub fn execute(&self) {
+    pub fn execute(&self) -> LoadAndExecuteSanitizedTransactionsOutput {
         let (transactions, check_results) = self.test_entry.prepare_transactions();
         let batch_output = self
             .batch_processor
@@ -276,6 +278,8 @@ impl SvmTestEnvironment<'_> {
         // merge new account states into the bank for multi-batch tests
         let mut mock_bank_accounts = self.mock_bank.account_shared_data.write().unwrap();
         mock_bank_accounts.extend(final_accounts_actual);
+
+        batch_output
     }
 }
 
@@ -2652,3 +2656,375 @@ fn svm_metrics_accumulation() {
         );
     }
 }
+
+// NOTE this could be moved to its own file in the future, but it requires a total refactor of the test runner
+mod balance_collector {
+    use {
+        super::*,
+        rand0_7::prelude::*,
+        solana_sdk::{bpf_loader, program_pack::Pack},
+        spl_generic_token::token_2022,
+        spl_token::state::{Account as TokenAccount, AccountState as TokenAccountState, Mint},
+        test_case::test_case,
+    };
+
+    // this could be part of mock_bank but so far nothing but this uses it
+    static SPL_TOKEN_BYTES: &[u8] =
+        include_bytes!("../../program-test/src/programs/spl_token-3.5.0.so");
+
+    const STARTING_BALANCE: u64 = LAMPORTS_PER_SOL * 100;
+
+    // a helper for constructing a transfer instruction, agnostic over system/token
+    // it also pulls double duty as a record of what the *result* of a transfer should be
+    // so we can instantiate a Transfer, gen the instruction, change it to fail, change the record to amount 0
+    // and then the final test confirms the pre/post balances are unchanged with no special casing
+    #[derive(Debug, Default)]
+    struct Transfer {
+        from: Pubkey,
+        to: Pubkey,
+        amount: u64,
+    }
+
+    impl Transfer {
+        // given a set of users, picks two randomly and does a random transfer between them
+        fn new_rand(users: &[Pubkey]) -> Self {
+            let mut rng = rand0_7::thread_rng();
+            let [from_idx, to_idx] = (0..users.len()).choose_multiple(&mut rng, 2)[..] else {
+                unreachable!()
+            };
+            let from = users[from_idx];
+            let to = users[to_idx];
+            let amount = rng.gen_range(1, STARTING_BALANCE / 100);
+
+            Self { from, to, amount }
+        }
+
+        fn to_system_instruction(&self) -> Instruction {
+            system_instruction::transfer(&self.from, &self.to, self.amount)
+        }
+
+        fn to_token_instruction(&self, fee_payer: &Pubkey) -> Instruction {
+            // true tokenkeg connoisseurs will note we shouldnt have to sign the sender
+            // we use a common account owner, the fee-payer, to conveniently reuse account state
+            // so why do we sign? to force the sender and receiver to be in a consistent order in account keys
+            // which means we can grab them by index in our final test instead of searching by key
+            let mut instruction = spl_token::instruction::transfer(
+                &spl_token::id(),
+                &self.from,
+                &self.to,
+                fee_payer,
+                &[],
+                self.amount,
+            )
+            .unwrap();
+            instruction.accounts[0].is_signer = true;
+
+            instruction
+        }
+
+        fn to_instruction(&self, fee_payer: &Pubkey, use_tokens: bool) -> Instruction {
+            if use_tokens {
+                self.to_token_instruction(fee_payer)
+            } else {
+                self.to_system_instruction()
+            }
+        }
+    }
+
+    #[test_case(false; "native")]
+    #[test_case(true; "token")]
+    fn svm_collect_balances(use_tokens: bool) {
+        let mut rng = rand0_7::thread_rng();
+
+        let fee_payer_keypair = Keypair::new();
+        let fake_fee_payer_keypair = Keypair::new();
+        let alice_keypair = Keypair::new();
+        let bob_keypair = Keypair::new();
+        let charlie_keypair = Keypair::new();
+
+        let fee_payer = fee_payer_keypair.pubkey();
+        let fake_fee_payer = fake_fee_payer_keypair.pubkey();
+        let mint = Pubkey::new_unique();
+        let alice = alice_keypair.pubkey();
+        let bob = bob_keypair.pubkey();
+        let charlie = charlie_keypair.pubkey();
+
+        let native_state = AccountSharedData::create(
+            STARTING_BALANCE,
+            vec![],
+            system_program::id(),
+            false,
+            u64::MAX,
+        );
+
+        let mut mint_buf = vec![0; Mint::get_packed_len()];
+        Mint {
+            decimals: 9,
+            is_initialized: true,
+            ..Mint::default()
+        }
+        .pack_into_slice(&mut mint_buf);
+
+        let mint_state =
+            AccountSharedData::create(LAMPORTS_PER_SOL, mint_buf, spl_token::id(), false, u64::MAX);
+
+        let token_account_for_tests = || TokenAccount {
+            mint,
+            owner: fee_payer,
+            amount: STARTING_BALANCE,
+            state: TokenAccountState::Initialized,
+            ..TokenAccount::default()
+        };
+
+        let mut token_buf = vec![0; TokenAccount::get_packed_len()];
+        token_account_for_tests().pack_into_slice(&mut token_buf);
+
+        let token_state = AccountSharedData::create(
+            LAMPORTS_PER_SOL,
+            token_buf,
+            spl_token::id(),
+            false,
+            u64::MAX,
+        );
+
+        let spl_token = AccountSharedData::create(
+            LAMPORTS_PER_SOL,
+            SPL_TOKEN_BYTES.to_vec(),
+            bpf_loader::id(),
+            true,
+            u64::MAX,
+        );
+
+        for _ in 0..100 {
+            let mut test_entry = SvmTestEntry::default();
+            test_entry.add_initial_account(fee_payer, &native_state.clone());
+
+            if use_tokens {
+                test_entry.add_initial_account(spl_token::id(), &spl_token);
+                test_entry.add_initial_account(mint, &mint_state);
+                test_entry.add_initial_account(alice, &token_state);
+                test_entry.add_initial_account(bob, &token_state);
+                test_entry.add_initial_account(charlie, &token_state);
+            } else {
+                test_entry.add_initial_account(alice, &native_state);
+                test_entry.add_initial_account(bob, &native_state);
+                test_entry.add_initial_account(charlie, &native_state);
+            }
+
+            // every time we perform a transfer, we mutate user_balances
+            // and then clone and push it into user_balance_history
+            // this lets us go through every svm balance record and confirm correctness
+            let mut user_balances = HashMap::new();
+            user_balances.insert(alice, STARTING_BALANCE);
+            user_balances.insert(bob, STARTING_BALANCE);
+            user_balances.insert(charlie, STARTING_BALANCE);
+            let mut user_balance_history = vec![(Transfer::default(), user_balances.clone())];
+
+            for _ in 0..50 {
+                // failures result in no balance changes (note we use a separate fee-payer)
+                // we mix some in with the successes to test that we never record changes for failures
+                let expected_status = match rng.gen::<f64>() {
+                    n if n < 0.85 => ExecutionStatus::Succeeded,
+                    n if n < 0.90 => ExecutionStatus::ExecutedFailed,
+                    n if n < 0.95 => ExecutionStatus::ProcessedFailed,
+                    _ => ExecutionStatus::Discarded,
+                };
+
+                let mut transfer = Transfer::new_rand(&[alice, bob, charlie]);
+                let from_signer = vec![&alice_keypair, &bob_keypair, &charlie_keypair]
+                    .into_iter()
+                    .find(|k| k.pubkey() == transfer.from)
+                    .unwrap();
+
+                let instructions = match expected_status {
+                    // a success results in balance changes and is a normal transaction
+                    ExecutionStatus::Succeeded => {
+                        user_balances
+                            .entry(transfer.from)
+                            .and_modify(|v| *v -= transfer.amount);
+                        user_balances
+                            .entry(transfer.to)
+                            .and_modify(|v| *v += transfer.amount);
+
+                        vec![transfer.to_instruction(&fee_payer, use_tokens)]
+                    }
+                    // transfer an unreasonable amount to fail execution
+                    ExecutionStatus::ExecutedFailed => {
+                        transfer.amount = u64::MAX / 2;
+                        let instruction = transfer.to_instruction(&fee_payer, use_tokens);
+                        transfer.amount = 0;
+
+                        vec![instruction]
+                    }
+                    // use a non-existant program to fail loading
+                    // token22 is very convenient because its presence ensures token bals are recorded
+                    // if we had to use a random program id we would need to push a token program onto account keys
+                    ExecutionStatus::ProcessedFailed => {
+                        let mut instruction = transfer.to_instruction(&fee_payer, use_tokens);
+                        instruction.program_id = token_2022::id();
+                        transfer.amount = 0;
+
+                        vec![instruction]
+                    }
+                    // use a non-existant fee-payer to trigger a discard
+                    ExecutionStatus::Discarded => {
+                        let mut instruction = transfer.to_instruction(&fee_payer, use_tokens);
+                        if use_tokens {
+                            instruction.accounts[2].pubkey = fake_fee_payer;
+                        }
+                        transfer.amount = 0;
+
+                        vec![instruction]
+                    }
+                };
+
+                let transaction = if expected_status.discarded() {
+                    Transaction::new_signed_with_payer(
+                        &instructions,
+                        Some(&fake_fee_payer),
+                        &[&fake_fee_payer_keypair, from_signer],
+                        Hash::default(),
+                    )
+                } else {
+                    test_entry.decrease_expected_lamports(&fee_payer, LAMPORTS_PER_SIGNATURE * 2);
+
+                    Transaction::new_signed_with_payer(
+                        &instructions,
+                        Some(&fee_payer),
+                        &[&fee_payer_keypair, from_signer],
+                        Hash::default(),
+                    )
+                };
+
+                test_entry.push_transaction_with_status(transaction, expected_status);
+                user_balance_history.push((transfer, user_balances.clone()));
+            }
+
+            // this block just updates the SvmTestEntry final account states to be accurate
+            // doing this instead of skipping it, we validate that user_balances is definitely correct
+            // because env.execute() will assert all these states match the final bank state
+            if use_tokens {
+                let mut token_account = token_account_for_tests();
+                let mut token_buf = vec![0; TokenAccount::get_packed_len()];
+
+                token_account.amount = *user_balances.get(&alice).unwrap();
+                token_account.pack_into_slice(&mut token_buf);
+                let final_token_state = AccountSharedData::create(
+                    LAMPORTS_PER_SOL,
+                    token_buf.clone(),
+                    spl_token::id(),
+                    false,
+                    u64::MAX,
+                );
+                test_entry.update_expected_account_data(alice, &final_token_state);
+
+                token_account.amount = *user_balances.get(&bob).unwrap();
+                token_account.pack_into_slice(&mut token_buf);
+                let final_token_state = AccountSharedData::create(
+                    LAMPORTS_PER_SOL,
+                    token_buf.clone(),
+                    spl_token::id(),
+                    false,
+                    u64::MAX,
+                );
+                test_entry.update_expected_account_data(bob, &final_token_state);
+
+                token_account.amount = *user_balances.get(&charlie).unwrap();
+                token_account.pack_into_slice(&mut token_buf);
+                let final_token_state = AccountSharedData::create(
+                    LAMPORTS_PER_SOL,
+                    token_buf.clone(),
+                    spl_token::id(),
+                    false,
+                    u64::MAX,
+                );
+                test_entry.update_expected_account_data(charlie, &final_token_state);
+            } else {
+                let mut alice_final_state = native_state.clone();
+                alice_final_state.set_lamports(*user_balances.get(&alice).unwrap());
+                test_entry.update_expected_account_data(alice, &alice_final_state);
+
+                let mut bob_final_state = native_state.clone();
+                bob_final_state.set_lamports(*user_balances.get(&bob).unwrap());
+                test_entry.update_expected_account_data(bob, &bob_final_state);
+
+                let mut charlie_final_state = native_state.clone();
+                charlie_final_state.set_lamports(*user_balances.get(&charlie).unwrap());
+                test_entry.update_expected_account_data(charlie, &charlie_final_state);
+            }
+
+            // turn on balance recording and run the batch
+            let mut env = SvmTestEnvironment::create(test_entry);
+            env.processing_config
+                .recording_config
+                .enable_transaction_balance_recording = true;
+            let batch_output = env.execute();
+
+            // thanks to execute() we know user_balances is correct
+            // now we test that every step in user_balance_history matches the svm recorded balances
+            // in other words, the test effectively has three balance trackers and we can test they *all* agree
+            // first get the collected balances in a manner that is system/token agnostic
+            let (batch_pre, batch_post) = if use_tokens {
+                let (_, _, pre_vecs, post_vecs) =
+                    batch_output.balance_collector.unwrap().into_vecs();
+
+                let pre_tupls: Vec<_> = pre_vecs
+                    .iter()
+                    .map(|bals| (bals[0].amount, bals[1].amount))
+                    .collect();
+
+                let post_tupls: Vec<_> = post_vecs
+                    .iter()
+                    .map(|bals| (bals[0].amount, bals[1].amount))
+                    .collect();
+
+                (pre_tupls, post_tupls)
+            } else {
+                let (pre_vecs, post_vecs, _, _) =
+                    batch_output.balance_collector.unwrap().into_vecs();
+
+                let pre_tupls: Vec<_> = pre_vecs.iter().map(|bals| (bals[1], bals[2])).collect();
+                let post_tupls: Vec<_> = post_vecs.iter().map(|bals| (bals[1], bals[2])).collect();
+
+                (pre_tupls, post_tupls)
+            };
+
+            // these two asserts are trivially true. we include them just to make it clearer what these vecs are
+            // for n transactions, we have n pre-balance sets and n post-balance sets from svm
+            // but we have *n+1* test balance sets: we push initial state, and then push post-tx bals once per tx
+            // this mismatch is not strange at all. we also only have n+1 distinct svm timesteps despite 2n records
+            // pre-balances: (0 1 2 3)
+            // post-balances:  (1 2 3 4)
+            // this does not mean time-overlapping svm records are equal. svm only captures the two accounts used by transfer
+            // whereas our test balances capture all three accounts at every timestep, so we require no pre/post separation
+            assert_eq!(user_balance_history.len(), batch_pre.len() + 1);
+            assert_eq!(user_balance_history.len(), batch_post.len() + 1);
+
+            // these are the real tests
+            for (i, (svm_pre_balances, svm_post_balances)) in
+                batch_pre.into_iter().zip(batch_post).enumerate()
+            {
+                let (_, ref expected_pre_balances) = user_balance_history[i];
+                let (ref transfer, ref expected_post_balances) = user_balance_history[i + 1];
+
+                assert_eq!(
+                    svm_pre_balances.0,
+                    *expected_pre_balances.get(&transfer.from).unwrap()
+                );
+                assert_eq!(
+                    svm_pre_balances.1,
+                    *expected_pre_balances.get(&transfer.to).unwrap()
+                );
+
+                assert_eq!(
+                    svm_post_balances.0,
+                    *expected_post_balances.get(&transfer.from).unwrap()
+                );
+                assert_eq!(
+                    svm_post_balances.1,
+                    *expected_post_balances.get(&transfer.to).unwrap()
+                );
+            }
+        }
+    }
+}

+ 8 - 1
svm/tests/mock_bank.rs

@@ -5,7 +5,8 @@ use solana_sdk::sysvar::recent_blockhashes::{Entry as BlockhashesEntry, RecentBl
 use {
     solana_bpf_loader_program::syscalls::{
         SyscallAbort, SyscallGetClockSysvar, SyscallGetRentSysvar, SyscallInvokeSignedRust,
-        SyscallLog, SyscallMemcpy, SyscallMemset, SyscallSetReturnData,
+        SyscallLog, SyscallMemcmp, SyscallMemcpy, SyscallMemmove, SyscallMemset,
+        SyscallSetReturnData,
     },
     solana_fee_structure::{FeeDetails, FeeStructure},
     solana_program_runtime::{
@@ -378,6 +379,12 @@ pub fn create_custom_loader<'a>() -> BuiltinProgram<InvokeContext<'a>> {
     loader
         .register_function("sol_memset_", SyscallMemset::vm)
         .expect("Registration failed");
+    loader
+        .register_function("sol_memcmp_", SyscallMemcmp::vm)
+        .expect("Registration failed");
+    loader
+        .register_function("sol_memmove_", SyscallMemmove::vm)
+        .expect("Registration failed");
     loader
         .register_function("sol_invoke_signed_rust", SyscallInvokeSignedRust::vm)
         .expect("Registration failed");

+ 8 - 0
timings/src/lib.rs

@@ -59,6 +59,7 @@ pub enum ExecuteTimingType {
     ProgramCacheUs,
     CheckBlockLimitsUs,
     FilterExecutableUs,
+    CollectBalancesUs,
 }
 
 pub struct Metrics([Saturating<u64>; ExecuteTimingType::CARDINALITY]);
@@ -137,6 +138,13 @@ eager_macro_rules! { $eager_1
                     .index(ExecuteTimingType::ExecuteUs).0,
                 i64
             ),
+            (
+                "collect_balances_us",
+                $self
+                    .metrics
+                    .index(ExecuteTimingType::CollectBalancesUs).0,
+                i64
+            ),
             (
                 "collect_logs_us",
                 $self