vc_processors/sys/
cgroup.rs

1//! This module provides cgroup helpers
2//!
3
4pub use imp::*;
5
6/// ENV name of cgroup name that we want to load
7pub const ENV_CGROUP_NAME: &str = "VC_CG_NAME";
8/// ENV name of cpuset that we want to load
9pub const ENV_CGROUP_CPUSET: &str = "VC_CG_CPUSET";
10
11#[cfg(not(target_os = "linux"))]
12mod imp {
13    use anyhow::Result;
14
15    /// try load cgroup from env
16    pub fn try_load_from_env() -> CtrlGroup {
17        CtrlGroup::new("", "").expect("never fail")
18    }
19
20    /// CgroupPid represents a pid
21    #[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
22    pub struct CgroupPid(u64);
23
24    impl From<u64> for CgroupPid {
25        fn from(pid: u64) -> Self {
26            Self(pid)
27        }
28    }
29
30    /// An Cgroup type
31    pub struct CtrlGroup {}
32
33    impl CtrlGroup {
34        /// Creates a new CtrlGroup with cgroup name `_cgname` and cpuset `_cpuset`.
35        pub fn new(_cgname: impl AsRef<str>, _cpuset: impl AsRef<str>) -> Result<Self> {
36            tracing::warn!("{} does not support cgroups.", std::env::consts::OS);
37            Ok(Self {})
38        }
39
40        /// Attach a task to this controller.
41        pub fn add_task(&mut self, _pid: CgroupPid) -> Result<()> {
42            Ok(())
43        }
44
45        /// Delete the controller.
46        #[allow(dead_code)]
47        pub fn delete(&mut self) {}
48    }
49}
50
51#[cfg(target_os = "linux")]
52mod imp {
53    use std::collections::HashSet;
54
55    use anyhow::{Context, Result};
56    use cgroups_rs::{hierarchies, Cgroup, CgroupPid, Controller, Subsystem};
57
58    use super::{ENV_CGROUP_CPUSET, ENV_CGROUP_NAME};
59
60    /// try load cgroup from env
61    pub fn try_load_from_env() -> CtrlGroup {
62        use std::env::var;
63
64        match (var(ENV_CGROUP_NAME), var(ENV_CGROUP_CPUSET)) {
65            (Ok(cgname), Ok(cpuset)) => match CtrlGroup::new(&cgname, &cpuset) {
66                Ok(mut cg) => {
67                    let pid = std::process::id() as u64;
68                    match cg.add_task(pid.into()) {
69                        Ok(_) => {
70                            tracing::info!(pid = pid, group = cgname.as_str(), "add into cgroup");
71                        }
72                        Err(e) => {
73                            tracing::error!(pid = pid, group = cgname.as_str(), err=?e, "failed to add cgroup");
74                        }
75                    }
76                    cg
77                }
78                Err(e) => {
79                    tracing::error!(err=?e, "failed to load cgroup cpuset from env. cgname: {}, cpuset: {}", cgname.as_str(), cpuset.as_str());
80                    CtrlGroup::empty()
81                }
82            },
83            _ => CtrlGroup::empty(),
84        }
85    }
86
87    /// An Cgroup type
88    pub struct CtrlGroup {
89        subsystems: Vec<Subsystem>,
90    }
91
92    impl CtrlGroup {
93        /// Returns an empty `CtrlGroup`
94        pub fn empty() -> Self {
95            Self { subsystems: vec![] }
96        }
97
98        /// Creates a new CtrlGroup with cgroup name `cgname` and cpuset.
99        pub fn new(cgname: impl AsRef<str>, cpuset: impl AsRef<str>) -> Result<Self> {
100            let mut wanted = HashSet::new();
101
102            let cg = Cgroup::load(hierarchies::auto(), cgname.as_ref());
103            let cgsubs = cg.subsystems();
104
105            for (sidx, sub) in cgsubs.iter().enumerate() {
106                if let Subsystem::CpuSet(ctrl) = sub {
107                    ctrl.create();
108                    ctrl.set_cpus(cpuset.as_ref())
109                        .with_context(|| format!("set cpuset to {}", cpuset.as_ref()))?;
110                    wanted.insert(sidx);
111                    break;
112                }
113            }
114
115            let mut subsystems = Vec::with_capacity(cgsubs.len());
116            for (sidx, sub) in cgsubs.iter().enumerate() {
117                if wanted.contains(&sidx) {
118                    subsystems.push(sub.clone());
119                }
120            }
121
122            Ok(CtrlGroup { subsystems })
123        }
124
125        /// Attach a task to this controller.
126        pub fn add_task(&mut self, pid: CgroupPid) -> Result<()> {
127            for sub in self.subsystems.iter() {
128                sub.to_controller()
129                    .add_task(&pid)
130                    .with_context(|| format!("subsystem {}", sub.controller_name()))?;
131            }
132
133            Ok(())
134        }
135
136        /// Delete the controller.
137        pub fn delete(&mut self) {
138            for sub in self.subsystems.iter() {
139                if let Err(e) = sub.to_controller().delete() {
140                    tracing::warn!(system = sub.controller_name().as_str(), "subsystem delete failed: {:?}", e);
141                };
142            }
143        }
144    }
145
146    impl Drop for CtrlGroup {
147        fn drop(&mut self) {
148            self.delete();
149        }
150    }
151}
vc_processors/sys/cgroup.rs

vc_processors/sys/
cgroup.rs
