diff --git a/polkadot/utils/staking-miner/src/main.rs b/polkadot/utils/staking-miner/src/main.rs index 7a6a0b71dbcb76d43615d34c1f20dcd52c2394f3..b10f87950445522a1eceb549e7a3b21603f5023d 100644 --- a/polkadot/utils/staking-miner/src/main.rs +++ b/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!( diff --git a/polkadot/utils/staking-miner/src/opts.rs b/polkadot/utils/staking-miner/src/opts.rs index 1e7b1f2ba294a0e9df11a371cbcf0de2419f118e..f43744ad45e8aa96ee48a4f1cb0191941ba5839e 100644 --- a/polkadot/utils/staking-miner/src/opts.rs +++ b/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 }) + } ); } diff --git a/polkadot/utils/staking-miner/src/rpc.rs b/polkadot/utils/staking-miner/src/rpc.rs index 4cf10f9ffef7886e0b62217dd18e37da92f89234..8929afcbe65656b6a212736a3dc4d5d586a2ea48 100644 --- a/polkadot/utils/staking-miner/src/rpc.rs +++ b/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)))