rpc: optimize `getTokenLargestAccounts` (#35315)

* rpc: optimize `getTokenLargestAccounts`

* use tuple instead of struct

* untuple

Co-authored-by: Tyera <teulberg@gmail.com>

---------

Co-authored-by: Tyera <teulberg@gmail.com>
This commit is contained in:
Kirill Fomichev 2024-02-27 02:08:29 -05:00 committed by GitHub
parent 09925a11eb
commit 8ad125d0c0
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 31 additions and 28 deletions

View File

@ -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![])?
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 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 token_balances = token_balances
.into_sorted_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 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 largest_token_balances = largest_token_balances
.iter()
.map(|(address, amount)| RpcTokenAccountBalance {
.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(