Skip to main content

code_executor/sandbox/
resource.rs

1use std::{io, time::Duration};
2
3use byte_unit::Byte;
4use cgroups_rs::fs::{Cgroup, cgroup_builder::CgroupBuilder, hierarchies};
5use uuid::Uuid;
6
7const PREFIX: &str = "judge";
8
9#[derive(Clone, Copy, Hash)]
10pub struct Resource {
11    pub memory: Byte,
12    pub cpu_quota: Duration,
13    pub cpu_period: Duration,
14}
15
16impl Default for Resource {
17    fn default() -> Self {
18        Resource {
19            memory: Byte::GIGABYTE,
20            cpu_quota: Duration::from_millis(100),
21            cpu_period: Duration::from_millis(100),
22        }
23    }
24}
25
26impl TryFrom<Resource> for Cgroup {
27    type Error = io::Error;
28
29    fn try_from(resource: Resource) -> Result<Self, Self::Error> {
30        let builder = CgroupBuilder::new(&format!("{}/{}", PREFIX, Uuid::new_v4()));
31
32        let memory = resource.memory.as_u64() as i64;
33        let builder = builder
34            .memory()
35            .memory_swap_limit(0)
36            .memory_soft_limit(memory)
37            .memory_hard_limit(memory)
38            .done();
39
40        let quota = resource.cpu_quota.as_micros() as i64;
41        let period = resource.cpu_period.as_micros() as u64;
42        let builder = builder.cpu().quota(quota).period(period).done();
43
44        let cgroup = builder
45            .build(hierarchies::auto())
46            .map_err(io::Error::other)?;
47        Ok(cgroup)
48    }
49}