authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-06-20 18:32:15-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-06-30 13:54:02-07:00
logc395df25aba0f1bdc1dd0cb6b9c7f14a90e72dae
treec79aa15fa1460b29b92986713991ab7242c38c38
parentcb308ba3ac2d7e3735d1cb42ef085edb1e6db723

std.Thread.Pool: implement jobserverv2 protocol

The host accepts N simultaneous connections and writes 1 byte to them each. Clients connect and read 1 byte in order to obtain a thread token. std.Thread.Pool now lazily spawns threads only when the work queue is non-empty. I think that was a bad idea and will revert it shortly. There is now a std.zig.initThreadPool wrapper that deals with: * Resolving a zig cache directory into a UNIX domain socket address. * Creating the "tmp" directory in .zig-cache but only if the listen failed due to ENOENT. * Deciding to connect to an existing jobserver, or become the host for child processes.

4 files changed, 312 insertions(+), 60 deletions(-)

lib/compiler/build_runner.zig+7-6
......@@ -19,13 +19,13 @@ pub fn main() !void {
1919 var single_threaded_arena = std.heap.ArenaAllocator.init(std.heap.page_allocator);
2020 defer single_threaded_arena.deinit();
2121
22 const args = try process.argsAlloc(single_threaded_arena.allocator());
23
2224 var thread_safe_arena: std.heap.ThreadSafeAllocator = .{
2325 .child_allocator = single_threaded_arena.allocator(),
2426 };
2527 const arena = thread_safe_arena.allocator();
2628
27 const args = try process.argsAlloc(arena);
28
2929 // skip my own exe name
3030 var arg_idx: usize = 1;
3131
......@@ -91,7 +91,9 @@ pub fn main() !void {
9191
9292 var targets = ArrayList([]const u8).init(arena);
9393 var debug_log_scopes = ArrayList([]const u8).init(arena);
94 var thread_pool_options: std.Thread.Pool.Options = .{ .allocator = arena };
94 var thread_pool_options: std.zig.ThreadPoolOptions = .{
95 .cache_directory = local_cache_directory,
96 };
9597
9698 var install_prefix: ?[]const u8 = null;
9799 var dir_list = std.Build.DirList{};
......@@ -387,7 +389,7 @@ fn runStepNames(
387389 b: *std.Build,
388390 step_names: []const []const u8,
389391 parent_prog_node: std.Progress.Node,
390 thread_pool_options: std.Thread.Pool.Options,
392 thread_pool_options: std.zig.ThreadPoolOptions,
391393 run: *Run,
392394 seed: u32,
393395) !void {
......@@ -446,8 +448,7 @@ fn runStepNames(
446448 }
447449 }
448450
449 var thread_pool: std.Thread.Pool = undefined;
450 try thread_pool.init(thread_pool_options);
451 var thread_pool = try std.zig.initThreadPool(gpa, thread_pool_options);
451452 defer thread_pool.deinit();
452453
453454 {
lib/std/Thread/Pool.zig+226-41
......@@ -2,13 +2,18 @@ const std = @import("std");
22const builtin = @import("builtin");
33const Pool = @This();
44const WaitGroup = @import("WaitGroup.zig");
5const assert = std.debug.assert;
56
6mutex: std.Thread.Mutex = .{},
7cond: std.Thread.Condition = .{},
8run_queue: RunQueue = .{},
9is_running: bool = true,
7mutex: std.Thread.Mutex,
8cond: std.Thread.Condition,
9run_queue: RunQueue,
10run_queue_len: usize,
11end_flag: bool,
1012allocator: std.mem.Allocator,
11threads: []std.Thread,
13threads_buffer: []std.Thread,
14threads_len: usize,
15job_server_options: Options.JobServer,
16job_server: ?*JobServer,
1217
1318const RunQueue = std.SinglyLinkedList(Runnable);
1419const Runnable = struct {
......@@ -18,63 +23,187 @@ const Runnable = struct {
1823const RunProto = *const fn (*Runnable) void;
1924
2025pub const Options = struct {
21 allocator: std.mem.Allocator,
26 /// Max number of threads to be actively working at the same time.
27 ///
28 /// `null` means to use the logical core count, leaving the main thread to
29 /// fill in the last slot.
30 ///
31 /// `0` is an illegal value.
2232 n_jobs: ?u32 = null,
23};
2433
25pub fn init(pool: *Pool, options: Options) !void {
26 const allocator = options.allocator;
34 /// For coordinating amongst an entire process tree.
35 job_server: Options.JobServer = .abstain,
36
37 pub const JobServer = union(enum) {
38 /// The thread pool neither hosts a jobserver nor connects to an existing one.
39 abstain,
40 /// The thread pool uses the Jobserver2 protocol to coordinate a global
41 /// thread pool across the entire process tree, avoiding cache
42 /// thrashing.
43 connect: std.net.Address,
44 /// The thread pool assumes the role of the root process and spawns a
45 /// dedicated thread for hosting the Jobserver2 protocol.
46 ///
47 /// Suggested to use a UNIX domain socket.
48 host: std.net.Address,
49 };
50};
2751
28 pool.* = .{
52/// After initializing the thread pool and spawning work, the main thread must
53/// call `waitAndWork`.
54pub fn init(
55 /// Not required to be thread-safe; protected by the pool's mutex.
56 allocator: std.mem.Allocator,
57 options: Options,
58) !Pool {
59 var pool: Pool = .{
60 .mutex = .{},
61 .cond = .{},
62 .run_queue = .{},
63 .run_queue_len = 0,
64 .end_flag = false,
2965 .allocator = allocator,
30 .threads = &[_]std.Thread{},
66 .threads_buffer = &.{},
67 .threads_len = 0,
68 .job_server_options = options.job_server,
69 .job_server = null,
3170 };
3271
33 if (builtin.single_threaded) {
72 if (builtin.single_threaded)
3473 return;
35 }
3674
3775 const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1);
76 assert(thread_count > 0);
77
78 pool.threads_buffer = try allocator.alloc(std.Thread, thread_count);
79 errdefer allocator.free(pool.threads_buffer);
80
81 switch (options.job_server) {
82 .abstain, .connect => {},
83 .host => |addr| {
84 var server = try addr.listen(.{});
85 errdefer server.deinit();
3886
39 // kill and join any threads we spawned and free memory on error.
40 pool.threads = try allocator.alloc(std.Thread, thread_count);
41 var spawned: usize = 0;
42 errdefer pool.join(spawned);
87 const pollfds = try allocator.alloc(std.posix.pollfd, thread_count);
88 errdefer allocator.free(pollfds);
4389
44 for (pool.threads) |*thread| {
45 thread.* = try std.Thread.spawn(.{}, worker, .{pool});
46 spawned += 1;
90 const job_server = try allocator.create(JobServer);
91 errdefer allocator.destroy(job_server);
92
93 job_server.* = .{
94 .server = server,
95 .pollfds = pollfds,
96 .thread = try std.Thread.spawn(.{}, JobServer.run, .{job_server}),
97 };
98
99 pool.job_server = job_server;
100 },
47101 }
48}
49102
50pub fn deinit(pool: *Pool) void {
51 pool.join(pool.threads.len); // kill and join all threads.
52 pool.* = undefined;
103 return pool;
53104}
54105
55fn join(pool: *Pool, spawned: usize) void {
56 if (builtin.single_threaded) {
106pub fn deinit(pool: *Pool) void {
107 if (builtin.single_threaded)
57108 return;
58 }
59109
60110 {
61111 pool.mutex.lock();
62112 defer pool.mutex.unlock();
63113
64 // ensure future worker threads exit the dequeue loop
65 pool.is_running = false;
114 // Ensure future worker threads exit the dequeue loop.
115 pool.end_flag = true;
66116 }
67117
68 // wake up any sleeping threads (this can be done outside the mutex)
69 // then wait for all the threads we know are spawned to complete.
118 // Wake up any sleeping threads (this can be done outside the mutex) then
119 // wait for all the threads we know are spawned to complete.
70120 pool.cond.broadcast();
71 for (pool.threads[0..spawned]) |thread| {
72 thread.join();
121
122 if (pool.job_server) |job_server| {
123 // Interrupt the jobserver thread from accepting connections.
124 // Since the server fd is also in the poll set, this handles both
125 // places where control flow could be blocked.
126 std.posix.shutdown(job_server.server.stream.handle, .both) catch {};
127 job_server.thread.join();
73128 }
74129
75 pool.allocator.free(pool.threads);
130 // Since we set end_flag with the mutex locked, no more threads could have
131 // been created.
132 const threads = pool.threads_buffer[0..pool.threads_len];
133
134 for (threads) |thread|
135 thread.join();
136
137 pool.allocator.free(pool.threads_buffer);
138 pool.* = undefined;
76139}
77140
141pub const JobServer = struct {
142 server: std.net.Server,
143 /// Has length n_jobs + 1. The first entry contains the server socket
144 /// itself, so that calling shutdown() in the other thread will both cause
145 /// the accept to return error.SocketNotListening and cause the poll() to
146 /// return.
147 pollfds: []std.posix.pollfd,
148 thread: std.Thread,
149
150 pub fn run(js: *JobServer) void {
151 @memset(js.pollfds, .{
152 .fd = -1,
153 // Only interested in errors and hangups.
154 .events = 0,
155 .revents = 0,
156 });
157
158 js.pollfds[0].fd = js.server.stream.handle;
159
160 main_loop: while (true) {
161 for (js.pollfds[1..]) |*pollfd| {
162 const err_event = (pollfd.revents & std.posix.POLL.ERR) != 0;
163 const hup_event = (pollfd.revents & std.posix.POLL.HUP) != 0;
164 if (err_event or hup_event) {
165 std.posix.close(pollfd.fd);
166 pollfd.fd = -1;
167 pollfd.revents = 0;
168 }
169
170 if (pollfd.fd >= 0) continue;
171
172 const connection = js.server.accept() catch |err| switch (err) {
173 error.SocketNotListening => break :main_loop, // Indicates a shutdown request.
174 else => |e| {
175 std.log.debug("job server accept failure: {s}", .{@errorName(e)});
176 continue;
177 },
178 };
179 _ = std.posix.send(connection.stream.handle, &.{0}, std.posix.MSG.NOSIGNAL) catch {
180 connection.stream.close();
181 continue;
182 };
183 pollfd.fd = connection.stream.handle;
184 }
185
186 _ = std.posix.poll(js.pollfds, -1) catch continue;
187 }
188
189 // Closes the active connections as well as the server itself.
190 for (js.pollfds) |pollfd| {
191 if (pollfd.fd >= 0) {
192 std.posix.close(pollfd.fd);
193 }
194 }
195
196 // Delete the UNIX domain socket.
197 switch (js.server.listen_address.any.family) {
198 std.posix.AF.UNIX => {
199 const path = std.mem.sliceTo(&js.server.listen_address.un.path, 0);
200 std.fs.cwd().deleteFile(path) catch {};
201 },
202 else => {},
203 }
204 }
205};
206
78207/// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and
79208/// `WaitGroup.finish` after it returns.
80209///
......@@ -127,6 +256,22 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args
127256 };
128257
129258 pool.run_queue.prepend(&closure.run_node);
259 pool.run_queue_len += 1;
260
261 // If there was already any queued work, spawn a new thread if we are
262 // under the max.
263 if (pool.run_queue_len > 1 and pool.threads_len < pool.threads_buffer.len) {
264 if (std.Thread.spawn(.{}, worker, .{pool})) |new_thread| {
265 pool.threads_buffer[pool.threads_len] = new_thread;
266 pool.threads_len += 1;
267 } else |_| if (pool.threads_len == 0) {
268 pool.mutex.unlock();
269 @call(.auto, func, args);
270 wait_group.finish();
271 return;
272 }
273 }
274
130275 pool.mutex.unlock();
131276 }
132277
......@@ -134,7 +279,7 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args
134279 pool.cond.signal();
135280}
136281
137pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void {
282pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void {
138283 if (builtin.single_threaded) {
139284 @call(.auto, func, args);
140285 return;
......@@ -162,15 +307,34 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void {
162307
163308 {
164309 pool.mutex.lock();
165 defer pool.mutex.unlock();
166310
167 const closure = try pool.allocator.create(Closure);
311 const closure = pool.allocator.create(Closure) catch {
312 pool.mutex.unlock();
313 @call(.auto, func, args);
314 return;
315 };
168316 closure.* = .{
169317 .arguments = args,
170318 .pool = pool,
171319 };
172320
173321 pool.run_queue.prepend(&closure.run_node);
322 pool.run_queue_len += 1;
323
324 // If there was already any queued work, spawn a new thread if we are
325 // under the max.
326 if (pool.run_queue_len > 1 and pool.threads_len < pool.threads_buffer.len) {
327 if (std.Thread.spawn(.{}, worker, .{pool})) |new_thread| {
328 pool.threads_buffer[pool.threads_len] = new_thread;
329 pool.threads_len += 1;
330 } else |_| if (pool.threads_len == 0) {
331 pool.mutex.unlock();
332 @call(.auto, func, args);
333 return;
334 }
335 }
336
337 pool.mutex.unlock();
174338 }
175339
176340 // 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 {
178342}
179343
180344fn worker(pool: *Pool) void {
345 var trash_buf: [1]u8 = undefined;
346 var connection: ?std.net.Stream = null;
347 defer if (connection) |stream| stream.close();
348
181349 pool.mutex.lock();
182350 defer pool.mutex.unlock();
183351
184352 while (true) {
185353 while (pool.run_queue.popFirst()) |run_node| {
186 // Temporarily unlock the mutex in order to execute the run_node
354 pool.run_queue_len -= 1;
355
356 // Temporarily unlock the mutex in order to execute the run_node.
187357 pool.mutex.unlock();
188358 defer pool.mutex.lock();
189359
360 if (connection == null) switch (pool.job_server_options) {
361 .abstain => {},
362 .connect, .host => |addr| {
363 if (std.net.tcpConnectToAddress(addr)) |stream| {
364 connection = stream;
365 _ = stream.readAll(&trash_buf) catch 1;
366 } else |_| {}
367 },
368 };
369
190370 const runFn = run_node.data.runFn;
191371 runFn(&run_node.data);
192372 }
193373
194374 // Stop executing instead of waiting if the thread pool is no longer running.
195 if (pool.is_running) {
196 pool.cond.wait(&pool.mutex);
197 } else {
375 if (pool.end_flag)
198376 break;
377
378 if (connection) |stream| {
379 stream.close();
380 connection = null;
199381 }
382
383 pool.cond.wait(&pool.mutex);
200384 }
201385}
202386
......@@ -207,6 +391,7 @@ pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void {
207391 defer pool.mutex.unlock();
208392 break :blk pool.run_queue.popFirst();
209393 }) |run_node| {
394 pool.run_queue_len -= 1;
210395 run_node.data.runFn(&run_node.data);
211396 continue;
212397 }
lib/std/zig.zig+63
......@@ -689,6 +689,7 @@ pub const EnvVar = enum {
689689 CLICOLOR_FORCE,
690690 XDG_CACHE_HOME,
691691 HOME,
692 JOBSERVER2,
692693
693694 pub fn isSet(comptime ev: EnvVar) bool {
694695 return std.process.hasEnvVarConstant(@tagName(ev));
......@@ -708,6 +709,68 @@ pub const EnvVar = enum {
708709 }
709710};
710711
712pub const ThreadPoolOptions = struct {
713 n_jobs: ?u32 = null,
714 cache_directory: std.Build.Cache.Directory,
715};
716
717pub const cache_tmp_basename = "tmp";
718
719pub fn initThreadPool(gpa: Allocator, options: ThreadPoolOptions) !std.Thread.Pool {
720 if (EnvVar.JOBSERVER2.getPosix()) |addr_string| {
721 return std.Thread.Pool.init(gpa, .{
722 .n_jobs = options.n_jobs,
723 .job_server = .{ .connect = try std.net.Address.initUnix(addr_string) },
724 });
725 }
726
727 const rand_int_string = hex64(std.crypto.random.int(u64));
728 const suffix = "/" ++ cache_tmp_basename ++ "/" ++ rand_int_string;
729
730 var addr: std.net.Address = .{
731 .un = .{
732 .family = std.posix.AF.UNIX,
733 .path = undefined,
734 },
735 };
736
737 const cache_dir = options.cache_directory.path orelse ".";
738
739 // Add 1 to ensure a terminating 0 is present in the path array for maximum portability.
740 if (cache_dir.len + suffix.len + 1 > addr.un.path.len)
741 return error.NameTooLong;
742
743 @memset(&addr.un.path, 0);
744 @memcpy(addr.un.path[0..cache_dir.len], cache_dir);
745 @memcpy(addr.un.path[cache_dir.len..][0..suffix.len], suffix);
746
747 return std.Thread.Pool.init(gpa, .{
748 .n_jobs = options.n_jobs,
749 .job_server = .{ .host = addr },
750 }) catch |err| switch (err) {
751 error.FileNotFound => {
752 try options.cache_directory.handle.makePath(cache_tmp_basename);
753 return std.Thread.Pool.init(gpa, .{
754 .n_jobs = options.n_jobs,
755 .job_server = .{ .host = addr },
756 });
757 },
758 else => |e| return e,
759 };
760}
761
762fn hex64(x: u64) [16]u8 {
763 const hex_charset = "0123456789abcdef";
764 var result: [16]u8 = undefined;
765 var i: usize = 0;
766 while (i < 8) : (i += 1) {
767 const byte = @as(u8, @truncate(x >> @as(u6, @intCast(8 * i))));
768 result[i * 2 + 0] = hex_charset[byte >> 4];
769 result[i * 2 + 1] = hex_charset[byte & 15];
770 }
771 return result;
772}
773
711774test {
712775 _ = Ast;
713776 _ = AstRlAnnotate;
src/main.zig+16-13
......@@ -10,7 +10,6 @@ const ArrayList = std.ArrayList;
1010const Ast = std.zig.Ast;
1111const Color = std.zig.Color;
1212const warn = std.log.warn;
13const ThreadPool = std.Thread.Pool;
1413const cleanExit = std.process.cleanExit;
1514const native_os = builtin.os.tag;
1615
......@@ -3093,10 +3092,6 @@ fn buildOutputType(
30933092 };
30943093 defer emit_implib_resolved.deinit();
30953094
3096 var thread_pool: ThreadPool = undefined;
3097 try thread_pool.init(.{ .allocator = gpa });
3098 defer thread_pool.deinit();
3099
31003095 var cleanup_local_cache_dir: ?fs.Dir = null;
31013096 defer if (cleanup_local_cache_dir) |*dir| dir.close();
31023097
......@@ -3141,6 +3136,11 @@ fn buildOutputType(
31413136 break :l global_cache_directory;
31423137 };
31433138
3139 var thread_pool = try std.zig.initThreadPool(gpa, .{
3140 .cache_directory = local_cache_directory,
3141 });
3142 defer thread_pool.deinit();
3143
31443144 for (create_module.c_source_files.items) |*src| {
31453145 if (!mem.eql(u8, src.src_path, "-")) continue;
31463146
......@@ -4896,8 +4896,9 @@ fn cmdBuild(gpa: Allocator, arena: Allocator, args: []const []const u8) !void {
48964896
48974897 child_argv.items[argv_index_cache_dir] = local_cache_directory.path orelse cwd_path;
48984898
4899 var thread_pool: ThreadPool = undefined;
4900 try thread_pool.init(.{ .allocator = gpa });
4899 var thread_pool = try std.zig.initThreadPool(gpa, .{
4900 .cache_directory = local_cache_directory,
4901 });
49014902 defer thread_pool.deinit();
49024903
49034904 // Dummy http client that is not actually used when only_core_functionality is enabled.
......@@ -5330,8 +5331,9 @@ fn jitCmd(
53305331 };
53315332 defer global_cache_directory.handle.close();
53325333
5333 var thread_pool: ThreadPool = undefined;
5334 try thread_pool.init(.{ .allocator = gpa });
5334 var thread_pool = try std.zig.initThreadPool(gpa, .{
5335 .cache_directory = global_cache_directory,
5336 });
53355337 defer thread_pool.deinit();
53365338
53375339 var child_argv: std.ArrayListUnmanaged([]const u8) = .{};
......@@ -6876,10 +6878,6 @@ fn cmdFetch(
68766878
68776879 const path_or_url = opt_path_or_url orelse fatal("missing url or path parameter", .{});
68786880
6879 var thread_pool: ThreadPool = undefined;
6880 try thread_pool.init(.{ .allocator = gpa });
6881 defer thread_pool.deinit();
6882
68836881 var http_client: std.http.Client = .{ .allocator = gpa };
68846882 defer http_client.deinit();
68856883
......@@ -6899,6 +6897,11 @@ fn cmdFetch(
68996897 };
69006898 defer global_cache_directory.handle.close();
69016899
6900 var thread_pool = try std.zig.initThreadPool(gpa, .{
6901 .cache_directory = global_cache_directory,
6902 });
6903 defer thread_pool.deinit();
6904
69026905 var job_queue: Package.Fetch.JobQueue = .{
69036906 .http_client = &http_client,
69046907 .thread_pool = &thread_pool,