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

code_executor/sandbox/
mod.rs
