feat: support setting the cores a container should use

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<usize>,
     cpu_limit: Option<f32>,
+    cpuset_cpus: Option<RangeInclusive<usize>>,
     workdir: Option<String>,
     user: Option<String>,
     cmd: Vec<String>,
@@ -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<RangeInclusive<usize>>) -> 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<usize>) -> 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");
+    }
+}

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| {