diff --git a/Cargo.lock b/Cargo.lock index c5912df..c0499b6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -244,6 +244,31 @@ dependencies = [ "piper", ] +[[package]] +name = "bon" +version = "3.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2529c31017402be841eb45892278a6c21a000c0a17643af326c73a73f83f0fb" +dependencies = [ + "bon-macros", + "rustversion", +] + +[[package]] +name = "bon-macros" +version = "3.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d82020dadcb845a345591863adb65d74fa8dc5c18a0b6d408470e13b7adc7005" +dependencies = [ + "darling", + "ident_case", + "prettyplease", + "proc-macro2", + "quote", + "rustversion", + "syn 2.0.93", +] + [[package]] name = "borsh" version = "1.5.7" @@ -347,6 +372,7 @@ dependencies = [ name = "code-executor" version = "2.6.1" dependencies = [ + "bon", "bstr", "byte-unit", "cgroups-rs", @@ -373,6 +399,41 @@ version = "0.8.21" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" +[[package]] +name = "darling" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9cdf337090841a411e2a7f3deb9187445851f91b309c0c0a29e05f74a00a48c0" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1247195ecd7e3c85f83c8d2a366e4210d588e802133e1e355180a9870b517ea4" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn 2.0.93", +] + +[[package]] +name = "darling_macro" +version = "0.21.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d38308df82d1080de0afee5d069fa14b0326a88c14f15c5ccda35b4a6c414c81" +dependencies = [ + "darling_core", + "quote", + "syn 2.0.93", +] + [[package]] name = "endi" version = "1.1.0" @@ -443,6 +504,12 @@ version = "2.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + [[package]] name = "funty" version = "2.0.0" @@ -573,6 +640,12 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + [[package]] name = "indexmap" version = "2.7.0" @@ -786,6 +859,16 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "prettyplease" +version = "0.2.25" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64d1ec885c64d0457d564db4ec299b2dae3f9c02808b8ad9c3a089c591b18033" +dependencies = [ + "proc-macro2", + "syn 2.0.93", +] + [[package]] name = "proc-macro-crate" version = "3.2.0" @@ -1059,6 +1142,12 @@ dependencies = [ "windows-sys 0.60.2", ] +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + [[package]] name = "ryu" version = "1.0.20" @@ -1172,6 +1261,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + [[package]] name = "syn" version = "1.0.109" diff --git a/Cargo.toml b/Cargo.toml index 32a9122..eb9df20 100644 --- a/Cargo.toml +++ b/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"] } diff --git a/examples/basic.rs b/examples/basic.rs index 9cb3ba5..f90634c 100644 --- a/examples/basic.rs +++ b/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(); diff --git a/examples/custom-language.rs b/examples/custom-language.rs index df3210b..b611b48 100644 --- a/examples/custom-language.rs +++ b/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(); diff --git a/src/cgroup.rs b/src/cgroup.rs index 7c47260..c81da42 100644 --- a/src/cgroup.rs +++ b/src/cgroup.rs @@ -1,12 +1,24 @@ +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, } @@ -14,21 +26,41 @@ impl TryFrom for Cgroup { type Error = cgroups_rs::fs::error::Error; fn try_from(config: ResourceConfig) -> Result { - 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 { + 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()) +} diff --git a/src/runner.rs b/src/runner.rs index 762ebe2..d641667 100644 --- a/src/runner.rs +++ b/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 { - 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();