diff --git a/src/cmd/sandbox.rs b/src/cmd/sandbox.rs index 2e92e63..e0b97d1 100644 --- a/src/cmd/sandbox.rs +++ b/src/cmd/sandbox.rs @@ -2,10 +2,13 @@ use crate::Workspace; use crate::cmd::{Command, CommandError, ProcessLinesActions, ProcessOutput, ProcessStatistics}; use log::{error, info}; use serde::Deserialize; -use std::error::Error; -use std::fmt; -use std::path::{Path, PathBuf}; -use std::time::Duration; +use std::{ + error::Error, + fmt, + ops::RangeInclusive, + path::{Path, PathBuf}, + time::Duration, +}; /// The Docker image used for sandboxing. pub struct SandboxImage { @@ -145,6 +148,7 @@ pub struct SandboxBuilder { env: Vec<(String, String)>, memory_limit: Option, cpu_limit: Option, + cpuset_cpus: Option>, workdir: Option, user: Option, cmd: Vec, @@ -160,6 +164,7 @@ impl SandboxBuilder { workdir: None, memory_limit: None, cpu_limit: None, + cpuset_cpus: None, user: None, cmd: Vec::new(), enable_networking: true, @@ -197,6 +202,15 @@ impl SandboxBuilder { self } + /// Restrict the sandbox to run on a specific inclusive range of CPU IDs. + /// + /// For example, `0..=1` will restrict the sandbox to CPUs 0 and 1 and translate to Docker's + /// `--cpuset-cpus 0-1`. + pub fn cpuset_cpus(mut self, cpus: Option>) -> Self { + self.cpuset_cpus = cpus; + self + } + /// Enable or disable the sandbox's networking. When it's disabled processes inside the sandbox /// won't be able to reach network service on the Internet or the host machine. /// @@ -255,6 +269,11 @@ impl SandboxBuilder { args.push(limit.to_string()); } + if let Some(cpus) = self.cpuset_cpus { + args.push("--cpuset-cpus".into()); + args.push(format_cpuset_cpus(&cpus)); + } + if !self.enable_networking { args.push("--network".into()); args.push("none".into()); @@ -538,3 +557,17 @@ pub fn docker_running(workspace: &Workspace) -> bool { .run() .is_ok() } + +fn format_cpuset_cpus(cpus: &RangeInclusive) -> String { + format!("{}-{}", cpus.start(), cpus.end()) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn formats_cpuset_cpus() { + assert_eq!(format_cpuset_cpus(&(2..=4)), "2-4"); + } +} diff --git a/tests/buildtest/mod.rs b/tests/buildtest/mod.rs index d27f9bb..c233982 100644 --- a/tests/buildtest/mod.rs +++ b/tests/buildtest/mod.rs @@ -213,6 +213,57 @@ fn test_sandbox_oom() { }); } +#[test] +#[cfg(not(windows))] +fn test_invalid_cpuset_cpus() { + use rustwide::cmd::CommandError; + + runner::run("hello-world", |run| { + let res = run.run( + SandboxBuilder::new() + .enable_networking(false) + .cpuset_cpus(Some(999_999..=999_999)), + |build| { + build.cmd("true").run()?; + Ok(()) + }, + ); + if let Some( + CommandError::SandboxContainerCreate(_) | CommandError::ExecutionFailed { .. }, + ) = res.err().and_then(|err| err.downcast().ok()) + { + // Everything is OK! + } else { + panic!( + "didn't get CommandError::SandboxContainerCreate or CommandError::ExecutionFailed" + ); + } + Ok(()) + }); +} + +#[test] +#[cfg(not(windows))] +fn test_cpuset_cpus_applied() { + runner::run("hello-world", |run| { + run.run( + SandboxBuilder::new() + .enable_networking(false) + .cpuset_cpus(Some(0..=1)), + |build| { + let output = build + .cmd("sh") + .args(&["-c", "grep '^Cpus_allowed_list:' /proc/self/status"]) + .run_capture()?; + + assert_eq!(output.stdout_lines(), ["Cpus_allowed_list:\t0-1"]); + Ok(()) + }, + )?; + Ok(()) + }); +} + #[test] fn test_override_files() { runner::run("cargo-config", |run| {