Add handling of SIGTERM, SIGINT and SIGQUIT

Fix #5649

utils/staking-miner/Cargo.toml

@@ -7,17 +7,16 @@ edition = "2021"
 [dependencies]
 codec = { package = "parity-scale-codec", version = "3.0.0" }
 clap = { version = "3.1", features = ["derive", "env"] }
-tracing-subscriber = { version = "0.3.11", features = ["env-filter"]   }
+tracing-subscriber = { version = "0.3.11", features = ["env-filter"] }
 jsonrpsee = { version = "0.14.0", features = ["ws-client", "macros"] }
 log = "0.4.17"
 paste = "1.0.7"
 serde = "1.0.137"
 serde_json = "1.0"
 thiserror = "1.0.31"
 tokio = { version = "1.18.2", features = ["macros", "rt-multi-thread", "sync"] }
-
 remote-externalities = { git = "https://github.com/paritytech/substrate", branch = "master" }
-
+signal-hook-tokio = { version = "0.3", features = ["futures-v0_3"] }
 sp-core = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-version = { git = "https://github.com/paritytech/substrate", branch = "master" }
 sp-io = { git = "https://github.com/paritytech/substrate", branch = "master" }
@@ -39,8 +38,11 @@ runtime-common = { package = "polkadot-runtime-common", path = "../../runtime/co
 polkadot-runtime = { path = "../../runtime/polkadot" }
 kusama-runtime = { path = "../../runtime/kusama" }
 westend-runtime = { path = "../../runtime/westend" }
+exitcode = "1.1"
 
 sub-tokens = { git = "https://github.com/paritytech/substrate-debug-kit", branch = "master" }
+signal-hook = "0.3"
+futures-util = "0.3"
 
 [dev-dependencies]
 assert_cmd = "2.0.4"

utils/staking-miner/src/main.rs

@@ -48,14 +48,16 @@ use crate::opts::*;
 use clap::Parser;
 use frame_election_provider_support::NposSolver;
 use frame_support::traits::Get;
+use futures_util::StreamExt;
 use jsonrpsee::ws_client::{WsClient, WsClientBuilder};
 use remote_externalities::{Builder, Mode, OnlineConfig};
 use rpc::{RpcApiClient, SharedRpcClient};
 use runtime_versions::RuntimeVersions;
 use sp_npos_elections::BalancingConfig;
 use sp_runtime::{traits::Block as BlockT, DeserializeOwned};
-use std::{ops::Deref, sync::Arc};
-use tracing_subscriber::{fmt, EnvFilter};
+use signal_hook::consts::signal::*;
+use signal_hook_tokio::Signals;
+
 pub(crate) enum AnyRuntime {
 	Polkadot,
 	Kusama,
@@ -437,13 +439,53 @@ pub(crate) async fn check_versions<T: frame_system::Config + EPM::Config>(
 	}
 }
 
+/// Control how we exit the application
+fn controlled_exit(code: i32) {
+	log::warn!(target: LOG_TARGET, "Exiting application");
+	std::process::exit(code);
+}
+
+/// Handles the various signal and exit the application
+/// when appropriate.
+async fn handle_signals(mut signals: Signals) {
+	let mut keyboard_sig_count: u8 = 0;
+	while let Some(signal) = signals.next().await {
+		match signal {
+			SIGTERM | SIGINT => {
+				if keyboard_sig_count >= 1 {
+					log::warn!(
+						target: LOG_TARGET,
+						"Received keyboard termination signal, quitting..."
+					);
+					controlled_exit(exitcode::OK);
+				}
+				keyboard_sig_count += 1;
+				log::warn!(
+					target: LOG_TARGET,
+					"Received keyboard termination signal #{}, if you keep doing that I will really quit",
+					keyboard_sig_count
+				);
+			},
+			SIGQUIT => {
+				log::warn!(target: LOG_TARGET, "Received SIGQUIT, quitting...");
+				controlled_exit(exitcode::OK);
+			},
+			_ => unreachable!(),
+		}
+	}
+}
+
 #[tokio::main]
 async fn main() {
 	fmt().with_env_filter(EnvFilter::from_default_env()).init();
 
 	let Opt { uri, command } = Opt::parse();
 	log::debug!(target: LOG_TARGET, "attempting to connect to {:?}", uri);
 
+	let signals = Signals::new(&[SIGTERM, SIGINT, SIGQUIT]).expect("Failed initializing Signals");
+	let handle = signals.handle();
+	let signals_task = tokio::spawn(handle_signals(signals));
+
 	let rpc = loop {
 		match SharedRpcClient::new(&uri).await {
 			Ok(client) => break client,
@@ -555,6 +597,9 @@ async fn main() {
 		}
 	};
 	log::info!(target: LOG_TARGET, "round of execution finished. outcome = {:?}", outcome);
+
+	handle.close();
+	let _ = signals_task.await;
 }
 
 #[cfg(test)]