build-manifest: allow configuring the number of threads

3 files changed, 14 insertions, 7 deletions

diff --git a/Cargo.lock b/Cargo.lock
index bfc2d3e066a..83ac8396a16 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -243,6 +243,7 @@ dependencies = [
  "anyhow",
  "flate2",
  "hex 0.4.2",
+ "num_cpus",
  "rayon",
  "serde",
  "serde_json",
diff --git a/src/tools/build-manifest/Cargo.toml b/src/tools/build-manifest/Cargo.toml
index 4ae4dbfc06e..4a2c710811f 100644
--- a/src/tools/build-manifest/Cargo.toml
+++ b/src/tools/build-manifest/Cargo.toml
@@ -14,3 +14,4 @@ tar = "0.4.29"
 sha2 = "0.9.1"
 rayon = "1.3.1"
 hex = "0.4.2"
+num_cpus = "1.13.0"
diff --git a/src/tools/build-manifest/src/main.rs b/src/tools/build-manifest/src/main.rs
index 0efebe363f8..6fda9f4e59f 100644
--- a/src/tools/build-manifest/src/main.rs
+++ b/src/tools/build-manifest/src/main.rs
@@ -207,13 +207,18 @@ fn main() {
     // related code in this tool and ./x.py dist hash-and-sign can be removed.
     let legacy = env::var("BUILD_MANIFEST_LEGACY").is_ok();
 
-    // Avoid overloading the old server in legacy mode.
-    if legacy {
-        rayon::ThreadPoolBuilder::new()
-            .num_threads(1)
-            .build_global()
-            .expect("failed to initialize Rayon");
-    }
+    let num_threads = if legacy {
+        // Avoid overloading the old server in legacy mode.
+        1
+    } else if let Ok(num) = env::var("BUILD_MANIFEST_NUM_THREADS") {
+        num.parse().expect("invalid number for BUILD_MANIFEST_NUM_THREADS")
+    } else {
+        num_cpus::get()
+    };
+    rayon::ThreadPoolBuilder::new()
+        .num_threads(num_threads)
+        .build_global()
+        .expect("failed to initialize Rayon");
 
     let mut args = env::args().skip(1);
     let input = PathBuf::from(args.next().unwrap());