Close on CTRL-BREAK

Cargo.toml

@@ -29,7 +29,7 @@ anyhow = "1.0"
 env_logger = "0.10"
 yansi = "0.5"
 chrono = "0.4"
-tokio = { version = "1.32", features = ["macros"] }
+tokio = { version = "1.32", features = ["macros", "signal"] }
 futures = "0.3"
 flexi_logger = { version = "0.27", features = ["colors"] }
 regex = "1"

src/main.rs

@@ -13,6 +13,7 @@ use chrono::Utc;
 use clap::Parser;
 use futures::prelude::*;
 
+use tokio::select;
 use ya_client_model::activity::activity_state::*;
 use ya_client_model::activity::ExeScriptCommand;
 use ya_client_model::activity::{ActivityUsage, CommandResult, ExeScriptCommandResult};
@@ -25,12 +26,14 @@ use crate::agreement::AgreementDesc;
 use crate::cli::*;
 use crate::logger::*;
 use crate::process::ProcessController;
+use crate::signal::SignalMonitor;
 
 mod agreement;
 mod cli;
 mod logger;
 mod offer_template;
 mod process;
+mod signal;
 
 async fn send_state<T>(ctx: &ExeUnitContext<T>, new_state: ActivityState) -> anyhow::Result<()> {
     Ok(gsb::service(ctx.report_url.clone())
@@ -122,6 +125,13 @@ async fn main() -> anyhow::Result<()> {
         }
     };
 
+    select! {
+        res = handle_cli(cli) => return res,
+        res = handle_signals() => return res,
+    };
+}
+
+async fn handle_cli(cli: Cli) -> anyhow::Result<()> {
     match cli.runtime.to_lowercase().as_str() {
         "dummy" => run::<process::dummy::Dummy>(cli).await,
         _ => {
@@ -132,6 +142,12 @@ async fn main() -> anyhow::Result<()> {
     }
 }
 
+async fn handle_signals() -> anyhow::Result<()> {
+    let signal = SignalMonitor::default().recv().await?;
+    log::info!("{} received, Shutting down runtime...", signal);
+    Ok(())
+}
+
 #[derive(Clone)]
 struct ExeUnitContext<T> {
     pub activity_id: String,

src/signal.rs

@@ -0,0 +1,66 @@
+pub(crate) type Signal = &'static str;
+
+use tokio::task::JoinHandle;
+use tokio::{
+    select,
+    sync::{
+        oneshot,
+        oneshot::{Receiver, Sender},
+    },
+};
+
+#[cfg(target_family = "unix")]
+use tokio::signal::unix;
+#[cfg(target_family = "windows")]
+use tokio::signal::windows;
+
+pub struct SignalMonitor {
+    stop_tx: Sender<Signal>,
+    stop_rx: Receiver<Signal>,
+}
+
+impl Default for SignalMonitor {
+    fn default() -> Self {
+        let (stop_tx, stop_rx) = oneshot::channel();
+        Self { stop_tx, stop_rx }
+    }
+}
+
+impl SignalMonitor {
+    pub async fn recv(self) -> anyhow::Result<Signal> {
+        Self::start(self.stop_tx)?;
+        Ok(self.stop_rx.await?)
+    }
+
+    #[cfg(target_family = "unix")]
+    fn start(stop_tx: Sender<Signal>) -> anyhow::Result<JoinHandle<()>> {
+        let mut sigterm = unix::signal(unix::SignalKind::terminate())?;
+        let mut sigint = unix::signal(unix::SignalKind::interrupt())?;
+        let mut sigquit = unix::signal(unix::SignalKind::quit())?;
+        Ok(tokio::spawn(async move {
+            select! {
+                _ = sigterm.recv() => stop_tx.send("SIGTERM").expect("Failed to handle SIGTERM event"),
+                _ = sigint.recv() => stop_tx.send("SIGINT").expect("Failed to handle SIGINT event"),
+                _ = sigquit.recv() => stop_tx.send("SIGQUIT").expect("Failed to handle SIGQUIT event"),
+            };
+        }))
+    }
+
+    #[cfg(target_family = "windows")]
+    fn start(stop_tx: Sender<Signal>) -> anyhow::Result<JoinHandle<()>> {
+        let mut ctrl_c = windows::ctrl_c()?;
+        let mut ctrl_close = windows::ctrl_close()?;
+        let mut ctrl_logoff = windows::ctrl_logoff()?;
+        let mut ctrl_shutdown = windows::ctrl_shutdown()?;
+        let mut ctrl_break = windows::ctrl_break()?;
+        Ok(tokio::spawn(async move {
+            select! {
+                _ = ctrl_c.recv() => stop_tx.send("CTRL-C").expect("Failed to handle CTRL-C event"),
+                _ = ctrl_close.recv() => stop_tx.send("CTRL-CLOSE").expect("Failed to handle CTRL-CLOSE event"),
+                _ = ctrl_logoff.recv() => stop_tx.send("CTRL-LOGOFF").expect("Failed to handle CTRL-LOGOFF event"),
+                _ = ctrl_shutdown.recv() => stop_tx.send("CTRL-SHUTDOWN").expect("Failed to handle CTRL-HUTDOWN event"),
+                _ = ctrl_break.recv() => stop_tx.send("CTRL-BREAK").expect("Failed to handle CTRL-BREAK event")
+            };
+        }))
+    }
+}