Catch signals to stop watchers

Cargo.lock

@@ -458,6 +458,7 @@ dependencies = [
  "fern",
  "ksni",
  "log",
+ "signal-hook",
  "smol",
  "toml 0.7.3",
  "watchers",

Cargo.toml

@@ -23,6 +23,7 @@ clap = { version = "4.2.1", features = ["string"] }
 fern = { version = "0.6.2", features = ["colored"] }
 log = { workspace = true }
 anyhow = { workspace = true }
+signal-hook = "0.3.15"
 
 ksni = {version = "0.2.0", optional = true}
 aw-server = { git = "https://github.com/ActivityWatch/aw-server-rust.git", optional = true }

src/main.rs

@@ -8,11 +8,16 @@ extern crate log;
 mod bundle;
 mod config;
 
+use std::sync::atomic::AtomicBool;
 use std::sync::Arc;
 use watchers::ConstructorFilter;
 use watchers::ReportClient;
 
 fn main() -> anyhow::Result<()> {
+    let is_stopped = Arc::new(AtomicBool::new(false));
+    signal_hook::flag::register(signal_hook::consts::SIGTERM, Arc::clone(&is_stopped))?;
+    signal_hook::flag::register(signal_hook::consts::SIGINT, Arc::clone(&is_stopped))?;
+
     let config = config::from_cli()?;
     let no_tray = config.no_tray;
     let config = config.watchers_config;
@@ -40,13 +45,16 @@ fn main() -> anyhow::Result<()> {
 
     let mut thread_handlers = Vec::new();
 
-    if let Some(idle_handler) = watchers::IDLE.run_first_supported(&client) {
+    if let Some(idle_handler) = watchers::IDLE.run_first_supported(&client, Arc::clone(&is_stopped))
+    {
         thread_handlers.push(idle_handler);
     } else {
         warn!("No supported idle handler is found");
     }
 
-    if let Some(active_window_handler) = watchers::ACTIVE_WINDOW.run_first_supported(&client) {
+    if let Some(active_window_handler) =
+        watchers::ACTIVE_WINDOW.run_first_supported(&client, is_stopped)
+    {
         thread_handlers.push(active_window_handler);
     } else {
         warn!("No supported active window handler is found");

watchers/src/watchers.rs

@@ -15,7 +15,7 @@ mod x11_window;
 
 use crate::report_client::ReportClient;
 use std::{
-    sync::Arc,
+    sync::{atomic::AtomicBool, Arc},
     thread::{self, JoinHandle},
 };
 
@@ -23,7 +23,7 @@ pub trait Watcher: Send {
     fn new() -> anyhow::Result<Self>
     where
         Self: Sized;
-    fn watch(&mut self, client: &Arc<ReportClient>);
+    fn watch(&mut self, client: &Arc<ReportClient>, is_stopped: Arc<AtomicBool>);
 }
 
 type BoxedWatcher = Box<dyn Watcher>;
@@ -34,7 +34,11 @@ type WatcherConstructors = [WatcherConstructor];
 pub trait ConstructorFilter {
     fn filter_first_supported(&self) -> Option<BoxedWatcher>;
 
-    fn run_first_supported(&self, client: &Arc<ReportClient>) -> Option<JoinHandle<()>>;
+    fn run_first_supported(
+        &self,
+        client: &Arc<ReportClient>,
+        is_stopped: Arc<AtomicBool>,
+    ) -> Option<JoinHandle<()>>;
 }
 
 impl ConstructorFilter for WatcherConstructors {
@@ -48,11 +52,15 @@ impl ConstructorFilter for WatcherConstructors {
         })
     }
 
-    fn run_first_supported(&self, client: &Arc<ReportClient>) -> Option<JoinHandle<()>> {
+    fn run_first_supported(
+        &self,
+        client: &Arc<ReportClient>,
+        is_stopped: Arc<AtomicBool>,
+    ) -> Option<JoinHandle<()>> {
         let idle_watcher = self.filter_first_supported();
         if let Some(mut watcher) = idle_watcher {
             let thread_client = Arc::clone(client);
-            let idle_handler = thread::spawn(move || watcher.watch(&thread_client));
+            let idle_handler = thread::spawn(move || watcher.watch(&thread_client, is_stopped));
             Some(idle_handler)
         } else {
             None

watchers/src/watchers/gnome_idle.rs

@@ -1,7 +1,9 @@
 use super::{idle, Watcher};
 use crate::report_client::ReportClient;
 use anyhow::Context;
-use std::{sync::Arc, thread};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::thread;
 use zbus::blocking::Connection;
 
 pub struct IdleWatcher {
@@ -34,10 +36,14 @@ impl Watcher for IdleWatcher {
         Ok(watcher)
     }
 
-    fn watch(&mut self, client: &Arc<ReportClient>) {
+    fn watch(&mut self, client: &Arc<ReportClient>, is_stopped: Arc<AtomicBool>) {
         let mut is_idle = false;
         info!("Starting idle watcher");
         loop {
+            if is_stopped.load(Ordering::Relaxed) {
+                warn!("Received an exit signal, shutting down");
+                break;
+            }
             match idle::ping_since_last_input(self, is_idle, client) {
                 Ok(is_idle_again) => {
                     is_idle = is_idle_again;

watchers/src/watchers/gnome_window.rs

@@ -1,7 +1,9 @@
 use crate::report_client::ReportClient;
 use anyhow::Context;
 use serde::Deserialize;
-use std::{sync::Arc, thread};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::thread;
 use zbus::blocking::Connection;
 
 use super::Watcher;
@@ -85,9 +87,13 @@ impl Watcher for WindowWatcher {
         Ok(watcher)
     }
 
-    fn watch(&mut self, client: &Arc<ReportClient>) {
+    fn watch(&mut self, client: &Arc<ReportClient>, is_stopped: Arc<AtomicBool>) {
         info!("Starting active window watcher");
         loop {
+            if is_stopped.load(Ordering::Relaxed) {
+                warn!("Received an exit signal, shutting down");
+                break;
+            }
             if let Err(error) = self.send_active_window(client) {
                 error!("Error on active window: {error}");
             }

watchers/src/watchers/kwin_window.rs

@@ -8,6 +8,7 @@ use crate::report_client::ReportClient;
 use anyhow::{anyhow, Context};
 use std::env::temp_dir;
 use std::path::Path;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{mpsc::channel, Arc, Mutex};
 use std::thread;
 use zbus::blocking::{Connection, ConnectionBuilder};
@@ -164,7 +165,7 @@ impl Watcher for WindowWatcher {
         Ok(Self { kwin_script })
     }
 
-    fn watch(&mut self, client: &Arc<ReportClient>) {
+    fn watch(&mut self, client: &Arc<ReportClient>, is_stopped: Arc<AtomicBool>) {
         self.kwin_script.load().unwrap();
 
         let active_window = Arc::new(Mutex::new(ActiveWindow {
@@ -200,6 +201,10 @@ impl Watcher for WindowWatcher {
 
         info!("Starting active window watcher");
         loop {
+            if is_stopped.load(Ordering::Relaxed) {
+                warn!("Received an exit signal, shutting down");
+                break;
+            }
             if let Err(error) = send_active_window(client, &active_window) {
                 error!("Error on sending active window: {error}");
             }

watchers/src/watchers/wl_foreign_toplevel.rs

@@ -9,6 +9,7 @@ use super::{wl_connection::subscribe_state, Watcher};
 use crate::report_client::ReportClient;
 use anyhow::{anyhow, Context};
 use std::collections::HashMap;
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::{sync::Arc, thread};
 use wayland_client::{
     event_created_child, globals::GlobalListContents, protocol::wl_registry, Connection, Dispatch,
@@ -141,7 +142,7 @@ impl Watcher for WindowWatcher {
         Ok(Self { connection })
     }
 
-    fn watch(&mut self, client: &Arc<ReportClient>) {
+    fn watch(&mut self, client: &Arc<ReportClient>, is_stopped: Arc<AtomicBool>) {
         let mut toplevel_state = ToplevelState::new(Arc::clone(client));
 
         self.connection
@@ -151,6 +152,10 @@ impl Watcher for WindowWatcher {
 
         info!("Starting wlr foreign toplevel watcher");
         loop {
+            if is_stopped.load(Ordering::Relaxed) {
+                warn!("Received an exit signal, shutting down");
+                break;
+            }
             if let Err(e) = self.connection.event_queue.roundtrip(&mut toplevel_state) {
                 error!("Event queue is not processed: {e}");
             } else if let Err(e) = toplevel_state.send_active_window() {

watchers/src/watchers/wl_kwin_idle.rs

@@ -3,6 +3,7 @@ use super::wl_connection::{subscribe_state, WlEventConnection};
 use super::Watcher;
 use crate::report_client::ReportClient;
 use chrono::{DateTime, Duration, Utc};
+use std::sync::atomic::{AtomicBool, Ordering};
 use std::{sync::Arc, thread};
 use wayland_client::{
     globals::GlobalListContents,
@@ -128,7 +129,7 @@ impl Watcher for IdleWatcher {
         Ok(Self { connection })
     }
 
-    fn watch(&mut self, client: &Arc<ReportClient>) {
+    fn watch(&mut self, client: &Arc<ReportClient>, is_stopped: Arc<AtomicBool>) {
         let timeout = u32::try_from(client.config.idle_timeout.as_secs() * 1000);
         let mut idle_state = IdleState::new(
             self.connection
@@ -143,6 +144,10 @@ impl Watcher for IdleWatcher {
 
         info!("Starting idle watcher");
         loop {
+            if is_stopped.load(Ordering::Relaxed) {
+                warn!("Received an exit signal, shutting down");
+                break;
+            }
             if let Err(e) = self.connection.event_queue.roundtrip(&mut idle_state) {
                 error!("Event queue is not processed: {e}");
             } else if let Err(e) = idle_state.send_ping() {

watchers/src/watchers/x11_screensaver_idle.rs

@@ -1,6 +1,8 @@
 use super::{idle, x11_connection::X11Client, Watcher};
 use crate::report_client::ReportClient;
-use std::{sync::Arc, thread};
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+use std::thread;
 
 pub struct IdleWatcher {
     client: X11Client,
@@ -22,10 +24,14 @@ impl Watcher for IdleWatcher {
         Ok(IdleWatcher { client })
     }
 
-    fn watch(&mut self, client: &Arc<ReportClient>) {
+    fn watch(&mut self, client: &Arc<ReportClient>, is_stopped: Arc<AtomicBool>) {
         info!("Starting idle watcher");
         let mut is_idle = false;
         loop {
+            if is_stopped.load(Ordering::Relaxed) {
+                warn!("Received an exit signal, shutting down");
+                break;
+            }
             match idle::ping_since_last_input(self, is_idle, client) {
                 Ok(is_idle_again) => {
                     is_idle = is_idle_again;

watchers/src/watchers/x11_window.rs

@@ -1,6 +1,8 @@
 use super::{x11_connection::X11Client, Watcher};
 use crate::report_client::ReportClient;
 use anyhow::Context;
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
 use std::thread;
 
 pub struct WindowWatcher {
@@ -40,9 +42,13 @@ impl Watcher for WindowWatcher {
         })
     }
 
-    fn watch(&mut self, client: &std::sync::Arc<crate::report_client::ReportClient>) {
+    fn watch(&mut self, client: &Arc<ReportClient>, is_stopped: Arc<AtomicBool>) {
         info!("Starting active window watcher");
         loop {
+            if is_stopped.load(Ordering::Relaxed) {
+                warn!("Received an exit signal, shutting down");
+                break;
+            }
             if let Err(error) = self.send_active_window(client) {
                 error!("Error on sending active window: {error}");
             }