obnam_benchmark/

daemon.rs

use log::{debug, error, info};
use nix::sys::signal::kill;
use nix::sys::signal::Signal;
use nix::unistd::Pid;
use std::ffi::OsStr;
use std::fs::read;
use std::os::unix::ffi::OsStrExt;
use std::path::{Path, PathBuf};
use std::process::Command;
use std::thread::sleep;
use std::time::{Duration, Instant};
use tempfile::NamedTempFile;

/// Possible errors from starting and stopping daemons.
#[derive(Debug, thiserror::Error)]
pub enum DaemonError {
    /// The daemon took too long to start. The timeout can be
    /// configured with [DaemonManager::timeout].
    #[error("daemon took longer than {0} ms to start: {1}\n{2}")]
    Timeout(u128, String, String),

    /// Something went wrong, when handling temporary files.
    #[error(transparent)]
    TempFile(#[from] std::io::Error),

    /// Something went wrong read standard output of daemon.
    #[error("failed to read daemon stdout: {0}")]
    Stdout(std::io::Error),

    /// Something went wrong read error output of daemon.
    #[error("failed to read daemon stderr: {0}")]
    Stderr(std::io::Error),

    /// Failed to kill a daemon.
    #[error("failed to kill process {0}: {1}")]
    Kill(i32, nix::Error),
}

/// Manage daemons.
///
/// A daemon is a process running in the background, doing useful
/// things. For Obnam benchmarks, it's the Obnam server, but this is a
/// generic manager. This version requires the `daemonize` helper
/// program to be available on $PATH.
pub struct DaemonManager {
    timeout: Duration,
}

impl Default for DaemonManager {
    fn default() -> Self {
        Self {
            timeout: Duration::from_millis(1000),
        }
    }
}

impl DaemonManager {
    /// Create a new manager instance, with default settings.
    pub fn new() -> Self {
        Self::default()
    }

    /// Set the timeout for waiting on a daemon to start, in
    /// milliseconds.
    pub fn timeout(mut self, millis: u64) -> Self {
        self.timeout = Duration::from_millis(millis);
        self
    }

    /// Start a daemon.
    ///
    /// The daemon is considered started if its process id (PID) is
    /// known. The daemon may take longer to actually be in a useful
    /// state, and it may fail after the PID is known, for example if
    /// it reads a configuration file and that has errors. This
    /// function won't wait for that to happen: it only cares about
    /// the PID.
    pub fn start(
        &self,
        argv: &[&OsStr],
        stdout: &Path,
        stderr: &Path,
    ) -> Result<Daemon, DaemonError> {
        info!("start daemon: {:?}", argv);
        let pid = NamedTempFile::new()?;
        let output = Command::new("daemonize")
            .args(&[
                "-c",
                "/",
                "-e",
                &stderr.display().to_string(),
                "-o",
                &stdout.display().to_string(),
                "-p",
                &pid.path().display().to_string(),
            ])
            .args(argv)
            .output()
            .unwrap();
        if output.status.code() != Some(0) {
            eprintln!("{}", String::from_utf8_lossy(&output.stdout));
            eprintln!("{}", String::from_utf8_lossy(&output.stderr));
            std::process::exit(1);
        }

        debug!("waiting for daemon to write PID file");
        let time = Instant::now();
        while time.elapsed() < self.timeout {
            // Do we have the pid file?
            if let Ok(pid) = std::fs::read(pid.path()) {
                // Parse it as a string. We don't mind if it's not purely UTF8.
                let pid = String::from_utf8_lossy(&pid).into_owned();
                // Strip newline, if any.
                if let Some(pid) = pid.strip_suffix('\n') {
                    // Parse as an integer, if possible.
                    if let Ok(pid) = pid.parse() {
                        // We have a pid, stop waiting.
                        info!("got pid for daemon: pid");
                        return Ok(Daemon::new(pid, stdout, stderr));
                    }
                }
                sleep_briefly();
            } else {
                sleep_briefly();
            }
        }

        error!(
            "no PID file within {} ms, giving up",
            self.timeout.as_millis()
        );
        let mut cmd = String::new();
        for arg in argv {
            if !cmd.is_empty() {
                cmd.push(' ');
            }
            cmd.push_str(
                &String::from_utf8_lossy(arg.as_bytes())
                    .to_owned()
                    .to_string(),
            );
        }
        let err = read(&stderr).map_err(DaemonError::Stderr)?;
        let err = String::from_utf8_lossy(&err).into_owned();
        Err(DaemonError::Timeout(self.timeout.as_millis(), cmd, err))
    }
}

/// A running daemon.
///
/// The daemon process is killed, when the `Daemon` struct is dropped.
#[derive(Debug)]
pub struct Daemon {
    pid: Option<i32>,
    stdout: PathBuf,
    stderr: PathBuf,
}

impl Daemon {
    fn new(pid: i32, stdout: &Path, stderr: &Path) -> Self {
        info!("started daemon with PID {}", pid);
        Self {
            pid: Some(pid),
            stdout: stdout.to_path_buf(),
            stderr: stderr.to_path_buf(),
        }
    }

    /// Explicitly stop a daemon.
    ///
    /// Calling this function is only useful if you want to handle
    /// errors. It can only be called once.
    pub fn stop(&mut self) -> Result<(), DaemonError> {
        if let Some(raw_pid) = self.pid.take() {
            info!("stopping daemon with PID {}", raw_pid);
            let pid = Pid::from_raw(raw_pid);
            kill(pid, Some(Signal::SIGKILL)).map_err(|e| DaemonError::Kill(raw_pid, e))?;
        }
        Ok(())
    }

    /// Return what the daemon has written to its stderr so far.
    pub fn stdout(&self) -> Result<Vec<u8>, DaemonError> {
        std::fs::read(&self.stdout).map_err(DaemonError::Stdout)
    }

    /// Return what the daemon has written to its stderr so far.
    pub fn stderr(&self) -> Result<Vec<u8>, DaemonError> {
        std::fs::read(&self.stderr).map_err(DaemonError::Stderr)
    }
}

impl Drop for Daemon {
    fn drop(&mut self) {
        if self.stop().is_err() {
            // Do nothing.
        }
    }
}

fn sleep_briefly() {
    sleep(Duration::from_millis(100));
}