From 63b8079aa819c6bc2b5172470b6630036964d92c Mon Sep 17 00:00:00 2001
From: Demmie <2e3s19@gmail.com>
Date: Mon, 5 Jun 2023 00:45:38 -0400
Subject: [PATCH] Shutdown the server more gracefully

---
 Cargo.lock           |  2 +-
 Cargo.toml           |  5 ++++-
 src/bundle.rs        | 10 +++++++---
 src/bundle/menu.rs   | 19 +++++++++++++++----
 src/bundle/server.rs | 28 ++++++++++++++++++++++++----
 src/main.rs          |  2 +-
 watchers/Cargo.toml  |  2 +-
 7 files changed, 53 insertions(+), 15 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index c91baab..e65c3df 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3775,7 +3775,7 @@ checksum = "0046fef7e28c3804e5e38bfa31ea2a0f73905319b677e57ebe37e49358989b5d"
 
 [[package]]
 name = "watchers"
-version = "0.0.1"
+version = "0.1.0"
 dependencies = [
  "anyhow",
  "aw-client-rust",
diff --git a/Cargo.toml b/Cargo.toml
index 3d825d1..c4994bf 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "awatcher"
-version = "0.1.0"
+version = { workspace = true }
 authors = ["Demmie <2e3s19@gmail.com>"]
 edition = "2021"
 license-file = "LICENSE"
@@ -16,6 +16,9 @@ image = { version = "0.24.6" }
 [workspace]
 members = ["watchers"]
 
+[workspace.package]
+version = "0.1.0"
+
 [workspace.dependencies]
 anyhow = "1.0.70"
 log = { version = "0.4.17", features = ["std"] }
diff --git a/src/bundle.rs b/src/bundle.rs
index 329bddd..eea164c 100644
--- a/src/bundle.rs
+++ b/src/bundle.rs
@@ -2,19 +2,23 @@ mod menu;
 mod server;
 
 pub use menu::Tray;
-use std::path::PathBuf;
+use std::{
+    path::PathBuf,
+    sync::{atomic::AtomicBool, Arc},
+};
 use watchers::config::Config;
 
-pub fn run(config: &Config, config_file: PathBuf, no_tray: bool) {
+pub fn run(config: &Config, config_file: PathBuf, no_tray: bool, is_stopped: Arc<AtomicBool>) {
     if !no_tray {
         let service = ksni::TrayService::new(Tray {
             server_host: config.host.clone(),
             server_port: config.port,
             config_file,
+            is_stopped: Arc::clone(&is_stopped),
         });
         service.spawn();
     }
 
     let port = config.port;
-    server::run(port);
+    server::run(port, is_stopped);
 }
diff --git a/src/bundle/menu.rs b/src/bundle/menu.rs
index 67a473d..986dd5c 100644
--- a/src/bundle/menu.rs
+++ b/src/bundle/menu.rs
@@ -1,10 +1,17 @@
-use std::path::PathBuf;
+use std::{
+    path::PathBuf,
+    sync::{
+        atomic::{AtomicBool, Ordering},
+        Arc,
+    },
+};
 
 #[derive(Debug)]
 pub struct Tray {
     pub server_host: String,
     pub server_port: u32,
     pub config_file: PathBuf,
+    pub is_stopped: Arc<AtomicBool>,
 }
 
 impl ksni::Tray for Tray {
@@ -51,9 +58,13 @@ impl ksni::Tray for Tray {
             ksni::menu::StandardItem {
                 label: "Exit".into(),
                 icon_name: "application-exit".into(),
-                activate: Box::new(|_| {
-                    std::process::exit(0);
-                }),
+                activate: {
+                    let is_stopped = Arc::clone(&self.is_stopped);
+
+                    Box::new(move |_| {
+                        is_stopped.store(true, Ordering::Relaxed);
+                    })
+                },
                 ..Default::default()
             }
             .into(),
diff --git a/src/bundle/server.rs b/src/bundle/server.rs
index 0d6d821..ac91880 100644
--- a/src/bundle/server.rs
+++ b/src/bundle/server.rs
@@ -1,8 +1,12 @@
 use anyhow::anyhow;
 use aw_server::endpoints::{build_rocket, embed_asset_resolver};
-use std::sync::Mutex;
+use std::sync::{
+    atomic::{AtomicBool, Ordering},
+    Arc, Mutex,
+};
+use tokio::time::{sleep, Duration};
 
-pub fn run(port: u32) {
+pub fn run(port: u32, is_stopped: Arc<AtomicBool>) {
     std::thread::spawn(move || {
         let db_path = aw_server::dirs::db_path(false)
             .map_err(|_| anyhow!("DB path is not found"))
@@ -21,12 +25,28 @@ pub fn run(port: u32) {
             asset_resolver: embed_asset_resolver!("$AW_WEBUI_DIST"),
             device_id,
         };
+        let server = build_rocket(server_state, config).launch();
+
+        let check = async {
+            loop {
+                if is_stopped.load(Ordering::Relaxed) {
+                    warn!("Received an exit signal, stopping the server");
+                    break;
+                }
+
+                sleep(Duration::from_secs(1)).await;
+            }
+        };
 
         tokio::runtime::Builder::new_current_thread()
             .enable_all()
             .build()
             .unwrap()
-            .block_on(build_rocket(server_state, config).launch())
-            .unwrap();
+            .block_on(async {
+                tokio::select! (
+                    r = server => {r.unwrap();},
+                    _ = check => {},
+                );
+            });
     });
 }
diff --git a/src/main.rs b/src/main.rs
index 5dcbead..f2b9b70 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -41,7 +41,7 @@ fn main() -> anyhow::Result<()> {
     );
 
     #[cfg(feature = "bundle")]
-    bundle::run(&config, config_file, no_tray);
+    bundle::run(&config, config_file, no_tray, Arc::clone(&is_stopped));
 
     let client = ReportClient::new(config)?;
     let client = Arc::new(client);
diff --git a/watchers/Cargo.toml b/watchers/Cargo.toml
index ad8a20c..581d812 100644
--- a/watchers/Cargo.toml
+++ b/watchers/Cargo.toml
@@ -1,6 +1,6 @@
 [package]
 name = "watchers"
-version = "0.0.1"
+version = { workspace = true }
 authors = ["Demmie <2e3s19@gmail.com>"]
 edition = "2021"