diff --git a/Cargo.lock b/Cargo.lock index 276b8ee..45c8dd4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -458,6 +458,7 @@ dependencies = [ "fern", "ksni", "log", + "signal-hook", "smol", "toml 0.7.3", "watchers", diff --git a/Cargo.toml b/Cargo.toml index ebb6e4a..53c39e8 100644 --- a/Cargo.toml +++ b/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 } diff --git a/src/main.rs b/src/main.rs index 76cdb9d..e281d95 100644 --- a/src/main.rs +++ b/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"); diff --git a/watchers/src/watchers.rs b/watchers/src/watchers.rs index 83c7b69..4ce5f33 100644 --- a/watchers/src/watchers.rs +++ b/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 where Self: Sized; - fn watch(&mut self, client: &Arc); + fn watch(&mut self, client: &Arc, is_stopped: Arc); } type BoxedWatcher = Box; @@ -34,7 +34,11 @@ type WatcherConstructors = [WatcherConstructor]; pub trait ConstructorFilter { fn filter_first_supported(&self) -> Option; - fn run_first_supported(&self, client: &Arc) -> Option>; + fn run_first_supported( + &self, + client: &Arc, + is_stopped: Arc, + ) -> Option>; } impl ConstructorFilter for WatcherConstructors { @@ -48,11 +52,15 @@ impl ConstructorFilter for WatcherConstructors { }) } - fn run_first_supported(&self, client: &Arc) -> Option> { + fn run_first_supported( + &self, + client: &Arc, + is_stopped: Arc, + ) -> Option> { 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 diff --git a/watchers/src/watchers/gnome_idle.rs b/watchers/src/watchers/gnome_idle.rs index 4361fa4..ad4b2c8 100644 --- a/watchers/src/watchers/gnome_idle.rs +++ b/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) { + fn watch(&mut self, client: &Arc, is_stopped: Arc) { 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; diff --git a/watchers/src/watchers/gnome_window.rs b/watchers/src/watchers/gnome_window.rs index ef7d771..89ab53d 100644 --- a/watchers/src/watchers/gnome_window.rs +++ b/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) { + fn watch(&mut self, client: &Arc, is_stopped: Arc) { 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}"); } diff --git a/watchers/src/watchers/kwin_window.rs b/watchers/src/watchers/kwin_window.rs index decc5bd..1b0ab72 100644 --- a/watchers/src/watchers/kwin_window.rs +++ b/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) { + fn watch(&mut self, client: &Arc, is_stopped: Arc) { 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}"); } diff --git a/watchers/src/watchers/wl_foreign_toplevel.rs b/watchers/src/watchers/wl_foreign_toplevel.rs index ed4c759..2867c47 100644 --- a/watchers/src/watchers/wl_foreign_toplevel.rs +++ b/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) { + fn watch(&mut self, client: &Arc, is_stopped: Arc) { 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() { diff --git a/watchers/src/watchers/wl_kwin_idle.rs b/watchers/src/watchers/wl_kwin_idle.rs index b823bf7..c2d1435 100644 --- a/watchers/src/watchers/wl_kwin_idle.rs +++ b/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) { + fn watch(&mut self, client: &Arc, is_stopped: Arc) { 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() { diff --git a/watchers/src/watchers/x11_screensaver_idle.rs b/watchers/src/watchers/x11_screensaver_idle.rs index 12632da..014af3a 100644 --- a/watchers/src/watchers/x11_screensaver_idle.rs +++ b/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) { + fn watch(&mut self, client: &Arc, is_stopped: Arc) { 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; diff --git a/watchers/src/watchers/x11_window.rs b/watchers/src/watchers/x11_window.rs index b0f842d..148b30a 100644 --- a/watchers/src/watchers/x11_window.rs +++ b/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) { + fn watch(&mut self, client: &Arc, is_stopped: Arc) { 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}"); }