Remove current node stake and total stake from cli

Instead request them from the network

bench-tps/src/cli.rs

@@ -69,9 +69,7 @@ pub struct Config {
     pub instruction_padding_config: Option<InstructionPaddingConfig>,
     pub num_conflict_groups: Option<usize>,
     pub bind_address: IpAddr,
-    pub client_node_id: Keypair,
-    pub client_node_stake: u64,
-    pub client_node_total_stake: u64,
+    pub client_node_id: Option<Keypair>,
 }
 
 impl Default for Config {
@@ -104,9 +102,7 @@ impl Default for Config {
             instruction_padding_config: None,
             num_conflict_groups: None,
             bind_address: IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
-            client_node_id: Keypair::new(),
-            client_node_stake: 0,
-            client_node_total_stake: 0,
+            client_node_id: None,
         }
     }
 }
@@ -367,33 +363,17 @@ pub fn build_args<'a>(version: &'_ str) -> App<'a, '_> {
                 .value_name("HOST")
                 .takes_value(true)
                 .validator(solana_net_utils::is_host)
+                .requires("client_node_id")
                 .help("IP address to use with connection cache"),
         )
         .arg(
             Arg::with_name("client_node_id")
                 .long("client-node-id")
                 .value_name("PATH")
                 .takes_value(true)
+                .requires("json_rpc_url")
                 .help("File containing a node id (keypair) to communicate with validators"),
         )
-        .arg(
-            Arg::with_name("client_node_stake")
-                .long("client-node-stake")
-                .value_name("LAMPORTS")
-                .takes_value(true)
-                .help(
-                    "Stake of the client_node_id",
-                ),
-        )
-        .arg(
-            Arg::with_name("client_node_total_stake")
-                .long("client-node-total-stake")
-                .value_name("LAMPORTS")
-                .takes_value(true)
-                .help(
-                    "Total stake of the client_node_id",
-                ),
-        )
 }
 
 /// Parses a clap `ArgMatches` structure into a `Config`
@@ -565,16 +545,10 @@ pub fn extract_args(matches: &ArgMatches) -> Config {
         &config.keypair_path,
     );
     if let Ok(node_id) = read_keypair_file(node_id_path) {
-        args.client_node_id = node_id;
+        args.client_node_id = Some(node_id);
     } else if matches.is_present("client_node_id") {
         panic!("could not parse identity path");
     }
-    if let Some(v) = matches.value_of("client-node-stake") {
-        args.client_node_stake = v.to_string().parse().expect("can't parse lamports");
-    }
-    if let Some(v) = matches.value_of("client-node-total-stake") {
-        args.client_node_total_stake = v.to_string().parse().expect("can't parse lamports");
-    }
 
     args
 }
@@ -594,20 +568,21 @@ mod tests {
     }
 
     #[test]
-    fn test_cli_parse() {
+    fn test_cli_parse_with_client_node_id() {
         let keypair = Keypair::new();
         let keypair_file_name = "./keypair.json";
         write_keypair_to_file(&keypair, keypair_file_name);
 
         let matches = build_args("1.0.0").get_matches_from(vec![
             "solana-bench-tps",
+            "-u http://192.0.0.1:8899",
             "--bind-address",
             "192.0.0.1",
             "--client-node-id",
             keypair_file_name,
         ]);
         let result = extract_args(&matches);
         assert_eq!(result.bind_address, IpAddr::V4(Ipv4Addr::new(192, 0, 0, 1)));
-        assert_eq!(result.client_node_id, keypair);
+        assert_eq!(result.client_node_id, Some(keypair));
     }
 }

bench-tps/src/main.rs

@@ -1,5 +1,4 @@
 #![allow(clippy::integer_arithmetic)]
-
 use {
     clap::value_t,
     log::*,
@@ -40,33 +39,76 @@ use {
 /// Number of signatures for all transactions in ~1 week at ~100K TPS
 pub const NUM_SIGNATURES_FOR_TXS: u64 = 100_000 * 60 * 60 * 24 * 7;
 
+fn find_node_activated_stake(
+    rpc_client: Arc<RpcClient>,
+    node_id: Pubkey,
+) -> Result<(u64, u64), ()> {
+    let vote_accounts = rpc_client.get_vote_accounts();
+    if let Err(error) = vote_accounts {
+        error!("Failed to get vote accounts, error: {}", error);
+        return Err(());
+    }
+
+    let vote_accounts = vote_accounts.unwrap();
+
+    let total_active_stake: u64 = vote_accounts
+        .current
+        .iter()
+        .chain(vote_accounts.delinquent.iter())
+        .map(|vote_account| vote_account.activated_stake)
+        .sum();
+
+    let node_id_as_str = node_id.to_string();
+    vote_accounts
+        .current
+        .iter()
+        .find(|&vote_account| vote_account.node_pubkey == node_id_as_str)
+        .map_or(Err(()), |value| {
+            Ok((value.activated_stake, total_active_stake))
+        })
+}
+
 fn create_connection_cache(
+    json_rpc_url: &str,
     tpu_connection_pool_size: usize,
     use_quic: bool,
     bind_address: IpAddr,
-    client_node_id: &Keypair,
-    stake: u64,
-    total_stake: u64,
+    client_node_id: Option<&Keypair>,
 ) -> ConnectionCache {
-    match use_quic {
-        true => {
-            let staked_nodes = Arc::new(RwLock::new(StakedNodes::default()));
-            staked_nodes.write().unwrap().total_stake = total_stake;
-            staked_nodes
-                .write()
-                .unwrap()
-                .pubkey_stake_map
-                .insert(client_node_id.pubkey(), stake);
-
-            ConnectionCache::new_with_client_options(
-                tpu_connection_pool_size,
-                None,
-                Some((client_node_id, bind_address)),
-                Some((&staked_nodes, &client_node_id.pubkey())),
-            )
-        }
-        false => ConnectionCache::with_udp(tpu_connection_pool_size),
+    if !use_quic {
+        return ConnectionCache::with_udp(tpu_connection_pool_size);
     }
+    if client_node_id.is_none() {
+        return ConnectionCache::new_with_client_options(
+            tpu_connection_pool_size,
+            None,
+            None,
+            None,
+        );
+    }
+
+    let rpc_client = Arc::new(RpcClient::new_with_commitment(
+        json_rpc_url.to_string(),
+        CommitmentConfig::confirmed(),
+    ));
+
+    let client_node_id = client_node_id.unwrap();
+    let (stake, total_stake) =
+        find_node_activated_stake(rpc_client, client_node_id.pubkey()).unwrap_or_default();
+    info!("Stake for specified client_node_id: {stake}, total stake: {total_stake}");
+    let staked_nodes = Arc::new(RwLock::new(StakedNodes::default()));
+    staked_nodes.write().unwrap().total_stake = total_stake;
+    staked_nodes
+        .write()
+        .unwrap()
+        .pubkey_stake_map
+        .insert(client_node_id.pubkey(), stake);
+    ConnectionCache::new_with_client_options(
+        tpu_connection_pool_size,
+        None,
+        Some((client_node_id, bind_address)),
+        Some((&staked_nodes, &client_node_id.pubkey())),
+    )
 }
 
 #[allow(clippy::too_many_arguments)]
@@ -198,8 +240,6 @@ fn main() {
         instruction_padding_config,
         bind_address,
         client_node_id,
-        client_node_stake,
-        client_node_total_stake,
         ..
     } = &cli_config;
 
@@ -258,12 +298,11 @@ fn main() {
         };
 
     let connection_cache = create_connection_cache(
+        json_rpc_url,
         *tpu_connection_pool_size,
         *use_quic,
         *bind_address,
-        client_node_id,
-        *client_node_stake,
-        *client_node_total_stake,
+        client_node_id.as_ref(),
     );
     let client = create_client(
         external_client_type,