Node identity for bench

bench-tps/src/cli.rs

@@ -1,7 +1,7 @@
 use {
     crate::spl_convert::FromOtherSolana,
     clap::{crate_description, crate_name, App, Arg, ArgMatches},
-    solana_clap_utils::input_validators::{is_url, is_url_or_moniker, is_within_range},
+    solana_clap_utils::input_validators::{is_keypair, is_url, is_url_or_moniker, is_within_range},
     solana_cli_config::{ConfigInput, CONFIG_FILE},
     solana_sdk::{
         fee_calculator::FeeRateGovernor,
@@ -10,7 +10,7 @@ use {
     },
     solana_tpu_client::tpu_client::{DEFAULT_TPU_CONNECTION_POOL_SIZE, DEFAULT_TPU_USE_QUIC},
     std::{
-        net::{Ipv4Addr, SocketAddr},
+        net::{IpAddr, Ipv4Addr, SocketAddr},
         process::exit,
         time::Duration,
     },
@@ -68,6 +68,8 @@ pub struct Config {
     pub use_durable_nonce: bool,
     pub instruction_padding_config: Option<InstructionPaddingConfig>,
     pub num_conflict_groups: Option<usize>,
+    pub bind_address: IpAddr,
+    pub client_node_id: Option<Keypair>,
 }
 
 impl Default for Config {
@@ -99,6 +101,8 @@ impl Default for Config {
             use_durable_nonce: false,
             instruction_padding_config: None,
             num_conflict_groups: None,
+            bind_address: IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0)),
+            client_node_id: None,
         }
     }
 }
@@ -353,6 +357,24 @@ pub fn build_args<'a>(version: &'_ str) -> App<'a, '_> {
                 .validator(|arg| is_within_range(arg, 1..))
                 .help("The number of unique destination accounts per transactions 'chunk'. Lower values will result in more transaction conflicts.")
         )
+        .arg(
+            Arg::with_name("bind_address")
+                .long("bind-address")
+                .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")
+                .validator(is_keypair)
+                .help("File containing the node identity (keypair) of a validator with active stake. This allows communicating with network using staked connection"),
+        )
 }
 
 /// Parses a clap `ArgMatches` structure into a `Config`
@@ -513,5 +535,54 @@ pub fn extract_args(matches: &ArgMatches) -> Config {
         );
     }
 
+    if let Some(addr) = matches.value_of("bind_address") {
+        args.bind_address = solana_net_utils::parse_host(addr).unwrap_or_else(|e| {
+            eprintln!("Failed to parse bind_address: {e}");
+            exit(1)
+        });
+    }
+    let (_, node_id_path) = ConfigInput::compute_keypair_path_setting(
+        matches.value_of("client_node_id").unwrap_or(""),
+        &config.keypair_path,
+    );
+    // error is checked by arg validator
+    if let Ok(node_id) = read_keypair_file(node_id_path) {
+        args.client_node_id = Some(node_id);
+    }
+
     args
 }
+
+#[cfg(test)]
+mod tests {
+    use {
+        super::*,
+        std::{fs::File, io::prelude::*, path::Path},
+    };
+
+    fn write_keypair_to_file(keypair: &Keypair, file_name: &str) {
+        let serialized = serde_json::to_string(&keypair.to_bytes().to_vec()).unwrap();
+        let path = Path::new(file_name);
+        let mut file = File::create(path).unwrap();
+        file.write_all(&serialized.into_bytes()).unwrap();
+    }
+
+    #[test]
+    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, Some(keypair));
+    }
+}

bench-tps/src/main.rs

@@ -1,5 +1,4 @@
 #![allow(clippy::integer_arithmetic)]
-
 use {
     clap::value_t,
     log::*,
@@ -19,43 +18,118 @@ use {
     solana_gossip::gossip_service::{discover_cluster, get_client, get_multi_client},
     solana_rpc_client::rpc_client::RpcClient,
     solana_sdk::{
-        commitment_config::CommitmentConfig, fee_calculator::FeeRateGovernor, pubkey::Pubkey,
+        commitment_config::CommitmentConfig,
+        fee_calculator::FeeRateGovernor,
+        pubkey::Pubkey,
+        signature::{Keypair, Signer},
         system_program,
     },
-    solana_streamer::socket::SocketAddrSpace,
+    solana_streamer::{socket::SocketAddrSpace, streamer::StakedNodes},
     std::{
-        collections::HashMap, fs::File, io::prelude::*, net::SocketAddr, path::Path, process::exit,
-        sync::Arc,
+        collections::HashMap,
+        fs::File,
+        io::prelude::*,
+        net::{IpAddr, SocketAddr},
+        path::Path,
+        process::exit,
+        sync::{Arc, RwLock},
     },
 };
 
 /// 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;
 
+/// Request information about node's stake
+/// If fail to get requested information, return error
+/// Otherwise return stake of the node
+/// along with total activated stake of the network
+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()
+        .map(|vote_account| vote_account.activated_stake)
+        .sum();
+
+    let node_id_as_str = node_id.to_string();
+    let find_result = vote_accounts
+        .current
+        .iter()
+        .find(|&vote_account| vote_account.node_pubkey == node_id_as_str);
+    match find_result {
+        Some(value) => Ok((value.activated_stake, total_active_stake)),
+        None => {
+            error!("failed to find stake for requested node");
+            Err(())
+        }
+    }
+}
+
+fn create_connection_cache(
+    json_rpc_url: &str,
+    tpu_connection_pool_size: usize,
+    use_quic: bool,
+    bind_address: IpAddr,
+    client_node_id: Option<&Keypair>,
+) -> ConnectionCache {
+    if !use_quic {
+        return ConnectionCache::with_udp(tpu_connection_pool_size);
+    }
+    if client_node_id.is_none() {
+        return ConnectionCache::new(tpu_connection_pool_size);
+    }
+
+    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 {
+        total_stake,
+        pubkey_stake_map: HashMap::from([(client_node_id.pubkey(), stake)]),
+        ..StakedNodes::default()
+    }));
+    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)]
 fn create_client(
     external_client_type: &ExternalClientType,
     entrypoint_addr: &SocketAddr,
     json_rpc_url: &str,
     websocket_url: &str,
     multi_client: bool,
-    use_quic: bool,
-    tpu_connection_pool_size: usize,
     rpc_tpu_sockets: Option<(SocketAddr, SocketAddr)>,
     num_nodes: usize,
     target_node: Option<Pubkey>,
+    connection_cache: ConnectionCache,
 ) -> Arc<dyn BenchTpsClient + Send + Sync> {
     match external_client_type {
         ExternalClientType::RpcClient => Arc::new(RpcClient::new_with_commitment(
             json_rpc_url.to_string(),
             CommitmentConfig::confirmed(),
         )),
         ExternalClientType::ThinClient => {
-            let connection_cache = match use_quic {
-                true => Arc::new(ConnectionCache::new(tpu_connection_pool_size)),
-                false => Arc::new(ConnectionCache::with_udp(tpu_connection_pool_size)),
-            };
-
+            let connection_cache = Arc::new(connection_cache);
             if let Some((rpc, tpu)) = rpc_tpu_sockets {
                 Arc::new(ThinClient::new(rpc, tpu, connection_cache))
             } else {
@@ -106,10 +180,6 @@ fn create_client(
                 json_rpc_url.to_string(),
                 CommitmentConfig::confirmed(),
             ));
-            let connection_cache = match use_quic {
-                true => ConnectionCache::new(tpu_connection_pool_size),
-                false => ConnectionCache::with_udp(tpu_connection_pool_size),
-            };
             match connection_cache {
                 ConnectionCache::Udp(cache) => Arc::new(
                     TpuClient::new_with_connection_cache(
@@ -168,6 +238,8 @@ fn main() {
         use_randomized_compute_unit_price,
         use_durable_nonce,
         instruction_padding_config,
+        bind_address,
+        client_node_id,
         ..
     } = &cli_config;
 
@@ -225,17 +297,23 @@ fn main() {
             None
         };
 
+    let connection_cache = create_connection_cache(
+        json_rpc_url,
+        *tpu_connection_pool_size,
+        *use_quic,
+        *bind_address,
+        client_node_id.as_ref(),
+    );
     let client = create_client(
         external_client_type,
         entrypoint_addr,
         json_rpc_url,
         websocket_url,
         *multi_client,
-        *use_quic,
-        *tpu_connection_pool_size,
         rpc_tpu_sockets,
         *num_nodes,
         *target_node,
+        connection_cache,
     );
     if let Some(instruction_padding_config) = instruction_padding_config {
         info!(