From f5b563bdc66d453703d2a35a62ec1ff15f187c5d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=81=E9=80=9F=E8=9C=97=E7=89=9B?= <31986081+Jisu-Woniu@users.noreply.github.com> Date: Fri, 10 Jan 2025 14:40:22 +0800 Subject: [PATCH 1/2] =?UTF-8?q?refactor(rsjudge-runner):=20=E2=99=BB?= =?UTF-8?q?=EF=B8=8F=20refactor=20resource-related=20code?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Now we can bind resource limit to a `Command` without `spawn`ing it. fix #208 --- Cargo.lock | 28 ++--- crates/rsjudge-runner/Cargo.toml | 2 +- crates/rsjudge-runner/examples/rusage_test.rs | 23 +++- crates/rsjudge-runner/src/error.rs | 5 +- .../rsjudge-runner/src/utils/resources/mod.rs | 119 ++++++++++++------ .../src/utils/resources/rusage.rs | 37 +++--- crates/rsjudge-traits/src/resource.rs | 21 ++-- 7 files changed, 142 insertions(+), 93 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 2d380eb..77e7e0b 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -332,9 +332,9 @@ checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" [[package]] name = "clap" -version = "4.5.24" +version = "4.5.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9560b07a799281c7e0958b9296854d6fafd4c5f31444a7e5bb1ad6dde5ccf1bd" +checksum = "a8eb5e908ef3a6efbe1ed62520fb7287959888c88485abe072543190ecc66783" dependencies = [ "clap_builder", "clap_derive", @@ -342,9 +342,9 @@ dependencies = [ [[package]] name = "clap_builder" -version = "4.5.24" +version = "4.5.26" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "874e0dd3eb68bf99058751ac9712f622e61e6f393a94f7128fa26e3f02f5c7cd" +checksum = "96b01801b5fc6a0a232407abc821660c9c6d25a1cafc0d4f85f29fb8d9afc121" dependencies = [ "anstream", "anstyle", @@ -354,9 +354,9 @@ dependencies = [ [[package]] name = "clap_complete" -version = "4.5.41" +version = "4.5.42" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "942dc5991a34d8cf58937ec33201856feba9cbceeeab5adf04116ec7c763bff1" +checksum = "33a7e468e750fa4b6be660e8b5651ad47372e8fb114030b594c2d75d48c5ffd0" dependencies = [ "clap", ] @@ -1541,18 +1541,18 @@ dependencies = [ [[package]] name = "thiserror" -version = "2.0.9" +version = "2.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f072643fd0190df67a8bab670c20ef5d8737177d6ac6b2e9a236cb096206b2cc" +checksum = "a3ac7f54ca534db81081ef1c1e7f6ea8a3ef428d2fc069097c079443d24124d3" dependencies = [ "thiserror-impl", ] [[package]] name = "thiserror-impl" -version = "2.0.9" +version = "2.0.10" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7b50fa271071aae2e6ee85f842e2e28ba8cd2c5fb67f11fcb1fd70b276f9e7d4" +checksum = "9e9465d30713b56a37ede7185763c3492a91be2f5fa68d958c44e41ab9248beb" dependencies = [ "proc-macro2", "quote", @@ -1561,9 +1561,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.42.0" +version = "1.43.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "5cec9b21b0450273377fc97bd4c33a8acffc8c996c987a7c5b319a0083707551" +checksum = "3d61fa4ffa3de412bfea335c6ecff681de2b609ba3c77ef3e00e521813a9ed9e" dependencies = [ "backtrace", "bytes", @@ -1579,9 +1579,9 @@ dependencies = [ [[package]] name = "tokio-macros" -version = "2.4.0" +version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "693d596312e88961bc67d7f1f97af8a70227d9f90c31bba5806eec004978d752" +checksum = "6e06d43f1345a3bcd39f6a56dbb7dcab2ba47e68e8ac134855e7e2bdbaf8cab8" dependencies = [ "proc-macro2", "quote", diff --git a/crates/rsjudge-runner/Cargo.toml b/crates/rsjudge-runner/Cargo.toml index 4e75852..94a4ae2 100644 --- a/crates/rsjudge-runner/Cargo.toml +++ b/crates/rsjudge-runner/Cargo.toml @@ -17,7 +17,7 @@ log.workspace = true nix = { version = "0.29.0", features = ["user", "resource", "process"] } rsjudge-traits.workspace = true rsjudge-utils.workspace = true -thiserror = "2.0.9" +thiserror = "2.0.10" tokio = { workspace = true, features = ["process", "sync", "time", "signal"] } tokio-util = "0.7.13" uzers = "0.12.1" diff --git a/crates/rsjudge-runner/examples/rusage_test.rs b/crates/rsjudge-runner/examples/rusage_test.rs index 26b85ee..7bfc2d3 100644 --- a/crates/rsjudge-runner/examples/rusage_test.rs +++ b/crates/rsjudge-runner/examples/rusage_test.rs @@ -1,10 +1,10 @@ // SPDX-License-Identifier: Apache-2.0 -use std::{os::unix::process::ExitStatusExt, path::PathBuf, time::Duration}; +use std::{num::NonZeroU64, os::unix::process::ExitStatusExt, path::PathBuf, time::Duration}; use anyhow::bail; use nix::{sys::wait::WaitStatus, unistd::Pid}; -use rsjudge_runner::utils::resources::{rusage::WaitForResourceUsage, RunWithResourceLimit}; +use rsjudge_runner::utils::resources::{rusage::WaitForResourceUsage, WithResourceLimit as _}; use rsjudge_traits::resource::ResourceLimit; use tokio::{process::Command, time::Instant}; @@ -28,7 +28,10 @@ async fn main() -> anyhow::Result<()> { .wait_for_resource_usage() .await? else { - bail!("Failed to get resource usage"); + bail!( + "Failed to get resource usage for `{}`", + stringify!(spin_lock) + ); }; dbg!(start_time.elapsed()); @@ -49,7 +52,7 @@ async fn main() -> anyhow::Result<()> { .wait_for_resource_usage() .await else { - bail!("Failed to get resource usage"); + bail!("Failed to get resource usage for `{}`", stringify!(sleep)); }; dbg!(start_time.elapsed()); @@ -59,11 +62,19 @@ async fn main() -> anyhow::Result<()> { eprintln!("Starting `large_alloc` with RAM limit of 1MB"); let Ok(Some((status, rusage))) = Command::new(large_alloc) - .spawn_with_resource_limit(ResourceLimit::new(None, None, Some(1 << 30), None))? + .spawn_with_resource_limit(ResourceLimit::new( + None, + None, + NonZeroU64::new(1 << 30), + None, + ))? .wait_for_resource_usage() .await else { - bail!("Failed to get resource usage"); + bail!( + "Failed to get resource usage for `{}`", + stringify!(large_alloc) + ); }; let status = WaitStatus::from_raw(Pid::from_raw(0), status.into_raw())?; diff --git a/crates/rsjudge-runner/src/error.rs b/crates/rsjudge-runner/src/error.rs index 81642ee..a315291 100644 --- a/crates/rsjudge-runner/src/error.rs +++ b/crates/rsjudge-runner/src/error.rs @@ -27,7 +27,7 @@ pub enum Error { TimeLimitExceeded(#[cfg(debug_assertions)] Option<(ExitStatus, ResourceUsage)>), #[error("Child process has exited with status: {0:?}")] - ChildExited(ExitStatus), + EarlyExited(ExitStatus), } /// Convert any error implementing [`Into`]`<`[`io::Error`]`>` into [`Error`]. @@ -37,4 +37,7 @@ impl> From for Error { } } +/// A specialized [`Result`] type for this crate. +/// +/// See the [`Error`] type for the error variants. pub type Result = StdResult; diff --git a/crates/rsjudge-runner/src/utils/resources/mod.rs b/crates/rsjudge-runner/src/utils/resources/mod.rs index 5ef4497..f4f464e 100644 --- a/crates/rsjudge-runner/src/utils/resources/mod.rs +++ b/crates/rsjudge-runner/src/utils/resources/mod.rs @@ -2,68 +2,85 @@ pub mod rusage; -use std::{ - future::Future, - process::ExitStatus, - time::{Duration, Instant}, -}; +use std::{future::Future, process::ExitStatus, time::Duration}; use nix::sys::resource::{setrlimit, Resource}; use rsjudge_traits::resource::ResourceLimit; -use tokio::process::{Child, Command}; +use tokio::{ + process::{Child, Command}, + time::Instant, +}; -use self::rusage::ResourceUsage; -use crate::{utils::resources::rusage::WaitForResourceUsage, Result}; +use self::rusage::{ResourceUsage, WaitForResourceUsage}; +use crate::Result; #[derive(Debug)] -pub struct ChildWithTimeout { - child: Child, - start: Instant, +pub struct CommandWithResourceLimit { + command: Command, timeout: Option, } -impl AsRef for ChildWithTimeout { - fn as_ref(&self) -> &Child { - &self.child +impl CommandWithResourceLimit { + /// Get a reference to the inner [`Command`]. + pub fn command(&self) -> &Command { + &self.command } -} -impl AsMut for ChildWithTimeout { - fn as_mut(&mut self) -> &mut Child { - &mut self.child + /// Get a mutable reference to the inner [`Command`]. + pub fn command_mut(&mut self) -> &mut Command { + &mut self.command + } + + /// Spawn the [`Command`] with the given resource limit. + /// + /// This function is synchronous and won't wait for the child to exit. + pub fn spawn(&mut self) -> Result { + Ok(ChildWithDeadline { + child: self.command.spawn()?, + deadline: self.timeout.map(|timeout| Instant::now() + timeout), + }) } } -pub trait RunWithResourceLimit { +/// Setting resource limits for a [`Command`]. +/// +/// This will take the [`Command`] by value and set the [`ResourceLimit`] for it. +pub trait WithResourceLimit { + /// Register resource limit for the command. + /// + /// Returns a [`CommandWithResourceLimit`] which can be spawned. + /// + /// You can also use [`command`][fn.command] or [`command_mut`][fn.command_mut] + /// to get the inner [`Command`][tokio::process::Command] object as needed. + /// + /// [fn.command]: CommandWithResourceLimit::command + /// [fn.command_mut]: CommandWithResourceLimit::command_mut + fn with_resource_limit(self, resource_limit: ResourceLimit) -> CommandWithResourceLimit; /// Spawn [`Self`] with optional resource limit. /// /// This function won't wait for the child to exit. /// Nor will it apply the [`ResourceLimit::wall_time_limit`] automatically. /// - /// However, the wall time limit can be applied by using [`WaitForResourceUsage::wait_for_resource_usage`]. + /// However, the wall time limit can be applied by using [`wait_for_resource_usage`]. /// /// This function is synchronous. /// /// # Errors /// /// This function will return an error if the child process cannot be spawned. - fn spawn_with_resource_limit( - &mut self, - resource_info: ResourceLimit, - ) -> Result; + /// + /// [`wait_for_resource_usage`]: WaitForResourceUsage::wait_for_resource_usage + fn spawn_with_resource_limit(self, resource_limit: ResourceLimit) -> Result; /// Run [`Self`] with given resource limit. fn wait_with_resource_limit( - &mut self, - resource_info: ResourceLimit, + self, + resource_limit: ResourceLimit, ) -> impl Future>> + Send; } -impl RunWithResourceLimit for Command { - fn spawn_with_resource_limit( - &mut self, - resource_info: ResourceLimit, - ) -> Result { +impl WithResourceLimit for Command { + fn with_resource_limit(mut self, resource_info: ResourceLimit) -> CommandWithResourceLimit { if let Some(cpu_time_limit) = resource_info.cpu_time_limit() { let set_cpu_limit = move || { setrlimit( @@ -78,6 +95,7 @@ impl RunWithResourceLimit for Command { self.pre_exec(set_cpu_limit); } } + if let Some(memory_limit) = resource_info.memory_limit() { let set_memory_limit = move || { setrlimit(Resource::RLIMIT_AS, memory_limit, memory_limit)?; @@ -104,15 +122,18 @@ impl RunWithResourceLimit for Command { } } - Ok(ChildWithTimeout { - child: self.spawn()?, - start: Instant::now(), + CommandWithResourceLimit { + command: self, timeout: resource_info.wall_time_limit(), - }) + } + } + + fn spawn_with_resource_limit(self, resource_limit: ResourceLimit) -> Result { + self.with_resource_limit(resource_limit).spawn() } async fn wait_with_resource_limit( - &mut self, + self, resource_limit: ResourceLimit, ) -> Result> { self.spawn_with_resource_limit(resource_limit)? @@ -121,6 +142,25 @@ impl RunWithResourceLimit for Command { } } +#[derive(Debug)] +pub struct ChildWithDeadline { + child: Child, + + deadline: Option, +} + +impl ChildWithDeadline { + /// Get a reference to the inner [`Child`]. + pub fn child(&self) -> &Child { + &self.child + } + + /// Get a mutable reference to the inner [`Child`]. + pub fn child_mut(&mut self) -> &mut Child { + &mut self.child + } +} + #[cfg(test)] mod tests { use std::time::{Duration, Instant}; @@ -128,14 +168,15 @@ mod tests { use rsjudge_traits::resource::ResourceLimit; use crate::{ - utils::resources::{rusage::WaitForResourceUsage as _, RunWithResourceLimit}, + utils::resources::{rusage::WaitForResourceUsage as _, WithResourceLimit as _}, Error, }; #[tokio::test] async fn test_wait_for_resource_usage() { - let mut child = tokio::process::Command::new("sleep") - .arg("10") + let mut command = tokio::process::Command::new("sleep"); + command.arg("10"); + let mut child = command .spawn_with_resource_limit(ResourceLimit::new( Some(Duration::from_secs(1)), Some(Duration::from_secs_f64(1.5)), diff --git a/crates/rsjudge-runner/src/utils/resources/rusage.rs b/crates/rsjudge-runner/src/utils/resources/rusage.rs index d830d41..522f9dc 100644 --- a/crates/rsjudge-runner/src/utils/resources/rusage.rs +++ b/crates/rsjudge-runner/src/utils/resources/rusage.rs @@ -15,12 +15,11 @@ use tokio::{ process::Child, select, signal::unix::{signal, SignalKind}, - task::spawn, - time::sleep, + time::sleep_until, }; -use tokio_util::sync::CancellationToken; -use crate::{utils::resources::ChildWithTimeout, Error, Result}; +// use tokio_util::sync::CancellationToken; +use crate::{utils::resources::ChildWithDeadline, Error, Result}; /// Resource usage of a process. /// @@ -36,9 +35,11 @@ impl From for ResourceUsage { fn from(rusage: rusage) -> Self { Self { cpu_time: Duration::new( + // User time rusage.ru_utime.tv_sec as u64, rusage.ru_utime.tv_usec as u32 * 1000, ) + Duration::new( + // System time rusage.ru_stime.tv_sec as u64, rusage.ru_stime.tv_usec as u32 * 1000, ), @@ -59,6 +60,7 @@ impl ResourceUsage { self.ram_usage } } + pub trait WaitForResourceUsage { /// Wait for the resource usage of the process. /// @@ -127,35 +129,24 @@ impl WaitForResourceUsage for Child { let exit_status = self .try_wait()? .ok_or_else(|| io::Error::other("Exit status not available"))?; - Err(Error::ChildExited(exit_status)) + if exit_status.success() { + Ok(None) + } else { + Err(Error::EarlyExited(exit_status)) + } } } } -impl WaitForResourceUsage for ChildWithTimeout { +impl WaitForResourceUsage for ChildWithDeadline { async fn wait_for_resource_usage(&mut self) -> Result> { - let Some(timeout) = self.timeout else { + let Some(deadline) = self.deadline else { return self.child.wait_for_resource_usage().await; }; - let cancellation_token = CancellationToken::new(); - let child_token = cancellation_token.child_token(); - - let start = self.start; - - spawn(async move { - loop { - if timeout <= start.elapsed() { - cancellation_token.cancel(); - break; - } - sleep(Duration::from_millis(10)).await; - } - }); - select! { res = self.child.wait_for_resource_usage() => res, - () = child_token.cancelled() => { + () = sleep_until(deadline) => { self.child.start_kill()?; return Err(Error::TimeLimitExceeded( #[cfg(debug_assertions)] diff --git a/crates/rsjudge-traits/src/resource.rs b/crates/rsjudge-traits/src/resource.rs index 745e255..66cd126 100644 --- a/crates/rsjudge-traits/src/resource.rs +++ b/crates/rsjudge-traits/src/resource.rs @@ -2,7 +2,7 @@ //! Resource limit for judging code. -use std::time::Duration; +use std::{num::NonZeroU64, time::Duration}; /// Resource limit for judging code. #[derive(Debug, Default, Clone, Copy)] @@ -16,9 +16,9 @@ pub struct ResourceLimit { /// Wall time limit may be inaccurate, due to the implementation of "wait-and-check" strategy. wall_time_limit: Option, /// The memory limit **in bytes**. - memory_limit: Option, + memory_limit: Option, /// Max file size limit **in bytes**. - max_file_size_limit: Option, + max_file_size_limit: Option, } impl ResourceLimit { @@ -27,8 +27,8 @@ impl ResourceLimit { pub fn new( cpu_time_limit: Option, wall_time_limit: Option, - memory_limit: Option, - max_file_size_limit: Option, + memory_limit: Option, + max_file_size_limit: Option, ) -> Self { Self { cpu_time_limit, @@ -53,13 +53,13 @@ impl ResourceLimit { /// Get the memory limit. #[must_use] pub fn memory_limit(&self) -> Option { - self.memory_limit + self.memory_limit.map(From::from) } /// Get the max file size limit. #[must_use] pub fn max_file_size_limit(&self) -> Option { - self.max_file_size_limit + self.max_file_size_limit.map(From::from) } /// Set the CPU time limit. @@ -75,13 +75,16 @@ impl ResourceLimit { } /// Set the memory limit. - pub fn set_memory_limit(&mut self, memory_limit: Option) -> &mut Self { + pub fn set_memory_limit(&mut self, memory_limit: Option) -> &mut Self { self.memory_limit = memory_limit; self } /// Set the max file size limit. - pub fn set_max_file_size_limit(&mut self, max_file_size_limit: Option) -> &mut Self { + pub fn set_max_file_size_limit( + &mut self, + max_file_size_limit: Option, + ) -> &mut Self { self.max_file_size_limit = max_file_size_limit; self } From f060f2648bc4052351dfed711352e90322b3dc5b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E6=9E=81=E9=80=9F=E8=9C=97=E7=89=9B?= <31986081+Jisu-Woniu@users.noreply.github.com> Date: Fri, 10 Jan 2025 14:42:59 +0800 Subject: [PATCH 2/2] =?UTF-8?q?build(deps):=20=E2=AC=86=EF=B8=8F=20update?= =?UTF-8?q?=20dependencies?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- Cargo.toml | 8 ++++---- crates/rsjudge-amqp/Cargo.toml | 2 +- crates/rsjudge-utils/Cargo.toml | 2 +- xtask/Cargo.toml | 2 +- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index f4deee4..4f0d739 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -27,7 +27,7 @@ log = "0.4.22" rsjudge-traits = { version = "0.1.0", path = "crates/rsjudge-traits" } rsjudge-utils = { version = "0.1.0", path = "crates/rsjudge-utils" } serde = { version = "1.0.217", features = ["derive"] } -tokio = "1.42.0" +tokio = "1.43.0" [package] name = "rsjudge" @@ -160,7 +160,7 @@ rsjudge-grpc = { version = "0.1.0", path = "crates/rsjudge-grpc", optional = tru rsjudge-rest = { version = "0.1.0", path = "crates/rsjudge-rest", optional = true } anyhow = "1.0.95" -clap = { version = "4.5.24", features = ["derive"] } +clap = { version = "4.5.26", features = ["derive"] } env_logger = "0.11.6" log.workspace = true mimalloc = { version = "0.1.43", optional = true } @@ -181,8 +181,8 @@ mimalloc = ["dep:mimalloc"] default = ["grpc", "rest", "mimalloc"] [build-dependencies] -clap = { version = "4.5.24", features = ["derive"] } -clap_complete = "4.5.41" +clap = { version = "4.5.26", features = ["derive"] } +clap_complete = "4.5.42" clap_mangen = "0.2.25" [profile.release] diff --git a/crates/rsjudge-amqp/Cargo.toml b/crates/rsjudge-amqp/Cargo.toml index e7b7227..88bd324 100644 --- a/crates/rsjudge-amqp/Cargo.toml +++ b/crates/rsjudge-amqp/Cargo.toml @@ -11,7 +11,7 @@ rust-version.workspace = true [dependencies] amqprs = { version = "2.1.0", features = ["urispec"] } -thiserror = "2.0.9" +thiserror = "2.0.10" tokio.workspace = true rsjudge-traits.workspace = true diff --git a/crates/rsjudge-utils/Cargo.toml b/crates/rsjudge-utils/Cargo.toml index ff2576c..5e3ce5f 100644 --- a/crates/rsjudge-utils/Cargo.toml +++ b/crates/rsjudge-utils/Cargo.toml @@ -12,7 +12,7 @@ rust-version.workspace = true [dependencies] log.workspace = true shell-words = "1.1.0" -thiserror = "2.0.9" +thiserror = "2.0.10" tokio = { workspace = true, features = ["process", "macros", "rt-multi-thread"] } [dev-dependencies] diff --git a/xtask/Cargo.toml b/xtask/Cargo.toml index 4a256ef..d878b0c 100644 --- a/xtask/Cargo.toml +++ b/xtask/Cargo.toml @@ -12,7 +12,7 @@ description = "A task runner for cargo workspaces" [dependencies] anyhow = "1.0.95" -clap = { version = "4.5.24", features = ["derive"] } +clap = { version = "4.5.26", features = ["derive"] } [features] dbg = []