builder: support cpu-core assignment for builders

Cargo.toml

@@ -40,6 +40,7 @@ http = "1.0.0"
 itertools = "0.14.0" 
 mime = "0.3.16"
 mockito = "1.0.2"
+num_cpus = "1.15.0"
 opentelemetry = "0.31.0"
 opentelemetry-otlp = { version = "0.31.0", features = ["grpc-tonic", "metrics"] }
 opentelemetry-resource-detectors = "0.10.0"

crates/bin/docs_rs_builder/Cargo.toml

@@ -27,9 +27,10 @@ docs_rs_types = { path = "../../lib/docs_rs_types" }
 docs_rs_utils = { path = "../../lib/docs_rs_utils" }
 docsrs-metadata = { path = "../../lib/metadata" }
 log = "0.4"
+num_cpus = { workspace = true }
 opentelemetry = { workspace = true }
 regex = { workspace = true }
-rustwide = { version = "0.22.0", features = ["unstable", "unstable-toolchain-ci"] }
+rustwide = { version = "0.23.0", features = ["unstable", "unstable-toolchain-ci"] }
 serde_json = { workspace = true }
 sqlx = { workspace = true }
 sysinfo = { version = "0.38.0", default-features = false, features = ["system"] }

crates/bin/docs_rs_builder/src/config.rs

@@ -1,7 +1,15 @@
-use anyhow::Result;
+use anyhow::{Result, bail};
 use docs_rs_config::AppConfig;
 use docs_rs_env_vars::{env, maybe_env, require_env};
-use std::{path::PathBuf, sync::Arc, time::Duration};
+use std::{
+    num::ParseIntError,
+    ops::{Deref, RangeInclusive},
+    path::PathBuf,
+    str::FromStr,
+    sync::Arc,
+    time::Duration,
+};
+use thiserror::Error;
 
 #[derive(Debug)]
 pub struct Config {
@@ -18,7 +26,10 @@ pub struct Config {
     pub rustwide_workspace: PathBuf,
     pub inside_docker: bool,
     pub docker_image: Option<String>,
+    /// Docker CPU quota / CPU count.
     pub build_cpu_limit: Option<u32>,
+    /// CPU cores the builder should use.
+    pub build_cpu_cores: Option<BuildCores>,
     pub include_default_targets: bool,
     pub disable_memory_limit: bool,
 
@@ -29,7 +40,7 @@ pub struct Config {
 impl AppConfig for Config {
     fn from_environment() -> Result<Self> {
         let prefix: PathBuf = require_env("DOCSRS_PREFIX")?;
-        Ok(Self {
+        let config = Self {
             temp_dir: prefix.join("tmp"),
             prefix,
             rustwide_workspace: env("DOCSRS_RUSTWIDE_WORKSPACE", PathBuf::from(".workspace"))?,
@@ -38,6 +49,7 @@ impl AppConfig for Config {
                 .or(maybe_env("DOCSRS_DOCKER_IMAGE")?),
 
             build_cpu_limit: maybe_env("DOCSRS_BUILD_CPU_LIMIT")?,
+            build_cpu_cores: maybe_env("DOCSRS_BUILD_CPU_CORES")?,
             include_default_targets: env("DOCSRS_INCLUDE_DEFAULT_TARGETS", true)?,
             disable_memory_limit: env("DOCSRS_DISABLE_MEMORY_LIMIT", false)?,
             build_workspace_reinitialization_interval: Duration::from_secs(env(
@@ -46,7 +58,13 @@ impl AppConfig for Config {
             )?),
             compiler_metrics_collection_path: maybe_env("DOCSRS_COMPILER_METRICS_PATH")?,
             build_limits: Arc::new(docs_rs_build_limits::Config::from_environment()?),
-        })
+        };
+
+        if config.build_cpu_limit.is_some() && config.build_cpu_cores.is_some() {
+            bail!("you only can define one of build_cpu_limit and build_cpu_cores");
+        }
+
+        Ok(config)
     }
 
     #[cfg(test)]
@@ -58,3 +76,196 @@ impl AppConfig for Config {
         Ok(config)
     }
 }
+
+impl Config {
+    /// The cargo job-limit we should set in builds.
+    ///
+    /// If we set either of the two CPU-limits, cargo should
+    /// limit itself automatically.
+    pub fn cargo_job_limit(&self) -> Option<usize> {
+        self.build_cpu_cores
+            .as_ref()
+            .map(|c| c.len())
+            .or(self.build_cpu_limit.map(|l| l as usize))
+    }
+}
+
+#[derive(Debug)]
+pub struct BuildCores(pub RangeInclusive<usize>);
+
+impl BuildCores {
+    pub fn len(&self) -> usize {
+        self.0.size_hint().0
+    }
+}
+
+impl Deref for BuildCores {
+    type Target = RangeInclusive<usize>;
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+#[derive(Debug, Error)]
+pub enum ParseBuildCoresError {
+    #[error("expected build core range in the form <start>-<end>")]
+    MissingSeparator,
+    #[error("invalid build core range start `{value}`: {source}")]
+    InvalidStart {
+        value: String,
+        #[source]
+        source: ParseIntError,
+    },
+    #[error("invalid build core range end `{value}`: {source}")]
+    InvalidEnd {
+        value: String,
+        #[source]
+        source: ParseIntError,
+    },
+    #[error("build core range start must be less than or equal to end")]
+    DescendingRange,
+    #[error("not enough cores, we only have {0}")]
+    NotEnoughCores(usize),
+}
+
+impl FromStr for BuildCores {
+    type Err = ParseBuildCoresError;
+
+    fn from_str(s: &str) -> std::result::Result<Self, Self::Err> {
+        let (start, end) = s
+            .split_once('-')
+            .ok_or(ParseBuildCoresError::MissingSeparator)?;
+
+        let start = start
+            .parse()
+            .map_err(|source| ParseBuildCoresError::InvalidStart {
+                value: start.to_string(),
+                source,
+            })?;
+
+        let end = end
+            .parse()
+            .map_err(|source| ParseBuildCoresError::InvalidEnd {
+                value: end.to_string(),
+                source,
+            })?;
+
+        if start > end {
+            return Err(ParseBuildCoresError::DescendingRange);
+        }
+
+        let cpus = num_cpus::get();
+
+        if end >= cpus {
+            // NOTE: docker counts the cores zero-based, so
+            // a core-number that is exactly the cpu-count is already
+            // too high.
+            return Err(ParseBuildCoresError::NotEnoughCores(cpus));
+        }
+
+        Ok(Self(start..=end))
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn parses_build_core_range() {
+        let build_cores: BuildCores = "2-3".parse().unwrap();
+
+        assert_eq!(build_cores.start(), &2);
+        assert_eq!(build_cores.end(), &3);
+        assert_eq!(build_cores.len(), 2);
+    }
+
+    #[test]
+    fn parses_single_build_core() {
+        let build_cores: BuildCores = "2-2".parse().unwrap();
+
+        assert_eq!(build_cores.start(), &2);
+        assert_eq!(build_cores.end(), &2);
+        assert_eq!(build_cores.len(), 1);
+    }
+
+    #[test]
+    fn rejects_build_core_range_without_separator() {
+        let err = "3".parse::<BuildCores>().unwrap_err();
+
+        assert!(
+            err.to_string()
+                .contains("expected build core range in the form <start>-<end>")
+        );
+    }
+
+    #[test]
+    fn rejects_build_core_range_with_descending_values() {
+        let err = "4-3".parse::<BuildCores>().unwrap_err();
+
+        assert!(
+            err.to_string()
+                .contains("build core range start must be less than or equal to end")
+        );
+    }
+
+    #[test]
+    fn rejects_build_core_range_with_invalid_end() {
+        let err = "3-a".parse::<BuildCores>().unwrap_err();
+
+        assert!(err.to_string().contains("invalid build core range end `a`"));
+    }
+
+    #[test]
+    fn rejects_build_core_range_with_invalid_core_number() {
+        let cpus = num_cpus::get();
+        let err = format!("0-{cpus}").parse::<BuildCores>().unwrap_err();
+
+        assert!(
+            err.to_string()
+                .contains(&format!("not enough cores, we only have {cpus}"))
+        );
+    }
+
+    #[test]
+    fn cargo_jobs_uses_core_range_length() {
+        let config = config_with_cpu_settings(Some(12), Some(BuildCores(3..=4)));
+
+        assert_eq!(config.cargo_job_limit(), Some(2));
+    }
+
+    #[test]
+    fn cargo_jobs_falls_back_to_cpu_limit() {
+        let config = config_with_cpu_settings(Some(12), None);
+
+        assert_eq!(config.cargo_job_limit(), Some(12));
+    }
+
+    #[test]
+    fn cargo_jobs_is_none_without_cpu_settings() {
+        let config = config_with_cpu_settings(None, None);
+
+        assert_eq!(config.cargo_job_limit(), None);
+    }
+
+    fn config_with_cpu_settings(
+        build_cpu_limit: Option<u32>,
+        build_cpu_cores: Option<BuildCores>,
+    ) -> Config {
+        Config {
+            prefix: PathBuf::new(),
+            temp_dir: PathBuf::new(),
+            compiler_metrics_collection_path: None,
+            build_workspace_reinitialization_interval: Duration::from_secs(0),
+            rustwide_workspace: PathBuf::new(),
+            inside_docker: false,
+            docker_image: None,
+            build_cpu_limit,
+            build_cpu_cores,
+            include_default_targets: true,
+            disable_memory_limit: false,
+            build_limits: Arc::new(docs_rs_build_limits::Config::default()),
+        }
+    }
+}

crates/bin/docs_rs_builder/src/docbuilder/rustwide_builder.rs

@@ -167,10 +167,17 @@ impl RustwideBuilder {
 
     #[instrument(skip(self))]
     fn prepare_sandbox(&self, limits: &Limits) -> SandboxBuilder {
-        SandboxBuilder::new()
-            .cpu_limit(self.config.build_cpu_limit.map(|limit| limit as f32))
+        let builder = SandboxBuilder::new()
             .memory_limit(Some(limits.memory()))
-            .enable_networking(limits.networking())
+            .enable_networking(limits.networking());
+
+        if let Some(cores) = &self.config.build_cpu_cores {
+            builder.cpuset_cpus(Some(cores.0.clone()))
+        } else if let Some(limit) = &self.config.build_cpu_limit {
+            builder.cpu_limit(Some(*limit as f32))
+        } else {
+            builder
+        }
     }
 
     pub fn purge_caches(&self) -> Result<()> {
@@ -1245,8 +1252,8 @@ impl RustwideBuilder {
             // docs.rs, but once it's stable we can remove this flag.
             "-Zrustdoc-scrape-examples".into(),
         ];
-        if let Some(cpu_limit) = self.config.build_cpu_limit {
-            cargo_args.push(format!("-j{cpu_limit}"));
+        if let Some(cargo_job_limit) = self.config.cargo_job_limit() {
+            cargo_args.push(format!("-j{cargo_job_limit}"));
         }
         // Cargo has a series of frightening bugs around cross-compiling proc-macros:
         // - Passing `--target` causes RUSTDOCFLAGS to fail to be passed 🤦
@@ -1369,7 +1376,10 @@ impl BuildResult {
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::testing::{TestEnvironment, TestEnvironmentExt as _};
+    use crate::{
+        config::BuildCores,
+        testing::{TestEnvironment, TestEnvironmentExt as _},
+    };
     use docs_rs_config::AppConfig as _;
     use docs_rs_utils::block_on_async_with_conn;
     // use crate::test::{AxumRouterTestExt, TestEnvironment};
@@ -2252,4 +2262,40 @@ mod tests {
 
         Ok(())
     }
+
+    #[test]
+    #[ignore]
+    fn test_build_with_cpu_limit() -> Result<()> {
+        let mut config = Config::test_config()?;
+        config.build_cpu_limit = Some(2);
+        let env = TestEnvironment::builder().config(config).build()?;
+
+        let mut builder = env.build_builder()?;
+        builder.update_toolchain()?;
+        assert!(
+            builder
+                .build_local_package(Path::new("tests/crates/hello-world"))?
+                .successful
+        );
+
+        Ok(())
+    }
+
+    #[test]
+    #[ignore]
+    fn test_build_with_cpu_cores() -> Result<()> {
+        let mut config = Config::test_config()?;
+        config.build_cpu_cores = Some(BuildCores(1..=2));
+        let env = TestEnvironment::builder().config(config).build()?;
+
+        let mut builder = env.build_builder()?;
+        builder.update_toolchain()?;
+        assert!(
+            builder
+                .build_local_package(Path::new("tests/crates/hello-world"))?
+                .successful
+        );
+
+        Ok(())
+    }
 }