cordon 0.3.0

//! Support for creating cgroups via systemd transient units.
//!
//! Note that the `dbus` create requires pkg-config and libdbus-1-dev to be installed. On
//! Debian-based systems:
//!
//! ```text
//! $ sudo apt install libdbus-1-dev pkg-config
//! ```

use dbus::blocking::Connection;
use std::sync::{Arc, Mutex};
use std::time::Duration;
use thiserror::Error;
use tracing::{debug, instrument};

// Pull in generated trait (from dbus-codegen-rust) to make calls to the systemd Manager object.
mod dbus_generated;
use dbus_generated::OrgFreedesktopSystemd1Manager;

#[derive(Error, Debug)]
pub enum Error {
    #[error("dbus error")]
    Dbus(#[from] dbus::Error),

    #[error("invalid scope name: {0:?}")]
    InvalidScopeName(String),
}

/// Configure a control group via systemd.
#[derive(Debug, Clone)]
pub struct ScopeParameters {
    /// Scope name, like `cordon-XXXXXXXXXXXXXXXX.scope`.
    /// Can be generated by [`unique_scope_name`].
    pub unit_name: String,
    pub description: Option<String>,
    pub cpu_weight: Option<u64>,
    pub memory_high: Option<u64>,
    pub memory_max: Option<u64>,
    pub tasks_max: Option<u64>,
}

impl ScopeParameters {
    pub fn with_unique_name() -> Self {
        Self {
            unit_name: unique_scope_name(),
            description: None,
            cpu_weight: None,
            memory_high: None,
            memory_max: None,
            tasks_max: None,
        }
    }
}

#[derive(Clone, Debug)]
pub struct ScopeHandle {
    unit_name: String,
}

impl ScopeHandle {
    pub fn unit_name(&self) -> &str {
        self.unit_name.as_str()
    }
}

/// Start a transient systemd unit containing the given pid.
/// `scope_name` must be a valid systemd scope name, ending in `.scope`.
#[instrument(level = "debug", skip_all, fields(unit = %params.unit_name))]
pub fn start_transient_unit(pid: u32, params: ScopeParameters) -> Result<ScopeHandle, Error> {
    // First open up a connection to the session bus.
    let conn = Connection::new_session()?;

    // Second, create a wrapper struct around the connection that makes it easy
    // to send method calls to a specific destination and path.
    let proxy = conn.with_proxy(
        "org.freedesktop.systemd1",
        "/org/freedesktop/systemd1",
        Duration::from_millis(5000),
    );

    // Name the unit after the pid.
    if !params
        .unit_name
        .chars()
        .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '.')
        || !params.unit_name.ends_with(".scope")
    {
        return Err(Error::InvalidScopeName(params.unit_name));
    }

    // Set up the dbus argument describing the unit.
    let pid_list = dbus::arg::Variant(Box::new(vec![pid]) as Box<dyn dbus::arg::RefArg>);
    let var_str =
        |s: &str| dbus::arg::Variant(Box::new(s.to_string()) as Box<dyn dbus::arg::RefArg>);
    let var_u64 = |u: u64| dbus::arg::Variant(Box::new(u) as Box<dyn dbus::arg::RefArg>);
    let mut props = vec![
        ("PIDs", pid_list),
        ("Delegate", dbus::arg::Variant(Box::new(true))),
    ];
    if let Some(desc) = &params.description {
        props.push(("Description", var_str(desc)));
    }
    if let Some(cpu_weight) = params.cpu_weight {
        props.push(("CPUWeight", var_u64(cpu_weight)));
    }
    if let Some(memory_high) = params.memory_high {
        props.push(("MemoryHigh", var_u64(memory_high)));
    }
    if let Some(memory_max) = params.memory_max {
        props.push(("MemoryMax", var_u64(memory_max)));
    }
    if let Some(tasks_max) = params.tasks_max {
        props.push(("TasksMax", var_u64(tasks_max)));
    }

    // Start the transient unit.
    let job_path = proxy.start_transient_unit(&params.unit_name, "fail", props, vec![])?;
    debug!("requested unit start");

    let _cgroup_wait_start = std::time::Instant::now();

    // Listen for the event marking the completion of the job.
    let job_done = Arc::new(Mutex::new(false));
    proxy.match_signal({
        let job_done = job_done.clone();
        move |msg: dbus_generated::OrgFreedesktopSystemd1ManagerJobRemoved,
              _: &Connection,
              _: &dbus::Message| {
            if msg.job != job_path {
                return true;
            }
            *job_done.lock().unwrap() = true;
            false
        }
    })?;

    // Block until the job completes.
    while !*job_done.lock().unwrap() {
        conn.process(Duration::from_millis(1000))?;
    }

    Ok(ScopeHandle {
        unit_name: params.unit_name,
    })
}

/// Generate a unique, random scope name, like `cordon-XXXXXXXXXXXXXXXX.scope`.
///
/// By naming scopes starting with `cordon-`, administrators can manage all cordon containers with
/// a drop-in file for `cordon-.scope`.
pub fn unique_scope_name() -> String {
    use rand::seq::IteratorRandom;

    const ALPHABET: &str = "abcdefghijklmnopqrstuvwxyz0123456789";

    let mut rng = rand::thread_rng();
    let mut name = String::from("cordon-");
    for _ in 0..16 {
        name.push(ALPHABET.chars().choose(&mut rng).unwrap());
    }
    name.push_str(".scope");

    name
}