use anyhow::{Context, Result};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::signal;
use tokio::sync::broadcast;
use tracing::{error, info, warn};
#[derive(Debug, Clone)]
pub struct DaemonConfig {
pub pid_file: Option<String>,
pub working_dir: Option<String>,
pub user: Option<String>,
pub group: Option<String>,
pub shutdown_timeout: u64,
}
impl Default for DaemonConfig {
fn default() -> Self {
Self {
pid_file: Some("/var/run/kindly-guard.pid".to_string()),
working_dir: None,
user: None,
group: None,
shutdown_timeout: 30,
}
}
}
pub struct DaemonHandle {
shutdown_tx: broadcast::Sender<()>,
running: Arc<AtomicBool>,
pid_file: Option<String>,
}
impl DaemonHandle {
pub fn new(config: DaemonConfig) -> Result<Self> {
let (shutdown_tx, _) = broadcast::channel(1);
let running = Arc::new(AtomicBool::new(true));
if let Some(ref pid_path) = config.pid_file {
write_pid_file(pid_path)?;
}
if let Some(ref dir) = config.working_dir {
std::env::set_current_dir(dir).context("Failed to change working directory")?;
}
#[cfg(unix)]
if config.user.is_some() || config.group.is_some() {
drop_privileges(config.user.as_deref(), config.group.as_deref())?;
}
Ok(Self {
shutdown_tx,
running,
pid_file: config.pid_file,
})
}
pub fn shutdown_receiver(&self) -> broadcast::Receiver<()> {
self.shutdown_tx.subscribe()
}
pub fn is_running(&self) -> bool {
self.running.load(Ordering::Relaxed)
}
pub fn shutdown(&self) {
info!("Daemon shutdown requested");
self.running.store(false, Ordering::Relaxed);
let _ = self.shutdown_tx.send(());
}
pub async fn setup_signal_handlers(self: Arc<Self>) {
let handle1 = self.clone();
tokio::spawn(async move {
match signal::ctrl_c().await {
Ok(()) => {
info!("Received SIGINT, shutting down gracefully");
handle1.shutdown();
},
Err(e) => {
error!("Failed to listen for SIGINT: {}", e);
},
}
});
#[cfg(unix)]
{
let handle2 = self.clone();
tokio::spawn(async move {
let mut stream = match signal::unix::signal(signal::unix::SignalKind::terminate()) {
Ok(s) => s,
Err(e) => {
error!("Failed to listen for SIGTERM: {}", e);
return;
},
};
loop {
stream.recv().await;
info!("Received SIGTERM, shutting down gracefully");
handle2.shutdown();
break;
}
});
let _handle3 = self;
tokio::spawn(async move {
let mut stream = match signal::unix::signal(signal::unix::SignalKind::hangup()) {
Ok(s) => s,
Err(e) => {
warn!("Failed to listen for SIGHUP: {}", e);
return;
},
};
loop {
stream.recv().await;
info!("Received SIGHUP, reloading configuration");
}
});
}
}
}
impl Drop for DaemonHandle {
fn drop(&mut self) {
if let Some(ref pid_path) = self.pid_file {
if let Err(e) = std::fs::remove_file(pid_path) {
warn!("Failed to remove PID file: {}", e);
}
}
}
}
fn write_pid_file(path: &str) -> Result<()> {
let pid = std::process::id();
std::fs::write(path, pid.to_string())
.with_context(|| format!("Failed to write PID file: {path}"))?;
info!("Wrote PID {} to {}", pid, path);
Ok(())
}
#[cfg(unix)]
fn drop_privileges(user: Option<&str>, group: Option<&str>) -> Result<()> {
if user.is_some() || group.is_some() {
warn!(
"Privilege dropping not yet implemented - would drop to user: {:?}, group: {:?}",
user, group
);
}
Ok(())
}
pub async fn run_with_daemon<F, Fut>(config: DaemonConfig, run_fn: F) -> Result<()>
where
F: FnOnce(broadcast::Receiver<()>) -> Fut,
Fut: std::future::Future<Output = Result<()>>,
{
let daemon = Arc::new(DaemonHandle::new(config)?);
daemon.clone().setup_signal_handlers().await;
let shutdown_rx = daemon.shutdown_receiver();
info!("Daemon started successfully");
tokio::select! {
result = run_fn(shutdown_rx) => {
match result {
Ok(()) => info!("Daemon function completed successfully"),
Err(e) => error!("Daemon function error: {}", e),
}
}
_ = tokio::signal::ctrl_c() => {
info!("Received interrupt signal");
}
}
info!("Daemon shutting down");
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::tempdir;
#[tokio::test]
async fn test_daemon_handle_creation() {
let config = DaemonConfig {
pid_file: None,
..Default::default()
};
let handle = DaemonHandle::new(config).unwrap();
assert!(handle.is_running());
handle.shutdown();
assert!(!handle.is_running());
}
#[tokio::test]
async fn test_pid_file_creation() {
let dir = tempdir().unwrap();
let pid_path = dir.path().join("test.pid");
let config = DaemonConfig {
pid_file: Some(pid_path.to_str().unwrap().to_string()),
..Default::default()
};
{
let _handle = DaemonHandle::new(config).unwrap();
assert!(pid_path.exists());
let content = std::fs::read_to_string(&pid_path).unwrap();
let pid: u32 = content.trim().parse().unwrap();
assert_eq!(pid, std::process::id());
}
assert!(!pid_path.exists());
}
#[tokio::test]
async fn test_shutdown_signal() {
let config = DaemonConfig {
pid_file: None,
..Default::default()
};
let handle = Arc::new(DaemonHandle::new(config).unwrap());
let mut rx = handle.shutdown_receiver();
handle.shutdown();
assert!(rx.recv().await.is_ok());
}
}