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(