feat(dkl rc)
This commit is contained in:
Mikaël Cluseau
@@ -1 +1,495 @@
|
||||
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
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user