Skip to content

Commit 0b0589e

Browse files
authored
Add custom rpc/tpu port options for ThinClient (solana-labs#24842)
1 parent 88c16c0 commit 0b0589e

File tree

2 files changed

+69
-29
lines changed

2 files changed

+69
-29
lines changed

bench-tps/src/cli.rs

+20
Original file line numberDiff line numberDiff line change
@@ -122,6 +122,26 @@ pub fn build_args<'a, 'b>(version: &'b str) -> App<'a, 'b> {
122122
.validator(is_url)
123123
.help("WebSocket URL for the solana cluster"),
124124
)
125+
.arg(
126+
Arg::with_name("rpc_addr")
127+
.long("rpc-addr")
128+
.value_name("HOST:PORT")
129+
.takes_value(true)
130+
.conflicts_with("tpu_client")
131+
.conflicts_with("rpc_client")
132+
.requires("tpu_addr")
133+
.help("Specify custom rpc_addr to create thin_client"),
134+
)
135+
.arg(
136+
Arg::with_name("tpu_addr")
137+
.long("tpu-addr")
138+
.value_name("HOST:PORT")
139+
.conflicts_with("tpu_client")
140+
.conflicts_with("rpc_client")
141+
.takes_value(true)
142+
.requires("rpc_addr")
143+
.help("Specify custom tpu_addr to create thin_client"),
144+
)
125145
.arg(
126146
Arg::with_name("entrypoint")
127147
.short("n")

bench-tps/src/main.rs

+49-29
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
#![allow(clippy::integer_arithmetic)]
22
use {
3+
clap::value_t,
34
log::*,
45
solana_bench_tps::{
56
bench::{do_bench_tps, generate_keypairs},
@@ -100,41 +101,60 @@ fn main() {
100101
do_bench_tps(client, cli_config, keypairs);
101102
}
102103
ExternalClientType::ThinClient => {
103-
let nodes = discover_cluster(entrypoint_addr, *num_nodes, SocketAddrSpace::Unspecified)
104-
.unwrap_or_else(|err| {
105-
eprintln!("Failed to discover {} nodes: {:?}", num_nodes, err);
104+
let client = if let Ok(rpc_addr) = value_t!(matches, "rpc_addr", String) {
105+
let rpc = rpc_addr.parse().unwrap_or_else(|e| {
106+
eprintln!("RPC address should parse as socketaddr {:?}", e);
106107
exit(1);
107108
});
108-
if *use_quic {
109-
connection_cache::set_use_quic(true);
110-
}
111-
let client = if *multi_client {
112-
let (client, num_clients) = get_multi_client(&nodes, &SocketAddrSpace::Unspecified);
113-
if nodes.len() < num_clients {
114-
eprintln!(
115-
"Error: Insufficient nodes discovered. Expecting {} or more",
116-
num_nodes
117-
);
118-
exit(1);
109+
let tpu = value_t!(matches, "tpu_addr", String)
110+
.unwrap()
111+
.parse()
112+
.unwrap_or_else(|e| {
113+
eprintln!("TPU address should parse to a socket: {:?}", e);
114+
exit(1);
115+
});
116+
117+
solana_client::thin_client::create_client(rpc, tpu)
118+
} else {
119+
let nodes =
120+
discover_cluster(entrypoint_addr, *num_nodes, SocketAddrSpace::Unspecified)
121+
.unwrap_or_else(|err| {
122+
eprintln!("Failed to discover {} nodes: {:?}", num_nodes, err);
123+
exit(1);
124+
});
125+
if *use_quic {
126+
connection_cache::set_use_quic(true);
119127
}
120-
Arc::new(client)
121-
} else if let Some(target_node) = target_node {
122-
info!("Searching for target_node: {:?}", target_node);
123-
let mut target_client = None;
124-
for node in nodes {
125-
if node.id == *target_node {
126-
target_client =
127-
Some(Arc::new(get_client(&[node], &SocketAddrSpace::Unspecified)));
128-
break;
128+
if *multi_client {
129+
let (client, num_clients) =
130+
get_multi_client(&nodes, &SocketAddrSpace::Unspecified);
131+
if nodes.len() < num_clients {
132+
eprintln!(
133+
"Error: Insufficient nodes discovered. Expecting {} or more",
134+
num_nodes
135+
);
136+
exit(1);
137+
}
138+
client
139+
} else if let Some(target_node) = target_node {
140+
info!("Searching for target_node: {:?}", target_node);
141+
let mut target_client = None;
142+
for node in nodes {
143+
if node.id == *target_node {
144+
target_client =
145+
Some(get_client(&[node], &SocketAddrSpace::Unspecified));
146+
break;
147+
}
129148
}
149+
target_client.unwrap_or_else(|| {
150+
eprintln!("Target node {} not found", target_node);
151+
exit(1);
152+
})
153+
} else {
154+
get_client(&nodes, &SocketAddrSpace::Unspecified)
130155
}
131-
target_client.unwrap_or_else(|| {
132-
eprintln!("Target node {} not found", target_node);
133-
exit(1);
134-
})
135-
} else {
136-
Arc::new(get_client(&nodes, &SocketAddrSpace::Unspecified))
137156
};
157+
let client = Arc::new(client);
138158
let keypairs = get_keypairs(
139159
client.clone(),
140160
id,

0 commit comments

Comments
 (0)