staking miner: CLI make request and connection timeout configurable (#6004)

* staking miner: CLI request and connection timeout

Make it possible to configure the `connection timeout` and `request timeout` via CLI flags.

* cargo fmt

polkadot/utils/staking-miner/src/main.rs

@@ -57,7 +57,7 @@ use signal_hook::consts::signal::*;
 use signal_hook_tokio::Signals;
 use sp_npos_elections::BalancingConfig;
 use sp_runtime::{traits::Block as BlockT, DeserializeOwned};
-use std::{ops::Deref, sync::Arc};
+use std::{ops::Deref, sync::Arc, time::Duration};
 use tracing_subscriber::{fmt, EnvFilter};
 
 pub(crate) enum AnyRuntime {
@@ -485,7 +485,7 @@ async fn handle_signals(mut signals: Signals) {
 async fn main() {
 	fmt().with_env_filter(EnvFilter::from_default_env()).init();
 
-	let Opt { uri, command } = Opt::parse();
+	let Opt { uri, command, connection_timeout, request_timeout } = Opt::parse();
 	log::debug!(target: LOG_TARGET, "attempting to connect to {:?}", uri);
 
 	let signals = Signals::new(&[SIGTERM, SIGINT, SIGQUIT]).expect("Failed initializing Signals");
@@ -493,7 +493,13 @@ async fn main() {
 	let signals_task = tokio::spawn(handle_signals(signals));
 
 	let rpc = loop {
-		match SharedRpcClient::new(&uri).await {
+		match SharedRpcClient::new(
+			&uri,
+			Duration::from_secs(connection_timeout as u64),
+			Duration::from_secs(request_timeout as u64),
+		)
+		.await
+		{
 			Ok(client) => break client,
 			Err(why) => {
 				log::warn!(

polkadot/utils/staking-miner/src/opts.rs

@@ -27,6 +27,14 @@ pub(crate) struct Opt {
 	#[clap(long, short, default_value = DEFAULT_URI, env = "URI", global = true)]
 	pub uri: String,
 
+	/// WS connection timeout in number of seconds.
+	#[clap(long, parse(try_from_str), default_value_t = 60)]
+	pub connection_timeout: usize,
+
+	/// WS request timeout in number of seconds.
+	#[clap(long, parse(try_from_str), default_value_t = 60 * 10)]
+	pub request_timeout: usize,
+
 	#[clap(subcommand)]
 	pub command: Command,
 }
@@ -223,6 +231,8 @@ mod test_super {
 			opt,
 			Opt {
 				uri: "hi".to_string(),
+				connection_timeout: 60,
+				request_timeout: 10 * 60,
 				command: Command::Monitor(MonitorConfig {
 					seed_or_path: "//Alice".to_string(),
 					listen: "head".to_string(),
@@ -251,6 +261,8 @@ mod test_super {
 			opt,
 			Opt {
 				uri: "hi".to_string(),
+				connection_timeout: 60,
+				request_timeout: 10 * 60,
 				command: Command::DryRun(DryRunConfig {
 					seed_or_path: "//Alice".to_string(),
 					at: None,
@@ -279,6 +291,8 @@ mod test_super {
 			opt,
 			Opt {
 				uri: "hi".to_string(),
+				connection_timeout: 60,
+				request_timeout: 10 * 60,
 				command: Command::EmergencySolution(EmergencySolutionConfig {
 					take: Some(99),
 					at: None,
@@ -294,7 +308,37 @@ mod test_super {
 
 		assert_eq!(
 			opt,
-			Opt { uri: "hi".to_string(), command: Command::Info(InfoOpts { json: false }) }
+			Opt {
+				uri: "hi".to_string(),
+				connection_timeout: 60,
+				request_timeout: 10 * 60,
+				command: Command::Info(InfoOpts { json: false })
+			}
+		);
+	}
+
+	#[test]
+	fn cli_request_conn_timeout_works() {
+		let opt = Opt::try_parse_from([
+			env!("CARGO_PKG_NAME"),
+			"--uri",
+			"hi",
+			"--request-timeout",
+			"10",
+			"--connection-timeout",
+			"9",
+			"info",
+		])
+		.unwrap();
+
+		assert_eq!(
+			opt,
+			Opt {
+				uri: "hi".to_string(),
+				connection_timeout: 9,
+				request_timeout: 10,
+				command: Command::Info(InfoOpts { json: false })
+			}
 		);
 	}
 

polkadot/utils/staking-miner/src/rpc.rs

@@ -27,9 +27,6 @@ use sp_core::{storage::StorageKey, Bytes};
 use sp_version::RuntimeVersion;
 use std::{future::Future, time::Duration};
 
-const MAX_CONNECTION_DURATION: Duration = Duration::from_secs(20);
-const MAX_REQUEST_DURATION: Duration = Duration::from_secs(60);
-
 #[derive(frame_support::DebugNoBound, thiserror::Error)]
 pub(crate) enum RpcHelperError {
 	JsonRpsee(#[from] jsonrpsee::core::Error),
@@ -125,11 +122,15 @@ impl SharedRpcClient {
 	}
 
 	/// Create a new shared JSON-RPC web-socket client.
-	pub(crate) async fn new(uri: &str) -> Result<Self, RpcError> {
+	pub(crate) async fn new(
+		uri: &str,
+		connection_timeout: Duration,
+		request_timeout: Duration,
+	) -> Result<Self, RpcError> {
 		let client = WsClientBuilder::default()
-			.connection_timeout(MAX_CONNECTION_DURATION)
+			.connection_timeout(connection_timeout)
 			.max_request_body_size(u32::MAX)
-			.request_timeout(MAX_REQUEST_DURATION)
+			.request_timeout(request_timeout)
 			.build(uri)
 			.await?;
 		Ok(Self(Arc::new(client)))