feature: cpu limit

Cargo.toml

@@ -11,6 +11,7 @@ A library designed for the backend of competitive programming platforms
 """
 
 [dependencies]
+bon = "3.7.2"
 byte-unit = "5.1.6"
 cgroups-rs = "0.4.0"
 nix = { version = "0.30.1", default-features = false, features = ["signal"] }

examples/basic.rs

@@ -22,10 +22,9 @@ async fn main() {
         CPP.runner_args,
         &project_path,
         Duration::from_secs(2),
-        ResourceConfig {
-            memory_limit: Byte::MEGABYTE,
-            process_count_limit: 1,
-        },
+        ResourceConfig::builder()
+            .memory_limit(Byte::GIBIBYTE)
+            .build(),
     )
     .unwrap();
     let metrics = runner.run(b"Hello").await.unwrap();

examples/custom-language.rs

@@ -36,10 +36,9 @@ async fn main() {
         CPP.runner_args,
         &project_path,
         Duration::from_secs(2),
-        ResourceConfig {
-            memory_limit: Byte::GIGABYTE,
-            process_count_limit: 1,
-        },
+        ResourceConfig::builder()
+            .memory_limit(Byte::GIBIBYTE)
+            .build(),
     )
     .unwrap();
     let metrics = runner.run(b"Hello").await.unwrap();

src/cgroup.rs

@@ -1,34 +1,66 @@
+use bon::Builder;
 use byte_unit::Byte;
-use cgroups_rs::fs::{Cgroup, MaxValue, cgroup_builder::CgroupBuilder, hierarchies};
+use cgroups_rs::fs::{
+    Cgroup, MaxValue, cgroup_builder::CgroupBuilder, cpu::CpuController, hierarchies,
+};
 use uuid::Uuid;
 
 const PREFIX: &str = "runner";
+const CPU_USAGE_PREFIX: &str = "usage_usec ";
 
-#[derive(Debug)]
+#[derive(Debug, Clone, Copy, Builder)]
 pub struct ResourceConfig {
     pub memory_limit: Byte,
+
+    #[builder(default = 100_000)]
+    pub quota: u64,
+
+    #[builder(default = 100_000)]
+    pub period: u64,
+
+    #[builder(default = 256)]
     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)
+
+        let builder = CgroupBuilder::new(&cgroup_name);
+
+        let memory_limit = config.memory_limit.as_u64() as i64;
+        let builder = builder
             .memory()
             .memory_swap_limit(0)
             .memory_soft_limit(memory_limit)
             .memory_hard_limit(memory_limit)
-            .done()
+            .done();
+
+        let builder = builder
+            .cpu()
+            .quota(config.quota as i64)
+            .period(config.period)
+            .done();
+
+        let process_count_limit = MaxValue::Value(config.process_count_limit as i64);
+        let builder = builder
             .pid()
             .maximum_number_of_processes(process_count_limit)
-            .done()
-            .build(hier)
+            .done();
+
+        builder.build(hier)
     }
 }
+
+pub(crate) fn get_cpu_usage(cgroup: &Cgroup) -> Option<u64> {
+    let cpu_controller: &CpuController = cgroup.controller_of()?;
+    let stats = cpu_controller.cpu().stat;
+
+    stats
+        .lines()
+        .find_map(|line| line.strip_prefix(CPU_USAGE_PREFIX))
+        .and_then(|x| x.parse().ok())
+}

src/runner.rs

@@ -1,19 +1,25 @@
 use std::{io, path::Path, process::Stdio, time::Duration};
 
 use cgroups_rs::{CgroupPid, fs::Cgroup};
+use nix::{
+    sys::signal::{self, Signal},
+    unistd::Pid,
+};
 use tokio::{
     io::{AsyncReadExt, AsyncWriteExt},
     process::Command,
+    task::JoinHandle,
     time::{Instant, sleep},
 };
 
-use crate::{CommandArgs, ExitStatus, ResourceConfig, Result, metrics::Metrics};
+use crate::{CommandArgs, ExitStatus, ResourceConfig, Result, cgroup, metrics::Metrics};
 
 #[derive(Debug)]
 pub struct Runner<'a> {
     pub args: CommandArgs<'a>,
     pub project_path: &'a Path,
     pub time_limit: Duration,
+    pub resource_config: ResourceConfig,
     pub cgroup: Cgroup,
 }
 
@@ -23,15 +29,49 @@ impl<'a> Runner<'a> {
         args: CommandArgs<'a>,
         project_path: &'a Path,
         time_limit: Duration,
-        config: ResourceConfig,
+        resource_config: ResourceConfig,
     ) -> Result<Self> {
-        let cgroup = config.try_into()?;
+        let cgroup = resource_config.try_into()?;
 
         Ok(Self {
             args,
             project_path,
-            cgroup,
             time_limit,
+            resource_config,
+            cgroup,
+        })
+    }
+
+    fn new_cpu_limiter(&self, id: i32) -> JoinHandle<()> {
+        let cg = self.cgroup.clone();
+        let quota = self.resource_config.quota;
+        let period = Duration::from_millis(self.resource_config.period);
+
+        tokio::spawn(async move {
+            let Some(mut prev_usage) = cgroup::get_cpu_usage(&cg) else {
+                return;
+            };
+
+            loop {
+                sleep(period).await;
+
+                let Some(cur_usage) = cgroup::get_cpu_usage(&cg) else {
+                    break;
+                };
+                if cur_usage - prev_usage <= quota {
+                    prev_usage = cur_usage;
+                    continue;
+                }
+
+                signal::kill(Pid::from_raw(id), Signal::SIGSTOP).unwrap();
+                sleep(period).await;
+                signal::kill(Pid::from_raw(id), Signal::SIGCONT).unwrap();
+
+                match cgroup::get_cpu_usage(&cg) {
+                    Some(v) => prev_usage = v,
+                    None => break,
+                }
+            }
         })
     }
 
@@ -49,8 +89,9 @@ impl<'a> Runner<'a> {
             .stderr(Stdio::piped())
             .spawn()?;
 
-        self.cgroup
-            .add_task_by_tgid(CgroupPid::from(child.id().unwrap() as u64))?;
+        let id = child.id().unwrap();
+        self.cgroup.add_task_by_tgid(CgroupPid::from(id as u64))?;
+        let cpu_limiter = self.new_cpu_limiter(id as i32);
 
         let mut stdin = child.stdin.take().unwrap();
         let mut stdout = child.stdout.take().unwrap();
@@ -86,6 +127,8 @@ impl<'a> Runner<'a> {
             }
         }?;
 
+        cpu_limiter.abort();
+
         let (stdout, stderr) = tokio::try_join!(stdout_observer, stderr_observer)?;
 
         Ok(Metrics {
@@ -137,10 +180,9 @@ mod test {
             language.runner_args,
             &project_path,
             Duration::from_secs(2),
-            ResourceConfig {
-                memory_limit: Byte::GIGABYTE,
-                process_count_limit: 256,
-            },
+            ResourceConfig::builder()
+                .memory_limit(Byte::GIBIBYTE)
+                .build(),
         )
         .unwrap();
         for (input, output) in test_cases {
@@ -161,10 +203,9 @@ mod test {
             language.runner_args,
             &project_path,
             Duration::from_secs(2),
-            ResourceConfig {
-                memory_limit: Byte::GIGABYTE,
-                process_count_limit: 256,
-            },
+            ResourceConfig::builder()
+                .memory_limit(Byte::GIBIBYTE)
+                .build(),
         )
         .unwrap();