code-executor 2.7.0

use bon::Builder;
use byte_unit::Byte;
use cgroups_rs::fs::{
    Cgroup, MaxValue, cgroup_builder::CgroupBuilder, cpu::CpuController, hierarchies,
};
use uuid::Uuid;

const PREFIX: &str = "runner";
const CPU_USAGE_PREFIX: &str = "usage_usec ";

#[derive(Debug, Clone, Copy, Builder)]
pub struct ResourceConfig {
    pub memory_limit: Byte,

    #[builder(default = 100_000)]
    pub quota: u64,

    #[builder(default = 100_000)]
    pub period: u64,

    #[builder(default = 256)]
    pub process_count_limit: usize,
}

impl TryFrom<ResourceConfig> for Cgroup {
    type Error = cgroups_rs::fs::error::Error;

    fn try_from(config: ResourceConfig) -> Result<Self, Self::Error> {
        let cgroup_name = format!("{PREFIX}/{}", Uuid::new_v4());
        let hier = hierarchies::auto();

        let builder = CgroupBuilder::new(&cgroup_name);

        let memory_limit = config.memory_limit.as_u64() as i64;
        let builder = builder
            .memory()
            .memory_swap_limit(0)
            .memory_soft_limit(memory_limit)
            .memory_hard_limit(memory_limit)
            .done();

        let builder = builder
            .cpu()
            .quota(config.quota as i64)
            .period(config.period)
            .done();

        let process_count_limit = MaxValue::Value(config.process_count_limit as i64);
        let builder = builder
            .pid()
            .maximum_number_of_processes(process_count_limit)
            .done();

        builder.build(hier)
    }
}

pub(crate) fn get_cpu_usage(cgroup: &Cgroup) -> Option<u64> {
    let cpu_controller: &CpuController = cgroup.controller_of()?;
    let stats = cpu_controller.cpu().stat;

    stats
        .lines()
        .find_map(|line| line.strip_prefix(CPU_USAGE_PREFIX))
        .and_then(|x| x.parse().ok())
}