From 9157241c5cd4f083604bbf508c92cdd063aaf5cf Mon Sep 17 00:00:00 2001
From: zleyyij <75810274+zleyyij@users.noreply.github.com>
Date: Thu, 6 Jun 2024 09:40:53 -0600
Subject: [PATCH] feat: added manual signal handler

---
 backend/Cargo.lock  | 10 ++++++++++
 backend/Cargo.toml  |  2 +-
 backend/src/main.rs | 32 +++++++++++++++++++++++++++-----
 3 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/backend/Cargo.lock b/backend/Cargo.lock
index d3a679dd..45c02c94 100644
--- a/backend/Cargo.lock
+++ b/backend/Cargo.lock
@@ -2082,6 +2082,15 @@ dependencies = [
  "lazy_static",
 ]
 
+[[package]]
+name = "signal-hook-registry"
+version = "1.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a9e9e0b4211b72e7b8b6e85c807d36c212bdb33ea8587f7569562a84df5465b1"
+dependencies = [
+ "libc",
+]
+
 [[package]]
 name = "signature"
 version = "2.2.0"
@@ -2536,6 +2545,7 @@ dependencies = [
  "mio",
  "num_cpus",
  "pin-project-lite",
+ "signal-hook-registry",
  "socket2",
  "tokio-macros",
  "windows-sys 0.48.0",
diff --git a/backend/Cargo.toml b/backend/Cargo.toml
index e1094654..0e23d98b 100644
--- a/backend/Cargo.toml
+++ b/backend/Cargo.toml
@@ -24,7 +24,7 @@ reqwest = { version = "0.12.4", features = ["stream"] }
 serde = { version = "1.0.203", features = ["derive"] }
 serde_json = "1.0.117"
 sqlx = { version = "0.7.4", features = ["sqlite", "runtime-tokio"] }
-tokio = { version = "1.38.0", features = ["macros", "rt-multi-thread"] }
+tokio = { version = "1.38.0", features = ["macros", "rt-multi-thread", "signal"] }
 tower = "0.4.13"
 tower-http = { version = "0.5.2", features = ["normalize-path", "fs", "cors"] }
 url = "2.5.0"
diff --git a/backend/src/main.rs b/backend/src/main.rs
index 895d5a8a..30e9f4f0 100644
--- a/backend/src/main.rs
+++ b/backend/src/main.rs
@@ -17,12 +17,15 @@ use color_eyre::Result;
 use db::DATABASE_URL;
 use gh::GithubAccessToken;
 use handlers_prelude::*;
-use log::{info, LevelFilter};
+use log::{error, debug, info, warn, LevelFilter};
 use oauth2::{basic::BasicClient, AuthUrl, ClientId, ClientSecret, TokenUrl};
 use reqwest::Client;
 use sqlx::SqlitePool;
 use std::env::{self, current_exe};
-use tokio::task;
+use tokio::{
+    signal::unix::{signal, SignalKind},
+    task,
+};
 use tower_http::cors::CorsLayer;
 use tower_http::{normalize_path::NormalizePathLayer, services::ServeDir};
 
@@ -55,12 +58,15 @@ async fn main() -> Result<()> {
     // Parse command line arguments
     let cli_args = Args::parse();
     // Read environment variables from dotenv file
-    dotenvy::from_path("cms-data/.env")
-        .context("No .env file was found in `cms-data/`, or failed to read from it.")?;
+    let dotenv_path = "cms-data/.env";
     // Initialize logging
     env_logger::builder()
-        .filter(None, log::LevelFilter::Info)
+        .filter(None, cli_args.logging_level)
         .init();
+    debug!("Initialized logging");
+    dotenvy::from_path(dotenv_path).unwrap_or_else(|_| {
+        warn!("Failed to read dotenv file located at {dotenv_path}, please ensure all config values are manually set");
+    });
     if cfg!(debug_assertions) {
         info!("Server running in development mode");
     } else {
@@ -70,6 +76,22 @@ async fn main() -> Result<()> {
     let state: AppState = init_state()
         .await
         .wrap_err("Failed to initialize app state")?;
+    debug!("Initialized app state");
+    // https://github.com/r-Techsupport/rts-cms/issues/27
+    // In docker, because the process is running with a PID of 1,
+    // we need to implement our own SIGINT/TERM handlers
+    #[cfg(target_family="unix")]
+    for sig in [SignalKind::interrupt(), SignalKind::terminate()] {
+        task::spawn(async move {
+            let mut listener = signal(sig)
+                .expect("Failed to initialize a signal handler");
+            listener.recv().await;
+            // At this point we've received SIGINT/SIGKILL and we can shut down
+            error!("SIGINT or SIGTERM received, terminating.");
+            std::process::exit(0);
+        });
+    }
+
     // files are served relative to the location of the executable, not where the
     // executable was run from
     let mut frontend_dir = current_exe()?;