Skip to content

Commit

Permalink
Improve cli args parsing in bench-tps (solana-labs#30307)
Browse files Browse the repository at this point in the history 
* Improve cli args parsing in bench-tps

* Rename parse funciton
* Get rid of panics, exit and println
* Add basic unit tests

* add allow dead_code and unused_imports for cli tests

* simplify code by using map_err instead of macro

* simplify data_size parsing

* use tempfile crate to create temp keypair files

* fix bug with reading client_id from config by default, re-added client_id tests

* minor fix of an error message

* add Eq to bench_tps::cli::Congig

* some minor error messages corrections
  • Loading branch information
@nickfrosty
kirill lykov authored and nickfrosty committed Mar 12, 2023
1 parent d376037 commit 81951e6
Show file tree
Hide file tree
Showing 4 changed files with 181 additions and 66 deletions.
1 change: 1 addition & 0 deletions Cargo.lock
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions bench-tps/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ thiserror = { workspace = true }
serial_test = { workspace = true }
solana-local-cluster = { workspace = true }
solana-test-validator = { workspace = true }
tempfile = "3.3.0"

[package.metadata.docs.rs]
targets = ["x86_64-unknown-linux-gnu"]
235 changes: 171 additions & 64 deletions bench-tps/src/cli.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,13 +11,13 @@ use {
solana_tpu_client::tpu_client::{DEFAULT_TPU_CONNECTION_POOL_SIZE, DEFAULT_TPU_USE_QUIC},
std::{
net::{IpAddr, Ipv4Addr, SocketAddr},
process::exit,
time::Duration,
},
};

const NUM_LAMPORTS_PER_ACCOUNT_DEFAULT: u64 = solana_sdk::native_token::LAMPORTS_PER_SOL;

#[derive(Eq, PartialEq, Debug)]
pub enum ExternalClientType {
// Submits transactions to an Rpc node using an RpcClient
RpcClient,
Expand All @@ -35,12 +35,14 @@ impl Default for ExternalClientType {
}
}

#[derive(Eq, PartialEq, Debug)]
pub struct InstructionPaddingConfig {
pub program_id: Pubkey,
pub data_size: u32,
}

/// Holds the configuration for a single run of the benchmark
#[derive(PartialEq, Debug)]
pub struct Config {
pub entrypoint_addr: SocketAddr,
pub json_rpc_url: String,
Expand Down Expand Up @@ -72,6 +74,8 @@ pub struct Config {
pub client_node_id: Option<Keypair>,
}

impl Eq for Config {}

impl Default for Config {
fn default() -> Config {
Config {
Expand Down Expand Up @@ -378,11 +382,7 @@ pub fn build_args<'a>(version: &'_ str) -> App<'a, '_> {
}

/// Parses a clap `ArgMatches` structure into a `Config`
/// # Arguments
/// * `matches` - command line arguments parsed by clap
/// # Panics
/// Panics if there is trouble parsing any of the arguments
pub fn extract_args(matches: &ArgMatches) -> Config {
pub fn parse_args(matches: &ArgMatches) -> Result<Config, &'static str> {
let mut args = Config::default();

let config = if let Some(config_file) = matches.value_of("config_file") {
Expand Down Expand Up @@ -411,7 +411,7 @@ pub fn extract_args(matches: &ArgMatches) -> Config {
if let Ok(id) = read_keypair_file(id_path) {
args.id = id;
} else if matches.is_present("identity") {
panic!("could not parse identity path");
return Err("could not parse identity path");
}

if matches.is_present("tpu_client") {
Expand All @@ -427,49 +427,50 @@ pub fn extract_args(matches: &ArgMatches) -> Config {
if let Some(v) = matches.value_of("tpu_connection_pool_size") {
args.tpu_connection_pool_size = v
.to_string()
.parse()
.expect("can't parse tpu_connection_pool_size");
.parse::<usize>()
.map_err(|_| "can't parse tpu-connection-pool-size")?;
}

if let Some(addr) = matches.value_of("entrypoint") {
args.entrypoint_addr = solana_net_utils::parse_host_port(addr).unwrap_or_else(|e| {
eprintln!("failed to parse entrypoint address: {e}");
exit(1)
});
args.entrypoint_addr = solana_net_utils::parse_host_port(addr)
.map_err(|_| "failed to parse entrypoint address")?;
}

if let Some(t) = matches.value_of("threads") {
args.threads = t.to_string().parse().expect("can't parse threads");
args.threads = t.to_string().parse().map_err(|_| "can't parse threads")?;
}

if let Some(n) = matches.value_of("num-nodes") {
args.num_nodes = n.to_string().parse().expect("can't parse num-nodes");
args.num_nodes = n.to_string().parse().map_err(|_| "can't parse num-nodes")?;
}

if let Some(duration) = matches.value_of("duration") {
args.duration = Duration::new(
duration.to_string().parse().expect("can't parse duration"),
0,
);
let seconds = duration
.to_string()
.parse()
.map_err(|_| "can't parse duration")?;
args.duration = Duration::new(seconds, 0);
}

if let Some(s) = matches.value_of("tx_count") {
args.tx_count = s.to_string().parse().expect("can't parse tx_count");
args.tx_count = s.to_string().parse().map_err(|_| "can't parse tx_count")?;
}

if let Some(s) = matches.value_of("keypair_multiplier") {
args.keypair_multiplier = s
.to_string()
.parse()
.expect("can't parse keypair-multiplier");
assert!(args.keypair_multiplier >= 2);
.map_err(|_| "can't parse keypair-multiplier")?;
if args.keypair_multiplier >= 2 {
return Err("args.keypair_multiplier must be lower than 2");
}
}

if let Some(t) = matches.value_of("thread-batch-sleep-ms") {
args.thread_batch_sleep_ms = t
.to_string()
.parse()
.expect("can't parse thread-batch-sleep-ms");
.map_err(|_| "can't parse thread-batch-sleep-ms")?;
}

args.sustained = matches.is_present("sustained");
Expand All @@ -486,23 +487,31 @@ pub fn extract_args(matches: &ArgMatches) -> Config {
}

if let Some(v) = matches.value_of("target_lamports_per_signature") {
args.target_lamports_per_signature = v.to_string().parse().expect("can't parse lamports");
args.target_lamports_per_signature = v
.to_string()
.parse()
.map_err(|_| "can't parse target-lamports-per-signature")?;
}

args.multi_client = !matches.is_present("no-multi-client");
args.target_node = matches
.value_of("target_node")
.map(|target_str| target_str.parse().unwrap());
.map(|target_str| target_str.parse::<Pubkey>())
.transpose()
.map_err(|_| "Failed to parse target-node")?;

if let Some(v) = matches.value_of("num_lamports_per_account") {
args.num_lamports_per_account = v.to_string().parse().expect("can't parse lamports");
args.num_lamports_per_account = v
.to_string()
.parse()
.map_err(|_| "can't parse num-lamports-per-account")?;
}

if let Some(t) = matches.value_of("target_slots_per_epoch") {
args.target_slots_per_epoch = t
.to_string()
.parse()
.expect("can't parse target slots per epoch");
.map_err(|_| "can't parse target-slots-per-epoch")?;
}

if matches.is_present("use_randomized_compute_unit_price") {
Expand All @@ -518,71 +527,169 @@ pub fn extract_args(matches: &ArgMatches) -> Config {
.value_of("instruction_padding_program_id")
.map(|target_str| target_str.parse().unwrap())
.unwrap_or_else(|| FromOtherSolana::from(spl_instruction_padding::ID));
let data_size = data_size
.parse()
.map_err(|_| "Can't parse padded instruction data size")?;
args.instruction_padding_config = Some(InstructionPaddingConfig {
program_id,
data_size: data_size
.to_string()
.parse()
.expect("Can't parse padded instruction data size"),
data_size,
});
}

if let Some(num_conflict_groups) = matches.value_of("num_conflict_groups") {
args.num_conflict_groups = Some(
num_conflict_groups
.parse()
.expect("Can't parse conflict groups"),
);
let parsed_num_conflict_groups = num_conflict_groups
.parse()
.map_err(|_| "Can't parse num-conflict-groups")?;
args.num_conflict_groups = Some(parsed_num_conflict_groups);
}

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)
});
args.bind_address =
solana_net_utils::parse_host(addr).map_err(|_| "Failed to parse bind-address")?;
}
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);

if let Some(client_node_id_filename) = matches.value_of("client_node_id") {
// error is checked by arg validator
let client_node_id = read_keypair_file(client_node_id_filename).map_err(|_| "")?;
args.client_node_id = Some(client_node_id);
}

args
Ok(args)
}

#[cfg(test)]
mod tests {
use {
super::*,
std::{fs::File, io::prelude::*, path::Path},
solana_sdk::signature::{read_keypair_file, write_keypair_file, Keypair, Signer},
std::{
net::{IpAddr, Ipv4Addr, SocketAddr},
time::Duration,
},
tempfile::{tempdir, TempDir},
};

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();
/// create a keypair and write it to json file in temporary directory
/// return both generated keypair and full file name
fn write_tmp_keypair(out_dir: &TempDir) -> (Keypair, String) {
let keypair = Keypair::new();
let file_path = out_dir
.path()
.join(format!("keypair_file-{}", keypair.pubkey()));
let keypair_file_name = file_path.into_os_string().into_string().unwrap();
write_keypair_file(&keypair, &keypair_file_name).unwrap();
(keypair, keypair_file_name)
}

#[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);
#[allow(clippy::cognitive_complexity)]
fn test_cli_parse() {
// create a directory inside of std::env::temp_dir(), removed when out_dir goes out of scope
let out_dir = tempdir().unwrap();
let (keypair, keypair_file_name) = write_tmp_keypair(&out_dir);

// parse provided rpc address, check that default ws address is correct
// always specify identity in these tests because otherwise a random one will be used
let matches = build_args("1.0.0").get_matches_from(vec![
"solana-bench-tps",
"--identity",
&keypair_file_name,
"-u",
"http://123.4.5.6:8899",
]);
let actual = parse_args(&matches).unwrap();
assert_eq!(
actual,
Config {
json_rpc_url: "http://123.4.5.6:8899".to_string(),
websocket_url: "ws://123.4.5.6:8900/".to_string(),
id: keypair,
..Config::default()
}
);

// parse cli args typical for private cluster tests
let keypair = read_keypair_file(&keypair_file_name).unwrap();
let matches = build_args("1.0.0").get_matches_from(vec![
"solana-bench-tps",
"-u http://192.0.0.1:8899",
"--identity",
&keypair_file_name,
"-u",
"http://123.4.5.6:8899",
"--duration",
"1000",
"--sustained",
"--threads",
"4",
"--read-client-keys",
"./client-accounts.yml",
"--entrypoint",
"192.1.2.3:8001",
]);
let actual = parse_args(&matches).unwrap();
assert_eq!(
actual,
Config {
json_rpc_url: "http://123.4.5.6:8899".to_string(),
websocket_url: "ws://123.4.5.6:8900/".to_string(),
id: keypair,
duration: Duration::new(1000, 0),
sustained: true,
threads: 4,
read_from_client_file: true,
client_ids_and_stake_file: "./client-accounts.yml".to_string(),
entrypoint_addr: SocketAddr::from((Ipv4Addr::new(192, 1, 2, 3), 8001)),
..Config::default()
}
);

// select different client type
let keypair = read_keypair_file(&keypair_file_name).unwrap();
let matches = build_args("1.0.0").get_matches_from(vec![
"solana-bench-tps",
"--identity",
&keypair_file_name,
"-u",
"http://123.4.5.6:8899",
"--use-tpu-client",
]);
let actual = parse_args(&matches).unwrap();
assert_eq!(
actual,
Config {
json_rpc_url: "http://123.4.5.6:8899".to_string(),
websocket_url: "ws://123.4.5.6:8900/".to_string(),
id: keypair,
external_client_type: ExternalClientType::TpuClient,
..Config::default()
}
);

// with client node id
let keypair = read_keypair_file(&keypair_file_name).unwrap();
let (client_id, client_id_file_name) = write_tmp_keypair(&out_dir);
let matches = build_args("1.0.0").get_matches_from(vec![
"solana-bench-tps",
"--identity",
&keypair_file_name,
"-u",
"http://192.0.0.1:8899",
"--bind-address",
"192.0.0.1",
"192.9.8.7",
"--client-node-id",
keypair_file_name,
&client_id_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));
let actual = parse_args(&matches).unwrap();
assert_eq!(
actual,
Config {
json_rpc_url: "http://192.0.0.1:8899".to_string(),
websocket_url: "ws://192.0.0.1:8900/".to_string(),
id: keypair,
bind_address: IpAddr::V4(Ipv4Addr::new(192, 9, 8, 7)),
client_node_id: Some(client_id),
..Config::default()
}
);
}
}
Loading

0 comments on commit 81951e6

Please sign in to comment.