authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-07-09 23:10:31-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2025-10-02 16:30:59-07:00
log9a62dc4d89e0874d3f5d2bebee11d57bdce8b0d5
tree24bdaaeb1481b2a08f8b2b0b9b27da20f1cd696e
parentfdd59f0fa73e5c2590b3bdd4a61e5e2e65557867

revert std.Thread.Pool for now

and move the Io impl to a separate file

3 files changed, 955 insertions(+), 628 deletions(-)

lib/std/Io.zig+1
......@@ -558,6 +558,7 @@ test {
558558const Io = @This();
559559
560560pub const EventLoop = @import("Io/EventLoop.zig");
561pub const ThreadPool = @import("Io/ThreadPool.zig");
561562
562563userdata: ?*anyopaque,
563564vtable: *const VTable,
lib/std/Io/ThreadPool.zig created+852
......@@ -0,0 +1,852 @@
1const builtin = @import("builtin");
2const std = @import("../std.zig");
3const Allocator = std.mem.Allocator;
4const assert = std.debug.assert;
5const WaitGroup = std.Thread.WaitGroup;
6const Io = std.Io;
7const Pool = @This();
8
9/// Must be a thread-safe allocator.
10allocator: std.mem.Allocator,
11mutex: std.Thread.Mutex = .{},
12cond: std.Thread.Condition = .{},
13run_queue: std.SinglyLinkedList = .{},
14is_running: bool = true,
15threads: std.ArrayListUnmanaged(std.Thread),
16ids: if (builtin.single_threaded) struct {
17 inline fn deinit(_: @This(), _: std.mem.Allocator) void {}
18 fn getIndex(_: @This(), _: std.Thread.Id) usize {
19 return 0;
20 }
21} else std.AutoArrayHashMapUnmanaged(std.Thread.Id, void),
22stack_size: usize,
23
24threadlocal var current_closure: ?*AsyncClosure = null;
25
26pub const Runnable = struct {
27 runFn: RunProto,
28 node: std.SinglyLinkedList.Node = .{},
29};
30
31pub const RunProto = *const fn (*Runnable, id: ?usize) void;
32
33pub const Options = struct {
34 allocator: std.mem.Allocator,
35 n_jobs: ?usize = null,
36 track_ids: bool = false,
37 stack_size: usize = std.Thread.SpawnConfig.default_stack_size,
38};
39
40pub fn init(pool: *Pool, options: Options) !void {
41 const gpa = options.allocator;
42 const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1);
43 const threads = try gpa.alloc(std.Thread, thread_count);
44 errdefer gpa.free(threads);
45
46 pool.* = .{
47 .allocator = gpa,
48 .threads = .initBuffer(threads),
49 .ids = .{},
50 .stack_size = options.stack_size,
51 };
52
53 if (builtin.single_threaded) return;
54
55 if (options.track_ids) {
56 try pool.ids.ensureTotalCapacity(gpa, 1 + thread_count);
57 pool.ids.putAssumeCapacityNoClobber(std.Thread.getCurrentId(), {});
58 }
59}
60
61pub fn deinit(pool: *Pool) void {
62 const gpa = pool.allocator;
63 pool.join();
64 pool.threads.deinit(gpa);
65 pool.ids.deinit(gpa);
66 pool.* = undefined;
67}
68
69fn join(pool: *Pool) void {
70 if (builtin.single_threaded) return;
71
72 {
73 pool.mutex.lock();
74 defer pool.mutex.unlock();
75
76 // ensure future worker threads exit the dequeue loop
77 pool.is_running = false;
78 }
79
80 // wake up any sleeping threads (this can be done outside the mutex)
81 // then wait for all the threads we know are spawned to complete.
82 pool.cond.broadcast();
83 for (pool.threads.items) |thread| thread.join();
84}
85
86/// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and
87/// `WaitGroup.finish` after it returns.
88///
89/// In the case that queuing the function call fails to allocate memory, or the
90/// target is single-threaded, the function is called directly.
91pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args: anytype) void {
92 wait_group.start();
93
94 if (builtin.single_threaded) {
95 @call(.auto, func, args);
96 wait_group.finish();
97 return;
98 }
99
100 const Args = @TypeOf(args);
101 const Closure = struct {
102 arguments: Args,
103 pool: *Pool,
104 runnable: Runnable = .{ .runFn = runFn },
105 wait_group: *WaitGroup,
106
107 fn runFn(runnable: *Runnable, _: ?usize) void {
108 const closure: *@This() = @alignCast(@fieldParentPtr("runnable", runnable));
109 @call(.auto, func, closure.arguments);
110 closure.wait_group.finish();
111 closure.pool.allocator.destroy(closure);
112 }
113 };
114
115 pool.mutex.lock();
116
117 const gpa = pool.allocator;
118 const closure = gpa.create(Closure) catch {
119 pool.mutex.unlock();
120 @call(.auto, func, args);
121 wait_group.finish();
122 return;
123 };
124 closure.* = .{
125 .arguments = args,
126 .pool = pool,
127 .wait_group = wait_group,
128 };
129
130 pool.run_queue.prepend(&closure.runnable.node);
131
132 if (pool.threads.items.len < pool.threads.capacity) {
133 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
134 .stack_size = pool.stack_size,
135 .allocator = gpa,
136 }, worker, .{pool}) catch t: {
137 pool.threads.items.len -= 1;
138 break :t undefined;
139 };
140 }
141
142 pool.mutex.unlock();
143 pool.cond.signal();
144}
145
146/// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and
147/// `WaitGroup.finish` after it returns.
148///
149/// The first argument passed to `func` is a dense `usize` thread id, the rest
150/// of the arguments are passed from `args`. Requires the pool to have been
151/// initialized with `.track_ids = true`.
152///
153/// In the case that queuing the function call fails to allocate memory, or the
154/// target is single-threaded, the function is called directly.
155pub fn spawnWgId(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args: anytype) void {
156 wait_group.start();
157
158 if (builtin.single_threaded) {
159 @call(.auto, func, .{0} ++ args);
160 wait_group.finish();
161 return;
162 }
163
164 const Args = @TypeOf(args);
165 const Closure = struct {
166 arguments: Args,
167 pool: *Pool,
168 runnable: Runnable = .{ .runFn = runFn },
169 wait_group: *WaitGroup,
170
171 fn runFn(runnable: *Runnable, id: ?usize) void {
172 const closure: *@This() = @alignCast(@fieldParentPtr("runnable", runnable));
173 @call(.auto, func, .{id.?} ++ closure.arguments);
174 closure.wait_group.finish();
175 closure.pool.allocator.destroy(closure);
176 }
177 };
178
179 pool.mutex.lock();
180
181 const gpa = pool.allocator;
182 const closure = gpa.create(Closure) catch {
183 const id: ?usize = pool.ids.getIndex(std.Thread.getCurrentId());
184 pool.mutex.unlock();
185 @call(.auto, func, .{id.?} ++ args);
186 wait_group.finish();
187 return;
188 };
189 closure.* = .{
190 .arguments = args,
191 .pool = pool,
192 .wait_group = wait_group,
193 };
194
195 pool.run_queue.prepend(&closure.runnable.node);
196
197 if (pool.threads.items.len < pool.threads.capacity) {
198 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
199 .stack_size = pool.stack_size,
200 .allocator = gpa,
201 }, worker, .{pool}) catch t: {
202 pool.threads.items.len -= 1;
203 break :t undefined;
204 };
205 }
206
207 pool.mutex.unlock();
208 pool.cond.signal();
209}
210
211pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void {
212 if (builtin.single_threaded) {
213 @call(.auto, func, args);
214 return;
215 }
216
217 const Args = @TypeOf(args);
218 const Closure = struct {
219 arguments: Args,
220 pool: *Pool,
221 runnable: Runnable = .{ .runFn = runFn },
222
223 fn runFn(runnable: *Runnable, _: ?usize) void {
224 const closure: *@This() = @alignCast(@fieldParentPtr("runnable", runnable));
225 @call(.auto, func, closure.arguments);
226 closure.pool.allocator.destroy(closure);
227 }
228 };
229
230 pool.mutex.lock();
231
232 const gpa = pool.allocator;
233 const closure = gpa.create(Closure) catch {
234 pool.mutex.unlock();
235 @call(.auto, func, args);
236 return;
237 };
238 closure.* = .{
239 .arguments = args,
240 .pool = pool,
241 };
242
243 pool.run_queue.prepend(&closure.runnable.node);
244
245 if (pool.threads.items.len < pool.threads.capacity) {
246 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
247 .stack_size = pool.stack_size,
248 .allocator = gpa,
249 }, worker, .{pool}) catch t: {
250 pool.threads.items.len -= 1;
251 break :t undefined;
252 };
253 }
254
255 pool.mutex.unlock();
256 pool.cond.signal();
257}
258
259test spawn {
260 const TestFn = struct {
261 fn checkRun(completed: *bool) void {
262 completed.* = true;
263 }
264 };
265
266 var completed: bool = false;
267
268 {
269 var pool: Pool = undefined;
270 try pool.init(.{
271 .allocator = std.testing.allocator,
272 });
273 defer pool.deinit();
274 pool.spawn(TestFn.checkRun, .{&completed});
275 }
276
277 try std.testing.expectEqual(true, completed);
278}
279
280fn worker(pool: *Pool) void {
281 pool.mutex.lock();
282 defer pool.mutex.unlock();
283
284 const id: ?usize = if (pool.ids.count() > 0) @intCast(pool.ids.count()) else null;
285 if (id) |_| pool.ids.putAssumeCapacityNoClobber(std.Thread.getCurrentId(), {});
286
287 while (true) {
288 while (pool.run_queue.popFirst()) |run_node| {
289 // Temporarily unlock the mutex in order to execute the run_node
290 pool.mutex.unlock();
291 defer pool.mutex.lock();
292
293 const runnable: *Runnable = @fieldParentPtr("node", run_node);
294 runnable.runFn(runnable, id);
295 }
296
297 // Stop executing instead of waiting if the thread pool is no longer running.
298 if (pool.is_running) {
299 pool.cond.wait(&pool.mutex);
300 } else {
301 break;
302 }
303 }
304}
305
306pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void {
307 var id: ?usize = null;
308
309 while (!wait_group.isDone()) {
310 pool.mutex.lock();
311 if (pool.run_queue.popFirst()) |run_node| {
312 id = id orelse pool.ids.getIndex(std.Thread.getCurrentId());
313 pool.mutex.unlock();
314 const runnable: *Runnable = @fieldParentPtr("node", run_node);
315 runnable.runFn(runnable, id);
316 continue;
317 }
318
319 pool.mutex.unlock();
320 wait_group.wait();
321 return;
322 }
323}
324
325pub fn getIdCount(pool: *Pool) usize {
326 return @intCast(1 + pool.threads.items.len);
327}
328
329pub fn io(pool: *Pool) Io {
330 return .{
331 .userdata = pool,
332 .vtable = &.{
333 .async = async,
334 .await = await,
335 .go = go,
336 .cancel = cancel,
337 .cancelRequested = cancelRequested,
338 .select = select,
339
340 .mutexLock = mutexLock,
341 .mutexUnlock = mutexUnlock,
342
343 .conditionWait = conditionWait,
344 .conditionWake = conditionWake,
345
346 .createFile = createFile,
347 .openFile = openFile,
348 .closeFile = closeFile,
349 .pread = pread,
350 .pwrite = pwrite,
351
352 .now = now,
353 .sleep = sleep,
354 },
355 };
356}
357
358const AsyncClosure = struct {
359 func: *const fn (context: *anyopaque, result: *anyopaque) void,
360 runnable: Runnable = .{ .runFn = runFn },
361 reset_event: std.Thread.ResetEvent,
362 select_condition: ?*std.Thread.ResetEvent,
363 cancel_tid: std.Thread.Id,
364 context_offset: usize,
365 result_offset: usize,
366
367 const done_reset_event: *std.Thread.ResetEvent = @ptrFromInt(@alignOf(std.Thread.ResetEvent));
368
369 const canceling_tid: std.Thread.Id = switch (@typeInfo(std.Thread.Id)) {
370 .int => |int_info| switch (int_info.signedness) {
371 .signed => -1,
372 .unsigned => std.math.maxInt(std.Thread.Id),
373 },
374 .pointer => @ptrFromInt(std.math.maxInt(usize)),
375 else => @compileError("unsupported std.Thread.Id: " ++ @typeName(std.Thread.Id)),
376 };
377
378 fn runFn(runnable: *Pool.Runnable, _: ?usize) void {
379 const closure: *AsyncClosure = @alignCast(@fieldParentPtr("runnable", runnable));
380 const tid = std.Thread.getCurrentId();
381 if (@cmpxchgStrong(
382 std.Thread.Id,
383 &closure.cancel_tid,
384 0,
385 tid,
386 .acq_rel,
387 .acquire,
388 )) |cancel_tid| {
389 assert(cancel_tid == canceling_tid);
390 return;
391 }
392 current_closure = closure;
393 closure.func(closure.contextPointer(), closure.resultPointer());
394 current_closure = null;
395 if (@cmpxchgStrong(
396 std.Thread.Id,
397 &closure.cancel_tid,
398 tid,
399 0,
400 .acq_rel,
401 .acquire,
402 )) |cancel_tid| assert(cancel_tid == canceling_tid);
403
404 if (@atomicRmw(
405 ?*std.Thread.ResetEvent,
406 &closure.select_condition,
407 .Xchg,
408 done_reset_event,
409 .release,
410 )) |select_reset| {
411 assert(select_reset != done_reset_event);
412 select_reset.set();
413 }
414 closure.reset_event.set();
415 }
416
417 fn contextOffset(context_alignment: std.mem.Alignment) usize {
418 return context_alignment.forward(@sizeOf(AsyncClosure));
419 }
420
421 fn resultOffset(
422 context_alignment: std.mem.Alignment,
423 context_len: usize,
424 result_alignment: std.mem.Alignment,
425 ) usize {
426 return result_alignment.forward(contextOffset(context_alignment) + context_len);
427 }
428
429 fn resultPointer(closure: *AsyncClosure) [*]u8 {
430 const base: [*]u8 = @ptrCast(closure);
431 return base + closure.result_offset;
432 }
433
434 fn contextPointer(closure: *AsyncClosure) [*]u8 {
435 const base: [*]u8 = @ptrCast(closure);
436 return base + closure.context_offset;
437 }
438
439 fn waitAndFree(closure: *AsyncClosure, gpa: Allocator, result: []u8) void {
440 closure.reset_event.wait();
441 const base: [*]align(@alignOf(AsyncClosure)) u8 = @ptrCast(closure);
442 @memcpy(result, closure.resultPointer()[0..result.len]);
443 gpa.free(base[0 .. closure.result_offset + result.len]);
444 }
445};
446
447fn async(
448 userdata: ?*anyopaque,
449 result: []u8,
450 result_alignment: std.mem.Alignment,
451 context: []const u8,
452 context_alignment: std.mem.Alignment,
453 start: *const fn (context: *const anyopaque, result: *anyopaque) void,
454) ?*Io.AnyFuture {
455 const pool: *Pool = @alignCast(@ptrCast(userdata));
456 pool.mutex.lock();
457
458 const gpa = pool.allocator;
459 const context_offset = context_alignment.forward(@sizeOf(AsyncClosure));
460 const result_offset = result_alignment.forward(context_offset + context.len);
461 const n = result_offset + result.len;
462 const closure: *AsyncClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(AsyncClosure), n) catch {
463 pool.mutex.unlock();
464 start(context.ptr, result.ptr);
465 return null;
466 }));
467 closure.* = .{
468 .func = start,
469 .context_offset = context_offset,
470 .result_offset = result_offset,
471 .reset_event = .{},
472 .cancel_tid = 0,
473 .select_condition = null,
474 };
475 @memcpy(closure.contextPointer()[0..context.len], context);
476 pool.run_queue.prepend(&closure.runnable.node);
477
478 if (pool.threads.items.len < pool.threads.capacity) {
479 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
480 .stack_size = pool.stack_size,
481 .allocator = gpa,
482 }, worker, .{pool}) catch t: {
483 pool.threads.items.len -= 1;
484 break :t undefined;
485 };
486 }
487
488 pool.mutex.unlock();
489 pool.cond.signal();
490
491 return @ptrCast(closure);
492}
493
494const DetachedClosure = struct {
495 pool: *Pool,
496 func: *const fn (context: *anyopaque) void,
497 run_node: Pool.RunQueue.Node = .{ .data = .{ .runFn = runFn } },
498 context_alignment: std.mem.Alignment,
499 context_len: usize,
500
501 fn runFn(runnable: *Pool.Runnable, _: ?usize) void {
502 const run_node: *Pool.RunQueue.Node = @fieldParentPtr("data", runnable);
503 const closure: *DetachedClosure = @alignCast(@fieldParentPtr("run_node", run_node));
504 closure.func(closure.contextPointer());
505 const gpa = closure.pool.allocator;
506 const base: [*]align(@alignOf(DetachedClosure)) u8 = @ptrCast(closure);
507 gpa.free(base[0..contextEnd(closure.context_alignment, closure.context_len)]);
508 }
509
510 fn contextOffset(context_alignment: std.mem.Alignment) usize {
511 return context_alignment.forward(@sizeOf(DetachedClosure));
512 }
513
514 fn contextEnd(context_alignment: std.mem.Alignment, context_len: usize) usize {
515 return contextOffset(context_alignment) + context_len;
516 }
517
518 fn contextPointer(closure: *DetachedClosure) [*]u8 {
519 const base: [*]u8 = @ptrCast(closure);
520 return base + contextOffset(closure.context_alignment);
521 }
522};
523
524fn go(
525 userdata: ?*anyopaque,
526 context: []const u8,
527 context_alignment: std.mem.Alignment,
528 start: *const fn (context: *const anyopaque) void,
529) void {
530 const pool: *Pool = @alignCast(@ptrCast(userdata));
531 pool.mutex.lock();
532
533 const gpa = pool.allocator;
534 const n = DetachedClosure.contextEnd(context_alignment, context.len);
535 const closure: *DetachedClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(DetachedClosure), n) catch {
536 pool.mutex.unlock();
537 start(context.ptr);
538 return;
539 }));
540 closure.* = .{
541 .pool = pool,
542 .func = start,
543 .context_alignment = context_alignment,
544 .context_len = context.len,
545 };
546 @memcpy(closure.contextPointer()[0..context.len], context);
547 pool.run_queue.prepend(&closure.run_node);
548
549 if (pool.threads.items.len < pool.threads.capacity) {
550 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
551 .stack_size = pool.stack_size,
552 .allocator = gpa,
553 }, worker, .{pool}) catch t: {
554 pool.threads.items.len -= 1;
555 break :t undefined;
556 };
557 }
558
559 pool.mutex.unlock();
560 pool.cond.signal();
561}
562
563fn await(
564 userdata: ?*anyopaque,
565 any_future: *std.Io.AnyFuture,
566 result: []u8,
567 result_alignment: std.mem.Alignment,
568) void {
569 _ = result_alignment;
570 const pool: *Pool = @alignCast(@ptrCast(userdata));
571 const closure: *AsyncClosure = @ptrCast(@alignCast(any_future));
572 closure.waitAndFree(pool.allocator, result);
573}
574
575fn cancel(
576 userdata: ?*anyopaque,
577 any_future: *Io.AnyFuture,
578 result: []u8,
579 result_alignment: std.mem.Alignment,
580) void {
581 _ = result_alignment;
582 const pool: *Pool = @alignCast(@ptrCast(userdata));
583 const closure: *AsyncClosure = @ptrCast(@alignCast(any_future));
584 switch (@atomicRmw(
585 std.Thread.Id,
586 &closure.cancel_tid,
587 .Xchg,
588 AsyncClosure.canceling_tid,
589 .acq_rel,
590 )) {
591 0, AsyncClosure.canceling_tid => {},
592 else => |cancel_tid| switch (builtin.os.tag) {
593 .linux => _ = std.os.linux.tgkill(
594 std.os.linux.getpid(),
595 @bitCast(cancel_tid),
596 std.posix.SIG.IO,
597 ),
598 else => {},
599 },
600 }
601 closure.waitAndFree(pool.allocator, result);
602}
603
604fn cancelRequested(userdata: ?*anyopaque) bool {
605 const pool: *Pool = @alignCast(@ptrCast(userdata));
606 _ = pool;
607 const closure = current_closure orelse return false;
608 return @atomicLoad(std.Thread.Id, &closure.cancel_tid, .acquire) == AsyncClosure.canceling_tid;
609}
610
611fn checkCancel(pool: *Pool) error{Canceled}!void {
612 if (cancelRequested(pool)) return error.Canceled;
613}
614
615fn mutexLock(userdata: ?*anyopaque, prev_state: Io.Mutex.State, mutex: *Io.Mutex) error{Canceled}!void {
616 _ = userdata;
617 if (prev_state == .contended) {
618 std.Thread.Futex.wait(@ptrCast(&mutex.state), @intFromEnum(Io.Mutex.State.contended));
619 }
620 while (@atomicRmw(
621 Io.Mutex.State,
622 &mutex.state,
623 .Xchg,
624 .contended,
625 .acquire,
626 ) != .unlocked) {
627 std.Thread.Futex.wait(@ptrCast(&mutex.state), @intFromEnum(Io.Mutex.State.contended));
628 }
629}
630fn mutexUnlock(userdata: ?*anyopaque, prev_state: Io.Mutex.State, mutex: *Io.Mutex) void {
631 _ = userdata;
632 _ = prev_state;
633 if (@atomicRmw(Io.Mutex.State, &mutex.state, .Xchg, .unlocked, .release) == .contended) {
634 std.Thread.Futex.wake(@ptrCast(&mutex.state), 1);
635 }
636}
637
638fn conditionWait(userdata: ?*anyopaque, cond: *Io.Condition, mutex: *Io.Mutex) Io.Cancelable!void {
639 const pool: *Pool = @alignCast(@ptrCast(userdata));
640 comptime assert(@TypeOf(cond.state) == u64);
641 const ints: *[2]std.atomic.Value(u32) = @ptrCast(&cond.state);
642 const cond_state = &ints[0];
643 const cond_epoch = &ints[1];
644 const one_waiter = 1;
645 const waiter_mask = 0xffff;
646 const one_signal = 1 << 16;
647 const signal_mask = 0xffff << 16;
648 // Observe the epoch, then check the state again to see if we should wake up.
649 // The epoch must be observed before we check the state or we could potentially miss a wake() and deadlock:
650 //
651 // - T1: s = LOAD(&state)
652 // - T2: UPDATE(&s, signal)
653 // - T2: UPDATE(&epoch, 1) + FUTEX_WAKE(&epoch)
654 // - T1: e = LOAD(&epoch) (was reordered after the state load)
655 // - T1: s & signals == 0 -> FUTEX_WAIT(&epoch, e) (missed the state update + the epoch change)
656 //
657 // Acquire barrier to ensure the epoch load happens before the state load.
658 var epoch = cond_epoch.load(.acquire);
659 var state = cond_state.fetchAdd(one_waiter, .monotonic);
660 assert(state & waiter_mask != waiter_mask);
661 state += one_waiter;
662
663 mutex.unlock(pool.io());
664 defer mutex.lock(pool.io()) catch @panic("TODO");
665
666 var futex_deadline = std.Thread.Futex.Deadline.init(null);
667
668 while (true) {
669 futex_deadline.wait(cond_epoch, epoch) catch |err| switch (err) {
670 error.Timeout => unreachable,
671 };
672
673 epoch = cond_epoch.load(.acquire);
674 state = cond_state.load(.monotonic);
675
676 // Try to wake up by consuming a signal and decremented the waiter we added previously.
677 // Acquire barrier ensures code before the wake() which added the signal happens before we decrement it and return.
678 while (state & signal_mask != 0) {
679 const new_state = state - one_waiter - one_signal;
680 state = cond_state.cmpxchgWeak(state, new_state, .acquire, .monotonic) orelse return;
681 }
682 }
683}
684
685fn conditionWake(userdata: ?*anyopaque, cond: *Io.Condition, wake: Io.Condition.Wake) void {
686 const pool: *Pool = @alignCast(@ptrCast(userdata));
687 _ = pool;
688 comptime assert(@TypeOf(cond.state) == u64);
689 const ints: *[2]std.atomic.Value(u32) = @ptrCast(&cond.state);
690 const cond_state = &ints[0];
691 const cond_epoch = &ints[1];
692 const one_waiter = 1;
693 const waiter_mask = 0xffff;
694 const one_signal = 1 << 16;
695 const signal_mask = 0xffff << 16;
696 var state = cond_state.load(.monotonic);
697 while (true) {
698 const waiters = (state & waiter_mask) / one_waiter;
699 const signals = (state & signal_mask) / one_signal;
700
701 // Reserves which waiters to wake up by incrementing the signals count.
702 // Therefore, the signals count is always less than or equal to the waiters count.
703 // We don't need to Futex.wake if there's nothing to wake up or if other wake() threads have reserved to wake up the current waiters.
704 const wakeable = waiters - signals;
705 if (wakeable == 0) {
706 return;
707 }
708
709 const to_wake = switch (wake) {
710 .one => 1,
711 .all => wakeable,
712 };
713
714 // Reserve the amount of waiters to wake by incrementing the signals count.
715 // Release barrier ensures code before the wake() happens before the signal it posted and consumed by the wait() threads.
716 const new_state = state + (one_signal * to_wake);
717 state = cond_state.cmpxchgWeak(state, new_state, .release, .monotonic) orelse {
718 // Wake up the waiting threads we reserved above by changing the epoch value.
719 // NOTE: a waiting thread could miss a wake up if *exactly* ((1<<32)-1) wake()s happen between it observing the epoch and sleeping on it.
720 // This is very unlikely due to how many precise amount of Futex.wake() calls that would be between the waiting thread's potential preemption.
721 //
722 // Release barrier ensures the signal being added to the state happens before the epoch is changed.
723 // If not, the waiting thread could potentially deadlock from missing both the state and epoch change:
724 //
725 // - T2: UPDATE(&epoch, 1) (reordered before the state change)
726 // - T1: e = LOAD(&epoch)
727 // - T1: s = LOAD(&state)
728 // - T2: UPDATE(&state, signal) + FUTEX_WAKE(&epoch)
729 // - T1: s & signals == 0 -> FUTEX_WAIT(&epoch, e) (missed both epoch change and state change)
730 _ = cond_epoch.fetchAdd(1, .release);
731 std.Thread.Futex.wake(cond_epoch, to_wake);
732 return;
733 };
734 }
735}
736
737fn createFile(
738 userdata: ?*anyopaque,
739 dir: Io.Dir,
740 sub_path: []const u8,
741 flags: Io.File.CreateFlags,
742) Io.File.OpenError!Io.File {
743 const pool: *Pool = @alignCast(@ptrCast(userdata));
744 try pool.checkCancel();
745 const fs_dir: std.fs.Dir = .{ .fd = dir.handle };
746 const fs_file = try fs_dir.createFile(sub_path, flags);
747 return .{ .handle = fs_file.handle };
748}
749
750fn openFile(
751 userdata: ?*anyopaque,
752 dir: Io.Dir,
753 sub_path: []const u8,
754 flags: Io.File.OpenFlags,
755) Io.File.OpenError!Io.File {
756 const pool: *Pool = @alignCast(@ptrCast(userdata));
757 try pool.checkCancel();
758 const fs_dir: std.fs.Dir = .{ .fd = dir.handle };
759 const fs_file = try fs_dir.openFile(sub_path, flags);
760 return .{ .handle = fs_file.handle };
761}
762
763fn closeFile(userdata: ?*anyopaque, file: Io.File) void {
764 const pool: *Pool = @alignCast(@ptrCast(userdata));
765 _ = pool;
766 const fs_file: std.fs.File = .{ .handle = file.handle };
767 return fs_file.close();
768}
769
770fn pread(userdata: ?*anyopaque, file: Io.File, buffer: []u8, offset: std.posix.off_t) Io.File.PReadError!usize {
771 const pool: *Pool = @alignCast(@ptrCast(userdata));
772 try pool.checkCancel();
773 const fs_file: std.fs.File = .{ .handle = file.handle };
774 return switch (offset) {
775 -1 => fs_file.read(buffer),
776 else => fs_file.pread(buffer, @bitCast(offset)),
777 };
778}
779
780fn pwrite(userdata: ?*anyopaque, file: Io.File, buffer: []const u8, offset: std.posix.off_t) Io.File.PWriteError!usize {
781 const pool: *Pool = @alignCast(@ptrCast(userdata));
782 try pool.checkCancel();
783 const fs_file: std.fs.File = .{ .handle = file.handle };
784 return switch (offset) {
785 -1 => fs_file.write(buffer),
786 else => fs_file.pwrite(buffer, @bitCast(offset)),
787 };
788}
789
790fn now(userdata: ?*anyopaque, clockid: std.posix.clockid_t) Io.ClockGetTimeError!Io.Timestamp {
791 const pool: *Pool = @alignCast(@ptrCast(userdata));
792 try pool.checkCancel();
793 const timespec = try std.posix.clock_gettime(clockid);
794 return @enumFromInt(@as(i128, timespec.sec) * std.time.ns_per_s + timespec.nsec);
795}
796
797fn sleep(userdata: ?*anyopaque, clockid: std.posix.clockid_t, deadline: Io.Deadline) Io.SleepError!void {
798 const pool: *Pool = @alignCast(@ptrCast(userdata));
799 const deadline_nanoseconds: i96 = switch (deadline) {
800 .duration => |duration| duration.nanoseconds,
801 .timestamp => |timestamp| @intFromEnum(timestamp),
802 };
803 var timespec: std.posix.timespec = .{
804 .sec = @intCast(@divFloor(deadline_nanoseconds, std.time.ns_per_s)),
805 .nsec = @intCast(@mod(deadline_nanoseconds, std.time.ns_per_s)),
806 };
807 while (true) {
808 try pool.checkCancel();
809 switch (std.os.linux.E.init(std.os.linux.clock_nanosleep(clockid, .{ .ABSTIME = switch (deadline) {
810 .duration => false,
811 .timestamp => true,
812 } }, &timespec, &timespec))) {
813 .SUCCESS => return,
814 .FAULT => unreachable,
815 .INTR => {},
816 .INVAL => return error.UnsupportedClock,
817 else => |err| return std.posix.unexpectedErrno(err),
818 }
819 }
820}
821
822fn select(userdata: ?*anyopaque, futures: []const *Io.AnyFuture) usize {
823 const pool: *Pool = @alignCast(@ptrCast(userdata));
824 _ = pool;
825
826 var reset_event: std.Thread.ResetEvent = .{};
827
828 for (futures, 0..) |future, i| {
829 const closure: *AsyncClosure = @ptrCast(@alignCast(future));
830 if (@atomicRmw(?*std.Thread.ResetEvent, &closure.select_condition, .Xchg, &reset_event, .seq_cst) == AsyncClosure.done_reset_event) {
831 for (futures[0..i]) |cleanup_future| {
832 const cleanup_closure: *AsyncClosure = @ptrCast(@alignCast(cleanup_future));
833 if (@atomicRmw(?*std.Thread.ResetEvent, &cleanup_closure.select_condition, .Xchg, null, .seq_cst) == AsyncClosure.done_reset_event) {
834 cleanup_closure.reset_event.wait(); // Ensure no reference to our stack-allocated reset_event.
835 }
836 }
837 return i;
838 }
839 }
840
841 reset_event.wait();
842
843 var result: ?usize = null;
844 for (futures, 0..) |future, i| {
845 const closure: *AsyncClosure = @ptrCast(@alignCast(future));
846 if (@atomicRmw(?*std.Thread.ResetEvent, &closure.select_condition, .Xchg, null, .seq_cst) == AsyncClosure.done_reset_event) {
847 closure.reset_event.wait(); // Ensure no reference to our stack-allocated reset_event.
848 if (result == null) result = i; // In case multiple are ready, return first.
849 }
850 }
851 return result.?;
852}
lib/std/Thread/Pool.zig+102-628
......@@ -1,34 +1,27 @@
1const builtin = @import("builtin");
21const std = @import("std");
3const Allocator = std.mem.Allocator;
4const assert = std.debug.assert;
5const WaitGroup = @import("WaitGroup.zig");
6const Io = std.Io;
2const builtin = @import("builtin");
73const Pool = @This();
4const WaitGroup = @import("WaitGroup.zig");
85
9/// Must be a thread-safe allocator.
10allocator: std.mem.Allocator,
116mutex: std.Thread.Mutex = .{},
127cond: std.Thread.Condition = .{},
138run_queue: std.SinglyLinkedList = .{},
149is_running: bool = true,
15threads: std.ArrayListUnmanaged(std.Thread),
10allocator: std.mem.Allocator,
11threads: if (builtin.single_threaded) [0]std.Thread else []std.Thread,
1612ids: if (builtin.single_threaded) struct {
1713 inline fn deinit(_: @This(), _: std.mem.Allocator) void {}
1814 fn getIndex(_: @This(), _: std.Thread.Id) usize {
1915 return 0;
2016 }
2117} else std.AutoArrayHashMapUnmanaged(std.Thread.Id, void),
22stack_size: usize,
2318
24threadlocal var current_closure: ?*AsyncClosure = null;
25
26pub const Runnable = struct {
19const Runnable = struct {
2720 runFn: RunProto,
2821 node: std.SinglyLinkedList.Node = .{},
2922};
3023
31pub const RunProto = *const fn (*Runnable, id: ?usize) void;
24const RunProto = *const fn (*Runnable, id: ?usize) void;
3225
3326pub const Options = struct {
3427 allocator: std.mem.Allocator,
......@@ -38,36 +31,48 @@ pub const Options = struct {
3831};
3932
4033pub fn init(pool: *Pool, options: Options) !void {
41 const gpa = options.allocator;
42 const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1);
43 const threads = try gpa.alloc(std.Thread, thread_count);
44 errdefer gpa.free(threads);
34 const allocator = options.allocator;
4535
4636 pool.* = .{
47 .allocator = gpa,
48 .threads = .initBuffer(threads),
37 .allocator = allocator,
38 .threads = if (builtin.single_threaded) .{} else &.{},
4939 .ids = .{},
50 .stack_size = options.stack_size,
5140 };
5241
53 if (builtin.single_threaded) return;
42 if (builtin.single_threaded) {
43 return;
44 }
5445
46 const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1);
5547 if (options.track_ids) {
56 try pool.ids.ensureTotalCapacity(gpa, 1 + thread_count);
48 try pool.ids.ensureTotalCapacity(allocator, 1 + thread_count);
5749 pool.ids.putAssumeCapacityNoClobber(std.Thread.getCurrentId(), {});
5850 }
51
52 // kill and join any threads we spawned and free memory on error.
53 pool.threads = try allocator.alloc(std.Thread, thread_count);
54 var spawned: usize = 0;
55 errdefer pool.join(spawned);
56
57 for (pool.threads) |*thread| {
58 thread.* = try std.Thread.spawn(.{
59 .stack_size = options.stack_size,
60 .allocator = allocator,
61 }, worker, .{pool});
62 spawned += 1;
63 }
5964}
6065
6166pub fn deinit(pool: *Pool) void {
62 const gpa = pool.allocator;
63 pool.join();
64 pool.threads.deinit(gpa);
65 pool.ids.deinit(gpa);
67 pool.join(pool.threads.len); // kill and join all threads.
68 pool.ids.deinit(pool.allocator);
6669 pool.* = undefined;
6770}
6871
69fn join(pool: *Pool) void {
70 if (builtin.single_threaded) return;
72fn join(pool: *Pool, spawned: usize) void {
73 if (builtin.single_threaded) {
74 return;
75 }
7176
7277 {
7378 pool.mutex.lock();
......@@ -80,7 +85,11 @@ fn join(pool: *Pool) void {
8085 // wake up any sleeping threads (this can be done outside the mutex)
8186 // then wait for all the threads we know are spawned to complete.
8287 pool.cond.broadcast();
83 for (pool.threads.items) |thread| thread.join();
88 for (pool.threads[0..spawned]) |thread| {
89 thread.join();
90 }
91
92 pool.allocator.free(pool.threads);
8493}
8594
8695/// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and
......@@ -108,38 +117,36 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args
108117 const closure: *@This() = @alignCast(@fieldParentPtr("runnable", runnable));
109118 @call(.auto, func, closure.arguments);
110119 closure.wait_group.finish();
111 closure.pool.allocator.destroy(closure);
112 }
113 };
114120
115 pool.mutex.lock();
121 // The thread pool's allocator is protected by the mutex.
122 const mutex = &closure.pool.mutex;
123 mutex.lock();
124 defer mutex.unlock();
116125
117 const gpa = pool.allocator;
118 const closure = gpa.create(Closure) catch {
119 pool.mutex.unlock();
120 @call(.auto, func, args);
121 wait_group.finish();
122 return;
123 };
124 closure.* = .{
125 .arguments = args,
126 .pool = pool,
127 .wait_group = wait_group,
126 closure.pool.allocator.destroy(closure);
127 }
128128 };
129129
130 pool.run_queue.prepend(&closure.runnable.node);
130 {
131 pool.mutex.lock();
131132
132 if (pool.threads.items.len < pool.threads.capacity) {
133 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
134 .stack_size = pool.stack_size,
135 .allocator = gpa,
136 }, worker, .{pool}) catch t: {
137 pool.threads.items.len -= 1;
138 break :t undefined;
133 const closure = pool.allocator.create(Closure) catch {
134 pool.mutex.unlock();
135 @call(.auto, func, args);
136 wait_group.finish();
137 return;
138 };
139 closure.* = .{
140 .arguments = args,
141 .pool = pool,
142 .wait_group = wait_group,
139143 };
144
145 pool.run_queue.prepend(&closure.runnable.node);
146 pool.mutex.unlock();
140147 }
141148
142 pool.mutex.unlock();
149 // Notify waiting threads outside the lock to try and keep the critical section small.
143150 pool.cond.signal();
144151}
145152
......@@ -172,43 +179,41 @@ pub fn spawnWgId(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, ar
172179 const closure: *@This() = @alignCast(@fieldParentPtr("runnable", runnable));
173180 @call(.auto, func, .{id.?} ++ closure.arguments);
174181 closure.wait_group.finish();
175 closure.pool.allocator.destroy(closure);
176 }
177 };
178182
179 pool.mutex.lock();
183 // The thread pool's allocator is protected by the mutex.
184 const mutex = &closure.pool.mutex;
185 mutex.lock();
186 defer mutex.unlock();
180187
181 const gpa = pool.allocator;
182 const closure = gpa.create(Closure) catch {
183 const id: ?usize = pool.ids.getIndex(std.Thread.getCurrentId());
184 pool.mutex.unlock();
185 @call(.auto, func, .{id.?} ++ args);
186 wait_group.finish();
187 return;
188 };
189 closure.* = .{
190 .arguments = args,
191 .pool = pool,
192 .wait_group = wait_group,
188 closure.pool.allocator.destroy(closure);
189 }
193190 };
194191
195 pool.run_queue.prepend(&closure.runnable.node);
192 {
193 pool.mutex.lock();
196194
197 if (pool.threads.items.len < pool.threads.capacity) {
198 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
199 .stack_size = pool.stack_size,
200 .allocator = gpa,
201 }, worker, .{pool}) catch t: {
202 pool.threads.items.len -= 1;
203 break :t undefined;
195 const closure = pool.allocator.create(Closure) catch {
196 const id: ?usize = pool.ids.getIndex(std.Thread.getCurrentId());
197 pool.mutex.unlock();
198 @call(.auto, func, .{id.?} ++ args);
199 wait_group.finish();
200 return;
201 };
202 closure.* = .{
203 .arguments = args,
204 .pool = pool,
205 .wait_group = wait_group,
204206 };
207
208 pool.run_queue.prepend(&closure.runnable.node);
209 pool.mutex.unlock();
205210 }
206211
207 pool.mutex.unlock();
212 // Notify waiting threads outside the lock to try and keep the critical section small.
208213 pool.cond.signal();
209214}
210215
211pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void {
216pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void {
212217 if (builtin.single_threaded) {
213218 @call(.auto, func, args);
214219 return;
......@@ -223,36 +228,30 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void {
223228 fn runFn(runnable: *Runnable, _: ?usize) void {
224229 const closure: *@This() = @alignCast(@fieldParentPtr("runnable", runnable));
225230 @call(.auto, func, closure.arguments);
226 closure.pool.allocator.destroy(closure);
227 }
228 };
229231
230 pool.mutex.lock();
232 // The thread pool's allocator is protected by the mutex.
233 const mutex = &closure.pool.mutex;
234 mutex.lock();
235 defer mutex.unlock();
231236
232 const gpa = pool.allocator;
233 const closure = gpa.create(Closure) catch {
234 pool.mutex.unlock();
235 @call(.auto, func, args);
236 return;
237 };
238 closure.* = .{
239 .arguments = args,
240 .pool = pool,
237 closure.pool.allocator.destroy(closure);
238 }
241239 };
242240
243 pool.run_queue.prepend(&closure.runnable.node);
241 {
242 pool.mutex.lock();
243 defer pool.mutex.unlock();
244244
245 if (pool.threads.items.len < pool.threads.capacity) {
246 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
247 .stack_size = pool.stack_size,
248 .allocator = gpa,
249 }, worker, .{pool}) catch t: {
250 pool.threads.items.len -= 1;
251 break :t undefined;
245 const closure = try pool.allocator.create(Closure);
246 closure.* = .{
247 .arguments = args,
248 .pool = pool,
252249 };
250
251 pool.run_queue.prepend(&closure.runnable.node);
253252 }
254253
255 pool.mutex.unlock();
254 // Notify waiting threads outside the lock to try and keep the critical section small.
256255 pool.cond.signal();
257256}
258257
......@@ -271,7 +270,7 @@ test spawn {
271270 .allocator = std.testing.allocator,
272271 });
273272 defer pool.deinit();
274 pool.spawn(TestFn.checkRun, .{&completed});
273 try pool.spawn(TestFn.checkRun, .{&completed});
275274 }
276275
277276 try std.testing.expectEqual(true, completed);
......@@ -323,530 +322,5 @@ pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void {
323322}
324323
325324pub fn getIdCount(pool: *Pool) usize {
326 return @intCast(1 + pool.threads.items.len);
327}
328
329pub fn io(pool: *Pool) Io {
330 return .{
331 .userdata = pool,
332 .vtable = &.{
333 .async = async,
334 .await = await,
335 .go = go,
336 .cancel = cancel,
337 .cancelRequested = cancelRequested,
338 .select = select,
339
340 .mutexLock = mutexLock,
341 .mutexUnlock = mutexUnlock,
342
343 .conditionWait = conditionWait,
344 .conditionWake = conditionWake,
345
346 .createFile = createFile,
347 .openFile = openFile,
348 .closeFile = closeFile,
349 .pread = pread,
350 .pwrite = pwrite,
351
352 .now = now,
353 .sleep = sleep,
354 },
355 };
356}
357
358const AsyncClosure = struct {
359 func: *const fn (context: *anyopaque, result: *anyopaque) void,
360 runnable: Runnable = .{ .runFn = runFn },
361 reset_event: std.Thread.ResetEvent,
362 select_condition: ?*std.Thread.ResetEvent,
363 cancel_tid: std.Thread.Id,
364 context_offset: usize,
365 result_offset: usize,
366
367 const done_reset_event: *std.Thread.ResetEvent = @ptrFromInt(@alignOf(std.Thread.ResetEvent));
368
369 const canceling_tid: std.Thread.Id = switch (@typeInfo(std.Thread.Id)) {
370 .int => |int_info| switch (int_info.signedness) {
371 .signed => -1,
372 .unsigned => std.math.maxInt(std.Thread.Id),
373 },
374 .pointer => @ptrFromInt(std.math.maxInt(usize)),
375 else => @compileError("unsupported std.Thread.Id: " ++ @typeName(std.Thread.Id)),
376 };
377
378 fn runFn(runnable: *std.Thread.Pool.Runnable, _: ?usize) void {
379 const closure: *AsyncClosure = @alignCast(@fieldParentPtr("runnable", runnable));
380 const tid = std.Thread.getCurrentId();
381 if (@cmpxchgStrong(
382 std.Thread.Id,
383 &closure.cancel_tid,
384 0,
385 tid,
386 .acq_rel,
387 .acquire,
388 )) |cancel_tid| {
389 assert(cancel_tid == canceling_tid);
390 return;
391 }
392 current_closure = closure;
393 closure.func(closure.contextPointer(), closure.resultPointer());
394 current_closure = null;
395 if (@cmpxchgStrong(
396 std.Thread.Id,
397 &closure.cancel_tid,
398 tid,
399 0,
400 .acq_rel,
401 .acquire,
402 )) |cancel_tid| assert(cancel_tid == canceling_tid);
403
404 if (@atomicRmw(
405 ?*std.Thread.ResetEvent,
406 &closure.select_condition,
407 .Xchg,
408 done_reset_event,
409 .release,
410 )) |select_reset| {
411 assert(select_reset != done_reset_event);
412 select_reset.set();
413 }
414 closure.reset_event.set();
415 }
416
417 fn contextOffset(context_alignment: std.mem.Alignment) usize {
418 return context_alignment.forward(@sizeOf(AsyncClosure));
419 }
420
421 fn resultOffset(
422 context_alignment: std.mem.Alignment,
423 context_len: usize,
424 result_alignment: std.mem.Alignment,
425 ) usize {
426 return result_alignment.forward(contextOffset(context_alignment) + context_len);
427 }
428
429 fn resultPointer(closure: *AsyncClosure) [*]u8 {
430 const base: [*]u8 = @ptrCast(closure);
431 return base + closure.result_offset;
432 }
433
434 fn contextPointer(closure: *AsyncClosure) [*]u8 {
435 const base: [*]u8 = @ptrCast(closure);
436 return base + closure.context_offset;
437 }
438
439 fn waitAndFree(closure: *AsyncClosure, gpa: Allocator, result: []u8) void {
440 closure.reset_event.wait();
441 const base: [*]align(@alignOf(AsyncClosure)) u8 = @ptrCast(closure);
442 @memcpy(result, closure.resultPointer()[0..result.len]);
443 gpa.free(base[0 .. closure.result_offset + result.len]);
444 }
445};
446
447fn async(
448 userdata: ?*anyopaque,
449 result: []u8,
450 result_alignment: std.mem.Alignment,
451 context: []const u8,
452 context_alignment: std.mem.Alignment,
453 start: *const fn (context: *const anyopaque, result: *anyopaque) void,
454) ?*Io.AnyFuture {
455 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
456 pool.mutex.lock();
457
458 const gpa = pool.allocator;
459 const context_offset = context_alignment.forward(@sizeOf(AsyncClosure));
460 const result_offset = result_alignment.forward(context_offset + context.len);
461 const n = result_offset + result.len;
462 const closure: *AsyncClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, @alignOf(AsyncClosure), n) catch {
463 pool.mutex.unlock();
464 start(context.ptr, result.ptr);
465 return null;
466 }));
467 closure.* = .{
468 .func = start,
469 .context_offset = context_offset,
470 .result_offset = result_offset,
471 .reset_event = .{},
472 .cancel_tid = 0,
473 .select_condition = null,
474 };
475 @memcpy(closure.contextPointer()[0..context.len], context);
476 pool.run_queue.prepend(&closure.runnable.node);
477
478 if (pool.threads.items.len < pool.threads.capacity) {
479 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
480 .stack_size = pool.stack_size,
481 .allocator = gpa,
482 }, worker, .{pool}) catch t: {
483 pool.threads.items.len -= 1;
484 break :t undefined;
485 };
486 }
487
488 pool.mutex.unlock();
489 pool.cond.signal();
490
491 return @ptrCast(closure);
492}
493
494const DetachedClosure = struct {
495 pool: *Pool,
496 func: *const fn (context: *anyopaque) void,
497 run_node: std.Thread.Pool.RunQueue.Node = .{ .data = .{ .runFn = runFn } },
498 context_alignment: std.mem.Alignment,
499 context_len: usize,
500
501 fn runFn(runnable: *std.Thread.Pool.Runnable, _: ?usize) void {
502 const run_node: *std.Thread.Pool.RunQueue.Node = @fieldParentPtr("data", runnable);
503 const closure: *DetachedClosure = @alignCast(@fieldParentPtr("run_node", run_node));
504 closure.func(closure.contextPointer());
505 const gpa = closure.pool.allocator;
506 const base: [*]align(@alignOf(DetachedClosure)) u8 = @ptrCast(closure);
507 gpa.free(base[0..contextEnd(closure.context_alignment, closure.context_len)]);
508 }
509
510 fn contextOffset(context_alignment: std.mem.Alignment) usize {
511 return context_alignment.forward(@sizeOf(DetachedClosure));
512 }
513
514 fn contextEnd(context_alignment: std.mem.Alignment, context_len: usize) usize {
515 return contextOffset(context_alignment) + context_len;
516 }
517
518 fn contextPointer(closure: *DetachedClosure) [*]u8 {
519 const base: [*]u8 = @ptrCast(closure);
520 return base + contextOffset(closure.context_alignment);
521 }
522};
523
524fn go(
525 userdata: ?*anyopaque,
526 context: []const u8,
527 context_alignment: std.mem.Alignment,
528 start: *const fn (context: *const anyopaque) void,
529) void {
530 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
531 pool.mutex.lock();
532
533 const gpa = pool.allocator;
534 const n = DetachedClosure.contextEnd(context_alignment, context.len);
535 const closure: *DetachedClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, @alignOf(DetachedClosure), n) catch {
536 pool.mutex.unlock();
537 start(context.ptr);
538 return;
539 }));
540 closure.* = .{
541 .pool = pool,
542 .func = start,
543 .context_alignment = context_alignment,
544 .context_len = context.len,
545 };
546 @memcpy(closure.contextPointer()[0..context.len], context);
547 pool.run_queue.prepend(&closure.run_node);
548
549 if (pool.threads.items.len < pool.threads.capacity) {
550 pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{
551 .stack_size = pool.stack_size,
552 .allocator = gpa,
553 }, worker, .{pool}) catch t: {
554 pool.threads.items.len -= 1;
555 break :t undefined;
556 };
557 }
558
559 pool.mutex.unlock();
560 pool.cond.signal();
561}
562
563fn await(
564 userdata: ?*anyopaque,
565 any_future: *std.Io.AnyFuture,
566 result: []u8,
567 result_alignment: std.mem.Alignment,
568) void {
569 _ = result_alignment;
570 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
571 const closure: *AsyncClosure = @ptrCast(@alignCast(any_future));
572 closure.waitAndFree(pool.allocator, result);
573}
574
575fn cancel(
576 userdata: ?*anyopaque,
577 any_future: *Io.AnyFuture,
578 result: []u8,
579 result_alignment: std.mem.Alignment,
580) void {
581 _ = result_alignment;
582 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
583 const closure: *AsyncClosure = @ptrCast(@alignCast(any_future));
584 switch (@atomicRmw(
585 std.Thread.Id,
586 &closure.cancel_tid,
587 .Xchg,
588 AsyncClosure.canceling_tid,
589 .acq_rel,
590 )) {
591 0, AsyncClosure.canceling_tid => {},
592 else => |cancel_tid| switch (builtin.os.tag) {
593 .linux => _ = std.os.linux.tgkill(
594 std.os.linux.getpid(),
595 @bitCast(cancel_tid),
596 std.posix.SIG.IO,
597 ),
598 else => {},
599 },
600 }
601 closure.waitAndFree(pool.allocator, result);
602}
603
604fn cancelRequested(userdata: ?*anyopaque) bool {
605 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
606 _ = pool;
607 const closure = current_closure orelse return false;
608 return @atomicLoad(std.Thread.Id, &closure.cancel_tid, .acquire) == AsyncClosure.canceling_tid;
609}
610
611fn checkCancel(pool: *Pool) error{Canceled}!void {
612 if (cancelRequested(pool)) return error.Canceled;
613}
614
615fn mutexLock(userdata: ?*anyopaque, prev_state: Io.Mutex.State, mutex: *Io.Mutex) error{Canceled}!void {
616 _ = userdata;
617 if (prev_state == .contended) {
618 std.Thread.Futex.wait(@ptrCast(&mutex.state), @intFromEnum(Io.Mutex.State.contended));
619 }
620 while (@atomicRmw(
621 Io.Mutex.State,
622 &mutex.state,
623 .Xchg,
624 .contended,
625 .acquire,
626 ) != .unlocked) {
627 std.Thread.Futex.wait(@ptrCast(&mutex.state), @intFromEnum(Io.Mutex.State.contended));
628 }
629}
630fn mutexUnlock(userdata: ?*anyopaque, prev_state: Io.Mutex.State, mutex: *Io.Mutex) void {
631 _ = userdata;
632 _ = prev_state;
633 if (@atomicRmw(Io.Mutex.State, &mutex.state, .Xchg, .unlocked, .release) == .contended) {
634 std.Thread.Futex.wake(@ptrCast(&mutex.state), 1);
635 }
636}
637
638fn conditionWait(userdata: ?*anyopaque, cond: *Io.Condition, mutex: *Io.Mutex) Io.Cancelable!void {
639 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
640 comptime assert(@TypeOf(cond.state) == u64);
641 const ints: *[2]std.atomic.Value(u32) = @ptrCast(&cond.state);
642 const cond_state = &ints[0];
643 const cond_epoch = &ints[1];
644 const one_waiter = 1;
645 const waiter_mask = 0xffff;
646 const one_signal = 1 << 16;
647 const signal_mask = 0xffff << 16;
648 // Observe the epoch, then check the state again to see if we should wake up.
649 // The epoch must be observed before we check the state or we could potentially miss a wake() and deadlock:
650 //
651 // - T1: s = LOAD(&state)
652 // - T2: UPDATE(&s, signal)
653 // - T2: UPDATE(&epoch, 1) + FUTEX_WAKE(&epoch)
654 // - T1: e = LOAD(&epoch) (was reordered after the state load)
655 // - T1: s & signals == 0 -> FUTEX_WAIT(&epoch, e) (missed the state update + the epoch change)
656 //
657 // Acquire barrier to ensure the epoch load happens before the state load.
658 var epoch = cond_epoch.load(.acquire);
659 var state = cond_state.fetchAdd(one_waiter, .monotonic);
660 assert(state & waiter_mask != waiter_mask);
661 state += one_waiter;
662
663 mutex.unlock(pool.io());
664 defer mutex.lock(pool.io()) catch @panic("TODO");
665
666 var futex_deadline = std.Thread.Futex.Deadline.init(null);
667
668 while (true) {
669 futex_deadline.wait(cond_epoch, epoch) catch |err| switch (err) {
670 error.Timeout => unreachable,
671 };
672
673 epoch = cond_epoch.load(.acquire);
674 state = cond_state.load(.monotonic);
675
676 // Try to wake up by consuming a signal and decremented the waiter we added previously.
677 // Acquire barrier ensures code before the wake() which added the signal happens before we decrement it and return.
678 while (state & signal_mask != 0) {
679 const new_state = state - one_waiter - one_signal;
680 state = cond_state.cmpxchgWeak(state, new_state, .acquire, .monotonic) orelse return;
681 }
682 }
683}
684
685fn conditionWake(userdata: ?*anyopaque, cond: *Io.Condition, wake: Io.Condition.Wake) void {
686 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
687 _ = pool;
688 comptime assert(@TypeOf(cond.state) == u64);
689 const ints: *[2]std.atomic.Value(u32) = @ptrCast(&cond.state);
690 const cond_state = &ints[0];
691 const cond_epoch = &ints[1];
692 const one_waiter = 1;
693 const waiter_mask = 0xffff;
694 const one_signal = 1 << 16;
695 const signal_mask = 0xffff << 16;
696 var state = cond_state.load(.monotonic);
697 while (true) {
698 const waiters = (state & waiter_mask) / one_waiter;
699 const signals = (state & signal_mask) / one_signal;
700
701 // Reserves which waiters to wake up by incrementing the signals count.
702 // Therefore, the signals count is always less than or equal to the waiters count.
703 // We don't need to Futex.wake if there's nothing to wake up or if other wake() threads have reserved to wake up the current waiters.
704 const wakeable = waiters - signals;
705 if (wakeable == 0) {
706 return;
707 }
708
709 const to_wake = switch (wake) {
710 .one => 1,
711 .all => wakeable,
712 };
713
714 // Reserve the amount of waiters to wake by incrementing the signals count.
715 // Release barrier ensures code before the wake() happens before the signal it posted and consumed by the wait() threads.
716 const new_state = state + (one_signal * to_wake);
717 state = cond_state.cmpxchgWeak(state, new_state, .release, .monotonic) orelse {
718 // Wake up the waiting threads we reserved above by changing the epoch value.
719 // NOTE: a waiting thread could miss a wake up if *exactly* ((1<<32)-1) wake()s happen between it observing the epoch and sleeping on it.
720 // This is very unlikely due to how many precise amount of Futex.wake() calls that would be between the waiting thread's potential preemption.
721 //
722 // Release barrier ensures the signal being added to the state happens before the epoch is changed.
723 // If not, the waiting thread could potentially deadlock from missing both the state and epoch change:
724 //
725 // - T2: UPDATE(&epoch, 1) (reordered before the state change)
726 // - T1: e = LOAD(&epoch)
727 // - T1: s = LOAD(&state)
728 // - T2: UPDATE(&state, signal) + FUTEX_WAKE(&epoch)
729 // - T1: s & signals == 0 -> FUTEX_WAIT(&epoch, e) (missed both epoch change and state change)
730 _ = cond_epoch.fetchAdd(1, .release);
731 std.Thread.Futex.wake(cond_epoch, to_wake);
732 return;
733 };
734 }
735}
736
737fn createFile(
738 userdata: ?*anyopaque,
739 dir: Io.Dir,
740 sub_path: []const u8,
741 flags: Io.File.CreateFlags,
742) Io.File.OpenError!Io.File {
743 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
744 try pool.checkCancel();
745 const fs_dir: std.fs.Dir = .{ .fd = dir.handle };
746 const fs_file = try fs_dir.createFile(sub_path, flags);
747 return .{ .handle = fs_file.handle };
748}
749
750fn openFile(
751 userdata: ?*anyopaque,
752 dir: Io.Dir,
753 sub_path: []const u8,
754 flags: Io.File.OpenFlags,
755) Io.File.OpenError!Io.File {
756 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
757 try pool.checkCancel();
758 const fs_dir: std.fs.Dir = .{ .fd = dir.handle };
759 const fs_file = try fs_dir.openFile(sub_path, flags);
760 return .{ .handle = fs_file.handle };
761}
762
763fn closeFile(userdata: ?*anyopaque, file: Io.File) void {
764 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
765 _ = pool;
766 const fs_file: std.fs.File = .{ .handle = file.handle };
767 return fs_file.close();
768}
769
770fn pread(userdata: ?*anyopaque, file: Io.File, buffer: []u8, offset: std.posix.off_t) Io.File.PReadError!usize {
771 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
772 try pool.checkCancel();
773 const fs_file: std.fs.File = .{ .handle = file.handle };
774 return switch (offset) {
775 -1 => fs_file.read(buffer),
776 else => fs_file.pread(buffer, @bitCast(offset)),
777 };
778}
779
780fn pwrite(userdata: ?*anyopaque, file: Io.File, buffer: []const u8, offset: std.posix.off_t) Io.File.PWriteError!usize {
781 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
782 try pool.checkCancel();
783 const fs_file: std.fs.File = .{ .handle = file.handle };
784 return switch (offset) {
785 -1 => fs_file.write(buffer),
786 else => fs_file.pwrite(buffer, @bitCast(offset)),
787 };
788}
789
790fn now(userdata: ?*anyopaque, clockid: std.posix.clockid_t) Io.ClockGetTimeError!Io.Timestamp {
791 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
792 try pool.checkCancel();
793 const timespec = try std.posix.clock_gettime(clockid);
794 return @enumFromInt(@as(i128, timespec.sec) * std.time.ns_per_s + timespec.nsec);
795}
796
797fn sleep(userdata: ?*anyopaque, clockid: std.posix.clockid_t, deadline: Io.Deadline) Io.SleepError!void {
798 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
799 const deadline_nanoseconds: i96 = switch (deadline) {
800 .duration => |duration| duration.nanoseconds,
801 .timestamp => |timestamp| @intFromEnum(timestamp),
802 };
803 var timespec: std.posix.timespec = .{
804 .sec = @intCast(@divFloor(deadline_nanoseconds, std.time.ns_per_s)),
805 .nsec = @intCast(@mod(deadline_nanoseconds, std.time.ns_per_s)),
806 };
807 while (true) {
808 try pool.checkCancel();
809 switch (std.os.linux.E.init(std.os.linux.clock_nanosleep(clockid, .{ .ABSTIME = switch (deadline) {
810 .duration => false,
811 .timestamp => true,
812 } }, &timespec, &timespec))) {
813 .SUCCESS => return,
814 .FAULT => unreachable,
815 .INTR => {},
816 .INVAL => return error.UnsupportedClock,
817 else => |err| return std.posix.unexpectedErrno(err),
818 }
819 }
820}
821
822fn select(userdata: ?*anyopaque, futures: []const *Io.AnyFuture) usize {
823 const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata));
824 _ = pool;
825
826 var reset_event: std.Thread.ResetEvent = .{};
827
828 for (futures, 0..) |future, i| {
829 const closure: *AsyncClosure = @ptrCast(@alignCast(future));
830 if (@atomicRmw(?*std.Thread.ResetEvent, &closure.select_condition, .Xchg, &reset_event, .seq_cst) == AsyncClosure.done_reset_event) {
831 for (futures[0..i]) |cleanup_future| {
832 const cleanup_closure: *AsyncClosure = @ptrCast(@alignCast(cleanup_future));
833 if (@atomicRmw(?*std.Thread.ResetEvent, &cleanup_closure.select_condition, .Xchg, null, .seq_cst) == AsyncClosure.done_reset_event) {
834 cleanup_closure.reset_event.wait(); // Ensure no reference to our stack-allocated reset_event.
835 }
836 }
837 return i;
838 }
839 }
840
841 reset_event.wait();
842
843 var result: ?usize = null;
844 for (futures, 0..) |future, i| {
845 const closure: *AsyncClosure = @ptrCast(@alignCast(future));
846 if (@atomicRmw(?*std.Thread.ResetEvent, &closure.select_condition, .Xchg, null, .seq_cst) == AsyncClosure.done_reset_event) {
847 closure.reset_event.wait(); // Ensure no reference to our stack-allocated reset_event.
848 if (result == null) result = i; // In case multiple are ready, return first.
849 }
850 }
851 return result.?;
325 return @intCast(1 + pool.threads.len);
852326}