From c9c620bd44a7fd84e90e16ca632335fea20e94f7 Mon Sep 17 00:00:00 2001
From: Lars Wirzenius <liw@liw.fi>
Date: Sat, 18 Dec 2021 10:53:46 +0200
Subject: feat: add abstraction for daemons and managing them

Sponsored-by: author
---
 src/daemon.rs | 156 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 src/lib.rs    |   1 +
 2 files changed, 157 insertions(+)
 create mode 100644 src/daemon.rs

diff --git a/src/daemon.rs b/src/daemon.rs
new file mode 100644
index 0000000..c6da5f5
--- /dev/null
+++ b/src/daemon.rs
@@ -0,0 +1,156 @@
+use nix::sys::signal::kill;
+use nix::sys::signal::Signal;
+use nix::unistd::Pid;
+use std::fs::read;
+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 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: &[&str]) -> Result<Daemon, DaemonError> {
+        let stdout = NamedTempFile::new()?;
+        let stderr = NamedTempFile::new()?;
+        let pid = NamedTempFile::new()?;
+        let output = Command::new("daemonize")
+            .args(&[
+                "-c",
+                "/",
+                "-e",
+                &stderr.path().display().to_string(),
+                "-o",
+                &stdout.path().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);
+        }
+
+        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.
+                        return Ok(Daemon::new(pid));
+                    }
+                }
+                sleep_briefly();
+            } else {
+                sleep_briefly();
+            }
+        }
+
+        let cmd = argv.join(" ");
+        let err = read(stderr.path()).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.
+pub struct Daemon {
+    pid: Option<i32>,
+}
+
+impl Daemon {
+    fn new(pid: i32) -> Self {
+        Self { pid: Some(pid) }
+    }
+
+    /// 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() {
+            let pid = Pid::from_raw(raw_pid);
+            kill(pid, Some(Signal::SIGKILL)).map_err(|e| DaemonError::Kill(raw_pid, e))?;
+        }
+        Ok(())
+    }
+}
+
+impl Drop for Daemon {
+    fn drop(&mut self) {
+        if self.stop().is_err() {
+            // Do nothing.
+        }
+    }
+}
+
+fn sleep_briefly() {
+    sleep(Duration::from_millis(100));
+}
diff --git a/src/lib.rs b/src/lib.rs
index db53aac..87c57d1 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -28,6 +28,7 @@
 //! This crate only collects data from a set of benchmarks. It does
 //! not analyze the data. The data can be stored for later analysis.
 
+pub mod daemon;
 pub mod obnam;
 pub mod result;
 pub mod specification;
-- 
cgit v1.2.1