code_executor/
cgroup.rs

1use bon::Builder;
2use byte_unit::Byte;
3use cgroups_rs::fs::{
4    Cgroup, MaxValue, cgroup_builder::CgroupBuilder, cpu::CpuController, hierarchies,
5};
6use uuid::Uuid;
7
8const PREFIX: &str = "runner";
9const CPU_USAGE_PREFIX: &str = "usage_usec ";
10
11#[derive(Debug, Clone, Copy, Builder)]
12pub struct ResourceConfig {
13    pub memory_limit: Byte,
14
15    #[builder(default = 100_000)]
16    pub quota: u64,
17
18    #[builder(default = 100_000)]
19    pub period: u64,
20
21    #[builder(default = 256)]
22    pub process_count_limit: usize,
23}
24
25impl TryFrom<ResourceConfig> for Cgroup {
26    type Error = cgroups_rs::fs::error::Error;
27
28    fn try_from(config: ResourceConfig) -> Result<Self, Self::Error> {
29        let cgroup_name = format!("{PREFIX}/{}", Uuid::new_v4());
30        let hier = hierarchies::auto();
31
32        let builder = CgroupBuilder::new(&cgroup_name);
33
34        let memory_limit = config.memory_limit.as_u64() as i64;
35        let builder = builder
36            .memory()
37            .memory_swap_limit(0)
38            .memory_soft_limit(memory_limit)
39            .memory_hard_limit(memory_limit)
40            .done();
41
42        let builder = builder
43            .cpu()
44            .quota(config.quota as i64)
45            .period(config.period)
46            .done();
47
48        let process_count_limit = MaxValue::Value(config.process_count_limit as i64);
49        let builder = builder
50            .pid()
51            .maximum_number_of_processes(process_count_limit)
52            .done();
53
54        builder.build(hier)
55    }
56}
57
58pub(crate) fn get_cpu_usage(cgroup: &Cgroup) -> Option<u64> {
59    let cpu_controller: &CpuController = cgroup.controller_of()?;
60    let stats = cpu_controller.cpu().stat;
61
62    stats
63        .lines()
64        .find_map(|line| line.strip_prefix(CPU_USAGE_PREFIX))
65        .and_then(|x| x.parse().ok())
66}
code_executor/cgroup.rs

code_executor/
cgroup.rs
