diff --git a/lib/compiler/build_runner.zig b/lib/compiler/build_runner.zig index 76fb9f22f8f3526ad4ce9848c8388450850fb1fb..ac0f0a370f3616264b7cd7c1ae802851455c41e5 100644 --- a/lib/compiler/build_runner.zig +++ b/lib/compiler/build_runner.zig @@ -19,13 +19,13 @@ pub fn main() !void { var single_threaded_arena = std.heap.ArenaAllocator.init(std.heap.page_allocator); defer single_threaded_arena.deinit(); + const args = try process.argsAlloc(single_threaded_arena.allocator()); + var thread_safe_arena: std.heap.ThreadSafeAllocator = .{ .child_allocator = single_threaded_arena.allocator(), }; const arena = thread_safe_arena.allocator(); - const args = try process.argsAlloc(arena); - // skip my own exe name var arg_idx: usize = 1; @@ -91,7 +91,9 @@ pub fn main() !void { var targets = ArrayList([]const u8).init(arena); var debug_log_scopes = ArrayList([]const u8).init(arena); - var thread_pool_options: std.Thread.Pool.Options = .{ .allocator = arena }; + var thread_pool_options: std.zig.ThreadPoolOptions = .{ + .cache_directory = local_cache_directory, + }; var install_prefix: ?[]const u8 = null; var dir_list = std.Build.DirList{}; @@ -387,7 +389,7 @@ fn runStepNames( b: *std.Build, step_names: []const []const u8, parent_prog_node: std.Progress.Node, - thread_pool_options: std.Thread.Pool.Options, + thread_pool_options: std.zig.ThreadPoolOptions, run: *Run, seed: u32, ) !void { @@ -446,8 +448,7 @@ fn runStepNames( } } - var thread_pool: std.Thread.Pool = undefined; - try thread_pool.init(thread_pool_options); + var thread_pool = try std.zig.initThreadPool(gpa, thread_pool_options); defer thread_pool.deinit(); { diff --git a/lib/std/Thread/Pool.zig b/lib/std/Thread/Pool.zig index 846c7035a75407fb54351605fe11847755d9ec8b..263a9aa96460cffc275563d9e8bc150c0eaab1e9 100644 --- a/lib/std/Thread/Pool.zig +++ b/lib/std/Thread/Pool.zig @@ -2,13 +2,18 @@ const std = @import("std"); const builtin = @import("builtin"); const Pool = @This(); const WaitGroup = @import("WaitGroup.zig"); +const assert = std.debug.assert; -mutex: std.Thread.Mutex = .{}, -cond: std.Thread.Condition = .{}, -run_queue: RunQueue = .{}, -is_running: bool = true, +mutex: std.Thread.Mutex, +cond: std.Thread.Condition, +run_queue: RunQueue, +run_queue_len: usize, +end_flag: bool, allocator: std.mem.Allocator, -threads: []std.Thread, +threads_buffer: []std.Thread, +threads_len: usize, +job_server_options: Options.JobServer, +job_server: ?*JobServer, const RunQueue = std.SinglyLinkedList(Runnable); const Runnable = struct { @@ -18,63 +23,187 @@ const Runnable = struct { const RunProto = *const fn (*Runnable) void; pub const Options = struct { - allocator: std.mem.Allocator, + /// Max number of threads to be actively working at the same time. + /// + /// `null` means to use the logical core count, leaving the main thread to + /// fill in the last slot. + /// + /// `0` is an illegal value. n_jobs: ?u32 = null, + + /// For coordinating amongst an entire process tree. + job_server: Options.JobServer = .abstain, + + pub const JobServer = union(enum) { + /// The thread pool neither hosts a jobserver nor connects to an existing one. + abstain, + /// The thread pool uses the Jobserver2 protocol to coordinate a global + /// thread pool across the entire process tree, avoiding cache + /// thrashing. + connect: std.net.Address, + /// The thread pool assumes the role of the root process and spawns a + /// dedicated thread for hosting the Jobserver2 protocol. + /// + /// Suggested to use a UNIX domain socket. + host: std.net.Address, + }; }; -pub fn init(pool: *Pool, options: Options) !void { - const allocator = options.allocator; - - pool.* = .{ +/// After initializing the thread pool and spawning work, the main thread must +/// call `waitAndWork`. +pub fn init( + /// Not required to be thread-safe; protected by the pool's mutex. + allocator: std.mem.Allocator, + options: Options, +) !Pool { + var pool: Pool = .{ + .mutex = .{}, + .cond = .{}, + .run_queue = .{}, + .run_queue_len = 0, + .end_flag = false, .allocator = allocator, - .threads = &[_]std.Thread{}, + .threads_buffer = &.{}, + .threads_len = 0, + .job_server_options = options.job_server, + .job_server = null, }; - if (builtin.single_threaded) { + if (builtin.single_threaded) return; - } const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1); + assert(thread_count > 0); - // kill and join any threads we spawned and free memory on error. - pool.threads = try allocator.alloc(std.Thread, thread_count); - var spawned: usize = 0; - errdefer pool.join(spawned); + pool.threads_buffer = try allocator.alloc(std.Thread, thread_count); + errdefer allocator.free(pool.threads_buffer); - for (pool.threads) |*thread| { - thread.* = try std.Thread.spawn(.{}, worker, .{pool}); - spawned += 1; + switch (options.job_server) { + .abstain, .connect => {}, + .host => |addr| { + var server = try addr.listen(.{}); + errdefer server.deinit(); + + const pollfds = try allocator.alloc(std.posix.pollfd, thread_count); + errdefer allocator.free(pollfds); + + const job_server = try allocator.create(JobServer); + errdefer allocator.destroy(job_server); + + job_server.* = .{ + .server = server, + .pollfds = pollfds, + .thread = try std.Thread.spawn(.{}, JobServer.run, .{job_server}), + }; + + pool.job_server = job_server; + }, } + + return pool; } pub fn deinit(pool: *Pool) void { - pool.join(pool.threads.len); // kill and join all threads. - pool.* = undefined; -} - -fn join(pool: *Pool, spawned: usize) void { - if (builtin.single_threaded) { + if (builtin.single_threaded) return; - } { pool.mutex.lock(); defer pool.mutex.unlock(); - // ensure future worker threads exit the dequeue loop - pool.is_running = false; + // Ensure future worker threads exit the dequeue loop. + pool.end_flag = true; } - // wake up any sleeping threads (this can be done outside the mutex) - // then wait for all the threads we know are spawned to complete. + // Wake up any sleeping threads (this can be done outside the mutex) then + // wait for all the threads we know are spawned to complete. pool.cond.broadcast(); - for (pool.threads[0..spawned]) |thread| { + + if (pool.job_server) |job_server| { + // Interrupt the jobserver thread from accepting connections. + // Since the server fd is also in the poll set, this handles both + // places where control flow could be blocked. + std.posix.shutdown(job_server.server.stream.handle, .both) catch {}; + job_server.thread.join(); + } + + // Since we set end_flag with the mutex locked, no more threads could have + // been created. + const threads = pool.threads_buffer[0..pool.threads_len]; + + for (threads) |thread| thread.join(); - } - pool.allocator.free(pool.threads); + pool.allocator.free(pool.threads_buffer); + pool.* = undefined; } +pub const JobServer = struct { + server: std.net.Server, + /// Has length n_jobs + 1. The first entry contains the server socket + /// itself, so that calling shutdown() in the other thread will both cause + /// the accept to return error.SocketNotListening and cause the poll() to + /// return. + pollfds: []std.posix.pollfd, + thread: std.Thread, + + pub fn run(js: *JobServer) void { + @memset(js.pollfds, .{ + .fd = -1, + // Only interested in errors and hangups. + .events = 0, + .revents = 0, + }); + + js.pollfds[0].fd = js.server.stream.handle; + + main_loop: while (true) { + for (js.pollfds[1..]) |*pollfd| { + const err_event = (pollfd.revents & std.posix.POLL.ERR) != 0; + const hup_event = (pollfd.revents & std.posix.POLL.HUP) != 0; + if (err_event or hup_event) { + std.posix.close(pollfd.fd); + pollfd.fd = -1; + pollfd.revents = 0; + } + + if (pollfd.fd >= 0) continue; + + const connection = js.server.accept() catch |err| switch (err) { + error.SocketNotListening => break :main_loop, // Indicates a shutdown request. + else => |e| { + std.log.debug("job server accept failure: {s}", .{@errorName(e)}); + continue; + }, + }; + _ = std.posix.send(connection.stream.handle, &.{0}, std.posix.MSG.NOSIGNAL) catch { + connection.stream.close(); + continue; + }; + pollfd.fd = connection.stream.handle; + } + + _ = std.posix.poll(js.pollfds, -1) catch continue; + } + + // Closes the active connections as well as the server itself. + for (js.pollfds) |pollfd| { + if (pollfd.fd >= 0) { + std.posix.close(pollfd.fd); + } + } + + // Delete the UNIX domain socket. + switch (js.server.listen_address.any.family) { + std.posix.AF.UNIX => { + const path = std.mem.sliceTo(&js.server.listen_address.un.path, 0); + std.fs.cwd().deleteFile(path) catch {}; + }, + else => {}, + } + } +}; + /// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and /// `WaitGroup.finish` after it returns. /// @@ -127,6 +256,22 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args }; pool.run_queue.prepend(&closure.run_node); + pool.run_queue_len += 1; + + // If there was already any queued work, spawn a new thread if we are + // under the max. + if (pool.run_queue_len > 1 and pool.threads_len < pool.threads_buffer.len) { + if (std.Thread.spawn(.{}, worker, .{pool})) |new_thread| { + pool.threads_buffer[pool.threads_len] = new_thread; + pool.threads_len += 1; + } else |_| if (pool.threads_len == 0) { + pool.mutex.unlock(); + @call(.auto, func, args); + wait_group.finish(); + return; + } + } + pool.mutex.unlock(); } @@ -134,7 +279,7 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args pool.cond.signal(); } -pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { +pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void { if (builtin.single_threaded) { @call(.auto, func, args); return; @@ -162,15 +307,34 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { { pool.mutex.lock(); - defer pool.mutex.unlock(); - const closure = try pool.allocator.create(Closure); + const closure = pool.allocator.create(Closure) catch { + pool.mutex.unlock(); + @call(.auto, func, args); + return; + }; closure.* = .{ .arguments = args, .pool = pool, }; pool.run_queue.prepend(&closure.run_node); + pool.run_queue_len += 1; + + // If there was already any queued work, spawn a new thread if we are + // under the max. + if (pool.run_queue_len > 1 and pool.threads_len < pool.threads_buffer.len) { + if (std.Thread.spawn(.{}, worker, .{pool})) |new_thread| { + pool.threads_buffer[pool.threads_len] = new_thread; + pool.threads_len += 1; + } else |_| if (pool.threads_len == 0) { + pool.mutex.unlock(); + @call(.auto, func, args); + return; + } + } + + pool.mutex.unlock(); } // Notify waiting threads outside the lock to try and keep the critical section small. @@ -178,25 +342,45 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { } fn worker(pool: *Pool) void { + var trash_buf: [1]u8 = undefined; + var connection: ?std.net.Stream = null; + defer if (connection) |stream| stream.close(); + pool.mutex.lock(); defer pool.mutex.unlock(); while (true) { while (pool.run_queue.popFirst()) |run_node| { - // Temporarily unlock the mutex in order to execute the run_node + pool.run_queue_len -= 1; + + // Temporarily unlock the mutex in order to execute the run_node. pool.mutex.unlock(); defer pool.mutex.lock(); + if (connection == null) switch (pool.job_server_options) { + .abstain => {}, + .connect, .host => |addr| { + if (std.net.tcpConnectToAddress(addr)) |stream| { + connection = stream; + _ = stream.readAll(&trash_buf) catch 1; + } else |_| {} + }, + }; + const runFn = run_node.data.runFn; runFn(&run_node.data); } // Stop executing instead of waiting if the thread pool is no longer running. - if (pool.is_running) { - pool.cond.wait(&pool.mutex); - } else { + if (pool.end_flag) break; + + if (connection) |stream| { + stream.close(); + connection = null; } + + pool.cond.wait(&pool.mutex); } } @@ -207,6 +391,7 @@ pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void { defer pool.mutex.unlock(); break :blk pool.run_queue.popFirst(); }) |run_node| { + pool.run_queue_len -= 1; run_node.data.runFn(&run_node.data); continue; } diff --git a/lib/std/zig.zig b/lib/std/zig.zig index 230dc45ddcbd450ae815176d4403843cfe49459f..b513c305eebd0dde1b988d1f2ab83543a2f2f712 100644 --- a/lib/std/zig.zig +++ b/lib/std/zig.zig @@ -689,6 +689,7 @@ pub const EnvVar = enum { CLICOLOR_FORCE, XDG_CACHE_HOME, HOME, + JOBSERVER2, pub fn isSet(comptime ev: EnvVar) bool { return std.process.hasEnvVarConstant(@tagName(ev)); @@ -708,6 +709,68 @@ pub const EnvVar = enum { } }; +pub const ThreadPoolOptions = struct { + n_jobs: ?u32 = null, + cache_directory: std.Build.Cache.Directory, +}; + +pub const cache_tmp_basename = "tmp"; + +pub fn initThreadPool(gpa: Allocator, options: ThreadPoolOptions) !std.Thread.Pool { + if (EnvVar.JOBSERVER2.getPosix()) |addr_string| { + return std.Thread.Pool.init(gpa, .{ + .n_jobs = options.n_jobs, + .job_server = .{ .connect = try std.net.Address.initUnix(addr_string) }, + }); + } + + const rand_int_string = hex64(std.crypto.random.int(u64)); + const suffix = "/" ++ cache_tmp_basename ++ "/" ++ rand_int_string; + + var addr: std.net.Address = .{ + .un = .{ + .family = std.posix.AF.UNIX, + .path = undefined, + }, + }; + + const cache_dir = options.cache_directory.path orelse "."; + + // Add 1 to ensure a terminating 0 is present in the path array for maximum portability. + if (cache_dir.len + suffix.len + 1 > addr.un.path.len) + return error.NameTooLong; + + @memset(&addr.un.path, 0); + @memcpy(addr.un.path[0..cache_dir.len], cache_dir); + @memcpy(addr.un.path[cache_dir.len..][0..suffix.len], suffix); + + return std.Thread.Pool.init(gpa, .{ + .n_jobs = options.n_jobs, + .job_server = .{ .host = addr }, + }) catch |err| switch (err) { + error.FileNotFound => { + try options.cache_directory.handle.makePath(cache_tmp_basename); + return std.Thread.Pool.init(gpa, .{ + .n_jobs = options.n_jobs, + .job_server = .{ .host = addr }, + }); + }, + else => |e| return e, + }; +} + +fn hex64(x: u64) [16]u8 { + const hex_charset = "0123456789abcdef"; + var result: [16]u8 = undefined; + var i: usize = 0; + while (i < 8) : (i += 1) { + const byte = @as(u8, @truncate(x >> @as(u6, @intCast(8 * i)))); + result[i * 2 + 0] = hex_charset[byte >> 4]; + result[i * 2 + 1] = hex_charset[byte & 15]; + } + return result; +} + test { _ = Ast; _ = AstRlAnnotate; diff --git a/src/main.zig b/src/main.zig index d9084686c885726d8f1b8dc6fcca0eb6021b0e1d..8b506366a3ea26f76c6059acbc8aa20ad685d0b0 100644 --- a/src/main.zig +++ b/src/main.zig @@ -10,7 +10,6 @@ const ArrayList = std.ArrayList; const Ast = std.zig.Ast; const Color = std.zig.Color; const warn = std.log.warn; -const ThreadPool = std.Thread.Pool; const cleanExit = std.process.cleanExit; const native_os = builtin.os.tag; @@ -3093,10 +3092,6 @@ fn buildOutputType( }; defer emit_implib_resolved.deinit(); - var thread_pool: ThreadPool = undefined; - try thread_pool.init(.{ .allocator = gpa }); - defer thread_pool.deinit(); - var cleanup_local_cache_dir: ?fs.Dir = null; defer if (cleanup_local_cache_dir) |*dir| dir.close(); @@ -3141,6 +3136,11 @@ fn buildOutputType( break :l global_cache_directory; }; + var thread_pool = try std.zig.initThreadPool(gpa, .{ + .cache_directory = local_cache_directory, + }); + defer thread_pool.deinit(); + for (create_module.c_source_files.items) |*src| { if (!mem.eql(u8, src.src_path, "-")) continue; @@ -4896,8 +4896,9 @@ fn cmdBuild(gpa: Allocator, arena: Allocator, args: []const []const u8) !void { child_argv.items[argv_index_cache_dir] = local_cache_directory.path orelse cwd_path; - var thread_pool: ThreadPool = undefined; - try thread_pool.init(.{ .allocator = gpa }); + var thread_pool = try std.zig.initThreadPool(gpa, .{ + .cache_directory = local_cache_directory, + }); defer thread_pool.deinit(); // Dummy http client that is not actually used when only_core_functionality is enabled. @@ -5330,8 +5331,9 @@ fn jitCmd( }; defer global_cache_directory.handle.close(); - var thread_pool: ThreadPool = undefined; - try thread_pool.init(.{ .allocator = gpa }); + var thread_pool = try std.zig.initThreadPool(gpa, .{ + .cache_directory = global_cache_directory, + }); defer thread_pool.deinit(); var child_argv: std.ArrayListUnmanaged([]const u8) = .{}; @@ -6876,10 +6878,6 @@ fn cmdFetch( const path_or_url = opt_path_or_url orelse fatal("missing url or path parameter", .{}); - var thread_pool: ThreadPool = undefined; - try thread_pool.init(.{ .allocator = gpa }); - defer thread_pool.deinit(); - var http_client: std.http.Client = .{ .allocator = gpa }; defer http_client.deinit(); @@ -6899,6 +6897,11 @@ fn cmdFetch( }; defer global_cache_directory.handle.close(); + var thread_pool = try std.zig.initThreadPool(gpa, .{ + .cache_directory = global_cache_directory, + }); + defer thread_pool.deinit(); + var job_queue: Package.Fetch.JobQueue = .{ .http_client = &http_client, .thread_pool = &thread_pool,