code-executor 1.1.0

A code runner library for online judge system
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

mod cgroup;
mod resource;

use std::{io, os::unix::process::ExitStatusExt, process, time::Duration};

use tokio::{
    process::{Child, Command},
    time::{Instant, interval},
};

use byte_unit::Byte;
use cgroups_rs::{
    CgroupPid,
    fs::{Cgroup, cpu::CpuController, memory::MemController},
};
pub use resource::Resource;

use crate::{
    Verdict,
    sandbox::cgroup::{CpuControllerExt, MemControllerExt},
};

// TODO: need further tuning
const POLL: Duration = Duration::from_millis(10);
const MIN_CPU_USAGE_PER_POLL: Duration = Duration::from_millis(1);

pub struct Sandbox {
    pub cgroup: Cgroup,
    pub cpu_usage_limit: Duration,
    pub wall_time_limit: Duration,
    pub idle_time_limit: Duration,
}

impl Sandbox {
    pub fn new(
        resource: Resource,
        time_limit: Duration,
        idle_time_limit: Duration,
    ) -> io::Result<Sandbox> {
        Ok(Sandbox {
            cgroup: resource.try_into()?,
            cpu_usage_limit: time_limit,
            wall_time_limit: Duration::max(time_limit * 2, time_limit + Duration::from_secs(2)),
            idle_time_limit,
        })
    }

    pub fn spawn(&self, mut command: Command) -> io::Result<Child> {
        let cgroup = self.cgroup.clone();

        unsafe {
            command
                .pre_exec(move || {
                    let id = process::id();

                    cgroup
                        .add_task_by_tgid(CgroupPid::from(id as u64))
                        .map_err(io::Error::other)
                })
                .spawn()
        }
    }

    pub async fn monitor(&self, mut child: Child) -> io::Result<(Option<Verdict>, Duration, Byte)> {
        let Some(id) = child.id() else {
            return Err(io::Error::other("Child exited"));
        };
        self.cgroup
            .add_task_by_tgid(CgroupPid::from(id as u64))
            .map_err(io::Error::other)?;
        let cpu: &CpuController = self
            .cgroup
            .controller_of()
            .ok_or(io::Error::other("Missing cpu controller"))?;
        let memory: &MemController = self
            .cgroup
            .controller_of()
            .ok_or(io::Error::other("Missing memory controller"))?;

        let start = Instant::now();
        let mut memory_usage = Byte::default();
        let mut prev_cpu_usage = cpu.usage();
        let mut idle_start: Option<Instant> = None;

        let mut interval = interval(POLL);

        while child.try_wait()?.is_none() {
            let cpu_usage = cpu.usage();
            memory_usage = memory_usage.max(memory.usage());

            if cpu_usage.abs_diff(prev_cpu_usage) <= MIN_CPU_USAGE_PER_POLL {
                match idle_start {
                    Some(idle_start) => {
                        if idle_start.elapsed() >= self.idle_time_limit {
                            return Ok((
                                Some(Verdict::IdleTimeLimitExceeded),
                                cpu_usage,
                                memory_usage,
                            ));
                        }
                    }
                    None => idle_start = Some(Instant::now()),
                }
            } else {
                idle_start = None;
            }

            if cpu_usage >= self.cpu_usage_limit || start.elapsed() >= self.wall_time_limit {
                return Ok((
                    Some(Verdict::TimeLimitExceeded),
                    self.cpu_usage_limit,
                    memory_usage,
                ));
            }

            prev_cpu_usage = cpu_usage;

            interval.tick().await;
        }

        let status = child.try_wait()?.unwrap();
        if status.success() {
            return Ok((None, prev_cpu_usage, memory_usage));
        }
        match status.signal() {
            // SIGKILL
            Some(9) => Ok((
                Some(Verdict::MemoryLimitExceeded),
                prev_cpu_usage,
                memory.limit(),
            )),
            _ => Ok((Some(Verdict::RuntimeError), prev_cpu_usage, memory_usage)),
        }
    }
}

impl Drop for Sandbox {
    fn drop(&mut self) {
        let _ = self.cgroup.kill();
        let _ = self.cgroup.delete();
    }
}