lib/choir/src/core/threading.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const sys = @import("sys");
3
4 pub const Options = struct {
5 max_threads: usize = 1,
6 worker_allocator: ?std.mem.Allocator = null,
7
8 pub fn workerCount(self: Options, item_count: usize) usize {
9 if (item_count == 0) return 1;
10 const configured = if (self.max_threads == 0)
11 sys.thread.cpuCount()
12 else
13 self.max_threads;
14 return @max(@as(usize, 1), @min(configured, item_count));
15 }
16
17 pub fn workerAllocator(self: Options, fallback: std.mem.Allocator) std.mem.Allocator {
18 return self.worker_allocator orelse fallback;
19 }
20
21 pub fn requestsParallelism(self: Options) bool {
22 return self.max_threads == 0 or self.max_threads > 1;
23 }
24 };
25
26 pub const ExecutionState = struct {
27 active: std.atomic.Value(usize),
28
29 pub fn init() ExecutionState {
30 return .{ .active = std.atomic.Value(usize).init(0) };
31 }
32
33 pub fn enter(self: *ExecutionState) ExecutionGuard {
34 _ = self.active.fetchAdd(1, .acq_rel);
35 return .{ .state = self };
36 }
37
38 pub fn isActive(self: *const ExecutionState) bool {
39 return self.active.load(.acquire) != 0;
40 }
41
42 fn exit(self: *ExecutionState) void {
43 const prior = self.active.fetchSub(1, .acq_rel);
44 std.debug.assert(prior != 0);
45 }
46 };
47
48 pub const ExecutionGuard = struct {
49 state: ?*ExecutionState,
50
51 pub fn deinit(self: *ExecutionGuard) void {
52 if (self.state) |state| {
53 state.exit();
54 self.state = null;
55 }
56 }
57 };
58
59 pub const ParallelError = std.mem.Allocator.Error || sys.thread.SpawnError;
60
61 pub fn parallelForEachIndex(
62 allocator: std.mem.Allocator,
63 options: Options,
64 item_count: usize,
65 context: anytype,
66 comptime run: fn (@TypeOf(context), usize) void,
67 ) ParallelError!void {
68 const worker_count = options.workerCount(item_count);
69 if (worker_count == 1) {
70 for (0..item_count) |index| {
71 run(context, index);
72 }
73 return;
74 }
75
76 const Context = @TypeOf(context);
77 const Batch = struct {
78 task_context: Context,
79 next_index: *std.atomic.Value(usize),
80 item_count: usize,
81 };
82 const Worker = struct {
83 fn work(batch: *Batch) void {
84 while (true) {
85 const index = batch.next_index.fetchAdd(1, .monotonic);
86 if (index >= batch.item_count) return;
87 run(batch.task_context, index);
88 }
89 }
90 };
91
92 var next_index = std.atomic.Value(usize).init(0);
93 var batch = Batch{
94 .task_context = context,
95 .next_index = &next_index,
96 .item_count = item_count,
97 };
98
99 const spawned_count = worker_count - 1;
100 const threads = try allocator.alloc(sys.thread.JoinHandle, spawned_count);
101 defer allocator.free(threads);
102
103 var spawned: usize = 0;
104 errdefer {
105 for (threads[0..spawned]) |thread| {
106 thread.join();
107 }
108 }
109
110 for (threads) |*thread| {
111 thread.* = try sys.thread.spawn(Worker.work, .{&batch});
112 spawned += 1;
113 }
114 Worker.work(&batch);
115 for (threads) |thread| {
116 thread.join();
117 }
118 }
119
120 test "ThreadingOptions counts workers with item and cpu limits" {
121 const testing = std.testing;
122
123 try testing.expectEqual(@as(usize, 1), (Options{}).workerCount(0));
124 try testing.expectEqual(@as(usize, 1), (Options{}).workerCount(8));
125 try testing.expectEqual(@as(usize, 2), (Options{ .max_threads = 2 }).workerCount(8));
126 try testing.expectEqual(@as(usize, 2), (Options{ .max_threads = 4 }).workerCount(2));
127 try testing.expect((Options{ .max_threads = 0 }).workerCount(8) >= 1);
128 }
129
130 test "ThreadingExecutionState tracks nested execution guards" {
131 const testing = std.testing;
132
133 var state = ExecutionState.init();
134 try testing.expect(!state.isActive());
135
136 var outer = state.enter();
137 defer outer.deinit();
138 try testing.expect(state.isActive());
139
140 {
141 var inner = state.enter();
142 defer inner.deinit();
143 try testing.expect(state.isActive());
144 }
145
146 try testing.expect(state.isActive());
147 outer.deinit();
148 try testing.expect(!state.isActive());
149 }
150
151 test "parallelForEachIndex visits every item" {
152 const testing = std.testing;
153 if (!sys.thread.threadsSupported()) return error.SkipZigTest;
154
155 var hits: [8]std.atomic.Value(usize) = undefined;
156 for (&hits) |*hit| hit.* = std.atomic.Value(usize).init(0);
157
158 const Task = struct {
159 values: []std.atomic.Value(usize),
160 };
161 const Runner = struct {
162 fn run(task: *Task, index: usize) void {
163 _ = task.values[index].fetchAdd(1, .acq_rel);
164 }
165 };
166
167 var task = Task{ .values = &hits };
168 try parallelForEachIndex(testing.allocator, .{ .max_threads = 2 }, hits.len, &task, Runner.run);
169
170 for (&hits) |*hit| {
171 try testing.expectEqual(@as(usize, 1), hit.load(.acquire));
172 }
173 }