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

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},
@@ -27,7 +29,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)]
@@ -115,6 +116,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 {
@@ -152,6 +178,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 {
@@ -163,6 +193,8 @@ impl Default for TransactionStatusMeta {
             post_balances: vec![],
             inner_instructions: None,
             log_messages: None,
+            pre_token_balances: None,
+            post_token_balances: None,
         }
     }
 }
@@ -178,6 +210,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 {
@@ -194,6 +228,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()),
         }
     }
 }
@@ -210,6 +250,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),
+    ))
+}