dkl/src/rc.rs

use eyre::format_err;
use log::{error, info, warn};
use nix::sys::signal::Signal;
use std::collections::{BTreeMap as Map, BTreeSet as Set};
use std::path::PathBuf;
use std::sync::LazyLock;
use tokio::{
    io::{copy, AsyncBufReadExt, AsyncReadExt, AsyncWriteExt, BufReader},
    net::{UnixListener, UnixStream},
    sync::{mpsc, watch, RwLock},
};

use crate::{cgroup, fs};

mod runner;

use runner::{Child, State};

const CFG_PATH: &str = "/etc/direktil/rc.yaml";
const SOCK_PATH: &str = "/run/dkl-rc/ctl.sock"; // Path::new when stable

#[derive(Default, serde::Serialize, serde::Deserialize)]
pub struct Config {
    #[serde(default, skip_serializing_if = "Map::is_empty")]
    pub cgroups: Map<String, CgroupConfig>,
}

#[derive(serde::Serialize, serde::Deserialize)]
pub struct CgroupConfig {
    pub controllers: String,
    #[serde(default, skip_serializing_if = "Map::is_empty")]
    pub settings: Map<String, String>,
    #[serde(default, skip_serializing_if = "Map::is_empty")]
    pub services: Map<String, Service>,
}

pub type Service = Vec<String>;

static MANAGER: LazyLock<RwLock<Manager>> = LazyLock::new(|| RwLock::new(Manager::default()));

type Result<T> = std::result::Result<T, Error>;

#[derive(Debug, thiserror::Error)]
enum Error {
    #[error("invalid command: {0:?}")]
    InvalidCommand(String),
    #[error("config read failed: {0}")]
    ConfigRead(fs::Error),
    #[error("config parse failed: {0}")]
    ConfigParse(serde_yaml::Error),
    #[error("cgroup setup failed: {0}")]
    CgroupSetup(fs::Error),
    #[error("invalid key (cgroup/service)")]
    InvalidKey,
    #[error("unknown cgroup: {0:?}")]
    UnknownCgroup(String),
    #[error("unknown service: {0:?}")]
    UnknownService(String),
    #[error("invalid signal: {0:?}")]
    InvalidSignal(String),
    #[error("process exited")]
    ProcessExited,
    #[error("nothing running under {0:?}")]
    NotRunning(String),
    #[error("kill failed: {0:?}")]
    KillFailed(nix::Error),
    #[error("service runner is dead")]
    RunnerDead,
}

pub async fn run() -> eyre::Result<()> {
    info!("starting");

    tokio::spawn(wait_terminate());

    let _ = reload_config().await;
    tokio::spawn(wait_reload());

    if let Some(sock_dir) = PathBuf::from(SOCK_PATH).parent() {
        let _ = tokio::fs::DirBuilder::new()
            .mode(0o700)
            .create(sock_dir)
            .await;
    }

    let _ = tokio::fs::remove_file(SOCK_PATH).await;
    let listener = UnixListener::bind(SOCK_PATH)?;

    loop {
        let Ok((conn, _)) = listener.accept().await else {
            warn!("listener closed");
            break;
        };

        tokio::spawn(async move { handle(conn).await });
    }

    cleanup().await;
    Ok(())
}

async fn cleanup() {
    let _ = tokio::fs::remove_file(SOCK_PATH).await;
}

pub async fn ctl<I, S>(args: I) -> eyre::Result<()>
where
    I: IntoIterator<Item = S>,
    S: Into<String>,
{
    let args: Vec<_> = args.into_iter().map(|s| s.into()).collect();
    let args = format!("{}\n", args.join(" "));

    match ctl_exec(args.as_bytes()).await {
        Ok(mut rd) => {
            copy(&mut rd, &mut tokio::io::stdout()).await?;
            std::process::exit(0);
        }
        Err(e) => {
            eprint!("{e}");
            std::process::exit(1);
        }
    }
}

async fn ctl_exec(request: &[u8]) -> eyre::Result<BufReader<UnixStream>> {
    let mut conn = UnixStream::connect(SOCK_PATH)
        .await
        .map_err(|e| format_err!("{SOCK_PATH}: {e}"))?;

    conn.write_all(request).await?;

    let mut rd = BufReader::with_capacity(64, conn);

    let mut code = String::new();
    rd.read_line(&mut code).await?;
    let code: i32 = code.trim_ascii().parse()?;

    if code != 0 {
        let mut err = String::new();
        rd.read_to_string(&mut err).await?;
        return Err(format_err!("{}", err.trim_ascii_end()));
    }

    Ok(rd)
}

async fn handle(mut conn: UnixStream) {
    let (rd, mut wr) = conn.split();
    let mut rd = BufReader::with_capacity(64, rd).lines();

    let Ok(Some(line)) = rd.next_line().await else {
        return;
    };

    let mut line = line.split_ascii_whitespace();

    macro_rules! next {
        () => {{
            match line.next() {
                Some(v) => v,
                None => return,
            }
        }};
    }

    let r = match next!() {
        "ls" => Ok(Some(ls().await)),
        "status" => Ok(Some(status().await)),
        "reload-config" => reload_config().await.map(|_| None),
        "start" => start(next!()).await.map(|_| None),
        "stop" => stop(next!()).await.map(|_| None),
        "reload" => reload(next!()).await.map(|_| None),
        "sig" => sig(next!(), next!()).await.map(|_| None),
        cmd => Err(Error::InvalidCommand(cmd.into())),
    };

    let _ = match r {
        Ok(None) => wr.write_all(b"0\n").await,
        Ok(Some(s)) => wr.write_all(format!("0\n{s}\n").as_bytes()).await,
        Err(e) => wr.write_all(format!("1\n{e}\n").as_bytes()).await,
    };

    let _ = wr.shutdown().await;
}

async fn wait_terminate() {
    use tokio::signal::unix::{signal, SignalKind};
    let Ok(mut sig) = signal(SignalKind::terminate())
        .inspect_err(|e| error!("failed to listen to SIGTERM (will be ignored): {e}"))
    else {
        return;
    };

    sig.recv().await;

    info!("SIGTERM received, terminating");
    MANAGER.write().await.terminate().await;

    cleanup().await;
    log::logger().flush();
    std::process::exit(0);
}

async fn wait_reload() {
    use tokio::signal::unix::{signal, SignalKind};
    let Ok(mut sig) = signal(SignalKind::hangup())
        .inspect_err(|e| error!("failed to listen to SIGHUP (will be ignored): {e}"))
    else {
        return;
    };

    loop {
        sig.recv().await;
        let _ = reload_config().await;
    }
}

async fn reload_config() -> Result<()> {
    let cfg = (fs::read(CFG_PATH).await)
        .map_err(Error::ConfigRead)
        .inspect_err(|e| error!("{e}"))?;

    let cfg = serde_yaml::from_slice::<Config>(&cfg)
        .map_err(Error::ConfigParse)
        .inspect_err(|e| error!("{CFG_PATH}: {e}"))?;

    info!("applying new config");
    let r = MANAGER.write().await.apply_config(cfg).await;
    match &r {
        Ok(_) => info!("applied new config"),
        Err(e) => info!("failed to apply new config: {e}"),
    }
    r
}

async fn ls() -> String {
    let mut keys = String::new();
    for (i, k) in MANAGER.read().await.runners.keys().enumerate() {
        if i != 0 {
            keys.push('\n');
        }
        keys.push_str(k);
    }
    keys
}

async fn status() -> String {
    let status = MANAGER.read().await.status();

    let mut table = tabled::builder::Builder::new();
    table.push_record(["cgroup", "service", "PID", "state", "msg"]);

    for (cg_svc, child) in status {
        let (cg, svc) = cg_svc.split_once('/').unwrap();
        let pid = child.pid.map(|p| p.to_string());
        table.push_record([
            cg,
            svc,
            pid.as_deref().unwrap_or("◌"),
            &format!("{:?}", child.state),
            child.msg.as_deref().unwrap_or("◌"),
        ]);
    }

    (table.build())
        .with(tabled::settings::Style::psql())
        .to_string()
}

async fn start(key: &str) -> Result<()> {
    MANAGER.write().await.start(key).await
}

async fn stop(key: &str) -> Result<()> {
    MANAGER.write().await.stop(key).await
}

async fn reload(key: &str) -> Result<()> {
    MANAGER.read().await.reload(key).await
}

async fn sig(key: &str, sig: &str) -> Result<()> {
    let sig: Signal = sig.parse().map_err(|_| Error::InvalidSignal(sig.into()))?;
    signal(key, sig).await
}

async fn child_for(key: &str) -> Result<Child> {
    MANAGER.read().await.child_for(key)
}

async fn signal(key: &str, sig: Signal) -> Result<()> {
    child_for(key).await?.kill(sig)
}

fn child_key(cg: &str, svc: &str) -> String {
    [cg, svc].join("/")
}
fn split_key(key: &str) -> Result<(&str, &str)> {
    key.split_once('/').ok_or(Error::InvalidKey)
}

#[derive(Default)]
struct Manager {
    cfg: Config,
    procs: Map<String, watch::Receiver<Child>>,
    runners: Map<String, mpsc::Sender<runner::Cmd>>,
}

impl Manager {
    fn status(&self) -> Vec<(String, Child)> {
        (self.procs.iter())
            .map(|(n, c)| (n.clone(), c.borrow().clone()))
            .collect()
    }

    fn child_for(&self, key: &str) -> Result<Child> {
        (self.procs.get(key))
            .map(|c| c.borrow().clone())
            .ok_or_else(|| Error::NotRunning(key.into()))
    }

    async fn apply_config(&mut self, new_cfg: Config) -> Result<()> {
        // create and configure cgroups
        for (name, cg) in &new_cfg.cgroups {
            let cg_path = PathBuf::from(cgroup::ROOT).join(name);
            fs::create_dir_all(&cg_path)
                .await
                .map_err(Error::CgroupSetup)?;

            fs::write(
                cg_path.join("cgroup.subtree_control"),
                cg.controllers.as_bytes(),
            )
            .await
            .map_err(Error::CgroupSetup)?;

            for (setting, value) in &cg.settings {
                fs::write(cg_path.join(setting), value.as_bytes())
                    .await
                    .map_err(Error::CgroupSetup)?;
            }
        }

        let new_svcs: Set<_> = new_cfg.service_keys().collect();

        // stop removed services
        let to_stop = Map::from_iter(self.runners.extract_if(.., |k, _| !new_svcs.contains(k)));
        let mut stopped = Set::new();
        for (key, runner_cmd) in to_stop {
            if runner_cmd.send(runner::Cmd::Stop).await.is_err() {
                // runner already dead
                continue;
            }
            stopped.insert(key);
        }

        // start added services
        for (key, cg, svc, service) in new_cfg.services() {
            if self.runners.contains_key(&key) {
                continue;
            };

            let cmd = self.spawn_runner(key, cg, svc, service.clone());
            if let Err(e) = cmd.send(runner::Cmd::Start).await {
                error!("runner instantly died: {e}");
            }
        }

        // wait & cleanup stopped
        for key in stopped {
            let Some(mut child_rx) = self.procs.remove(&key) else {
                continue;
            };
            let _ = child_rx
                .wait_for(|c| matches!(c.state, State::Finalized))
                .await;
        }

        self.cfg = new_cfg;
        Ok(())
    }

    async fn terminate(&mut self) {
        self.runners.clear();

        for child in self.procs.values_mut() {
            let _ = child
                .wait_for(|c| matches!(c.state, State::Finalized))
                .await;
        }

        self.procs.clear();
    }

    fn runner(&mut self, key: &str) -> Result<mpsc::Sender<runner::Cmd>> {
        if let Some(c) = self.runners.get(key) {
            return Ok(c.clone());
        }

        let (cg, svc) = split_key(key)?;
        let service = self.cfg.service(key)?;

        Ok(self.spawn_runner(key.into(), cg, svc, service.clone()))
    }

    fn spawn_runner(
        &mut self,
        key: String,
        cg: &str,
        svc: &str,
        service: Service,
    ) -> mpsc::Sender<runner::Cmd> {
        let (runner, child_rx, cmds_tx) = runner::new(cg, svc, service);

        tokio::spawn(runner.run());

        self.procs.insert(key.clone(), child_rx);
        self.runners.insert(key, cmds_tx.clone());

        cmds_tx
    }

    async fn cmd(&mut self, key: &str, cmd: runner::Cmd) -> Result<()> {
        if self.runner(key)?.send(cmd).await.is_err() {
            // runner died
            self.runners.remove(key);
            return Err(Error::RunnerDead);
        }
        Ok(())
    }

    async fn start(&mut self, key: &str) -> Result<()> {
        self.cmd(key, runner::Cmd::Start).await
    }

    async fn stop(&mut self, key: &str) -> Result<()> {
        self.cmd(key, runner::Cmd::Stop).await
    }

    async fn reload(&self, key: &str) -> Result<()> {
        let proc = (self.procs.get(key)) //
            .ok_or_else(|| Error::UnknownService(key.into()))?;
        proc.borrow().reload()
    }
}

impl Config {
    fn cgroup(&self, cg: &str) -> Result<&CgroupConfig> {
        self.cgroups
            .get(cg)
            .ok_or_else(|| Error::UnknownCgroup(cg.into()))
    }

    fn service(&self, key: &str) -> Result<&Service> {
        let (cg, svc) = split_key(key)?;
        self.cgroup(cg)?.service(svc)
    }

    fn service_keys(&self) -> impl Iterator<Item = String> {
        (self.cgroups.iter())
            .map(|(cg_name, cg)| cg.services.keys().map(move |n| child_key(cg_name, n)))
            .flatten()
    }

    fn services(&self) -> impl Iterator<Item = (String, &String, &String, &Service)> {
        (self.cgroups.iter())
            .map(|(cg_name, cg)| {
                cg.services
                    .iter()
                    .map(move |(n, service)| (child_key(cg_name, n), cg_name, n, service))
            })
            .flatten()
    }
}

impl CgroupConfig {
    fn service(&self, svc: &str) -> Result<&Vec<String>> {
        self.services
            .get(svc)
            .ok_or_else(|| Error::UnknownService(svc.into()))
    }
}

pub async fn complete() -> Vec<String> {
    let mut r = vec![];
    let Ok(rd) = ctl_exec(b"ls\n").await else {
        return r;
    };
    let mut rd = rd.lines();
    while let Some(line) = rd.next_line().await.ok().flatten() {
        r.push(line);
    }
    r
}