From 46f63674604a42f6f8496e6dd86ca500b94caf88 Mon Sep 17 00:00:00 2001 From: Kirill Fomichev Date: Sun, 25 Feb 2024 19:03:07 -0500 Subject: [PATCH 1/3] rpc: optimize `getTokenLargestAccounts` --- rpc/src/rpc.rs | 75 +++++++++++++++++++++++++++++++++----------------- 1 file changed, 49 insertions(+), 26 deletions(-) diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs index 82eda9489ef247..6d160d62d145b2 100644 --- a/rpc/src/rpc.rs +++ b/rpc/src/rpc.rs @@ -101,7 +101,7 @@ use { std::{ any::type_name, cmp::{max, min}, - collections::{HashMap, HashSet}, + collections::{BinaryHeap, HashMap, HashSet}, convert::TryFrom, net::SocketAddr, str::FromStr, @@ -1861,36 +1861,59 @@ 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::::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); + #[derive(Debug, PartialEq, Eq)] + struct TokenAccountBalance { + address: Pubkey, + amount: u64, + } + impl PartialOrd for TokenAccountBalance { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.amount.cmp(&other.amount).reverse()) + } + } + impl Ord for TokenAccountBalance { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.amount.cmp(&other.amount).reverse() + } + } - 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 mut token_balances = + BinaryHeap::::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::::unpack(account.data()) + .map(|account| account.base.amount) + .unwrap_or(0); - let largest_token_balances = largest_token_balances - .iter() - .map(|(address, amount)| RpcTokenAccountBalance { - address: address.to_string(), - amount: token_amount_to_ui_amount(*amount, decimals), + let requires_insert = if token_balances.len() < NUM_LARGEST_ACCOUNTS { + true + } else { + if let Some(peek) = token_balances.peek() { + peek.amount < amount + } else { + false + } + }; + + if requires_insert { + if token_balances.len() >= NUM_LARGEST_ACCOUNTS { + token_balances.pop(); + } + token_balances.push(TokenAccountBalance { address, amount }); + } + } + let token_balances = token_balances + .into_sorted_vec() + .into_iter() + .map(|token| RpcTokenAccountBalance { + address: token.address.to_string(), + amount: token_amount_to_ui_amount(token.amount, decimals), }) .collect(); - Ok(new_response(&bank, largest_token_balances)) + + Ok(new_response(&bank, token_balances)) } pub fn get_token_accounts_by_owner( From ca64bce4d490f5005989f037daa2bcab029fa623 Mon Sep 17 00:00:00 2001 From: Kirill Fomichev Date: Mon, 26 Feb 2024 21:01:44 -0500 Subject: [PATCH 2/3] use tuple instead of struct --- rpc/src/rpc.rs | 50 +++++++++++++++----------------------------------- 1 file changed, 15 insertions(+), 35 deletions(-) diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs index 6d160d62d145b2..155ffe4ada77fc 100644 --- a/rpc/src/rpc.rs +++ b/rpc/src/rpc.rs @@ -100,7 +100,7 @@ use { }, std::{ any::type_name, - cmp::{max, min}, + cmp::{max, min, Reverse}, collections::{BinaryHeap, HashMap, HashSet}, convert::TryFrom, net::SocketAddr, @@ -1862,24 +1862,8 @@ impl JsonRpcRequestProcessor { )); } - #[derive(Debug, PartialEq, Eq)] - struct TokenAccountBalance { - address: Pubkey, - amount: u64, - } - impl PartialOrd for TokenAccountBalance { - fn partial_cmp(&self, other: &Self) -> Option { - Some(self.amount.cmp(&other.amount).reverse()) - } - } - impl Ord for TokenAccountBalance { - fn cmp(&self, other: &Self) -> std::cmp::Ordering { - self.amount.cmp(&other.amount).reverse() - } - } - let mut token_balances = - BinaryHeap::::with_capacity(NUM_LARGEST_ACCOUNTS); + BinaryHeap::>::with_capacity(NUM_LARGEST_ACCOUNTS); for (address, account) in self.get_filtered_spl_token_accounts_by_mint(&bank, &mint_owner, mint, vec![])? { @@ -1887,29 +1871,25 @@ impl JsonRpcRequestProcessor { .map(|account| account.base.amount) .unwrap_or(0); - let requires_insert = if token_balances.len() < NUM_LARGEST_ACCOUNTS { - true - } else { - if let Some(peek) = token_balances.peek() { - peek.amount < amount - } else { - false - } - }; - - if requires_insert { - if token_balances.len() >= NUM_LARGEST_ACCOUNTS { - token_balances.pop(); + 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.push(TokenAccountBalance { address, amount }); + token_balances.pop(); } + token_balances.push(Reverse(new_entry)); } + let token_balances = token_balances .into_sorted_vec() .into_iter() - .map(|token| RpcTokenAccountBalance { - address: token.address.to_string(), - amount: token_amount_to_ui_amount(token.amount, decimals), + .map(|Reverse(token)| RpcTokenAccountBalance { + address: token.1.to_string(), + amount: token_amount_to_ui_amount(token.0, decimals), }) .collect(); From 4e96622b2f9cad395fb98fe0688478c66283ab6c Mon Sep 17 00:00:00 2001 From: Kirill Fomichev Date: Mon, 26 Feb 2024 21:13:34 -0500 Subject: [PATCH 3/3] untuple Co-authored-by: Tyera --- rpc/src/rpc.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/rpc/src/rpc.rs b/rpc/src/rpc.rs index 155ffe4ada77fc..7bde6b837f2a13 100644 --- a/rpc/src/rpc.rs +++ b/rpc/src/rpc.rs @@ -1887,9 +1887,9 @@ impl JsonRpcRequestProcessor { let token_balances = token_balances .into_sorted_vec() .into_iter() - .map(|Reverse(token)| RpcTokenAccountBalance { - address: token.1.to_string(), - amount: token_amount_to_ui_amount(token.0, decimals), + .map(|Reverse((amount, address))| RpcTokenAccountBalance { + address: address.to_string(), + amount: token_amount_to_ui_amount(amount, decimals), }) .collect();