cordon/
systemd.rs

1//! Support for creating cgroups via systemd transient units.
2//!
3//! Note that the `dbus` create requires pkg-config and libdbus-1-dev to be installed. On
4//! Debian-based systems:
5//!
6//! ```text
7//! $ sudo apt install libdbus-1-dev pkg-config
8//! ```
9
10use dbus::blocking::Connection;
11use std::sync::{Arc, Mutex};
12use std::time::Duration;
13use thiserror::Error;
14use tracing::{debug, instrument};
15
16// Pull in generated trait (from dbus-codegen-rust) to make calls to the systemd Manager object.
17mod dbus_generated;
18use dbus_generated::OrgFreedesktopSystemd1Manager;
19
20#[derive(Error, Debug)]
21pub enum Error {
22    #[error("dbus error")]
23    Dbus(#[from] dbus::Error),
24
25    #[error("invalid scope name: {0:?}")]
26    InvalidScopeName(String),
27}
28
29/// Configure a control group via systemd.
30#[derive(Debug, Clone)]
31pub struct ScopeParameters {
32    /// Scope name, like `cordon-XXXXXXXXXXXXXXXX.scope`.
33    /// Can be generated by [`unique_scope_name`].
34    pub unit_name: String,
35    pub description: Option<String>,
36    pub cpu_weight: Option<u64>,
37    pub memory_high: Option<u64>,
38    pub memory_max: Option<u64>,
39    pub tasks_max: Option<u64>,
40}
41
42impl ScopeParameters {
43    pub fn with_unique_name() -> Self {
44        Self {
45            unit_name: unique_scope_name(),
46            description: None,
47            cpu_weight: None,
48            memory_high: None,
49            memory_max: None,
50            tasks_max: None,
51        }
52    }
53}
54
55#[derive(Clone, Debug)]
56pub struct ScopeHandle {
57    unit_name: String,
58}
59
60impl ScopeHandle {
61    pub fn unit_name(&self) -> &str {
62        self.unit_name.as_str()
63    }
64}
65
66/// Start a transient systemd unit containing the given pid.
67/// `scope_name` must be a valid systemd scope name, ending in `.scope`.
68#[instrument(level = "debug", skip_all, fields(unit = %params.unit_name))]
69pub fn start_transient_unit(pid: u32, params: ScopeParameters) -> Result<ScopeHandle, Error> {
70    // First open up a connection to the session bus.
71    let conn = Connection::new_session()?;
72
73    // Second, create a wrapper struct around the connection that makes it easy
74    // to send method calls to a specific destination and path.
75    let proxy = conn.with_proxy(
76        "org.freedesktop.systemd1",
77        "/org/freedesktop/systemd1",
78        Duration::from_millis(5000),
79    );
80
81    // Name the unit after the pid.
82    if !params
83        .unit_name
84        .chars()
85        .all(|c| c.is_ascii_alphanumeric() || c == '-' || c == '.')
86        || !params.unit_name.ends_with(".scope")
87    {
88        return Err(Error::InvalidScopeName(params.unit_name));
89    }
90
91    // Set up the dbus argument describing the unit.
92    let pid_list = dbus::arg::Variant(Box::new(vec![pid]) as Box<dyn dbus::arg::RefArg>);
93    let var_str =
94        |s: &str| dbus::arg::Variant(Box::new(s.to_string()) as Box<dyn dbus::arg::RefArg>);
95    let var_u64 = |u: u64| dbus::arg::Variant(Box::new(u) as Box<dyn dbus::arg::RefArg>);
96    let mut props = vec![
97        ("PIDs", pid_list),
98        ("Delegate", dbus::arg::Variant(Box::new(true))),
99    ];
100    if let Some(desc) = &params.description {
101        props.push(("Description", var_str(desc)));
102    }
103    if let Some(cpu_weight) = params.cpu_weight {
104        props.push(("CPUWeight", var_u64(cpu_weight)));
105    }
106    if let Some(memory_high) = params.memory_high {
107        props.push(("MemoryHigh", var_u64(memory_high)));
108    }
109    if let Some(memory_max) = params.memory_max {
110        props.push(("MemoryMax", var_u64(memory_max)));
111    }
112    if let Some(tasks_max) = params.tasks_max {
113        props.push(("TasksMax", var_u64(tasks_max)));
114    }
115
116    // Start the transient unit.
117    let job_path = proxy.start_transient_unit(&params.unit_name, "fail", props, vec![])?;
118    debug!("requested unit start");
119
120    let _cgroup_wait_start = std::time::Instant::now();
121
122    // Listen for the event marking the completion of the job.
123    let job_done = Arc::new(Mutex::new(false));
124    proxy.match_signal({
125        let job_done = job_done.clone();
126        move |msg: dbus_generated::OrgFreedesktopSystemd1ManagerJobRemoved,
127              _: &Connection,
128              _: &dbus::Message| {
129            if msg.job != job_path {
130                return true;
131            }
132            *job_done.lock().unwrap() = true;
133            false
134        }
135    })?;
136
137    // Block until the job completes.
138    while !*job_done.lock().unwrap() {
139        conn.process(Duration::from_millis(1000))?;
140    }
141
142    Ok(ScopeHandle {
143        unit_name: params.unit_name,
144    })
145}
146
147/// Generate a unique, random scope name, like `cordon-XXXXXXXXXXXXXXXX.scope`.
148///
149/// By naming scopes starting with `cordon-`, administrators can manage all cordon containers with
150/// a drop-in file for `cordon-.scope`.
151pub fn unique_scope_name() -> String {
152    use rand::seq::IteratorRandom;
153
154    const ALPHABET: &str = "abcdefghijklmnopqrstuvwxyz0123456789";
155
156    let mut rng = rand::thread_rng();
157    let mut name = String::from("cordon-");
158    for _ in 0..16 {
159        name.push(ALPHABET.chars().choose(&mut rng).unwrap());
160    }
161    name.push_str(".scope");
162
163    name
164}
cordon/systemd.rs

cordon/
systemd.rs
