code_executor/sandbox/
resource.rs1use 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}