liboj_cgroups/
lib.rs

1//! High-level APIs for cgroups v1.
2//!
3//! The APIs do not support all feature that cgroup can provide. If some features is not used when
4//! implement an online judge or the kernel document doesn't provide enough information, the feature
5//! will not be included in the APIs.
6#[macro_use]
7mod attr_file;
8mod hierarchy;
9pub mod subsystem;
10
11#[cfg(test)]
12mod tests;
13
14use std::collections::HashSet;
15use std::fs;
16use std::io;
17use std::path::{Path, PathBuf};
18
19use lazy_static::lazy_static;
20use nix::unistd::Pid;
21
22use hierarchy::{Hierarchy, HierarchyNode};
23use subsystem::{prelude::*, Subsystem};
24
25lazy_static! {
26    static ref SYSTEM_CGROUP_ROOT: &'static Path = Path::new("/sys/fs/cgroup/");
27    static ref GLOBAL_CGROUP: GlobalContext =
28        GlobalContext::new().expect("failed to initialize cgroup context");
29}
30
31struct GlobalContext;
32
33impl GlobalContext {
34    const GLOBAL_NAME: &'static str = "liboj";
35
36    fn new() -> io::Result<GlobalContext> {
37        let res = GlobalContext {};
38        for controller in Subsystem::all() {
39            fs::create_dir_all(res.subsystem_path(controller))?;
40        }
41        Ok(res)
42    }
43
44    fn subsystem_path(&self, sub: Subsystem) -> PathBuf {
45        SYSTEM_CGROUP_ROOT.join(sub.name()).join(Self::GLOBAL_NAME)
46    }
47
48    fn subsystem(&self, sub: Subsystem) -> Hierarchy {
49        Hierarchy::new(self.subsystem_path(sub))
50    }
51}
52
53impl Drop for GlobalContext {
54    fn drop(&mut self) {
55        for sub in Subsystem::all() {
56            let _ = fs::remove_dir_all(GLOBAL_CGROUP.subsystem_path(sub));
57        }
58    }
59}
60
61/// Cgroup context.
62pub struct Context {
63    name: String,
64}
65
66impl Context {
67    fn get_controller(&self, sub: Subsystem) -> HierarchyNode {
68        match GLOBAL_CGROUP.subsystem(sub).get_child(&self.name) {
69            Some(controller) => controller,
70            None => panic!(
71                "Cgroup {} does not support {} subsystem",
72                self.name,
73                sub.name()
74            ),
75        }
76    }
77
78    pub fn add_process(&self, pid: Pid) -> io::Result<()> {
79        for sub in Subsystem::all() {
80            if let Some(controller) = GLOBAL_CGROUP.subsystem(sub).get_child(&self.name) {
81                controller.procs().write(&pid)?;
82            }
83        }
84        Ok(())
85    }
86
87    pub fn remove_process(&self, pid: Pid) -> io::Result<()> {
88        for sub in Subsystem::all().map(|sub| GLOBAL_CGROUP.subsystem(sub).root()) {
89            sub.procs().write(&pid)?;
90        }
91        Ok(())
92    }
93
94    pub fn cpu_controller(&self) -> Box<dyn '_ + CpuController> {
95        Box::new(self.get_controller(Subsystem::Cpu))
96    }
97
98    pub fn cpuacct_controller(&self) -> Box<dyn '_ + CpuAcctController> {
99        Box::new(self.get_controller(Subsystem::Cpuacct))
100    }
101
102    pub fn freezer_controller(&self) -> Box<dyn '_ + FreezerController> {
103        Box::new(self.get_controller(Subsystem::Freezer))
104    }
105
106    pub fn memory_controller(&self) -> Box<dyn '_ + MemoryController> {
107        Box::new(self.get_controller(Subsystem::Memory))
108    }
109
110    pub fn pids_controller(&self) -> Box<dyn '_ + PidsController> {
111        Box::new(self.get_controller(Subsystem::Pids))
112    }
113}
114
115pub struct Builder {
116    name: Option<String>,
117    subsystems: HashSet<Subsystem>,
118}
119
120impl Default for Builder {
121    fn default() -> Builder {
122        Builder::new().enable_all_subsystems()
123    }
124}
125
126impl Builder {
127    pub fn new() -> Builder {
128        Builder {
129            name: None,
130            subsystems: HashSet::new(),
131        }
132    }
133
134    pub fn set_name(mut self, name: String) -> Builder {
135        self.name = Some(name);
136        self
137    }
138
139    pub fn add_subsystem(mut self, sub: Subsystem) -> Builder {
140        self.subsystems.insert(sub);
141        self
142    }
143
144    pub fn enable_all_subsystems(mut self) -> Builder {
145        for sub in Subsystem::all() {
146            self.subsystems.insert(sub);
147        }
148        self
149    }
150
151    pub fn build(self) -> io::Result<Context> {
152        use uuid::Uuid;
153        let name = self.name.unwrap_or_else(|| format!("{:x}", Uuid::new_v4()));
154        for controller in self.subsystems.into_iter() {
155            GLOBAL_CGROUP.subsystem(controller).create(&name)?;
156        }
157        Ok(Context { name })
158    }
159}
liboj_cgroups/lib.rs

liboj_cgroups/
lib.rs
