From 3a04d4428e4bd6ca12115242a1714f2f3799a3eb Mon Sep 17 00:00:00 2001
From: Serban Iorga <serban@parity.io>
Date: Wed, 28 Dec 2022 15:47:09 +0200
Subject: [PATCH] Handle `SIGTERM` for the docker containers + relay (#1735)

* Handle SIGTERM for some docker containers

* Implement SIGTERM handling for the relay
---
 bridges/relays/bin-substrate/Cargo.toml     |  2 ++
 bridges/relays/bin-substrate/src/cli/mod.rs | 36 +++++++++++++++++++--
 bridges/relays/bin-substrate/src/main.rs    |  5 +--
 3 files changed, 37 insertions(+), 6 deletions(-)

diff --git a/bridges/relays/bin-substrate/Cargo.toml b/bridges/relays/bin-substrate/Cargo.toml
index 9ae6a98868e39..de1807cdc9708 100644
--- a/bridges/relays/bin-substrate/Cargo.toml
+++ b/bridges/relays/bin-substrate/Cargo.toml
@@ -16,6 +16,8 @@ log = "0.4.17"
 num-format = "0.4"
 num-traits = "0.2"
 structopt = "0.3"
+signal-hook = "0.3.14"
+signal-hook-async-std = "0.2.2"
 strum = { version = "0.21.0", features = ["derive"] }
 
 # Bridge dependencies
diff --git a/bridges/relays/bin-substrate/src/cli/mod.rs b/bridges/relays/bin-substrate/src/cli/mod.rs
index 5733aa411d8e7..4bd735c6334ad 100644
--- a/bridges/relays/bin-substrate/src/cli/mod.rs
+++ b/bridges/relays/bin-substrate/src/cli/mod.rs
@@ -18,7 +18,11 @@
 
 use std::convert::TryInto;
 
+use async_std::prelude::*;
 use codec::{Decode, Encode};
+use futures::{select, FutureExt};
+use signal_hook::consts::*;
+use signal_hook_async_std::Signals;
 use structopt::{clap::arg_enum, StructOpt};
 use strum::{EnumString, EnumVariantNames};
 
@@ -37,6 +41,9 @@ mod relay_messages;
 mod relay_parachains;
 mod resubmit_transactions;
 
+/// The target that will be used when publishing logs related to this pallet.
+pub const LOG_TARGET: &str = "bridge";
+
 /// Parse relay CLI args.
 pub fn parse_args() -> Command {
 	Command::from_args()
@@ -100,8 +107,7 @@ impl Command {
 	}
 
 	/// Run the command.
-	pub async fn run(self) -> anyhow::Result<()> {
-		self.init_logger();
+	async fn do_run(self) -> anyhow::Result<()> {
 		match self {
 			Self::RelayHeaders(arg) => arg.run().await?,
 			Self::RelayMessages(arg) => arg.run().await?,
@@ -114,6 +120,32 @@ impl Command {
 		}
 		Ok(())
 	}
+
+	/// Run the command.
+	pub async fn run(self) {
+		self.init_logger();
+
+		let exit_signals = match Signals::new([SIGINT, SIGTERM]) {
+			Ok(signals) => signals,
+			Err(e) => {
+				log::error!(target: LOG_TARGET, "Could not register exit signals: {}", e);
+				return
+			},
+		};
+		let run = self.do_run().fuse();
+		futures::pin_mut!(exit_signals, run);
+
+		select! {
+			signal = exit_signals.next().fuse() => {
+				log::info!(target: LOG_TARGET, "Received exit signal {:?}", signal);
+			},
+			result = run => {
+				if let Err(e) = result {
+					log::error!(target: LOG_TARGET, "substrate-relay: {}", e);
+				}
+			},
+		}
+	}
 }
 
 arg_enum! {
diff --git a/bridges/relays/bin-substrate/src/main.rs b/bridges/relays/bin-substrate/src/main.rs
index bc84786ee27f2..8bd89a4193627 100644
--- a/bridges/relays/bin-substrate/src/main.rs
+++ b/bridges/relays/bin-substrate/src/main.rs
@@ -24,8 +24,5 @@ mod cli;
 fn main() {
 	let command = cli::parse_args();
 	let run = command.run();
-	let result = async_std::task::block_on(run);
-	if let Err(error) = result {
-		log::error!(target: "bridge", "substrate-relay: {}", error);
-	}
+	async_std::task::block_on(run);
 }