code_executor/sandbox/
mod.rs

1mod cgroup;
2mod resource;
3
4use std::{io, os::unix::process::ExitStatusExt, process, time::Duration};
5
6use tokio::{
7    process::{Child, Command},
8    time::{Instant, interval},
9};
10
11use byte_unit::Byte;
12use cgroups_rs::{
13    CgroupPid,
14    fs::{Cgroup, cpu::CpuController, memory::MemController},
15};
16pub use resource::Resource;
17
18use crate::{
19    Verdict,
20    sandbox::cgroup::{CpuControllerExt, MemControllerExt},
21};
22
23// TODO: need further tuning
24const POLL: Duration = Duration::from_millis(10);
25const MIN_CPU_USAGE_PER_POLL: Duration = Duration::from_millis(1);
26const IDLE_TIME_LIMIT: Duration = Duration::from_millis(100);
27
28pub struct Sandbox {
29    pub cgroup: Cgroup,
30    pub cpu_usage_limit: Duration,
31    pub wall_time_limit: Duration,
32}
33
34impl Sandbox {
35    pub fn new(resource: Resource, time_limit: Duration) -> io::Result<Sandbox> {
36        Ok(Sandbox {
37            cgroup: resource.try_into()?,
38            cpu_usage_limit: time_limit,
39            wall_time_limit: Duration::max(time_limit * 2, time_limit + Duration::from_secs(2)),
40        })
41    }
42
43    pub fn spawn(&self, mut command: Command) -> io::Result<Child> {
44        let cgroup = self.cgroup.clone();
45
46        unsafe {
47            command
48                .pre_exec(move || {
49                    let id = process::id();
50
51                    cgroup
52                        .add_task_by_tgid(CgroupPid::from(id as u64))
53                        .map_err(io::Error::other)
54                })
55                .spawn()
56        }
57    }
58
59    pub async fn monitor(&self, mut child: Child) -> io::Result<(Option<Verdict>, Duration, Byte)> {
60        let Some(id) = child.id() else {
61            return Err(io::Error::other("Child exited"));
62        };
63        self.cgroup
64            .add_task_by_tgid(CgroupPid::from(id as u64))
65            .map_err(io::Error::other)?;
66        let cpu: &CpuController = self
67            .cgroup
68            .controller_of()
69            .ok_or(io::Error::other("Missing cpu controller"))?;
70        let memory: &MemController = self
71            .cgroup
72            .controller_of()
73            .ok_or(io::Error::other("Missing memory controller"))?;
74
75        let start = Instant::now();
76        let mut memory_usage = Byte::default();
77        let mut prev_cpu_usage = cpu.usage();
78        let mut idle_start: Option<Instant> = None;
79
80        let mut interval = interval(POLL);
81
82        while child.try_wait()?.is_none() {
83            let cpu_usage = cpu.usage();
84            memory_usage = memory_usage.max(memory.usage());
85
86            if cpu_usage.abs_diff(prev_cpu_usage) <= MIN_CPU_USAGE_PER_POLL {
87                match idle_start {
88                    Some(idle_start) => {
89                        if idle_start.elapsed() >= IDLE_TIME_LIMIT {
90                            return Ok((
91                                Some(Verdict::IdleTimeLimitExceeded),
92                                cpu_usage,
93                                memory_usage,
94                            ));
95                        }
96                    }
97                    None => idle_start = Some(Instant::now()),
98                }
99            } else {
100                idle_start = None;
101            }
102
103            if cpu_usage >= self.cpu_usage_limit || start.elapsed() >= self.wall_time_limit {
104                return Ok((
105                    Some(Verdict::TimeLimitExceeded),
106                    self.cpu_usage_limit,
107                    memory_usage,
108                ));
109            }
110
111            prev_cpu_usage = cpu_usage;
112
113            interval.tick().await;
114        }
115
116        let status = child.try_wait()?.unwrap();
117        if status.success() {
118            return Ok((None, prev_cpu_usage, memory_usage));
119        }
120        match status.signal() {
121            // SIGKILL
122            Some(9) => Ok((
123                Some(Verdict::MemoryLimitExceeded),
124                prev_cpu_usage,
125                memory.limit(),
126            )),
127            _ => Ok((Some(Verdict::RuntimeError), prev_cpu_usage, memory_usage)),
128        }
129    }
130}
131
132impl Drop for Sandbox {
133    fn drop(&mut self) {
134        let _ = self.cgroup.kill();
135        let _ = self.cgroup.delete();
136    }
137}
code_executor/sandbox/mod.rs

code_executor/sandbox/
mod.rs
