Include token owners in TransactionTokenBalances (#20642)

* Cache owners in TransactionTokenBalances

* Light cleanup

* Use return struct, and remove pub

* Single-use statements

* Why not, just do the whole crate

* Add metrics

* Make datapoint_debug to prevent spam unless troubleshooting

Cargo.lock

@@ -5758,10 +5758,13 @@ dependencies = [
  "bincode",
  "bs58 0.4.0",
  "lazy_static",
+ "log 0.4.14",
  "serde",
  "serde_derive",
  "serde_json",
  "solana-account-decoder",
+ "solana-measure",
+ "solana-metrics",
  "solana-runtime",
  "solana-sdk",
  "solana-vote-program",

ledger/src/blockstore.rs

@@ -8361,6 +8361,7 @@ pub mod tests {
                     amount: "11".to_string(),
                     ui_amount_string: "1.1".to_string(),
                 },
+                owner: Pubkey::new_unique().to_string(),
             }]),
             post_token_balances: Some(vec![TransactionTokenBalance {
                 account_index: 0,
@@ -8371,6 +8372,7 @@ pub mod tests {
                     amount: "11".to_string(),
                     ui_amount_string: "1.1".to_string(),
                 },
+                owner: Pubkey::new_unique().to_string(),
             }]),
             rewards: Some(vec![Reward {
                 pubkey: "My11111111111111111111111111111111111111111".to_string(),

programs/bpf/Cargo.lock

@@ -3394,10 +3394,13 @@ dependencies = [
  "bincode",
  "bs58 0.4.0",
  "lazy_static",
+ "log",
  "serde",
  "serde_derive",
  "serde_json",
  "solana-account-decoder",
+ "solana-measure",
+ "solana-metrics",
  "solana-runtime",
  "solana-sdk",
  "solana-vote-program",

storage-proto/proto/confirmed_block.proto

@@ -66,6 +66,7 @@ message TokenBalance {
     uint32 account_index = 1;
     string mint = 2;
     UiTokenAmount ui_token_amount = 3;
+    string owner = 4;
 }
 
 message UiTokenAmount {

storage-proto/src/convert.rs

@@ -408,6 +408,7 @@ impl From<TransactionTokenBalance> for generated::TokenBalance {
                 amount: value.ui_token_amount.amount,
                 ui_amount_string: value.ui_token_amount.ui_amount_string,
             }),
+            owner: value.owner,
         }
     }
 }
@@ -435,6 +436,7 @@ impl From<generated::TokenBalance> for TransactionTokenBalance {
                     )
                 },
             },
+            owner: value.owner,
         }
     }
 }

storage-proto/src/lib.rs

@@ -111,6 +111,8 @@ pub struct StoredTransactionTokenBalance {
     pub account_index: u8,
     pub mint: String,
     pub ui_token_amount: StoredTokenAmount,
+    #[serde(deserialize_with = "default_on_eof")]
+    pub owner: String,
 }
 
 impl From<StoredTransactionTokenBalance> for TransactionTokenBalance {
@@ -119,11 +121,13 @@ impl From<StoredTransactionTokenBalance> for TransactionTokenBalance {
             account_index,
             mint,
             ui_token_amount,
+            owner,
         } = value;
         Self {
             account_index,
             mint,
             ui_token_amount: ui_token_amount.into(),
+            owner,
         }
     }
 }
@@ -134,11 +138,13 @@ impl From<TransactionTokenBalance> for StoredTransactionTokenBalance {
             account_index,
             mint,
             ui_token_amount,
+            owner,
         } = value;
         Self {
             account_index,
             mint,
             ui_token_amount: ui_token_amount.into(),
+            owner,
         }
     }
 }

transaction-status/Cargo.toml

@@ -15,10 +15,13 @@ bincode = "1.3.3"
 bs58 = "0.4.0"
 Inflector = "0.11.4"
 lazy_static = "1.4.0"
+log = "0.4.14"
 serde = "1.0.130"
 serde_derive = "1.0.103"
 serde_json = "1.0.68"
 solana-account-decoder = { path = "../account-decoder", version = "=1.9.0" }
+solana-measure = { path = "../measure", version = "=1.9.0" }
+solana-metrics = { path = "../metrics", version = "=1.9.0" }
 solana-runtime = { path = "../runtime", version = "=1.9.0" }
 solana-sdk = { path = "../sdk", version = "=1.9.0" }
 solana-vote-program = { path = "../programs/vote", version = "=1.9.0" }

transaction-status/src/lib.rs

@@ -15,14 +15,14 @@ pub mod parse_token;
 pub mod parse_vote;
 pub mod token_balances;
 
-pub use crate::extract_memos::extract_and_fmt_memos;
+pub use {crate::extract_memos::extract_and_fmt_memos, solana_runtime::bank::RewardType};
-use crate::{
+use {
+    crate::{
         parse_accounts::{parse_accounts, ParsedAccount},
         parse_instruction::{parse, ParsedInstruction},
-};
+    },
-use solana_account_decoder::parse_token::UiTokenAmount;
+    solana_account_decoder::parse_token::UiTokenAmount,
-pub use solana_runtime::bank::RewardType;
+    solana_sdk::{
-use solana_sdk::{
         clock::{Slot, UnixTimestamp},
         commitment_config::CommitmentConfig,
         deserialize_utils::default_on_eof,
@@ -32,8 +32,9 @@ use solana_sdk::{
         sanitize::Sanitize,
         signature::Signature,
         transaction::{Result, Transaction, TransactionError},
+    },
+    std::fmt,
 };
-use std::fmt;
 /// A duplicate representation of an Instruction for pretty JSON serialization
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
 #[serde(rename_all = "camelCase", untagged)]
@@ -126,6 +127,7 @@ pub struct TransactionTokenBalance {
     pub account_index: u8,
     pub mint: String,
     pub ui_token_amount: UiTokenAmount,
+    pub owner: String,
 }
 
 #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
@@ -134,6 +136,8 @@ pub struct UiTransactionTokenBalance {
     pub account_index: u8,
     pub mint: String,
     pub ui_token_amount: UiTokenAmount,
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub owner: Option<String>,
 }
 
 impl From<TransactionTokenBalance> for UiTransactionTokenBalance {
@@ -142,6 +146,11 @@ impl From<TransactionTokenBalance> for UiTransactionTokenBalance {
             account_index: token_balance.account_index,
             mint: token_balance.mint,
             ui_token_amount: token_balance.ui_token_amount,
+            owner: if !token_balance.owner.is_empty() {
+                Some(token_balance.owner)
+            } else {
+                None
+            },
         }
     }
 }

transaction-status/src/parse_accounts.rs

@@ -22,8 +22,7 @@ pub fn parse_accounts(message: &Message) -> Vec<ParsedAccount> {
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {super::*, solana_sdk::message::MessageHeader};
-    use solana_sdk::message::MessageHeader;
 
     #[test]
     fn test_parse_accounts() {

transaction-status/src/parse_associated_token.rs

@@ -1,8 +1,10 @@
-use crate::parse_instruction::{
+use {
+    crate::parse_instruction::{
         check_num_accounts, ParsableProgram, ParseInstructionError, ParsedInstructionEnum,
+    },
+    serde_json::json,
+    solana_sdk::{instruction::CompiledInstruction, pubkey::Pubkey},
 };
-use serde_json::json;
-use solana_sdk::{instruction::CompiledInstruction, pubkey::Pubkey};
 
 // A helper function to convert spl_associated_token_account_v1_0::id() as spl_sdk::pubkey::Pubkey
 // to solana_sdk::pubkey::Pubkey
@@ -47,13 +49,15 @@ fn check_num_associated_token_accounts(
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {
-    use spl_associated_token_account_v1_0::{
+        super::*,
+        spl_associated_token_account_v1_0::{
             create_associated_token_account,
             solana_program::{
                 instruction::CompiledInstruction as SplAssociatedTokenCompiledInstruction,
                 message::Message, pubkey::Pubkey as SplAssociatedTokenPubkey,
             },
+        },
     };
 
     fn convert_pubkey(pubkey: Pubkey) -> SplAssociatedTokenPubkey {

transaction-status/src/parse_bpf_loader.rs

@@ -1,11 +1,13 @@
-use crate::parse_instruction::{
+use {
+    crate::parse_instruction::{
         check_num_accounts, ParsableProgram, ParseInstructionError, ParsedInstructionEnum,
-};
+    },
-use bincode::deserialize;
+    bincode::deserialize,
-use serde_json::json;
+    serde_json::json,
-use solana_sdk::{
+    solana_sdk::{
         instruction::CompiledInstruction, loader_instruction::LoaderInstruction,
         loader_upgradeable_instruction::UpgradeableLoaderInstruction, pubkey::Pubkey,
+    },
 };
 
 pub fn parse_bpf_loader(
@@ -154,9 +156,11 @@ fn check_num_bpf_upgradeable_loader_accounts(
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {
-    use serde_json::Value;
+        super::*,
-    use solana_sdk::{message::Message, pubkey};
+        serde_json::Value,
+        solana_sdk::{message::Message, pubkey},
+    };
 
     #[test]
     fn test_parse_bpf_loader_instructions() {

transaction-status/src/parse_instruction.rs

@@ -1,4 +1,5 @@
-use crate::{
+use {
+    crate::{
         extract_memos::{spl_memo_id_v1, spl_memo_id_v3},
         parse_associated_token::{parse_associated_token, spl_associated_token_id_v1_0},
         parse_bpf_loader::{parse_bpf_loader, parse_bpf_upgradeable_loader},
@@ -6,16 +7,17 @@ use crate::{
         parse_system::parse_system,
         parse_token::parse_token,
         parse_vote::parse_vote,
-};
+    },
-use inflector::Inflector;
+    inflector::Inflector,
-use serde_json::Value;
+    serde_json::Value,
-use solana_account_decoder::parse_token::spl_token_id_v2_0;
+    solana_account_decoder::parse_token::spl_token_id_v2_0,
-use solana_sdk::{instruction::CompiledInstruction, pubkey::Pubkey, stake, system_program};
+    solana_sdk::{instruction::CompiledInstruction, pubkey::Pubkey, stake, system_program},
-use std::{
+    std::{
         collections::HashMap,
         str::{from_utf8, Utf8Error},
+    },
+    thiserror::Error,
 };
-use thiserror::Error;
 
 lazy_static! {
     static ref ASSOCIATED_TOKEN_PROGRAM_ID: Pubkey = spl_associated_token_id_v1_0();
@@ -150,8 +152,7 @@ pub(crate) fn check_num_accounts(
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {super::*, serde_json::json};
-    use serde_json::json;
 
     #[test]
     fn test_parse() {

transaction-status/src/parse_stake.rs

@@ -1,10 +1,12 @@
-use crate::parse_instruction::{
+use {
+    crate::parse_instruction::{
         check_num_accounts, ParsableProgram, ParseInstructionError, ParsedInstructionEnum,
-};
+    },
-use bincode::deserialize;
+    bincode::deserialize,
-use serde_json::{json, Map};
+    serde_json::{json, Map},
-use solana_sdk::{
+    solana_sdk::{
         instruction::CompiledInstruction, pubkey::Pubkey, stake::instruction::StakeInstruction,
+    },
 };
 
 pub fn parse_stake(
@@ -275,14 +277,16 @@ fn check_num_stake_accounts(accounts: &[u8], num: usize) -> Result<(), ParseInst
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {
-    use solana_sdk::{
+        super::*,
+        solana_sdk::{
             message::Message,
             pubkey::Pubkey,
             stake::{
                 instruction::{self, LockupArgs},
                 state::{Authorized, Lockup, StakeAuthorize},
             },
+        },
     };
 
     #[test]

transaction-status/src/parse_system.rs

@@ -1,10 +1,12 @@
-use crate::parse_instruction::{
+use {
+    crate::parse_instruction::{
         check_num_accounts, ParsableProgram, ParseInstructionError, ParsedInstructionEnum,
-};
+    },
-use bincode::deserialize;
+    bincode::deserialize,
-use serde_json::json;
+    serde_json::json,
-use solana_sdk::{
+    solana_sdk::{
         instruction::CompiledInstruction, pubkey::Pubkey, system_instruction::SystemInstruction,
+    },
 };
 
 pub fn parse_system(
@@ -197,8 +199,10 @@ fn check_num_system_accounts(accounts: &[u8], num: usize) -> Result<(), ParseIns
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {
-    use solana_sdk::{message::Message, pubkey::Pubkey, system_instruction};
+        super::*,
+        solana_sdk::{message::Message, pubkey::Pubkey, system_instruction},
+    };
 
     #[test]
     #[allow(clippy::same_item_push)]

transaction-status/src/parse_token.rs

@@ -1,15 +1,19 @@
-use crate::parse_instruction::{
+use {
+    crate::parse_instruction::{
         check_num_accounts, ParsableProgram, ParseInstructionError, ParsedInstructionEnum,
-};
+    },
-use serde_json::{json, Map, Value};
+    serde_json::{json, Map, Value},
-use solana_account_decoder::parse_token::{pubkey_from_spl_token_v2_0, token_amount_to_ui_amount};
+    solana_account_decoder::parse_token::{pubkey_from_spl_token_v2_0, token_amount_to_ui_amount},
-use solana_sdk::{
+    solana_sdk::{
         instruction::{AccountMeta, CompiledInstruction, Instruction},
         pubkey::Pubkey,
-};
+    },
-use spl_token_v2_0::{
+    spl_token_v2_0::{
         instruction::{AuthorityType, TokenInstruction},
-    solana_program::{instruction::Instruction as SplTokenInstruction, program_option::COption},
+        solana_program::{
+            instruction::Instruction as SplTokenInstruction, program_option::COption,
+        },
+    },
 };
 
 pub fn parse_token(
@@ -452,16 +456,18 @@ pub fn spl_token_v2_0_instruction(instruction: SplTokenInstruction) -> Instructi
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {
-    use solana_sdk::instruction::CompiledInstruction;
+        super::*,
-    use spl_token_v2_0::{
+        solana_sdk::instruction::CompiledInstruction,
+        spl_token_v2_0::{
             instruction::*,
             solana_program::{
                 instruction::CompiledInstruction as SplTokenCompiledInstruction, message::Message,
                 pubkey::Pubkey as SplTokenPubkey,
             },
+        },
+        std::str::FromStr,
     };
-    use std::str::FromStr;
 
     fn convert_pubkey(pubkey: Pubkey) -> SplTokenPubkey {
         SplTokenPubkey::from_str(&pubkey.to_string()).unwrap()

transaction-status/src/parse_vote.rs

@@ -1,10 +1,12 @@
-use crate::parse_instruction::{
+use {
+    crate::parse_instruction::{
         check_num_accounts, ParsableProgram, ParseInstructionError, ParsedInstructionEnum,
+    },
+    bincode::deserialize,
+    serde_json::json,
+    solana_sdk::{instruction::CompiledInstruction, pubkey::Pubkey},
+    solana_vote_program::vote_instruction::VoteInstruction,
 };
-use bincode::deserialize;
-use serde_json::json;
-use solana_sdk::{instruction::CompiledInstruction, pubkey::Pubkey};
-use solana_vote_program::vote_instruction::VoteInstruction;
 
 pub fn parse_vote(
     instruction: &CompiledInstruction,
@@ -143,11 +145,13 @@ fn check_num_vote_accounts(accounts: &[u8], num: usize) -> Result<(), ParseInstr
 
 #[cfg(test)]
 mod test {
-    use super::*;
+    use {
-    use solana_sdk::{hash::Hash, message::Message, pubkey::Pubkey};
+        super::*,
-    use solana_vote_program::{
+        solana_sdk::{hash::Hash, message::Message, pubkey::Pubkey},
+        solana_vote_program::{
             vote_instruction,
             vote_state::{Vote, VoteAuthorize, VoteInit},
+        },
     };
 
     #[test]

transaction-status/src/token_balances.rs

@@ -1,14 +1,19 @@
-use crate::TransactionTokenBalance;
+use {
-use solana_account_decoder::parse_token::{
+    crate::TransactionTokenBalance,
-    spl_token_id_v2_0, spl_token_v2_0_native_mint, token_amount_to_ui_amount, UiTokenAmount,
+    solana_account_decoder::parse_token::{
-};
+        pubkey_from_spl_token_v2_0, spl_token_id_v2_0, spl_token_v2_0_native_mint,
-use solana_runtime::{bank::Bank, transaction_batch::TransactionBatch};
+        token_amount_to_ui_amount, UiTokenAmount,
-use solana_sdk::{account::ReadableAccount, pubkey::Pubkey};
+    },
-use spl_token_v2_0::{
+    solana_measure::measure::Measure,
+    solana_metrics::datapoint_debug,
+    solana_runtime::{bank::Bank, transaction_batch::TransactionBatch},
+    solana_sdk::{account::ReadableAccount, pubkey::Pubkey},
+    spl_token_v2_0::{
         solana_program::program_pack::Pack,
         state::{Account as TokenAccount, Mint},
+    },
+    std::collections::HashMap,
 };
-use std::{collections::HashMap, str::FromStr};
 
 pub type TransactionTokenBalances = Vec<Vec<TransactionTokenBalance>>;
 
@@ -54,6 +59,7 @@ pub fn collect_token_balances(
     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 has_token_program = transaction
@@ -68,41 +74,56 @@ pub fn collect_token_balances(
                     continue;
                 }
 
-                if let Some((mint, ui_token_amount)) =
+                if let Some(TokenBalanceData {
-                    collect_token_balance_from_account(bank, account_id, mint_decimals)
+                    mint,
+                    ui_token_amount,
+                    owner,
+                }) = collect_token_balance_from_account(bank, account_id, mint_decimals)
                 {
                     transaction_balances.push(TransactionTokenBalance {
                         account_index: index as u8,
                         mint,
                         ui_token_amount,
+                        owner,
                     });
                 }
             }
         }
         balances.push(transaction_balances);
     }
+    collect_time.stop();
+    datapoint_debug!(
+        "collect_token_balances",
+        ("collect_time_us", collect_time.as_us(), i64),
+    );
     balances
 }
 
-pub fn collect_token_balance_from_account(
+struct TokenBalanceData {
+    mint: String,
+    owner: String,
+    ui_token_amount: UiTokenAmount,
+}
+
+fn collect_token_balance_from_account(
     bank: &Bank,
     account_id: &Pubkey,
     mint_decimals: &mut HashMap<Pubkey, u8>,
-) -> Option<(String, UiTokenAmount)> {
+) -> Option<TokenBalanceData> {
     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_spl_token_v2_0(&token_account.mint);
-    let mint = &Pubkey::from_str(mint_string).unwrap_or_default();
 
-    let decimals = mint_decimals.get(mint).cloned().or_else(|| {
+    let decimals = mint_decimals.get(&mint).cloned().or_else(|| {
-        let decimals = get_mint_decimals(bank, mint)?;
+        let decimals = get_mint_decimals(bank, &mint)?;
-        mint_decimals.insert(*mint, decimals);
+        mint_decimals.insert(mint, decimals);
         Some(decimals)
     })?;
 
-    Some((
+    Some(TokenBalanceData {
-        mint_string.to_string(),
+        mint: token_account.mint.to_string(),
-        token_amount_to_ui_amount(token_account.amount, decimals),
+        owner: token_account.owner.to_string(),
-    ))
+        ui_token_amount: token_amount_to_ui_amount(token_account.amount, decimals),
+    })
 }