refactor(rsjudge-runner): ♻️ refactor resource-related code

Now we can bind resource limit to a `Command` without `spawn`ing it.

fix #208

Author: Jisu-Woniu

Cargo.lock

@@ -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"

crates/rsjudge-runner/Cargo.toml

@@ -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())?;

crates/rsjudge-runner/examples/rusage_test.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<E: Into<io::Error>> From<E> for Error {
     }
 }
 
+/// A specialized [`Result`] type for this crate.
+///
+/// See the [`Error`] type for the error variants.
 pub type Result<T, E = Error> = StdResult<T, E>;

crates/rsjudge-runner/src/error.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<Duration>,
 }
 
-impl AsRef<Child> 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<Child> 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<ChildWithDeadline> {
+        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<ChildWithTimeout>;
+    ///
+    /// [`wait_for_resource_usage`]: WaitForResourceUsage::wait_for_resource_usage
+    fn spawn_with_resource_limit(self, resource_limit: ResourceLimit) -> Result<ChildWithDeadline>;
 
     /// Run [`Self`] with given resource limit.
     fn wait_with_resource_limit(
-        &mut self,
-        resource_info: ResourceLimit,
+        self,
+        resource_limit: ResourceLimit,
     ) -> impl Future<Output = Result<Option<(ExitStatus, ResourceUsage)>>> + Send;
 }
 
-impl RunWithResourceLimit for Command {
-    fn spawn_with_resource_limit(
-        &mut self,
-        resource_info: ResourceLimit,
-    ) -> Result<ChildWithTimeout> {
+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<ChildWithDeadline> {
+        self.with_resource_limit(resource_limit).spawn()
     }
 
     async fn wait_with_resource_limit(
-        &mut self,
+        self,
         resource_limit: ResourceLimit,
     ) -> Result<Option<(ExitStatus, ResourceUsage)>> {
         self.spawn_with_resource_limit(resource_limit)?
@@ -121,21 +142,41 @@ impl RunWithResourceLimit for Command {
     }
 }
 
+#[derive(Debug)]
+pub struct ChildWithDeadline {
+    child: Child,
+
+    deadline: Option<Instant>,
+}
+
+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};
 
     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)),