zeph_core/

daemon.rs

// SPDX-FileCopyrightText: 2026 Andrei G <bug-ops>
// SPDX-License-Identifier: MIT OR Apache-2.0

//! Daemon supervisor for component lifecycle management.

use std::time::Duration;

use tokio::sync::watch;
use tokio::task::JoinHandle;

use crate::config::DaemonConfig;

#[non_exhaustive]
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ComponentStatus {
    Running,
    Failed(String),
    Stopped,
}

#[non_exhaustive]
/// Error type for daemon component task failures and pid file guard acquisition.
#[derive(Debug, thiserror::Error)]
pub enum DaemonError {
    #[error("task error: {0}")]
    Task(String),
    #[error("shutdown error: {0}")]
    Shutdown(String),
    #[error(
        "another daemon instance is already running (PID {pid}); stop it before starting a new one"
    )]
    AlreadyRunning { pid: u32 },
    #[error("pid file error: {0}")]
    PidFile(#[from] std::io::Error),
}

pub struct ComponentHandle {
    pub name: String,
    handle: JoinHandle<Result<(), DaemonError>>,
    pub status: ComponentStatus,
    pub restart_count: u32,
}

impl ComponentHandle {
    #[must_use]
    pub fn new(name: impl Into<String>, handle: JoinHandle<Result<(), DaemonError>>) -> Self {
        Self {
            name: name.into(),
            handle,
            status: ComponentStatus::Running,
            restart_count: 0,
        }
    }

    #[must_use]
    pub fn is_finished(&self) -> bool {
        self.handle.is_finished()
    }
}

pub struct DaemonSupervisor {
    components: Vec<ComponentHandle>,
    health_interval: Duration,
    _max_backoff: Duration,
    shutdown_rx: watch::Receiver<bool>,
}

impl DaemonSupervisor {
    #[must_use]
    pub fn new(config: &DaemonConfig, shutdown_rx: watch::Receiver<bool>) -> Self {
        Self {
            components: Vec::new(),
            health_interval: Duration::from_secs(config.health_interval_secs),
            _max_backoff: Duration::from_secs(config.max_restart_backoff_secs),
            shutdown_rx,
        }
    }

    pub fn add_component(&mut self, handle: ComponentHandle) {
        self.components.push(handle);
    }

    #[must_use]
    pub fn component_count(&self) -> usize {
        self.components.len()
    }

    /// Run the health monitoring loop until shutdown signal.
    pub async fn run(&mut self) {
        let mut interval = tokio::time::interval(self.health_interval);
        loop {
            tokio::select! {
                _ = interval.tick() => {
                    self.check_health();
                }
                _ = self.shutdown_rx.changed() => {
                    if *self.shutdown_rx.borrow() {
                        tracing::info!("daemon supervisor shutting down");
                        break;
                    }
                }
            }
        }
    }

    fn check_health(&mut self) {
        for component in &mut self.components {
            if component.status == ComponentStatus::Running && component.is_finished() {
                component.status = ComponentStatus::Failed("task exited".into());
                component.restart_count += 1;
                tracing::warn!(
                    component = %component.name,
                    restarts = component.restart_count,
                    "component exited unexpectedly"
                );
            }
        }
    }

    #[must_use]
    pub fn component_statuses(&self) -> Vec<(&str, &ComponentStatus)> {
        self.components
            .iter()
            .map(|c| (c.name.as_str(), &c.status))
            .collect()
    }
}

/// Check whether a process with the given PID is currently alive.
///
/// On Unix, uses `kill -0` which returns success if the process exists and the current user
/// has permission to signal it.
/// On Windows, uses `tasklist /FI "PID eq <pid>"` and checks for the PID in the output.
#[must_use]
pub fn is_process_alive(pid: u32) -> bool {
    #[cfg(unix)]
    {
        // PIDs on Unix are signed (pid_t = i32); u32::MAX wraps to -1 which would
        // signal every process, so reject anything that does not fit in a positive i32.
        let Ok(signed) = i32::try_from(pid) else {
            return false;
        };
        if signed <= 0 {
            return false;
        }
        std::process::Command::new("kill")
            .args(["-0", &signed.to_string()])
            .output()
            .is_ok_and(|o| o.status.success())
    }
    #[cfg(windows)]
    {
        std::process::Command::new("tasklist")
            .args(["/FI", &format!("PID eq {pid}"), "/NH", "/FO", "CSV"])
            .output()
            .map(|o| {
                let stdout = String::from_utf8_lossy(&o.stdout);
                // tasklist outputs lines like: "process.exe","PID","..."
                // We look for the PID appearing as a quoted field.
                stdout.contains(&format!("\"{pid}\""))
            })
            .unwrap_or(false)
    }
    #[cfg(not(any(unix, windows)))]
    {
        let _ = pid;
        false
    }
}

/// Write a PID file atomically using `O_CREAT | O_EXCL` to prevent TOCTOU races.
///
/// # Errors
///
/// Returns an error if the PID file directory cannot be created, the file already exists,
/// or the file cannot be written.
pub fn write_pid_file(path: &str) -> std::io::Result<()> {
    use std::io::Write as _;
    let expanded = expand_tilde(path);
    let path = std::path::Path::new(&expanded);
    if let Some(parent) = path.parent() {
        std::fs::create_dir_all(parent)?;
    }
    let mut file = std::fs::OpenOptions::new()
        .write(true)
        .create_new(true)
        .open(path)?;
    file.write_all(std::process::id().to_string().as_bytes())
}

/// Read the PID from a PID file.
///
/// # Errors
///
/// Returns an error if the file cannot be read or the content is not a valid PID.
pub fn read_pid_file(path: &str) -> std::io::Result<u32> {
    let expanded = expand_tilde(path);
    let content = std::fs::read_to_string(&expanded)?;
    content
        .trim()
        .parse::<u32>()
        .map_err(|e| std::io::Error::new(std::io::ErrorKind::InvalidData, e))
}

/// Remove the PID file.
///
/// # Errors
///
/// Returns an error if the file cannot be removed.
pub fn remove_pid_file(path: &str) -> std::io::Result<()> {
    let expanded = expand_tilde(path);
    match std::fs::remove_file(&expanded) {
        Ok(()) => Ok(()),
        Err(e) if e.kind() == std::io::ErrorKind::NotFound => Ok(()),
        Err(e) => Err(e),
    }
}

/// Exclusive, `flock(2)`-backed guard on the daemon's pid file.
///
/// Replaces the previous check-then-act sequence (read pid file, check liveness, remove if
/// stale, then write) with a single atomic lock acquisition: only one process can ever hold the
/// advisory lock on the pid file, so a second concurrent `--daemon` invocation fails immediately
/// with [`DaemonError::AlreadyRunning`] instead of racing the first to write the file. Dropping
/// the guard unlinks the pid file and releases the lock.
///
/// # Examples
///
/// ```no_run
/// use zeph_core::daemon::PidGuard;
///
/// let guard = PidGuard::acquire("/tmp/zeph-daemon.pid").expect("no other instance running");
/// // ... run the daemon ...
/// drop(guard); // releases the lock and removes the pid file
/// ```
// Wraps the shared guard purely for its `Drop` impl (unlinks the pid file on release).
#[cfg(unix)]
#[derive(Debug)]
pub struct PidGuard(#[allow(dead_code)] zeph_common::pidfile::PidLockGuard);

#[cfg(unix)]
impl PidGuard {
    /// Acquire an exclusive lock on the pid file at `path`, creating it if necessary, and write
    /// the current process id into it.
    ///
    /// # Errors
    ///
    /// Returns [`DaemonError::AlreadyRunning`] if another process already holds the lock, or
    /// [`DaemonError::PidFile`] for filesystem failures.
    pub fn acquire(path: &str) -> Result<Self, DaemonError> {
        use zeph_common::pidfile::{PidLockError, PidLockGuard};

        let expanded = expand_tilde(path);
        let path = std::path::Path::new(&expanded);

        PidLockGuard::acquire(path).map(Self).map_err(|e| match e {
            PidLockError::AlreadyRunning { pid } => DaemonError::AlreadyRunning { pid },
            PidLockError::Io(err) => DaemonError::PidFile(err),
        })
    }
}

/// Exclusive guard on the daemon's pid file (non-Unix fallback).
///
/// `flock(2)` is unavailable outside Unix, so this falls back to the create-then-check
/// sequence: [`write_pid_file`] already uses an atomic exclusive create, so two processes
/// racing to start fresh are still serialized correctly; only recovery from a stale pid file
/// left by a crashed process re-opens a (narrow) TOCTOU window between the liveness check and
/// the write.
#[cfg(not(unix))]
#[derive(Debug)]
pub struct PidGuard {
    path: String,
}

#[cfg(not(unix))]
impl PidGuard {
    /// Acquire the pid file guard at `path`, removing a stale (dead-process) file first.
    ///
    /// # Errors
    ///
    /// Returns [`DaemonError::AlreadyRunning`] if the existing pid file names a live process, or
    /// [`DaemonError::PidFile`] for filesystem failures.
    pub fn acquire(path: &str) -> Result<Self, DaemonError> {
        if let Ok(existing_pid) = read_pid_file(path) {
            if is_process_alive(existing_pid) {
                return Err(DaemonError::AlreadyRunning { pid: existing_pid });
            }
            remove_pid_file(path)?;
        }
        write_pid_file(path)?;
        Ok(Self {
            path: path.to_owned(),
        })
    }
}

#[cfg(not(unix))]
impl Drop for PidGuard {
    fn drop(&mut self) {
        let _ = remove_pid_file(&self.path);
    }
}

fn expand_tilde(path: &str) -> String {
    if let Some(rest) = path.strip_prefix("~/")
        && let Some(home) = std::env::var_os("HOME").or_else(|| std::env::var_os("USERPROFILE"))
    {
        return format!("{}/{rest}", home.to_string_lossy());
    }
    path.to_owned()
}

#[cfg(test)]
mod tests {
    #![allow(clippy::field_reassign_with_default)]
    use std::assert_matches;

    use super::*;

    #[test]
    fn expand_tilde_with_home() {
        let result = expand_tilde("~/test/file.pid");
        assert!(!result.starts_with("~/"));
    }

    #[test]
    fn expand_tilde_absolute_unchanged() {
        assert_eq!(expand_tilde("/tmp/zeph.pid"), "/tmp/zeph.pid");
    }

    #[test]
    fn pid_file_roundtrip() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("test.pid");
        let path_str = path.to_string_lossy().to_string();

        write_pid_file(&path_str).unwrap();
        let pid = read_pid_file(&path_str).unwrap();
        assert_eq!(pid, std::process::id());
        remove_pid_file(&path_str).unwrap();
        assert!(!path.exists());
    }

    #[test]
    fn remove_nonexistent_pid_file_ok() {
        assert!(remove_pid_file("/tmp/nonexistent_zeph_test.pid").is_ok());
    }

    #[cfg(unix)]
    #[test]
    fn pid_guard_acquire_writes_pid_and_removes_on_drop() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("guard.pid");
        let path_str = path.to_string_lossy().to_string();

        let guard = PidGuard::acquire(&path_str).expect("acquire should succeed");
        let pid = read_pid_file(&path_str).expect("pid file must exist");
        assert_eq!(pid, std::process::id());
        drop(guard);
        assert!(!path.exists(), "pid file must be removed on drop");
    }

    /// Regression test for #5679: a second concurrent acquisition attempt on a pid file
    /// already locked by a live guard must fail closed with `AlreadyRunning`, not silently
    /// overwrite the file or proceed.
    #[cfg(unix)]
    #[test]
    fn pid_guard_second_acquire_fails_with_already_running() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("guard-race.pid");
        let path_str = path.to_string_lossy().to_string();

        let _first = PidGuard::acquire(&path_str).expect("first acquire must succeed");
        let err = PidGuard::acquire(&path_str).expect_err("second acquire must fail");
        assert_matches!(err, DaemonError::AlreadyRunning { .. });
    }

    #[test]
    fn read_invalid_pid_file() {
        let dir = tempfile::tempdir().unwrap();
        let path = dir.path().join("bad.pid");
        std::fs::write(&path, "not_a_number").unwrap();
        assert!(read_pid_file(&path.to_string_lossy()).is_err());
    }

    #[tokio::test]
    async fn supervisor_tracks_components() {
        let config = DaemonConfig::default();
        let (_tx, rx) = watch::channel(false);
        let mut supervisor = DaemonSupervisor::new(&config, rx);

        let handle = tokio::spawn(async { Ok::<(), DaemonError>(()) });
        supervisor.add_component(ComponentHandle::new("test", handle));
        assert_eq!(supervisor.component_count(), 1);
    }

    #[tokio::test]
    async fn supervisor_detects_finished_component() {
        let config = DaemonConfig::default();
        let (_tx, rx) = watch::channel(false);
        let mut supervisor = DaemonSupervisor::new(&config, rx);

        let handle = tokio::spawn(async { Ok::<(), DaemonError>(()) });
        tokio::time::sleep(Duration::from_millis(10)).await;
        supervisor.add_component(ComponentHandle::new("finished", handle));
        supervisor.check_health();

        let statuses = supervisor.component_statuses();
        assert_eq!(statuses.len(), 1);
        assert_matches!(statuses[0].1, ComponentStatus::Failed(_));
    }

    #[tokio::test]
    async fn supervisor_shutdown() {
        let config = DaemonConfig {
            health_interval_secs: 1,
            ..DaemonConfig::default()
        };
        let (tx, rx) = watch::channel(false);
        let mut supervisor = DaemonSupervisor::new(&config, rx);

        let run_handle = tokio::spawn(async move { supervisor.run().await });
        tokio::time::sleep(Duration::from_millis(50)).await;
        let _ = tx.send(true);
        tokio::time::timeout(Duration::from_secs(2), run_handle)
            .await
            .expect("supervisor should stop on shutdown")
            .expect("task should complete");
    }

    #[test]
    fn component_status_eq() {
        assert_eq!(ComponentStatus::Running, ComponentStatus::Running);
        assert_eq!(ComponentStatus::Stopped, ComponentStatus::Stopped);
        assert_ne!(ComponentStatus::Running, ComponentStatus::Stopped);
    }

    #[test]
    fn is_process_alive_current_process() {
        let pid = std::process::id();
        assert!(is_process_alive(pid), "current process must be alive");
    }

    #[test]
    fn is_process_alive_nonexistent_pid() {
        // u32::MAX is effectively guaranteed to not be a valid running PID.
        assert!(
            !is_process_alive(u32::MAX),
            "PID u32::MAX must not be alive"
        );
    }
}