sausage/host/src/process/services.rs

//! Supervises the world's services list, live edits, crash limits, and graceful stops.
use super::{GroupGuard, Request, Runtime, Started, Stdio};
use serde::Deserialize;
use std::{
    collections::{BTreeMap, VecDeque},
    sync::Arc,
    time::Duration,
};
use tokio::time::Instant;

#[derive(Clone, Deserialize, PartialEq, Eq)]
struct Spec {
    name: String,
    argv: Vec<String>,
    #[serde(default = "root")]
    cwd: String,
    #[serde(default)]
    env: BTreeMap<String, String>,
}
fn root() -> String {
    "/".into()
}
struct Live {
    child: Started,
    _guard: GroupGuard,
}
struct Service {
    spec: Spec,
    live: Option<Live>,
    crashes: VecDeque<Instant>,
    stopped: bool,
}
impl Service {
    async fn start(&mut self, runtime: &Arc<Runtime>) {
        match runtime
            .start(
                Request {
                    program: self.spec.argv.first().cloned().unwrap_or_default(),
                    argv: self.spec.argv.clone(),
                    cwd: self.spec.cwd.clone(),
                    env: self
                        .spec
                        .env
                        .iter()
                        .map(|(k, v)| (k.clone(), v.clone()))
                        .collect(),
                    group: None,
                    stdin: Stdio::Pipe,
                    stdout: Stdio::Null,
                    stderr: Stdio::Null,
                },
                0,
                0,
                0,
                None,
                None,
            )
            .await
        {
            Ok(child) => {
                let guard = GroupGuard {
                    runtime: runtime.clone(),
                    group: child.pid,
                };
                self.live = Some(Live {
                    child,
                    _guard: guard,
                });
            }
            Err(_) => self.crashed(),
        }
    }
    fn crashed(&mut self) {
        let now = Instant::now();
        while self
            .crashes
            .front()
            .is_some_and(|t| now.duration_since(*t) >= Duration::from_secs(10))
        {
            self.crashes.pop_front();
        }
        self.crashes.push_back(now);
        self.stopped = self.crashes.len() > 5;
    }
    async fn stop(&mut self) {
        if let Some(mut live) = self.live.take() {
            drop(live.child.input.take());
            let _ = tokio::time::timeout(Duration::from_secs(5), async {
                loop {
                    if live.child.code.borrow().is_some() {
                        break;
                    }
                    if live.child.code.changed().await.is_err() {
                        break;
                    }
                }
            })
            .await;
            // The guard stops remaining descendants, or the process after its grace period.
        }
    }
}
/// Start the supervisor. Aborting its task stops its service processes.
/// The YAML file contains name, argv, optional cwd, and an optional map of environment variables.
pub fn supervise(runtime: Arc<Runtime>) -> tokio::task::JoinHandle<()> {
    tokio::spawn(async move {
        let mut services: BTreeMap<String, Service> = BTreeMap::new();
        let mut previous = Vec::new();
        let mut interval = tokio::time::interval(Duration::from_millis(100));
        loop {
            interval.tick().await;
            let content = match runtime
                .world
                .root
                .open("/config/services.yaml", Default::default())
                .await
            {
                Ok(file) => match file.stat().await {
                    Ok(stat) => file.read(0, stat.size as usize).await.ok(),
                    Err(_) => None,
                },
                Err(e) if e.kind() == std::io::ErrorKind::NotFound => Some(b"[]".to_vec()),
                Err(_) => None,
            };
            if let Some(content) = content.filter(|c| *c != previous)
                && let Ok(specs) = serde_yaml_ng::from_slice::<Vec<Spec>>(&content)
            {
                let specs: BTreeMap<_, _> =
                    specs.into_iter().map(|s| (s.name.clone(), s)).collect();
                let removed: Vec<_> = services
                    .keys()
                    .filter(|name| !specs.contains_key(*name))
                    .cloned()
                    .collect();
                for name in removed {
                    services.remove(&name).unwrap().stop().await;
                }
                for (name, spec) in specs {
                    let service = services.entry(name).or_insert_with(|| Service {
                        spec: spec.clone(),
                        live: None,
                        crashes: VecDeque::new(),
                        stopped: false,
                    });
                    if service.spec != spec {
                        service.stop().await;
                    }
                    service.spec = spec;
                    service.crashes.clear();
                    service.stopped = false;
                }
                previous = content;
            }
            for service in services.values_mut() {
                if service
                    .live
                    .as_ref()
                    .is_some_and(|live| live.child.code.borrow().is_some())
                {
                    service.live = None;
                    service.crashed();
                }
                if service.live.is_none() && !service.stopped {
                    service.start(&runtime).await;
                }
            }
        }
    })
}