nordisk/monetize

crates/monetize-admin/src/supervise.rs

//! **A child process kept running** — PID 1's two children (the server and the
//! console). A child that dies is started again after a backoff (one second,
//! doubling to thirty, back to one after a long run); [`Supervisor::restart`]
//! stops the running one and starts it again at once; dropping the supervisor
//! stops the child. The child's lines are said again in PID 1's log (the
//! installer collects that one).
//!
//! As PID 1 the reaper (`appliance_pid1::runtime::reaper`) may take a child's
//! exit status first: a wait that finds no child (ECHILD) is an exit too.

use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::Arc;
use std::time::{Duration, Instant};

use tokio::io::AsyncBufReadExt as _;
use tokio::sync::Notify;

const BACKOFF_FIRST: Duration = Duration::from_secs(1);
const BACKOFF_MAX: Duration = Duration::from_secs(30);
/// A run longer than this resets the backoff.
const LONG_RUN: Duration = Duration::from_secs(60);

/// **One child, kept running.**
pub struct Supervisor {
    starts: Arc<AtomicU64>,
    restart: Arc<Notify>,
    task: tokio::task::JoinHandle<()>,
}

impl Supervisor {
    /// Start `program args` now, and again whenever it ends. Must run inside
    /// the tokio runtime.
    pub fn spawn(name: &'static str, program: PathBuf, args: Vec<String>) -> Supervisor {
        let starts = Arc::new(AtomicU64::new(0));
        let restart = Arc::new(Notify::new());
        // A32-14, stays a task: the supervisor of one child; its handle is kept and aborted at the stop.
        let task = tokio::spawn(keep(name, program, args, Arc::clone(&starts), Arc::clone(&restart)));
        Supervisor { starts, restart, task }
    }

    /// Stop the running child and start it again at once.
    pub fn restart(&self) {
        self.restart.notify_one();
    }

    /// How many times the child was started.
    pub fn starts(&self) -> u64 {
        self.starts.load(Ordering::SeqCst)
    }
}

impl Drop for Supervisor {
    fn drop(&mut self) {
        // The task's child is `kill_on_drop`: aborting the task kills it.
        self.task.abort();
    }
}

async fn keep(name: &'static str, program: PathBuf, args: Vec<String>, starts: Arc<AtomicU64>, restart: Arc<Notify>) {
    let mut backoff = BACKOFF_FIRST;
    loop {
        // A32-14, a process by design: PID 1 supervises its own children (this binary's `server` and
        // `ui`), restarting one that stops; a box's processes are its design, never cargo or sh.
        let mut cmd = tokio::process::Command::new(&program);
        cmd.args(&args).stdin(std::process::Stdio::null()).stdout(std::process::Stdio::piped()).stderr(std::process::Stdio::piped()).kill_on_drop(true);
        let began = Instant::now();
        // A32-14, a process by design: the child process itself (see above): PID 1 supervises it.
        let mut child = match cmd.spawn() {
            Ok(c) => c,
            Err(e) => {
                tracing::error!(child = name, error = %e, "supervise: cannot start {}", program.display());
                tokio::time::sleep(backoff).await;
                backoff = (backoff * 2).min(BACKOFF_MAX);
                continue;
            }
        };
        starts.fetch_add(1, Ordering::SeqCst);
        tracing::info!(child = name, pid = child.id(), "supervise: started");
        if let Some(out) = child.stdout.take() {
            // A32-14, stays a task: the child's stdout reader, ends with the child.
            tokio::spawn(say(name, tokio::io::BufReader::new(out)));
        }
        if let Some(err) = child.stderr.take() {
            // A32-14, stays a task: the child's stderr reader, ends with the child.
            tokio::spawn(say(name, tokio::io::BufReader::new(err)));
        }
        let asked = tokio::select! {
            ended = child.wait() => {
                match ended {
                    Ok(st) => tracing::warn!(child = name, status = %st, "supervise: ended"),
                    // The reaper took the status first: it ended all the same.
                    Err(e) => tracing::warn!(child = name, error = %e, "supervise: ended (status taken by the reaper)"),
                }
                false
            }
            _ = restart.notified() => {
                tracing::info!(child = name, "supervise: restart asked");
                let _ = child.kill().await;
                true
            }
        };
        if asked {
            backoff = BACKOFF_FIRST;
            continue;
        }
        if began.elapsed() > LONG_RUN {
            backoff = BACKOFF_FIRST;
        }
        tokio::select! {
            _ = tokio::time::sleep(backoff) => {}
            _ = restart.notified() => {}
        }
        backoff = (backoff * 2).min(BACKOFF_MAX);
    }
}

async fn say<R: tokio::io::AsyncRead + Unpin>(name: &'static str, r: tokio::io::BufReader<R>) {
    let mut lines = r.lines();
    while let Ok(Some(l)) = lines.next_line().await {
        tracing::info!(child = name, "{l}");
    }
}