staking-miner: Add handling of SIGTERM, SIGKILL, SIGQUIT and SIGINT (#5780)

* Add handling of SIGTERM, SIGINT and SIGQUIT

Fix #5649

* Fix signals and log levels

* Add missing imports

* Add sig count to the logs

polkadot/Cargo.lock

@@ -278,9 +278,9 @@ dependencies = [
 
 [[package]]
 name = "async-std"
-version = "1.11.0"
+version = "1.12.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "52580991739c5cdb36cde8b2a516371c0a3b70dda36d916cc08b82372916808c"
+checksum = "62565bb4402e926b29953c785397c6dc0391b7b446e45008b0049eb43cec6f5d"
 dependencies = [
  "async-attributes",
  "async-channel",
@@ -297,7 +297,6 @@ dependencies = [
  "kv-log-macro",
  "log",
  "memchr",
- "num_cpus",
  "once_cell",
  "pin-project-lite 0.2.7",
  "pin-utils",
@@ -1155,9 +1154,9 @@ dependencies = [
 
 [[package]]
 name = "crossbeam-channel"
-version = "0.5.1"
+version = "0.5.5"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "06ed27e177f16d65f0f0c22a213e17c696ace5dd64b14258b52f9417ccb52db4"
+checksum = "4c02a4d71819009c192cf4872265391563fd6a84c81ff2c0f2a7026ca4c1d85c"
 dependencies = [
  "cfg-if 1.0.0",
  "crossbeam-utils",
@@ -1732,6 +1731,12 @@ dependencies = [
  "futures",
 ]
 
+[[package]]
+name = "exitcode"
+version = "1.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "de853764b47027c2e862a995c34978ffa63c1501f2e15f987ba11bd4f9bba193"
+
 [[package]]
 name = "expander"
 version = "0.0.4"
@@ -9715,9 +9720,9 @@ checksum = "43b2853a4d09f215c24cc5489c992ce46052d359b5109343cbafbf26bc62f8a3"
 
 [[package]]
 name = "signal-hook"
-version = "0.3.10"
+version = "0.3.14"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "9c98891d737e271a2954825ef19e46bd16bdb98e2746f2eec4f7a4ef7946efd1"
+checksum = "a253b5e89e2698464fc26b545c9edceb338e18a89effeeecfea192c3025be29d"
 dependencies = [
  "libc",
  "signal-hook-registry",
@@ -9732,6 +9737,18 @@ dependencies = [
  "libc",
 ]
 
+[[package]]
+name = "signal-hook-tokio"
+version = "0.3.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "213241f76fb1e37e27de3b6aa1b068a2c333233b59cca6634f634b80a27ecf1e"
+dependencies = [
+ "futures-core",
+ "libc",
+ "signal-hook",
+ "tokio",
+]
+
 [[package]]
 name = "signature"
 version = "1.4.0"
@@ -10571,9 +10588,11 @@ version = "0.9.26"
 dependencies = [
  "assert_cmd",
  "clap",
+ "exitcode",
  "frame-election-provider-support",
  "frame-support",
  "frame-system",
+ "futures-util",
  "jsonrpsee",
  "kusama-runtime",
  "log",
@@ -10590,6 +10609,8 @@ dependencies = [
  "sc-transaction-pool-api",
  "serde",
  "serde_json",
+ "signal-hook",
+ "signal-hook-tokio",
  "sp-core",
  "sp-io",
  "sp-npos-elections",

polkadot/utils/staking-miner/Cargo.toml

@@ -15,9 +15,8 @@ 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"

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

@@ -48,14 +48,18 @@ 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 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 tracing_subscriber::{fmt, EnvFilter};
+
 pub(crate) enum AnyRuntime {
 	Polkadot,
 	Kusama,
@@ -437,6 +441,46 @@ pub(crate) async fn check_versions<T: frame_system::Config + EPM::Config>(
 	}
 }
 
+/// Control how we exit the application
+fn controlled_exit(code: i32) {
+	log::info!(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 {
+			// Interrupts come from the keyboard
+			SIGQUIT | SIGINT => {
+				if keyboard_sig_count >= 1 {
+					log::info!(
+						target: LOG_TARGET,
+						"Received keyboard termination signal #{}/{}, quitting...",
+						keyboard_sig_count + 1,
+						2
+					);
+					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
+				);
+			},
+
+			SIGKILL | SIGTERM => {
+				log::info!(target: LOG_TARGET, "Received SIGKILL | SIGTERM, quitting...");
+				controlled_exit(exitcode::OK);
+			},
+			_ => unreachable!(),
+		}
+	}
+}
+
 #[tokio::main]
 async fn main() {
 	fmt().with_env_filter(EnvFilter::from_default_env()).init();
@@ -444,6 +488,10 @@ async fn main() {
 	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 +603,9 @@ async fn main() {
 		}
 	};
 	log::info!(target: LOG_TARGET, "round of execution finished. outcome = {:?}", outcome);
+
+	handle.close();
+	let _ = signals_task.await;
 }
 
 #[cfg(test)]