svm: collect transaction balances

Author: 2501babe

Cargo.lock

@@ -577,6 +577,7 @@ solana-zk-sdk = { path = "zk-sdk", version = "=2.3.0" }
 solana-zk-token-proof-program = { path = "programs/zk-token-proof", version = "=2.3.0" }
 solana-zk-token-sdk = { path = "zk-token-sdk", version = "=2.3.0" }
 spl-associated-token-account = "=6.0.0"
+spl-generic-token = "1.0.1"
 spl-instruction-padding = "0.3"
 spl-memo = "=6.0.0"
 spl-pod = "=0.5.1"

Cargo.toml

@@ -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,7 @@ 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 (balances, token_balances) = compile_collected_balances(balance_collector);
             let mut transaction_index = starting_transaction_index.unwrap_or_default();
             let (batch_transaction_indexes, tx_costs): (Vec<_>, Vec<_>) = commit_results
                 .iter()
@@ -177,14 +165,8 @@ impl Committer {
                 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/committer.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/consume_worker.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 {
@@ -388,7 +376,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/consumer.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;
@@ -33,7 +31,6 @@ impl LeaderExecuteAndCommitTimings {
             "banking_stage-leader_slot_execute_and_commit_timings",
             "id" => id,
             ("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/src/banking_stage/leader_slot_timing_metrics.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,

core/tests/scheduler_cost_adjustment.rs

@@ -74,8 +74,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 }

ledger/Cargo.toml

@@ -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 {
@@ -222,12 +214,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,
@@ -259,6 +250,7 @@ pub fn execute_batch<'a>(
         check_block_costs_elapsed.as_us(),
     );
     let tx_costs = tx_costs?;
+    let (balances, token_balances) = compile_collected_balances(balance_collector);
 
     bank_utils::find_and_send_votes(
         batch.sanitized_transactions(),
@@ -278,15 +270,6 @@ 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);
-
         // The length of costs vector needs to be consistent with all other
         // vectors that are sent over (such as `transactions`). So, replace the
         // None elements with Some(0)
@@ -4416,7 +4399,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/blockstore_processor.rs

@@ -30,8 +30,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]