From cb094700631ea1ae238ea678c192ce4f85fbecc0 Mon Sep 17 00:00:00 2001
From: Andrew Kelley <andrew@ziglang.org>
Date: Sun, 19 Feb 2023 16:21:36 -0700
Subject: [PATCH] zig build: add a -j<N> option for limiting concurrency

---
 lib/build_runner.zig    | 20 ++++++++++++++++++--
 lib/std/Thread/Pool.zig | 11 +++++++++--
 src/main.zig            |  4 ++--
 3 files changed, 29 insertions(+), 6 deletions(-)

diff --git a/lib/build_runner.zig b/lib/build_runner.zig
index 24fd3440a3..e12d46f03d 100644
--- a/lib/build_runner.zig
+++ b/lib/build_runner.zig
@@ -87,6 +87,7 @@ pub fn main() !void {
 
     var targets = ArrayList([]const u8).init(allocator);
     var debug_log_scopes = ArrayList([]const u8).init(allocator);
+    var thread_pool_options: std.Thread.Pool.Options = .{ .allocator = allocator };
 
     const stderr_stream = io.getStdErr().writer();
     const stdout_stream = io.getStdOut().writer();
@@ -231,6 +232,19 @@ pub fn main() !void {
                 };
             } else if (mem.eql(u8, arg, "-fno-reference-trace")) {
                 builder.reference_trace = null;
+            } else if (mem.startsWith(u8, arg, "-j")) {
+                const num = arg["-j".len..];
+                const n_jobs = std.fmt.parseUnsigned(u32, num, 10) catch |err| {
+                    std.debug.print("unable to parse jobs count '{s}': {s}", .{
+                        num, @errorName(err),
+                    });
+                    process.exit(1);
+                };
+                if (n_jobs < 1) {
+                    std.debug.print("number of jobs must be at least 1\n", .{});
+                    process.exit(1);
+                }
+                thread_pool_options.n_jobs = n_jobs;
             } else if (mem.eql(u8, arg, "--")) {
                 builder.args = argsRest(args, arg_idx);
                 break;
@@ -258,7 +272,7 @@ pub fn main() !void {
     if (builder.validateUserInputDidItFail())
         usageAndErr(builder, true, stderr_stream);
 
-    runStepNames(builder, targets.items, main_progress_node) catch |err| {
+    runStepNames(builder, targets.items, main_progress_node, thread_pool_options) catch |err| {
         switch (err) {
             error.UncleanExit => process.exit(1),
             else => return err,
@@ -270,6 +284,7 @@ fn runStepNames(
     b: *std.Build,
     step_names: []const []const u8,
     parent_prog_node: *std.Progress.Node,
+    thread_pool_options: std.Thread.Pool.Options,
 ) !void {
     var step_stack = ArrayList(*Step).init(b.allocator);
     defer step_stack.deinit();
@@ -297,7 +312,7 @@ fn runStepNames(
     }
 
     var thread_pool: std.Thread.Pool = undefined;
-    try thread_pool.init(b.allocator);
+    try thread_pool.init(thread_pool_options);
     defer thread_pool.deinit();
 
     {
@@ -523,6 +538,7 @@ fn usage(builder: *std.Build, already_ran_build: bool, out_stream: anytype) !voi
         \\  --verbose                    Print commands before executing them
         \\  --color [auto|off|on]        Enable or disable colored error messages
         \\  --prominent-compile-errors   Output compile errors formatted for a human to read
+        \\  -j<N>                        Limit concurrent jobs (default is to use all CPU cores)
         \\
         \\Project-Specific Options:
         \\
diff --git a/lib/std/Thread/Pool.zig b/lib/std/Thread/Pool.zig
index 930befbac5..ed1a4dc052 100644
--- a/lib/std/Thread/Pool.zig
+++ b/lib/std/Thread/Pool.zig
@@ -17,7 +17,14 @@ const Runnable = struct {
 
 const RunProto = *const fn (*Runnable) void;
 
-pub fn init(pool: *Pool, allocator: std.mem.Allocator) !void {
+pub const Options = struct {
+    allocator: std.mem.Allocator,
+    n_jobs: ?u32 = null,
+};
+
+pub fn init(pool: *Pool, options: Options) !void {
+    const allocator = options.allocator;
+
     pool.* = .{
         .allocator = allocator,
         .threads = &[_]std.Thread{},
@@ -27,7 +34,7 @@ pub fn init(pool: *Pool, allocator: std.mem.Allocator) !void {
         return;
     }
 
-    const thread_count = std.math.max(1, std.Thread.getCpuCount() catch 1);
+    const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1);
     pool.threads = try allocator.alloc(std.Thread, thread_count);
     errdefer allocator.free(pool.threads);
 
diff --git a/src/main.zig b/src/main.zig
index dd0faa628c..3a06d0272f 100644
--- a/src/main.zig
+++ b/src/main.zig
@@ -2999,7 +2999,7 @@ fn buildOutputType(
     defer zig_lib_directory.handle.close();
 
     var thread_pool: ThreadPool = undefined;
-    try thread_pool.init(gpa);
+    try thread_pool.init(.{ .allocator = gpa });
     defer thread_pool.deinit();
 
     var libc_installation: ?LibCInstallation = null;
@@ -4201,7 +4201,7 @@ pub fn cmdBuild(gpa: Allocator, arena: Allocator, args: []const []const u8) !voi
             .basename = exe_basename,
         };
         var thread_pool: ThreadPool = undefined;
-        try thread_pool.init(gpa);
+        try thread_pool.init(.{ .allocator = gpa });
         defer thread_pool.deinit();
 
         var cleanup_build_runner_dir: ?fs.Dir = null;