From 2a51098dafdc8a4bcdcf1c4beea499b2ab4762a3 Mon Sep 17 00:00:00 2001
From: Dmitry <b4tm4n@mail.ru>
Date: Thu, 12 Sep 2024 20:43:07 +0300
Subject: [PATCH] stop on ctrl+c

---
 Cargo.lock  | 29 +++++++++++++++++++++++++++++
 Cargo.toml  |  1 +
 src/main.rs | 26 +++++++++++++++++---------
 3 files changed, 47 insertions(+), 9 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index fb16623..fa9ad73 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -207,6 +207,12 @@ version = "1.0.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "baf1de4339761588bc0619e3cbc0120ee582ebb74b53b4efbf79117bd2da40fd"
 
+[[package]]
+name = "cfg_aliases"
+version = "0.2.1"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724"
+
 [[package]]
 name = "chrono"
 version = "0.4.38"
@@ -282,6 +288,16 @@ version = "0.8.7"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b"
 
+[[package]]
+name = "ctrlc"
+version = "3.4.5"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "90eeab0aa92f3f9b4e87f258c72b139c207d251f9cbc1080a0086b86a8870dd3"
+dependencies = [
+ "nix",
+ "windows-sys 0.59.0",
+]
+
 [[package]]
 name = "deranged"
 version = "0.3.11"
@@ -300,6 +316,7 @@ dependencies = [
  "bollard",
  "clap",
  "clap_derive",
+ "ctrlc",
  "flexi_logger",
  "futures-util",
  "log",
@@ -732,6 +749,18 @@ dependencies = [
  "windows-sys 0.52.0",
 ]
 
+[[package]]
+name = "nix"
+version = "0.29.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "71e2746dc3a24dd78b3cfcb7be93368c6de9963d30f43a6a73998a9cf4b17b46"
+dependencies = [
+ "bitflags",
+ "cfg-if",
+ "cfg_aliases",
+ "libc",
+]
+
 [[package]]
 name = "nu-ansi-term"
 version = "0.50.1"
diff --git a/Cargo.toml b/Cargo.toml
index a5b1fd1..be298f4 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -11,6 +11,7 @@ anyhow = "1.0.88"
 bollard = { version = "0.17.1", features = ["ssl", "json_data_content", "tokio-stream"] }
 clap = { version = "4.5.17", features = ["derive"] }
 clap_derive = "4.5.13"
+ctrlc = "3.4.5"
 flexi_logger = "0.29.0"
 futures-util = "0.3.30"
 log = "0.4.22"
diff --git a/src/main.rs b/src/main.rs
index 6d1beef..67920f8 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -160,13 +160,25 @@ async fn restart_task(connection: Arc<Docker>, mut rx: mpsc::Receiver<Containers
             log::warn!("restart -> container: {}...", &container_id);
             let res = restart_container(&connection, container_id.as_str()).await;
             match res {
-                Ok(_) => log::info!("ok\n"),
-                Err(e) => log::error!("error: \n{e:?}\n"),
+                Ok(_) => log::info!("ok"),
+                Err(e) => log::error!("error: \n{e:?}"),
             }
         }
     }
 }
 
+fn shutdown_control(shutdown: Option<mpsc::Sender<()>>) {
+    let mut shutdown = shutdown;
+    let res = ctrlc::set_handler(move || {
+        log::info!("recieved Ctrl-C");
+        shutdown.take();
+    });
+
+    if res.is_ok() {
+        log::info!("Press Ctrl-C to stop");
+    }
+}
+
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
     let logger = create_logger()?;
@@ -190,6 +202,8 @@ async fn main() -> anyhow::Result<()> {
     let interval = cli.interval;
     let unhealthy_timeout = cli.unhealthy_timeout;
 
+    shutdown_control(Some(shutdown_tx));
+
     tokio::try_join!(
         tokio::spawn(query_task(
             query_connection,
@@ -199,13 +213,7 @@ async fn main() -> anyhow::Result<()> {
             cli
         )),
         tokio::spawn(filter_task(unhealthy_timeout, filter_rx, filter_tx)),
-        tokio::spawn(restart_task(restart_connection, restart_rx)),
-        tokio::spawn(async {
-            log::debug!("shutdown -> sleep");
-            tokio::time::sleep(Duration::from_secs(80)).await;
-            log::warn!("shutdown -> drop");
-            drop(shutdown_tx);
-        })
+        tokio::spawn(restart_task(restart_connection, restart_rx))
     )?;
 
     drop(logger);