From 7a4bfa367c8cb37f19ba85c4d1cc4f062a938df3 Mon Sep 17 00:00:00 2001
From: Steven Luscher <steven.luscher@anza.xyz>
Date: Wed, 11 Dec 2024 00:16:45 +0000
Subject: [PATCH] A benchmark suite for the `getTokenSupply` RPC call

---
 accounts-cluster-bench/src/main.rs | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/accounts-cluster-bench/src/main.rs b/accounts-cluster-bench/src/main.rs
index f7b89eecf48e95..8b1c215cdc6f9a 100644
--- a/accounts-cluster-bench/src/main.rs
+++ b/accounts-cluster-bench/src/main.rs
@@ -318,6 +318,7 @@ pub enum RpcBench {
     Transaction,
     TransactionParsed,
     FirstAvailableBlock,
+    TokenSupply,
 }
 
 #[derive(Debug)]
@@ -339,6 +340,7 @@ impl FromStr for RpcBench {
             "multiple-accounts" => Ok(RpcBench::MultipleAccounts),
             "token-accounts-by-delegate" => Ok(RpcBench::TokenAccountsByDelegate),
             "token-accounts-by-owner" => Ok(RpcBench::TokenAccountsByOwner),
+            "token-supply" => Ok(RpcBench::TokenSupply),
             "transaction" => Ok(RpcBench::Transaction),
             "transaction-parsed" => Ok(RpcBench::TransactionParsed),
             "version" => Ok(RpcBench::Version),
@@ -701,6 +703,27 @@ fn run_rpc_bench_loop(
                     }
                 }
             }
+            RpcBench::TokenSupply => {
+                let mut rpc_time = Measure::start("rpc-get-token-supply");
+                match client
+                    .get_token_supply_with_commitment(&mint.unwrap(), CommitmentConfig::confirmed())
+                {
+                    Ok(_ui_token_amount) => {
+                        rpc_time.stop();
+                        stats.success += 1;
+                        stats.total_success_time_us += rpc_time.as_us();
+                    }
+                    Err(e) => {
+                        rpc_time.stop();
+                        stats.errors += 1;
+                        stats.total_errors_time_us += rpc_time.as_us();
+                        if last_error.elapsed().as_secs() > 2 {
+                            info!("get-token-supply error: {:?}", e);
+                            last_error = Instant::now();
+                        }
+                    }
+                }
+            }
             RpcBench::Transaction => {
                 process_get_transaction(
                     "rpc-get-transaction-base64",
@@ -1264,6 +1287,7 @@ fn main() {
                 .requires_ifs(&[
                     ("supply", "mint"),
                     ("token-accounts-by-owner", "mint"),
+                    ("token-supply", "mint"),
                 ])
                 .help("Spawn a thread which calls a specific RPC method in a loop to benchmark it"),
         )