Shutdown the server more gracefully

Cargo.lock

@@ -3775,7 +3775,7 @@ checksum = "0046fef7e28c3804e5e38bfa31ea2a0f73905319b677e57ebe37e49358989b5d"
 
 [[package]]
 name = "watchers"
-version = "0.0.1"
+version = "0.1.0"
 dependencies = [
  "anyhow",
  "aw-client-rust",

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"] }

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);
 }

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(),

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 => {},
+                );
+            });
     });
 }

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);

watchers/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "watchers"
-version = "0.0.1"
+version = { workspace = true }
 authors = ["Demmie <2e3s19@gmail.com>"]
 edition = "2021"