Feature/enhance cgroup

Cargo.toml

@@ -11,13 +11,14 @@ A library designed for the backend of competitive programming platforms
 """
 
 [dependencies]
-cached = "0.55.1"
-cgroups-rs = "0.3.4"
+byte-unit = "5.1.6"
+cgroups-rs = "0.4.0"
+nix = { version = "0.30.1", default-features = false, features = ["signal"] }
 thiserror = "2.0.12"
 tokio = { version = "1.45.0", features = ["full"] }
 tracing = "0.1.41"
 uuid = { version = "1.16.0", features = ["v4", "fast-rng"] }
 
 [dev-dependencies]
 bstr = "1.12.0"
-rstest = "0.25.0"
+rstest = "0.26.1"

examples/basic.rs

@@ -1,6 +1,7 @@
 use std::time::Duration;
 
-use code_executor::{CPP, Runner};
+use byte_unit::Byte;
+use code_executor::{CPP, ResourceConfig, Runner};
 
 #[tokio::main]
 async fn main() {
@@ -21,8 +22,10 @@ async fn main() {
         CPP.runner_args,
         &project_path,
         Duration::from_secs(2),
-        i64::MAX,
-        256,
+        ResourceConfig {
+            memory_limit: Byte::MEGABYTE,
+            process_count_limit: 1,
+        },
     )
     .unwrap();
     let metrics = runner.run(b"Hello").await.unwrap();

examples/custom-language.rs

@@ -1,6 +1,7 @@
 use std::time::Duration;
 
-use code_executor::{CommandArgs, Compiler, Language, Runner};
+use byte_unit::Byte;
+use code_executor::{CommandArgs, Compiler, Language, ResourceConfig, Runner};
 
 pub const CPP: Language = Language {
     compiler: Compiler {
@@ -35,8 +36,10 @@ async fn main() {
         CPP.runner_args,
         &project_path,
         Duration::from_secs(2),
-        i64::MAX,
-        512,
+        ResourceConfig {
+            memory_limit: Byte::GIGABYTE,
+            process_count_limit: 1,
+        },
     )
     .unwrap();
     let metrics = runner.run(b"Hello").await.unwrap();

src/cgroup.rs

@@ -0,0 +1,34 @@
+use byte_unit::Byte;
+use cgroups_rs::fs::{Cgroup, MaxValue, cgroup_builder::CgroupBuilder, hierarchies};
+use uuid::Uuid;
+
+const PREFIX: &str = "runner";
+
+#[derive(Debug)]
+pub struct ResourceConfig {
+    pub memory_limit: Byte,
+    pub process_count_limit: usize,
+}
+
+impl TryFrom<ResourceConfig> for Cgroup {
+    type Error = cgroups_rs::fs::error::Error;
+
+    fn try_from(config: ResourceConfig) -> Result<Self, Self::Error> {
+        let memory_limit = config.memory_limit.as_u64() as i64;
+        let process_count_limit = MaxValue::Value(config.process_count_limit as i64);
+
+        let cgroup_name = format!("{PREFIX}/{}", Uuid::new_v4());
+
+        let hier = hierarchies::auto();
+        CgroupBuilder::new(&cgroup_name)
+            .memory()
+            .memory_swap_limit(0)
+            .memory_soft_limit(memory_limit)
+            .memory_hard_limit(memory_limit)
+            .done()
+            .pid()
+            .maximum_number_of_processes(process_count_limit)
+            .done()
+            .build(hier)
+    }
+}

src/error.rs

@@ -9,7 +9,7 @@ pub enum Error {
     IO(#[from] std::io::Error),
 
     #[error("Failed to create cgroup: {0}")]
-    Cgroup(#[from] cgroups_rs::error::Error),
+    Cgroup(#[from] cgroups_rs::fs::error::Error),
 }
 
 impl From<std::string::FromUtf8Error> for Error {

src/lib.rs

@@ -1,10 +1,12 @@
+mod cgroup;
 mod compiler;
 mod error;
 mod language;
 mod metrics;
 mod runner;
 mod util;
 
+pub use cgroup::*;
 pub use compiler::*;
 pub use error::*;
 pub use language::*;

src/metrics.rs

@@ -1,18 +1,24 @@
-use std::time::Duration;
+use std::{os::unix::process::ExitStatusExt, time::Duration};
+
+use nix::sys::signal::Signal;
 
 #[derive(Debug, PartialEq, Eq)]
 pub enum ExitStatus {
     Success,
     RuntimeError,
-    Timeout,
+    TimeLimitExceeded,
+    ResourceLimitExceeded,
 }
 
 impl From<std::process::ExitStatus> for ExitStatus {
     fn from(raw: std::process::ExitStatus) -> Self {
         if raw.success() {
-            Self::Success
-        } else {
-            Self::RuntimeError
+            return Self::Success;
+        }
+
+        match raw.signal().and_then(|raw| Signal::try_from(raw).ok()) {
+            Some(Signal::SIGKILL) => ExitStatus::ResourceLimitExceeded,
+            _ => Self::RuntimeError,
         }
     }
 }

src/runner.rs

@@ -1,45 +1,20 @@
-use std::{
-    io,
-    path::Path,
-    process::{self, Stdio},
-    sync::Arc,
-    time::Duration,
-};
+use std::{io, path::Path, process::Stdio, time::Duration};
 
-use cached::proc_macro::cached;
-use cgroups_rs::{Cgroup, CgroupPid, cgroup_builder::CgroupBuilder, hierarchies};
+use cgroups_rs::{CgroupPid, fs::Cgroup};
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     process::Command,
     time::{Instant, sleep},
 };
 
-use crate::{CommandArgs, ExitStatus, Result, metrics::Metrics};
-
-#[cached(result = true)]
-fn create_cgroup(memory_limit: i64, process_count_limit: usize) -> Result<Cgroup> {
-    let cgroup_name = format!("runner/{}-{}", memory_limit, process_count_limit);
-    let hier = hierarchies::auto();
-    let cgroup = CgroupBuilder::new(&cgroup_name)
-        .memory()
-        .memory_swap_limit(memory_limit)
-        .memory_soft_limit(memory_limit)
-        .memory_hard_limit(memory_limit)
-        .done()
-        .pid()
-        .maximum_number_of_processes(cgroups_rs::MaxValue::Value(process_count_limit as i64))
-        .done()
-        .build(hier)?;
-
-    Ok(cgroup)
-}
+use crate::{CommandArgs, ExitStatus, ResourceConfig, Result, metrics::Metrics};
 
 #[derive(Debug)]
 pub struct Runner<'a> {
     pub args: CommandArgs<'a>,
     pub project_path: &'a Path,
     pub time_limit: Duration,
-    pub cgroup: Arc<Cgroup>,
+    pub cgroup: Cgroup,
 }
 
 impl<'a> Runner<'a> {
@@ -48,15 +23,14 @@ impl<'a> Runner<'a> {
         args: CommandArgs<'a>,
         project_path: &'a Path,
         time_limit: Duration,
-        memory_limit: i64,
-        process_count_limit: usize,
+        config: ResourceConfig,
     ) -> Result<Self> {
-        let cgroup = create_cgroup(memory_limit, process_count_limit)?;
+        let cgroup = config.try_into()?;
 
         Ok(Self {
             args,
             project_path,
-            cgroup: Arc::new(cgroup),
+            cgroup,
             time_limit,
         })
     }
@@ -65,24 +39,19 @@ impl<'a> Runner<'a> {
     pub async fn run(&self, input: &[u8]) -> Result<Metrics> {
         let CommandArgs { binary, args } = self.args;
 
-        let cgroup = self.cgroup.clone();
+        let start = Instant::now();
 
-        let mut child = Command::new(binary);
-        let child = child
+        let mut child = Command::new(binary)
             .current_dir(self.project_path)
             .args(args)
             .stdin(Stdio::piped())
             .stdout(Stdio::piped())
-            .stderr(Stdio::piped());
-        let child = unsafe {
-            child.pre_exec(move || {
-                cgroup
-                    .add_task_by_tgid(CgroupPid::from(process::id() as u64))
-                    .map_err(std::io::Error::other)
-            })
-        };
-        let start = Instant::now();
-        let mut child = child.spawn()?;
+            .stderr(Stdio::piped())
+            .spawn()?;
+
+        self.cgroup
+            .add_task_by_tgid(CgroupPid::from(child.id().unwrap() as u64))?;
+
         let mut stdin = child.stdin.take().unwrap();
         let mut stdout = child.stdout.take().unwrap();
         let mut stderr = child.stderr.take().unwrap();
@@ -113,7 +82,7 @@ impl<'a> Runner<'a> {
                 child.kill().await?;
                 child.wait().await?;
 
-                Ok(ExitStatus::Timeout)
+                Ok(ExitStatus::TimeLimitExceeded)
             }
         }?;
 
@@ -128,6 +97,12 @@ impl<'a> Runner<'a> {
     }
 }
 
+impl<'a> Drop for Runner<'a> {
+    fn drop(&mut self) {
+        let _ = self.cgroup.delete();
+    }
+}
+
 #[cfg(test)]
 mod test {
     use std::{
@@ -136,18 +111,22 @@ mod test {
     };
 
     use bstr::ByteSlice;
+    use byte_unit::Byte;
     use rstest::rstest;
 
     use crate::{
-        CPP, ExitStatus, JAVA, Language, PYTHON, RUST, Runner,
+        CPP, ExitStatus, JAVA, Language, PYTHON, RUST, ResourceConfig, Runner,
         test::{read_code, read_test_cases},
     };
 
     #[rstest]
     #[tokio::test]
     async fn should_output_correct(
         #[values(CPP, RUST, JAVA, PYTHON)] language: Language<'static>,
-        #[files("tests/data/problem/*")] problem_path: PathBuf,
+
+        #[dirs]
+        #[files("tests/data/problem/*")]
+        problem_path: PathBuf,
     ) {
         let test_cases = read_test_cases(&problem_path);
 
@@ -158,8 +137,10 @@ mod test {
             language.runner_args,
             &project_path,
             Duration::from_secs(2),
-            i64::MAX,
-            512,
+            ResourceConfig {
+                memory_limit: Byte::GIGABYTE,
+                process_count_limit: 256,
+            },
         )
         .unwrap();
         for (input, output) in test_cases {
@@ -180,13 +161,15 @@ mod test {
             language.runner_args,
             &project_path,
             Duration::from_secs(2),
-            i64::MAX,
-            512,
+            ResourceConfig {
+                memory_limit: Byte::GIGABYTE,
+                process_count_limit: 256,
+            },
         )
         .unwrap();
 
         let metrics = runner.run(b"").await.unwrap();
 
-        assert_eq!(metrics.exit_status, ExitStatus::Timeout)
+        assert_eq!(metrics.exit_status, ExitStatus::TimeLimitExceeded)
     }
 }