perf(sandbox): ext for controllers instead of getting controller everytime

Author: akagiyuu

src/language/mod.rs

@@ -26,7 +26,6 @@ impl Language {
         let args = self.compile_args?;
         let args = args.replace("{main}", main);
         let mut args = args.split_whitespace();
-        // SAFETY: there is always at least 1 element
         let binary = args.next().unwrap();
 
         let mut command = Command::new(binary);
@@ -37,7 +36,6 @@ impl Language {
     pub fn get_run_command(&self, main: &str) -> Command {
         let args = self.run_args.replace("{main}", main);
         let mut args = args.split_whitespace();
-        // SAFETY: there is always at least 1 element
         let binary = args.next().unwrap();
 
         let mut command = Command::new(binary);

src/sandbox/cgroup.rs

@@ -1,43 +1,41 @@
 use std::time::Duration;
 
 use byte_unit::Byte;
-use cgroups_rs::fs::{Cgroup, cpu::CpuController, memory::MemController};
+use cgroups_rs::fs::{cpu::CpuController, memory::MemController};
 
 const CPU_USAGE_PREFIX: &str = "usage_usec ";
 
-pub trait CgroupExt {
-    fn get_cpu_time(&self) -> Duration;
-    fn get_memory_usage(&self) -> Byte;
-    fn get_memory_limit(&self) -> Byte;
+pub trait CpuControllerExt {
+    fn usage(&self) -> Duration;
 }
 
-impl CgroupExt for Cgroup {
-    fn get_cpu_time(&self) -> Duration {
-        let cpu_controller: &CpuController = self.controller_of().unwrap();
-        let stats = cpu_controller.cpu().stat;
+impl CpuControllerExt for CpuController {
+    fn usage(&self) -> Duration {
+        let stats = self.cpu().stat;
 
-        // SAFETY: there must be cpu usage for valid cgroup
         let usage = stats
             .lines()
             .find_map(|line| line.strip_prefix(CPU_USAGE_PREFIX))
             .unwrap();
-        // SAFETY: cpu usage must be duration in microsecond
         let usage = usage.parse().unwrap();
         Duration::from_micros(usage)
     }
+}
+
+pub trait MemControllerExt {
+    fn usage(&self) -> Byte;
+    fn limit(&self) -> Byte;
+}
 
-    fn get_memory_usage(&self) -> Byte {
-        // SAFETY: there must be memory controller for cgroup v2
-        let memory_controller: &MemController = self.controller_of().unwrap();
-        let stats = memory_controller.memory_stat();
+impl MemControllerExt for MemController {
+    fn usage(&self) -> Byte {
+        let stats = self.memory_stat();
 
         Byte::from_u64(stats.usage_in_bytes)
     }
 
-    fn get_memory_limit(&self) -> Byte {
-        // SAFETY: there must be memory controller for cgroup v2
-        let memory_controller: &MemController = self.controller_of().unwrap();
-        let stats = memory_controller.memory_stat();
+    fn limit(&self) -> Byte {
+        let stats = self.memory_stat();
 
         Byte::from_u64(stats.limit_in_bytes.max(0) as u64)
     }

src/sandbox/mod.rs

@@ -10,28 +10,34 @@ use std::{
 };
 
 use byte_unit::Byte;
-use cgroups_rs::{CgroupPid, fs::Cgroup};
+use cgroups_rs::{
+    CgroupPid,
+    fs::{Cgroup, cpu::CpuController, memory::MemController},
+};
 use nix::{libc::getpid, sys::signal::Signal};
 pub use resource::Resource;
 
-use crate::{Verdict, sandbox::cgroup::CgroupExt};
+use crate::{
+    Verdict,
+    sandbox::cgroup::{CpuControllerExt, MemControllerExt},
+};
 
 // TODO: need further tuning
 const POLL: Duration = Duration::from_millis(10);
-const MIN_CPU_TIME_PER_POLL: Duration = Duration::from_millis(1);
+const MIN_CPU_USAGE_PER_POLL: Duration = Duration::from_millis(1);
 const IDLE_TIME_LIMIT: Duration = Duration::from_millis(100);
 
 pub struct Sandbox {
     pub cgroup: Cgroup,
-    pub cpu_time_limit: Duration,
+    pub cpu_usage_limit: Duration,
     pub wall_time_limit: Duration,
 }
 
 impl Sandbox {
     pub fn new(resource: Resource, time_limit: Duration) -> io::Result<Sandbox> {
         Ok(Sandbox {
             cgroup: resource.try_into()?,
-            cpu_time_limit: time_limit,
+            cpu_usage_limit: time_limit,
             wall_time_limit: Duration::max(time_limit * 2, time_limit + Duration::from_secs(2)),
         })
     }
@@ -56,23 +62,31 @@ impl Sandbox {
         self.cgroup
             .add_task_by_tgid(CgroupPid::from(child.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_time = self.cgroup.get_cpu_time();
+        let mut prev_cpu_usage = cpu.usage();
         let mut idle_start: Option<Instant> = None;
 
         while child.try_wait()?.is_none() {
-            let cpu_time = self.cgroup.get_cpu_time();
-            memory_usage = memory_usage.max(self.cgroup.get_memory_usage());
+            let cpu_usage = cpu.usage();
+            memory_usage = memory_usage.max(memory.usage());
 
-            if cpu_time.abs_diff(prev_cpu_time) <= MIN_CPU_TIME_PER_POLL {
+            if cpu_usage.abs_diff(prev_cpu_usage) <= MIN_CPU_USAGE_PER_POLL {
                 match idle_start {
                     Some(idle_start) => {
                         if idle_start.elapsed() >= IDLE_TIME_LIMIT {
                             return Ok((
                                 Some(Verdict::IdleTimeLimitExceeded),
-                                cpu_time,
+                                cpu_usage,
                                 memory_usage,
                             ));
                         }
@@ -83,41 +97,37 @@ impl Sandbox {
                 idle_start = None;
             }
 
-            if cpu_time >= self.cpu_time_limit || start.elapsed() >= self.wall_time_limit {
+            if cpu_usage >= self.cpu_usage_limit || start.elapsed() >= self.wall_time_limit {
                 return Ok((
                     Some(Verdict::TimeLimitExceeded),
-                    self.cpu_time_limit,
+                    self.cpu_usage_limit,
                     memory_usage,
                 ));
             }
 
-            prev_cpu_time = cpu_time;
+            prev_cpu_usage = cpu_usage;
 
             sleep(POLL);
         }
 
-        // SAFETY: child must be finished at this point to exit the previous loop
         let status = child.try_wait()?.unwrap();
         if status.success() {
-            return Ok((None, prev_cpu_time, memory_usage));
+            return Ok((None, prev_cpu_usage, memory_usage));
         }
         match status.signal().and_then(|x| Signal::try_from(x).ok()) {
             Some(Signal::SIGKILL) => Ok((
                 Some(Verdict::MemoryLimitExceeded),
-                prev_cpu_time,
-                self.cgroup.get_memory_limit(),
+                prev_cpu_usage,
+                memory.limit(),
             )),
-            _ => Ok((Some(Verdict::RuntimeError), prev_cpu_time, memory_usage)),
+            _ => Ok((Some(Verdict::RuntimeError), prev_cpu_usage, memory_usage)),
         }
     }
 }
 
 impl Drop for Sandbox {
     fn drop(&mut self) {
-        // SAFETY: always be used with stable version of linux kernel
         let _ = self.cgroup.kill();
-
-        // SAFETY: no descendant is created previously by judge
         let _ = self.cgroup.delete();
     }
 }