authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-05-03 17:57:48-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-05-03 20:58:02-07:00
loga72292513e378159c542a785ca69f8111bacbdcf
treeaac9e20a88f90e9117b338bcc0e925c7ffa90182
parenta96b78c170ef0464e51a1c2fa226c51d49cfde04

add std.Thread.Pool.spawnWg

This function accepts a WaitGroup parameter and manages the reference counting therein. It also is infallible. The existing `spawn` function is still handy when the job wants to further schedule more tasks.

6 files changed, 88 insertions(+), 86 deletions(-)

lib/compiler/build_runner.zig+6-11
......@@ -466,10 +466,9 @@ fn runStepNames(
466466 const step = steps_slice[steps_slice.len - i - 1];
467467 if (step.state == .skipped_oom) continue;
468468
469 wait_group.start();
470 thread_pool.spawn(workerMakeOneStep, .{
469 thread_pool.spawnWg(&wait_group, workerMakeOneStep, .{
471470 &wait_group, &thread_pool, b, step, &step_prog, run,
472 }) catch @panic("OOM");
471 });
473472 }
474473 }
475474 assert(run.memory_blocked_steps.items.len == 0);
......@@ -895,8 +894,6 @@ fn workerMakeOneStep(
895894 prog_node: *std.Progress.Node,
896895 run: *Run,
897896) void {
898 defer wg.finish();
899
900897 // First, check the conditions for running this step. If they are not met,
901898 // then we return without doing the step, relying on another worker to
902899 // queue this step up again when dependencies are met.
......@@ -976,10 +973,9 @@ fn workerMakeOneStep(
976973
977974 // Successful completion of a step, so we queue up its dependants as well.
978975 for (s.dependants.items) |dep| {
979 wg.start();
980 thread_pool.spawn(workerMakeOneStep, .{
976 thread_pool.spawnWg(wg, workerMakeOneStep, .{
981977 wg, thread_pool, b, dep, prog_node, run,
982 }) catch @panic("OOM");
978 });
983979 }
984980 }
985981
......@@ -1002,10 +998,9 @@ fn workerMakeOneStep(
1002998 if (dep.max_rss <= remaining) {
1003999 remaining -= dep.max_rss;
10041000
1005 wg.start();
1006 thread_pool.spawn(workerMakeOneStep, .{
1001 thread_pool.spawnWg(wg, workerMakeOneStep, .{
10071002 wg, thread_pool, b, dep, prog_node, run,
1008 }) catch @panic("OOM");
1003 });
10091004 } else {
10101005 run.memory_blocked_steps.items[i] = dep;
10111006 i += 1;
lib/std/Thread/Pool.zig+59
......@@ -75,6 +75,65 @@ fn join(pool: *Pool, spawned: usize) void {
7575 pool.allocator.free(pool.threads);
7676}
7777
78/// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and
79/// `WaitGroup.finish` after it returns.
80///
81/// In the case that queuing the function call fails to allocate memory, or the
82/// target is single-threaded, the function is called directly.
83pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args: anytype) void {
84 wait_group.start();
85
86 if (builtin.single_threaded) {
87 @call(.auto, func, args);
88 wait_group.finish();
89 return;
90 }
91
92 const Args = @TypeOf(args);
93 const Closure = struct {
94 arguments: Args,
95 pool: *Pool,
96 run_node: RunQueue.Node = .{ .data = .{ .runFn = runFn } },
97 wait_group: *WaitGroup,
98
99 fn runFn(runnable: *Runnable) void {
100 const run_node: *RunQueue.Node = @fieldParentPtr("data", runnable);
101 const closure: *@This() = @alignCast(@fieldParentPtr("run_node", run_node));
102 @call(.auto, func, closure.arguments);
103 closure.wait_group.finish();
104
105 // The thread pool's allocator is protected by the mutex.
106 const mutex = &closure.pool.mutex;
107 mutex.lock();
108 defer mutex.unlock();
109
110 closure.pool.allocator.destroy(closure);
111 }
112 };
113
114 {
115 pool.mutex.lock();
116
117 const closure = pool.allocator.create(Closure) catch {
118 pool.mutex.unlock();
119 @call(.auto, func, args);
120 wait_group.finish();
121 return;
122 };
123 closure.* = .{
124 .arguments = args,
125 .pool = pool,
126 .wait_group = wait_group,
127 };
128
129 pool.run_queue.prepend(&closure.run_node);
130 pool.mutex.unlock();
131 }
132
133 // Notify waiting threads outside the lock to try and keep the critical section small.
134 pool.cond.signal();
135}
136
78137pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void {
79138 if (builtin.single_threaded) {
80139 @call(.auto, func, args);
src/Compilation.zig+14-46
......@@ -3273,7 +3273,7 @@ pub fn performAllTheWork(
32733273
32743274 if (!build_options.only_c and !build_options.only_core_functionality) {
32753275 if (comp.docs_emit != null) {
3276 try taskDocsCopy(comp, &comp.work_queue_wait_group);
3276 comp.thread_pool.spawnWg(&comp.work_queue_wait_group, workerDocsCopy, .{comp});
32773277 comp.work_queue_wait_group.spawnManager(workerDocsWasm, .{ comp, &wasm_prog_node });
32783278 }
32793279 }
......@@ -3305,39 +3305,34 @@ pub fn performAllTheWork(
33053305
33063306 const file = mod.builtin_file orelse continue;
33073307
3308 comp.astgen_wait_group.start();
3309 try comp.thread_pool.spawn(workerUpdateBuiltinZigFile, .{
3310 comp, mod, file, &comp.astgen_wait_group,
3308 comp.thread_pool.spawnWg(&comp.astgen_wait_group, workerUpdateBuiltinZigFile, .{
3309 comp, mod, file,
33113310 });
33123311 }
33133312 }
33143313
33153314 while (comp.astgen_work_queue.readItem()) |file| {
3316 comp.astgen_wait_group.start();
3317 try comp.thread_pool.spawn(workerAstGenFile, .{
3315 comp.thread_pool.spawnWg(&comp.astgen_wait_group, workerAstGenFile, .{
33183316 comp, file, &zir_prog_node, &comp.astgen_wait_group, .root,
33193317 });
33203318 }
33213319
33223320 while (comp.embed_file_work_queue.readItem()) |embed_file| {
3323 comp.astgen_wait_group.start();
3324 try comp.thread_pool.spawn(workerCheckEmbedFile, .{
3325 comp, embed_file, &comp.astgen_wait_group,
3321 comp.thread_pool.spawnWg(&comp.astgen_wait_group, workerCheckEmbedFile, .{
3322 comp, embed_file,
33263323 });
33273324 }
33283325
33293326 while (comp.c_object_work_queue.readItem()) |c_object| {
3330 comp.work_queue_wait_group.start();
3331 try comp.thread_pool.spawn(workerUpdateCObject, .{
3332 comp, c_object, &c_obj_prog_node, &comp.work_queue_wait_group,
3327 comp.thread_pool.spawnWg(&comp.work_queue_wait_group, workerUpdateCObject, .{
3328 comp, c_object, &c_obj_prog_node,
33333329 });
33343330 }
33353331
33363332 if (!build_options.only_core_functionality) {
33373333 while (comp.win32_resource_work_queue.readItem()) |win32_resource| {
3338 comp.work_queue_wait_group.start();
3339 try comp.thread_pool.spawn(workerUpdateWin32Resource, .{
3340 comp, win32_resource, &win32_resource_prog_node, &comp.work_queue_wait_group,
3334 comp.thread_pool.spawnWg(&comp.work_queue_wait_group, workerUpdateWin32Resource, .{
3335 comp, win32_resource, &win32_resource_prog_node,
33413336 });
33423337 }
33433338 }
......@@ -3680,14 +3675,7 @@ fn processOneJob(comp: *Compilation, job: Job, prog_node: *std.Progress.Node) !v
36803675 }
36813676}
36823677
3683fn taskDocsCopy(comp: *Compilation, wg: *WaitGroup) !void {
3684 wg.start();
3685 errdefer wg.finish();
3686 try comp.thread_pool.spawn(workerDocsCopy, .{ comp, wg });
3687}
3688
3689fn workerDocsCopy(comp: *Compilation, wg: *WaitGroup) void {
3690 defer wg.finish();
3678fn workerDocsCopy(comp: *Compilation) void {
36913679 docsCopyFallible(comp) catch |err| {
36923680 return comp.lockAndSetMiscFailure(
36933681 .docs_copy,
......@@ -3965,8 +3953,6 @@ fn workerAstGenFile(
39653953 wg: *WaitGroup,
39663954 src: AstGenSrc,
39673955) void {
3968 defer wg.finish();
3969
39703956 var child_prog_node = prog_node.start(file.sub_file_path, 0);
39713957 child_prog_node.activate();
39723958 defer child_prog_node.end();
......@@ -4025,13 +4011,9 @@ fn workerAstGenFile(
40254011 .importing_file = file,
40264012 .import_tok = item.data.token,
40274013 } };
4028 wg.start();
4029 comp.thread_pool.spawn(workerAstGenFile, .{
4014 comp.thread_pool.spawnWg(wg, workerAstGenFile, .{
40304015 comp, import_result.file, prog_node, wg, sub_src,
4031 }) catch {
4032 wg.finish();
4033 continue;
4034 };
4016 });
40354017 }
40364018 }
40374019 }
......@@ -4041,9 +4023,7 @@ fn workerUpdateBuiltinZigFile(
40414023 comp: *Compilation,
40424024 mod: *Package.Module,
40434025 file: *Module.File,
4044 wg: *WaitGroup,
40454026) void {
4046 defer wg.finish();
40474027 Builtin.populateFile(comp, mod, file) catch |err| {
40484028 comp.mutex.lock();
40494029 defer comp.mutex.unlock();
......@@ -4054,13 +4034,7 @@ fn workerUpdateBuiltinZigFile(
40544034 };
40554035}
40564036
4057fn workerCheckEmbedFile(
4058 comp: *Compilation,
4059 embed_file: *Module.EmbedFile,
4060 wg: *WaitGroup,
4061) void {
4062 defer wg.finish();
4063
4037fn workerCheckEmbedFile(comp: *Compilation, embed_file: *Module.EmbedFile) void {
40644038 comp.detectEmbedFileUpdate(embed_file) catch |err| {
40654039 comp.reportRetryableEmbedFileError(embed_file, err) catch |oom| switch (oom) {
40664040 // Swallowing this error is OK because it's implied to be OOM when
......@@ -4289,10 +4263,7 @@ fn workerUpdateCObject(
42894263 comp: *Compilation,
42904264 c_object: *CObject,
42914265 progress_node: *std.Progress.Node,
4292 wg: *WaitGroup,
42934266) void {
4294 defer wg.finish();
4295
42964267 comp.updateCObject(c_object, progress_node) catch |err| switch (err) {
42974268 error.AnalysisFail => return,
42984269 else => {
......@@ -4309,10 +4280,7 @@ fn workerUpdateWin32Resource(
43094280 comp: *Compilation,
43104281 win32_resource: *Win32Resource,
43114282 progress_node: *std.Progress.Node,
4312 wg: *WaitGroup,
43134283) void {
4314 defer wg.finish();
4315
43164284 comp.updateWin32Resource(win32_resource, progress_node) catch |err| switch (err) {
43174285 error.AnalysisFail => return,
43184286 else => {
src/Package/Fetch.zig+5-22
......@@ -722,14 +722,7 @@ fn queueJobsForDeps(f: *Fetch) RunError!void {
722722 const thread_pool = f.job_queue.thread_pool;
723723
724724 for (new_fetches, prog_names) |*new_fetch, prog_name| {
725 f.job_queue.wait_group.start();
726 thread_pool.spawn(workerRun, .{ new_fetch, prog_name }) catch |err| switch (err) {
727 error.OutOfMemory => {
728 new_fetch.oom_flag = true;
729 f.job_queue.wait_group.finish();
730 continue;
731 },
732 };
725 thread_pool.spawnWg(&f.job_queue.wait_group, workerRun, .{ new_fetch, prog_name });
733726 }
734727}
735728
......@@ -750,8 +743,6 @@ pub fn relativePathDigest(
750743}
751744
752745pub fn workerRun(f: *Fetch, prog_name: []const u8) void {
753 defer f.job_queue.wait_group.finish();
754
755746 var prog_node = f.prog_node.start(prog_name, 0);
756747 defer prog_node.end();
757748 prog_node.activate();
......@@ -1477,10 +1468,7 @@ fn computeHash(
14771468 .fs_path = fs_path,
14781469 .failure = undefined, // to be populated by the worker
14791470 };
1480 wait_group.start();
1481 try thread_pool.spawn(workerDeleteFile, .{
1482 root_dir, deleted_file, &wait_group,
1483 });
1471 thread_pool.spawnWg(&wait_group, workerDeleteFile, .{ root_dir, deleted_file });
14841472 try deleted_files.append(deleted_file);
14851473 continue;
14861474 }
......@@ -1507,10 +1495,7 @@ fn computeHash(
15071495 .hash = undefined, // to be populated by the worker
15081496 .failure = undefined, // to be populated by the worker
15091497 };
1510 wait_group.start();
1511 try thread_pool.spawn(workerHashFile, .{
1512 root_dir, hashed_file, &wait_group,
1513 });
1498 thread_pool.spawnWg(&wait_group, workerHashFile, .{ root_dir, hashed_file });
15141499 try all_files.append(hashed_file);
15151500 }
15161501 }
......@@ -1602,13 +1587,11 @@ fn dumpHashInfo(all_files: []const *const HashedFile) !void {
16021587 try bw.flush();
16031588}
16041589
1605fn workerHashFile(dir: fs.Dir, hashed_file: *HashedFile, wg: *WaitGroup) void {
1606 defer wg.finish();
1590fn workerHashFile(dir: fs.Dir, hashed_file: *HashedFile) void {
16071591 hashed_file.failure = hashFileFallible(dir, hashed_file);
16081592}
16091593
1610fn workerDeleteFile(dir: fs.Dir, deleted_file: *DeletedFile, wg: *WaitGroup) void {
1611 defer wg.finish();
1594fn workerDeleteFile(dir: fs.Dir, deleted_file: *DeletedFile) void {
16121595 deleted_file.failure = deleteFileFallible(dir, deleted_file);
16131596}
16141597
src/link/MachO/hasher.zig+1-5
......@@ -36,14 +36,12 @@ pub fn ParallelHasher(comptime Hasher: type) type {
3636 file_size - fstart
3737 else
3838 chunk_size;
39 wg.start();
40 try self.thread_pool.spawn(worker, .{
39 self.thread_pool.spawnWg(&wg, worker, .{
4140 file,
4241 fstart,
4342 buffer[fstart..][0..fsize],
4443 &(out_buf.*),
4544 &(result.*),
46 &wg,
4745 });
4846 }
4947 }
......@@ -56,9 +54,7 @@ pub fn ParallelHasher(comptime Hasher: type) type {
5654 buffer: []u8,
5755 out: *[hash_size]u8,
5856 err: *fs.File.PReadError!usize,
59 wg: *WaitGroup,
6057 ) void {
61 defer wg.finish();
6258 err.* = file.preadAll(buffer, fstart);
6359 Hasher.hash(buffer, out, .{});
6460 }
src/main.zig+3-2
......@@ -5109,8 +5109,9 @@ fn cmdBuild(gpa: Allocator, arena: Allocator, args: []const []const u8) !void {
51095109 &fetch,
51105110 );
51115111
5112 job_queue.wait_group.start();
5113 try job_queue.thread_pool.spawn(Package.Fetch.workerRun, .{ &fetch, "root" });
5112 job_queue.thread_pool.spawnWg(&job_queue.wait_group, Package.Fetch.workerRun, .{
5113 &fetch, "root",
5114 });
51145115 job_queue.wait_group.wait();
51155116
51165117 try job_queue.consolidateErrors();