svm: collect transaction balances

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,
             );

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),

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,
@@ -239,18 +238,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()
@@ -323,6 +310,7 @@ impl Consumer {
         let LoadAndExecuteTransactionsOutput {
             processing_results,
             processed_counts,
+            balance_collector,
         } = load_and_execute_transactions_output;
 
         let transaction_counts = LeaderProcessedTransactionCounts {
@@ -392,7 +380,7 @@ impl Consumer {
                     processing_results,
                     starting_transaction_index,
                     bank,
-                    &mut pre_balance_info,
+                    balance_collector,
                     &mut execute_and_commit_timings,
                     &processed_counts,
                 )

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),

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,

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 }
 

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,

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]