rpc: optimize getTokenLargestAccounts (#35315)

* rpc: optimize `getTokenLargestAccounts` * use tuple instead of struct * untuple Co-authored-by: Tyera <[email protected]> --------- Co-authored-by: Tyera <[email protected]>
solana-labs · Feb 27, 2024 · 8ad125d · 8ad125d
1 parent 09925a1
commit 8ad125d
Showing 1 changed file with 29 additions and 26 deletions.
diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs
@@ -100,8 +100,8 @@ use {
  },
  std::{
  any::type_name,
- cmp::{max, min},
- collections::{HashMap, HashSet},
+ cmp::{max, min, Reverse},
+ collections::{BinaryHeap, HashMap, HashSet},
  convert::TryFrom,
  net::SocketAddr,
  str::FromStr,
@@ -1861,36 +1861,39 @@ impl JsonRpcRequestProcessor {
  "Invalid param: not a Token mint".to_string(),
  ));
  }
- let mut token_balances: Vec<_> = self
- .get_filtered_spl_token_accounts_by_mint(&bank, &mint_owner, mint, vec![])?
- .into_iter()
- .map(|(address, account)| {
- let amount = StateWithExtensions::<TokenAccount>::unpack(account.data())
- .map(|account| account.base.amount)
- .unwrap_or(0);
- (address, amount)
- })
- .collect();
 
- let sort_largest = |a: &(_, u64), b: &(_, u64)| b.1.cmp(&a.1);
+ let mut token_balances =
+ BinaryHeap::<Reverse<(u64, Pubkey)>>::with_capacity(NUM_LARGEST_ACCOUNTS);
+ for (address, account) in
+ self.get_filtered_spl_token_accounts_by_mint(&bank, &mint_owner, mint, vec![])?
+ {
+ let amount = StateWithExtensions::<TokenAccount>::unpack(account.data())
+ .map(|account| account.base.amount)
+ .unwrap_or(0);
 
- let largest_token_balances = if token_balances.len() > NUM_LARGEST_ACCOUNTS {
- token_balances
- .select_nth_unstable_by(NUM_LARGEST_ACCOUNTS, sort_largest)
- .0
- } else {
- token_balances.as_mut_slice()
- };
- largest_token_balances.sort_unstable_by(sort_largest);
+ let new_entry = (amount, address);
+ if token_balances.len() >= NUM_LARGEST_ACCOUNTS {
+ let Reverse(entry) = token_balances
+ .peek()
+ .expect("BinaryHeap::peek should succeed when len > 0");
+ if *entry >= new_entry {
+ continue;
+ }
+ token_balances.pop();
+ }
+ token_balances.push(Reverse(new_entry));
+ }
 
- let largest_token_balances = largest_token_balances
- .iter()
- .map(|(address, amount)| RpcTokenAccountBalance {
+ let token_balances = token_balances
+ .into_sorted_vec()
+ .into_iter()
+ .map(|Reverse((amount, address))| RpcTokenAccountBalance {
  address: address.to_string(),
- amount: token_amount_to_ui_amount(*amount, decimals),
+ amount: token_amount_to_ui_amount(amount, decimals),
  })
  .collect();
- Ok(new_response(&bank, largest_token_balances))
+
+ Ok(new_response(&bank, token_balances))
  }
 
  pub fn get_token_accounts_by_owner(