SPL token balance in transaction metadata (bp #13673) (#14752)

* SPL token balance in transaction metadata  (#13673)

* feat: store pre / post token balances

* move helper functions into separate include

* move token balance functionality to transaction-status crate

* fix blockstore processor test

* fix bigtable legacy test

* add caching to decimals

(cherry picked from commit 13db3eca9f)

# Conflicts:
#	core/src/banking_stage.rs
#	storage-proto/Cargo.toml
#	storage-proto/src/convert.rs

* resolve merges storage-proto banking_stage

* fix bpf tests with token deltas

Co-authored-by: Josh <josh.hundley@gmail.com>

Cargo.lock

@@ -4929,6 +4929,7 @@ dependencies = [
  "prost",
  "serde",
  "serde_derive",
+ "solana-account-decoder",
  "solana-sdk",
  "solana-transaction-status",
 ]

core/src/banking_stage.rs

@@ -41,8 +41,13 @@ use solana_sdk::{
     timing::{duration_as_ms, timestamp},
     transaction::{self, Transaction, TransactionError},
 };
+use solana_transaction_status::token_balances::{
+    collect_token_balances, TransactionTokenBalancesSet,
+};
 use std::{
-    cmp, env,
+    cmp,
+    collections::HashMap,
+    env,
     net::UdpSocket,
     sync::atomic::AtomicBool,
     sync::mpsc::Receiver,
@@ -535,6 +540,14 @@ impl BankingStage {
         };
 
         let mut execute_timings = ExecuteTimings::default();
+        let mut mint_decimals: HashMap<Pubkey, u8> = HashMap::new();
+
+        let pre_token_balances = if transaction_status_sender.is_some() {
+            collect_token_balances(&bank, &batch, &mut mint_decimals)
+        } else {
+            vec![]
+        };
+
         let (
             mut loaded_accounts,
             results,
@@ -581,12 +594,14 @@ impl BankingStage {
             bank_utils::find_and_send_votes(txs, &tx_results, Some(gossip_vote_sender));
             if let Some(sender) = transaction_status_sender {
                 let post_balances = bank.collect_balances(batch);
+                let post_token_balances = collect_token_balances(&bank, &batch, &mut mint_decimals);
                 send_transaction_status_batch(
                     bank.clone(),
                     batch.transactions(),
                     batch.iteration_order_vec(),
                     tx_results.execution_results,
                     TransactionBalancesSet::new(pre_balances, post_balances),
+                    TransactionTokenBalancesSet::new(pre_token_balances, post_token_balances),
                     inner_instructions,
                     transaction_logs,
                     sender,

core/src/transaction_status_service.rs

@@ -54,6 +54,7 @@ impl TransactionStatusService {
             iteration_order,
             statuses,
             balances,
+            token_balances,
             inner_instructions,
             transaction_logs,
         } = write_transaction_status_receiver.recv_timeout(Duration::from_secs(1))?;
@@ -64,6 +65,8 @@ impl TransactionStatusService {
             (status, nonce_rollback),
             pre_balances,
             post_balances,
+            pre_token_balances,
+            post_token_balances,
             inner_instructions,
             log_messages,
         ) in izip!(
@@ -71,6 +74,8 @@ impl TransactionStatusService {
             statuses,
             balances.pre_balances,
             balances.post_balances,
+            token_balances.pre_token_balances,
+            token_balances.post_token_balances,
             inner_instructions,
             transaction_logs
         ) {
@@ -98,6 +103,8 @@ impl TransactionStatusService {
                 });
 
                 let log_messages = Some(log_messages);
+                let pre_token_balances = Some(pre_token_balances);
+                let post_token_balances = Some(post_token_balances);
 
                 blockstore
                     .write_transaction_status(
@@ -112,6 +119,8 @@ impl TransactionStatusService {
                             post_balances,
                             inner_instructions,
                             log_messages,
+                            pre_token_balances,
+                            post_token_balances,
                         },
                     )
                     .expect("Expect database write to succeed");

ledger/src/blockstore.rs

@@ -5875,6 +5875,8 @@ pub mod tests {
                             post_balances: post_balances.clone(),
                             inner_instructions: Some(vec![]),
                             log_messages: Some(vec![]),
+                            pre_token_balances: Some(vec![]),
+                            post_token_balances: Some(vec![]),
                         },
                     )
                     .unwrap();
@@ -5889,6 +5891,8 @@ pub mod tests {
                             post_balances: post_balances.clone(),
                             inner_instructions: Some(vec![]),
                             log_messages: Some(vec![]),
+                            pre_token_balances: Some(vec![]),
+                            post_token_balances: Some(vec![]),
                         },
                     )
                     .unwrap();
@@ -5901,6 +5905,8 @@ pub mod tests {
                         post_balances,
                         inner_instructions: Some(vec![]),
                         log_messages: Some(vec![]),
+                        pre_token_balances: Some(vec![]),
+                        post_token_balances: Some(vec![]),
                     }),
                 }
             })
@@ -6091,6 +6097,8 @@ pub mod tests {
                 instructions: vec![CompiledInstruction::new(1, &(), vec![0])],
             }];
             let log_messages_vec = vec![String::from("Test message\n")];
+            let pre_token_balances_vec = vec![];
+            let post_token_balances_vec = vec![];
 
             // result not found
             assert!(transaction_status_cf
@@ -6111,6 +6119,8 @@ pub mod tests {
                         post_balances: post_balances_vec.clone(),
                         inner_instructions: Some(inner_instructions_vec.clone()),
                         log_messages: Some(log_messages_vec.clone()),
+                        pre_token_balances: Some(pre_token_balances_vec.clone()),
+                        post_token_balances: Some(post_token_balances_vec.clone())
                     },
                 )
                 .is_ok());
@@ -6123,6 +6133,8 @@ pub mod tests {
                 post_balances,
                 inner_instructions,
                 log_messages,
+                pre_token_balances,
+                post_token_balances,
             } = transaction_status_cf
                 .get((0, Signature::default(), 0))
                 .unwrap()
@@ -6133,6 +6145,8 @@ pub mod tests {
             assert_eq!(post_balances, post_balances_vec);
             assert_eq!(inner_instructions.unwrap(), inner_instructions_vec);
             assert_eq!(log_messages.unwrap(), log_messages_vec);
+            assert_eq!(pre_token_balances.unwrap(), pre_token_balances_vec);
+            assert_eq!(post_token_balances.unwrap(), post_token_balances_vec);
 
             // insert value
             assert!(transaction_status_cf
@@ -6145,6 +6159,8 @@ pub mod tests {
                         post_balances: post_balances_vec.clone(),
                         inner_instructions: Some(inner_instructions_vec.clone()),
                         log_messages: Some(log_messages_vec.clone()),
+                        pre_token_balances: Some(pre_token_balances_vec.clone()),
+                        post_token_balances: Some(post_token_balances_vec.clone())
                     },
                 )
                 .is_ok());
@@ -6157,6 +6173,8 @@ pub mod tests {
                 post_balances,
                 inner_instructions,
                 log_messages,
+                pre_token_balances,
+                post_token_balances,
             } = transaction_status_cf
                 .get((0, Signature::new(&[2u8; 64]), 9))
                 .unwrap()
@@ -6169,6 +6187,8 @@ pub mod tests {
             assert_eq!(post_balances, post_balances_vec);
             assert_eq!(inner_instructions.unwrap(), inner_instructions_vec);
             assert_eq!(log_messages.unwrap(), log_messages_vec);
+            assert_eq!(pre_token_balances.unwrap(), pre_token_balances_vec);
+            assert_eq!(post_token_balances.unwrap(), post_token_balances_vec);
         }
         Blockstore::destroy(&blockstore_path).expect("Expected successful database destruction");
     }
@@ -6397,6 +6417,8 @@ pub mod tests {
                 post_balances: post_balances_vec,
                 inner_instructions: Some(vec![]),
                 log_messages: Some(vec![]),
+                pre_token_balances: Some(vec![]),
+                post_token_balances: Some(vec![]),
             };
 
             let signature1 = Signature::new(&[1u8; 64]);
@@ -6531,6 +6553,8 @@ pub mod tests {
                     instructions: vec![CompiledInstruction::new(1, &(), vec![0])],
                 }]);
                 let log_messages = Some(vec![String::from("Test message\n")]);
+                let pre_token_balances = Some(vec![]);
+                let post_token_balances = Some(vec![]);
                 let signature = transaction.signatures[0];
                 blockstore
                     .transaction_status_cf
@@ -6543,6 +6567,8 @@ pub mod tests {
                             post_balances: post_balances.clone(),
                             inner_instructions: inner_instructions.clone(),
                             log_messages: log_messages.clone(),
+                            pre_token_balances: pre_token_balances.clone(),
+                            post_token_balances: post_token_balances.clone(),
                         },
                     )
                     .unwrap();
@@ -6555,6 +6581,8 @@ pub mod tests {
                         post_balances,
                         inner_instructions,
                         log_messages,
+                        pre_token_balances,
+                        post_token_balances,
                     }),
                 }
             })
@@ -7000,6 +7028,8 @@ pub mod tests {
                             post_balances: vec![],
                             inner_instructions: Some(vec![]),
                             log_messages: Some(vec![]),
+                            pre_token_balances: Some(vec![]),
+                            post_token_balances: Some(vec![]),
                         },
                     )
                     .unwrap();

ledger/src/blockstore_processor.rs

@@ -37,6 +37,10 @@ use solana_sdk::{
     signature::{Keypair, Signature},
     transaction::{Result, Transaction, TransactionError},
 };
+use solana_transaction_status::token_balances::{
+    collect_token_balances, TransactionTokenBalancesSet,
+};
+
 use std::{
     cell::RefCell,
     collections::{HashMap, HashSet},
@@ -104,6 +108,16 @@ fn execute_batch(
     replay_vote_sender: Option<&ReplayVoteSender>,
     timings: &mut ExecuteTimings,
 ) -> Result<()> {
+    let record_token_balances = transaction_status_sender.is_some();
+
+    let mut mint_decimals: HashMap<Pubkey, u8> = HashMap::new();
+
+    let pre_token_balances = if record_token_balances {
+        collect_token_balances(&bank, &batch, &mut mint_decimals)
+    } else {
+        vec![]
+    };
+
     let (tx_results, balances, inner_instructions, transaction_logs) =
         batch.bank().load_execute_and_commit_transactions(
             batch,
@@ -123,12 +137,22 @@ fn execute_batch(
     } = tx_results;
 
     if let Some(sender) = transaction_status_sender {
+        let post_token_balances = if record_token_balances {
+            collect_token_balances(&bank, &batch, &mut mint_decimals)
+        } else {
+            vec![]
+        };
+
+        let token_balances =
+            TransactionTokenBalancesSet::new(pre_token_balances, post_token_balances);
+
         send_transaction_status_batch(
             bank.clone(),
             batch.transactions(),
             batch.iteration_order_vec(),
             execution_results,
             balances,
+            token_balances,
             inner_instructions,
             transaction_logs,
             sender,
@@ -1061,6 +1085,7 @@ pub struct TransactionStatusBatch {
     pub iteration_order: Option<Vec<usize>>,
     pub statuses: Vec<TransactionExecutionResult>,
     pub balances: TransactionBalancesSet,
+    pub token_balances: TransactionTokenBalancesSet,
     pub inner_instructions: Vec<Option<InnerInstructionsList>>,
     pub transaction_logs: Vec<TransactionLogMessages>,
 }
@@ -1073,6 +1098,7 @@ pub fn send_transaction_status_batch(
     iteration_order: Option<Vec<usize>>,
     statuses: Vec<TransactionExecutionResult>,
     balances: TransactionBalancesSet,
+    token_balances: TransactionTokenBalancesSet,
     inner_instructions: Vec<Option<InnerInstructionsList>>,
     transaction_logs: Vec<TransactionLogMessages>,
     transaction_status_sender: TransactionStatusSender,
@@ -1084,6 +1110,7 @@ pub fn send_transaction_status_batch(
         iteration_order,
         statuses,
         balances,
+        token_balances,
         inner_instructions,
         transaction_logs,
     }) {

programs/bpf/tests/programs.rs

@@ -36,10 +36,10 @@ use solana_sdk::{
     transaction::{Transaction, TransactionError},
 };
 use solana_transaction_status::{
-    ConfirmedTransaction, InnerInstructions, TransactionStatusMeta, TransactionWithStatusMeta,
+    token_balances::collect_token_balances, ConfirmedTransaction, InnerInstructions,
-    UiTransactionEncoding,
+    TransactionStatusMeta, TransactionWithStatusMeta, UiTransactionEncoding,
 };
-use std::{cell::RefCell, env, fs::File, io::Read, path::PathBuf, sync::Arc};
+use std::{cell::RefCell, collections::HashMap, env, fs::File, io::Read, path::PathBuf, sync::Arc};
 
 /// BPF program file extension
 const PLATFORM_FILE_EXTENSION_BPF: &str = "so";
@@ -222,6 +222,8 @@ fn process_transaction_and_record_inner(
 fn execute_transactions(bank: &Bank, txs: &[Transaction]) -> Vec<ConfirmedTransaction> {
     let batch = bank.prepare_batch(txs, None);
     let mut timings = ExecuteTimings::default();
+    let mut mint_decimals = HashMap::new();
+    let tx_pre_token_balances = collect_token_balances(&bank, &batch, &mut mint_decimals);
     let (
         TransactionResults {
             execution_results, ..
@@ -241,6 +243,7 @@ fn execute_transactions(bank: &Bank, txs: &[Transaction]) -> Vec<ConfirmedTransa
         true,
         &mut timings,
     );
+    let tx_post_token_balances = collect_token_balances(&bank, &batch, &mut mint_decimals);
 
     for _ in 0..(txs.len() - transaction_logs.len()) {
         transaction_logs.push(vec![]);
@@ -255,6 +258,8 @@ fn execute_transactions(bank: &Bank, txs: &[Transaction]) -> Vec<ConfirmedTransa
         inner_instructions.into_iter(),
         pre_balances.into_iter(),
         post_balances.into_iter(),
+        tx_pre_token_balances.into_iter(),
+        tx_post_token_balances.into_iter(),
         transaction_logs.into_iter(),
     )
     .map(
@@ -264,6 +269,8 @@ fn execute_transactions(bank: &Bank, txs: &[Transaction]) -> Vec<ConfirmedTransa
             inner_instructions,
             pre_balances,
             post_balances,
+            pre_token_balances,
+            post_token_balances,
             log_messages,
         )| {
             let fee_calculator = nonce_rollback
@@ -289,6 +296,8 @@ fn execute_transactions(bank: &Bank, txs: &[Transaction]) -> Vec<ConfirmedTransa
                 fee,
                 pre_balances,
                 post_balances,
+                pre_token_balances: Some(pre_token_balances),
+                post_token_balances: Some(post_token_balances),
                 inner_instructions,
                 log_messages: Some(log_messages),
             };

storage-bigtable/src/bigtable.rs

@@ -687,6 +687,8 @@ mod tests {
                 post_balances: vec![0, 42, 1],
                 inner_instructions: Some(vec![]),
                 log_messages: Some(vec![]),
+                pre_token_balances: Some(vec![]),
+                post_token_balances: Some(vec![]),
             }),
         };
         let block = ConfirmedBlock {
@@ -736,6 +738,8 @@ mod tests {
             if let Some(meta) = &mut block.transactions[0].meta {
                 meta.inner_instructions = None; // Legacy bincode implementation does not support inner_instructions
                 meta.log_messages = None; // Legacy bincode implementation does not support log_messages
+                meta.pre_token_balances = None; // Legacy bincode implementation does not support token balances
+                meta.post_token_balances = None; // Legacy bincode implementation does not support token balances
             }
             assert_eq!(block, bincode_block.into());
         } else {

storage-bigtable/src/lib.rs

@@ -185,6 +185,8 @@ impl From<StoredConfirmedBlockTransactionStatusMeta> for TransactionStatusMeta {
             post_balances,
             inner_instructions: None,
             log_messages: None,
+            pre_token_balances: None,
+            post_token_balances: None,
         }
     }
 }

storage-proto/Cargo.toml

@@ -14,6 +14,7 @@ bs58 = "0.3.1"
 prost = "0.6.1"
 serde = "1.0.112"
 serde_derive = "1.0.103"
+solana-account-decoder = { path = "../account-decoder", version = "1.4.25" }
 solana-sdk = { path = "../sdk", version = "1.4.25" }
 solana-transaction-status = { path = "../transaction-status", version = "1.4.25" }
 

storage-proto/proto/solana.storage.confirmed_block.rs

@@ -61,6 +61,10 @@ pub struct TransactionStatusMeta {
     pub inner_instructions: ::std::vec::Vec<InnerInstructions>,
     #[prost(string, repeated, tag = "6")]
     pub log_messages: ::std::vec::Vec<std::string::String>,
+    #[prost(message, repeated, tag = "7")]
+    pub pre_token_balances: ::std::vec::Vec<TokenBalance>,
+    #[prost(message, repeated, tag = "8")]
+    pub post_token_balances: ::std::vec::Vec<TokenBalance>,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
 pub struct TransactionError {
@@ -84,6 +88,24 @@ pub struct CompiledInstruction {
     pub data: std::vec::Vec<u8>,
 }
 #[derive(Clone, PartialEq, ::prost::Message)]
+pub struct TokenBalance {
+    #[prost(uint32, tag = "1")]
+    pub account_index: u32,
+    #[prost(string, tag = "2")]
+    pub mint: std::string::String,
+    #[prost(message, optional, tag = "3")]
+    pub ui_token_amount: ::std::option::Option<UiTokenAmount>,
+}
+#[derive(Clone, PartialEq, ::prost::Message)]
+pub struct UiTokenAmount {
+    #[prost(double, tag = "1")]
+    pub ui_amount: f64,
+    #[prost(uint32, tag = "2")]
+    pub decimals: u32,
+    #[prost(string, tag = "3")]
+    pub amount: std::string::String,
+}
+#[derive(Clone, PartialEq, ::prost::Message)]
 pub struct Reward {
     #[prost(string, tag = "1")]
     pub pubkey: std::string::String,

storage-proto/src/confirmed_block.proto

@@ -41,6 +41,8 @@ message TransactionStatusMeta {
     repeated uint64 post_balances = 4;
     repeated InnerInstructions inner_instructions = 5;
     repeated string log_messages = 6;
+    repeated TokenBalance pre_token_balances = 7;
+    repeated TokenBalance post_token_balances = 8;
 }
 
 message TransactionError {
@@ -58,6 +60,18 @@ message CompiledInstruction {
     bytes data = 3;
 }
 
+message TokenBalance {
+    uint32 account_index = 1;
+    string mint = 2;
+    UiTokenAmount ui_token_amount = 3;
+}
+
+message UiTokenAmount {
+    double ui_amount = 1;
+    uint32 decimals = 2;
+    string amount = 3;
+}
+
 enum RewardType {
     Unspecified = 0;
     Fee = 1;

storage-proto/src/convert.rs

@@ -1,4 +1,5 @@
 use crate::StoredExtendedRewards;
+use solana_account_decoder::parse_token::UiTokenAmount;
 use solana_sdk::{
     hash::Hash,
     instruction::CompiledInstruction,
@@ -11,7 +12,7 @@ use solana_sdk::{
 };
 use solana_transaction_status::{
     ConfirmedBlock, InnerInstructions, Reward, RewardType, TransactionByAddrInfo,
-    TransactionStatusMeta, TransactionWithStatusMeta,
+    TransactionStatusMeta, TransactionTokenBalance, TransactionWithStatusMeta,
 };
 use std::convert::{TryFrom, TryInto};
 
@@ -269,6 +270,8 @@ impl From<TransactionStatusMeta> for generated::TransactionStatusMeta {
             post_balances,
             inner_instructions,
             log_messages,
+            pre_token_balances,
+            post_token_balances,
         } = value;
         let err = match status {
             Ok(()) => None,
@@ -282,6 +285,17 @@ impl From<TransactionStatusMeta> for generated::TransactionStatusMeta {
             .map(|ii| ii.into())
             .collect();
         let log_messages = log_messages.unwrap_or_default();
+        let pre_token_balances = pre_token_balances
+            .unwrap_or_default()
+            .into_iter()
+            .map(|balance| balance.into())
+            .collect();
+        let post_token_balances = post_token_balances
+            .unwrap_or_default()
+            .into_iter()
+            .map(|balance| balance.into())
+            .collect();
+
         Self {
             err,
             fee,
@@ -289,6 +303,8 @@ impl From<TransactionStatusMeta> for generated::TransactionStatusMeta {
             post_balances,
             inner_instructions,
             log_messages,
+            pre_token_balances,
+            post_token_balances,
         }
     }
 }
@@ -304,6 +320,8 @@ impl TryFrom<generated::TransactionStatusMeta> for TransactionStatusMeta {
             post_balances,
             inner_instructions,
             log_messages,
+            pre_token_balances,
+            post_token_balances,
         } = value;
         let status = match &err {
             None => Ok(()),
@@ -316,6 +334,18 @@ impl TryFrom<generated::TransactionStatusMeta> for TransactionStatusMeta {
                 .collect(),
         );
         let log_messages = Some(log_messages);
+        let pre_token_balances = Some(
+            pre_token_balances
+                .into_iter()
+                .map(|balance| balance.into())
+                .collect(),
+        );
+        let post_token_balances = Some(
+            post_token_balances
+                .into_iter()
+                .map(|balance| balance.into())
+                .collect(),
+        );
         Ok(Self {
             status,
             fee,
@@ -323,6 +353,8 @@ impl TryFrom<generated::TransactionStatusMeta> for TransactionStatusMeta {
             post_balances,
             inner_instructions,
             log_messages,
+            pre_token_balances,
+            post_token_balances,
         })
     }
 }
@@ -345,6 +377,35 @@ impl From<generated::InnerInstructions> for InnerInstructions {
     }
 }
 
+impl From<TransactionTokenBalance> for generated::TokenBalance {
+    fn from(value: TransactionTokenBalance) -> Self {
+        Self {
+            account_index: value.account_index as u32,
+            mint: value.mint,
+            ui_token_amount: Some(generated::UiTokenAmount {
+                ui_amount: value.ui_token_amount.ui_amount,
+                decimals: value.ui_token_amount.decimals as u32,
+                amount: value.ui_token_amount.amount,
+            }),
+        }
+    }
+}
+
+impl From<generated::TokenBalance> for TransactionTokenBalance {
+    fn from(value: generated::TokenBalance) -> Self {
+        let ui_token_amount = value.ui_token_amount.unwrap_or_default();
+        Self {
+            account_index: value.account_index as u8,
+            mint: value.mint,
+            ui_token_amount: UiTokenAmount {
+                ui_amount: ui_token_amount.ui_amount,
+                decimals: ui_token_amount.decimals as u8,
+                amount: ui_token_amount.amount,
+            },
+        }
+    }
+}
+
 impl From<CompiledInstruction> for generated::CompiledInstruction {
     fn from(value: CompiledInstruction) -> Self {
         Self {

transaction-status/src/lib.rs

@@ -10,11 +10,13 @@ pub mod parse_stake;
 pub mod parse_system;
 pub mod parse_token;
 pub mod parse_vote;
+pub mod token_balances;
 
 use crate::{
     parse_accounts::{parse_accounts, ParsedAccount},
     parse_instruction::{parse, ParsedInstruction},
 };
+use solana_account_decoder::parse_token::UiTokenAmount;
 pub use solana_runtime::bank::RewardType;
 use solana_sdk::{
     clock::{Slot, UnixTimestamp},
@@ -28,7 +30,6 @@ use solana_sdk::{
     transaction::{Result, Transaction, TransactionError},
 };
 use std::fmt;
-
 /// A duplicate representation of an Instruction for pretty JSON serialization
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase", untagged)]
@@ -116,6 +117,31 @@ pub struct UiInnerInstructions {
     pub instructions: Vec<UiInstruction>,
 }
 
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+pub struct TransactionTokenBalance {
+    pub account_index: u8,
+    pub mint: String,
+    pub ui_token_amount: UiTokenAmount,
+}
+
+#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
+#[serde(rename_all = "camelCase")]
+pub struct UiTransactionTokenBalance {
+    pub account_index: u8,
+    pub mint: String,
+    pub ui_token_amount: UiTokenAmount,
+}
+
+impl From<TransactionTokenBalance> for UiTransactionTokenBalance {
+    fn from(token_balance: TransactionTokenBalance) -> Self {
+        Self {
+            account_index: token_balance.account_index,
+            mint: token_balance.mint,
+            ui_token_amount: token_balance.ui_token_amount,
+        }
+    }
+}
+
 impl UiInnerInstructions {
     fn parse(inner_instructions: InnerInstructions, message: &Message) -> Self {
         Self {
@@ -153,6 +179,10 @@ pub struct TransactionStatusMeta {
     pub inner_instructions: Option<Vec<InnerInstructions>>,
     #[serde(deserialize_with = "default_on_eof")]
     pub log_messages: Option<Vec<String>>,
+    #[serde(deserialize_with = "default_on_eof")]
+    pub pre_token_balances: Option<Vec<TransactionTokenBalance>>,
+    #[serde(deserialize_with = "default_on_eof")]
+    pub post_token_balances: Option<Vec<TransactionTokenBalance>>,
 }
 
 impl Default for TransactionStatusMeta {
@@ -164,6 +194,8 @@ impl Default for TransactionStatusMeta {
             post_balances: vec![],
             inner_instructions: None,
             log_messages: None,
+            pre_token_balances: None,
+            post_token_balances: None,
         }
     }
 }
@@ -179,6 +211,8 @@ pub struct UiTransactionStatusMeta {
     pub post_balances: Vec<u64>,
     pub inner_instructions: Option<Vec<UiInnerInstructions>>,
     pub log_messages: Option<Vec<String>>,
+    pub pre_token_balances: Option<Vec<UiTransactionTokenBalance>>,
+    pub post_token_balances: Option<Vec<UiTransactionTokenBalance>>,
 }
 
 impl UiTransactionStatusMeta {
@@ -195,6 +229,12 @@ impl UiTransactionStatusMeta {
                     .collect()
             }),
             log_messages: meta.log_messages,
+            pre_token_balances: meta
+                .pre_token_balances
+                .map(|balance| balance.into_iter().map(|balance| balance.into()).collect()),
+            post_token_balances: meta
+                .post_token_balances
+                .map(|balance| balance.into_iter().map(|balance| balance.into()).collect()),
         }
     }
 }
@@ -211,6 +251,12 @@ impl From<TransactionStatusMeta> for UiTransactionStatusMeta {
                 .inner_instructions
                 .map(|ixs| ixs.into_iter().map(|ix| ix.into()).collect()),
             log_messages: meta.log_messages,
+            pre_token_balances: meta
+                .pre_token_balances
+                .map(|balance| balance.into_iter().map(|balance| balance.into()).collect()),
+            post_token_balances: meta
+                .post_token_balances
+                .map(|balance| balance.into_iter().map(|balance| balance.into()).collect()),
         }
     }
 }

transaction-status/src/token_balances.rs

@@ -0,0 +1,113 @@
+use crate::TransactionTokenBalance;
+use solana_account_decoder::parse_token::{
+    spl_token_id_v2_0, spl_token_v2_0_native_mint, token_amount_to_ui_amount, UiTokenAmount,
+};
+use solana_runtime::{
+    bank::Bank, transaction_batch::TransactionBatch, transaction_utils::OrderedIterator,
+};
+use solana_sdk::pubkey::Pubkey;
+use spl_token_v2_0::{
+    solana_program::program_pack::Pack,
+    state::{Account as TokenAccount, Mint},
+};
+use std::{collections::HashMap, str::FromStr};
+
+pub type TransactionTokenBalances = Vec<Vec<TransactionTokenBalance>>;
+
+pub struct TransactionTokenBalancesSet {
+    pub pre_token_balances: TransactionTokenBalances,
+    pub post_token_balances: TransactionTokenBalances,
+}
+
+impl TransactionTokenBalancesSet {
+    pub fn new(
+        pre_token_balances: TransactionTokenBalances,
+        post_token_balances: TransactionTokenBalances,
+    ) -> Self {
+        assert_eq!(pre_token_balances.len(), post_token_balances.len());
+        Self {
+            pre_token_balances,
+            post_token_balances,
+        }
+    }
+}
+
+fn is_token_program(program_id: &Pubkey) -> bool {
+    program_id == &spl_token_id_v2_0()
+}
+
+fn get_mint_decimals(bank: &Bank, mint: &Pubkey) -> Option<u8> {
+    if mint == &spl_token_v2_0_native_mint() {
+        Some(spl_token_v2_0::native_mint::DECIMALS)
+    } else {
+        let mint_account = bank.get_account(mint)?;
+
+        let decimals = Mint::unpack(&mint_account.data)
+            .map(|mint| mint.decimals)
+            .ok()?;
+
+        Some(decimals)
+    }
+}
+
+pub fn collect_token_balances(
+    bank: &Bank,
+    batch: &TransactionBatch,
+    mut mint_decimals: &mut HashMap<Pubkey, u8>,
+) -> TransactionTokenBalances {
+    let mut balances: TransactionTokenBalances = vec![];
+
+    for (_, transaction) in OrderedIterator::new(batch.transactions(), batch.iteration_order()) {
+        let account_keys = &transaction.message.account_keys;
+        let mut fetch_account_hash: HashMap<u8, bool> = HashMap::new();
+        for instruction in transaction.message.instructions.iter() {
+            if let Some(program_id) = account_keys.get(instruction.program_id_index as usize) {
+                if is_token_program(&program_id) {
+                    for account in &instruction.accounts {
+                        fetch_account_hash.insert(*account, true);
+                    }
+                }
+            }
+        }
+
+        let mut transaction_balances: Vec<TransactionTokenBalance> = vec![];
+        for index in fetch_account_hash.keys() {
+            if let Some(account_id) = account_keys.get(*index as usize) {
+                if let Some((mint, ui_token_amount)) =
+                    collect_token_balance_from_account(&bank, account_id, &mut mint_decimals)
+                {
+                    transaction_balances.push(TransactionTokenBalance {
+                        account_index: *index,
+                        mint,
+                        ui_token_amount,
+                    });
+                }
+            }
+        }
+        balances.push(transaction_balances);
+    }
+    balances
+}
+
+pub fn collect_token_balance_from_account(
+    bank: &Bank,
+    account_id: &Pubkey,
+    mint_decimals: &mut HashMap<Pubkey, u8>,
+) -> Option<(String, UiTokenAmount)> {
+    let account = bank.get_account(account_id)?;
+
+    let token_account = TokenAccount::unpack(&account.data).ok()?;
+    let mint_string = &token_account.mint.to_string();
+    let mint = &Pubkey::from_str(&mint_string).unwrap_or_default();
+
+    let decimals = mint_decimals.get(&mint).cloned().or_else(|| {
+        let decimals = get_mint_decimals(bank, &mint)?;
+        mint_decimals.insert(*mint, decimals);
+        Some(decimals)
+    })?;
+
+    Some((
+        mint_string.to_string(),
+        token_amount_to_ui_amount(token_account.amount, decimals),
+    ))
+}