| ... | ... | @@ -1,22 +1,27 @@ |
| 1 | 1 | const builtin = @import("builtin"); |
| 2 | 2 | const std = @import("std"); |
| 3 | const Allocator = std.mem.Allocator; |
| 3 | 4 | const assert = std.debug.assert; |
| 4 | 5 | const WaitGroup = @import("WaitGroup.zig"); |
| 6 | const Io = std.Io; |
| 5 | 7 | const Pool = @This(); |
| 6 | 8 | |
| 9 | /// Must be a thread-safe allocator. |
| 10 | allocator: std.mem.Allocator, |
| 7 | 11 | mutex: std.Thread.Mutex = .{}, |
| 8 | 12 | cond: std.Thread.Condition = .{}, |
| 9 | 13 | run_queue: std.SinglyLinkedList = .{}, |
| 10 | 14 | is_running: bool = true, |
| 11 | | /// Must be a thread-safe allocator. |
| 12 | | allocator: std.mem.Allocator, |
| 13 | | threads: if (builtin.single_threaded) [0]std.Thread else []std.Thread, |
| 15 | threads: std.ArrayListUnmanaged(std.Thread), |
| 14 | 16 | ids: if (builtin.single_threaded) struct { |
| 15 | 17 | inline fn deinit(_: @This(), _: std.mem.Allocator) void {} |
| 16 | 18 | fn getIndex(_: @This(), _: std.Thread.Id) usize { |
| 17 | 19 | return 0; |
| 18 | 20 | } |
| 19 | 21 | } else std.AutoArrayHashMapUnmanaged(std.Thread.Id, void), |
| 22 | stack_size: usize, |
| 23 | |
| 24 | threadlocal var current_closure: ?*AsyncClosure = null; |
| 20 | 25 | |
| 21 | 26 | pub const Runnable = struct { |
| 22 | 27 | runFn: RunProto, |
| ... | ... | @@ -33,48 +38,36 @@ pub const Options = struct { |
| 33 | 38 | }; |
| 34 | 39 | |
| 35 | 40 | pub fn init(pool: *Pool, options: Options) !void { |
| 36 | | const allocator = options.allocator; |
| 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); |
| 37 | 45 | |
| 38 | 46 | pool.* = .{ |
| 39 | | .allocator = allocator, |
| 40 | | .threads = if (builtin.single_threaded) .{} else &.{}, |
| 47 | .allocator = gpa, |
| 48 | .threads = .initBuffer(threads), |
| 41 | 49 | .ids = .{}, |
| 50 | .stack_size = options.stack_size, |
| 42 | 51 | }; |
| 43 | 52 | |
| 44 | | if (builtin.single_threaded) { |
| 45 | | return; |
| 46 | | } |
| 53 | if (builtin.single_threaded) return; |
| 47 | 54 | |
| 48 | | const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1); |
| 49 | 55 | if (options.track_ids) { |
| 50 | | try pool.ids.ensureTotalCapacity(allocator, 1 + thread_count); |
| 56 | try pool.ids.ensureTotalCapacity(gpa, 1 + thread_count); |
| 51 | 57 | pool.ids.putAssumeCapacityNoClobber(std.Thread.getCurrentId(), {}); |
| 52 | 58 | } |
| 53 | | |
| 54 | | // kill and join any threads we spawned and free memory on error. |
| 55 | | pool.threads = try allocator.alloc(std.Thread, thread_count); |
| 56 | | var spawned: usize = 0; |
| 57 | | errdefer pool.join(spawned); |
| 58 | | |
| 59 | | for (pool.threads) |*thread| { |
| 60 | | thread.* = try std.Thread.spawn(.{ |
| 61 | | .stack_size = options.stack_size, |
| 62 | | .allocator = allocator, |
| 63 | | }, worker, .{pool}); |
| 64 | | spawned += 1; |
| 65 | | } |
| 66 | 59 | } |
| 67 | 60 | |
| 68 | 61 | pub fn deinit(pool: *Pool) void { |
| 69 | | pool.join(pool.threads.len); // kill and join all threads. |
| 70 | | pool.ids.deinit(pool.allocator); |
| 62 | const gpa = pool.allocator; |
| 63 | pool.join(); |
| 64 | pool.threads.deinit(gpa); |
| 65 | pool.ids.deinit(gpa); |
| 71 | 66 | pool.* = undefined; |
| 72 | 67 | } |
| 73 | 68 | |
| 74 | | fn join(pool: *Pool, spawned: usize) void { |
| 75 | | if (builtin.single_threaded) { |
| 76 | | return; |
| 77 | | } |
| 69 | fn join(pool: *Pool) void { |
| 70 | if (builtin.single_threaded) return; |
| 78 | 71 | |
| 79 | 72 | { |
| 80 | 73 | pool.mutex.lock(); |
| ... | ... | @@ -87,11 +80,7 @@ fn join(pool: *Pool, spawned: usize) void { |
| 87 | 80 | // wake up any sleeping threads (this can be done outside the mutex) |
| 88 | 81 | // then wait for all the threads we know are spawned to complete. |
| 89 | 82 | pool.cond.broadcast(); |
| 90 | | for (pool.threads[0..spawned]) |thread| { |
| 91 | | thread.join(); |
| 92 | | } |
| 93 | | |
| 94 | | pool.allocator.free(pool.threads); |
| 83 | for (pool.threads.items) |thread| thread.join(); |
| 95 | 84 | } |
| 96 | 85 | |
| 97 | 86 | /// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and |
| ... | ... | @@ -123,26 +112,34 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args |
| 123 | 112 | } |
| 124 | 113 | }; |
| 125 | 114 | |
| 126 | | { |
| 127 | | pool.mutex.lock(); |
| 128 | | |
| 129 | | const closure = pool.allocator.create(Closure) catch { |
| 130 | | pool.mutex.unlock(); |
| 131 | | @call(.auto, func, args); |
| 132 | | wait_group.finish(); |
| 133 | | return; |
| 134 | | }; |
| 135 | | closure.* = .{ |
| 136 | | .arguments = args, |
| 137 | | .pool = pool, |
| 138 | | .wait_group = wait_group, |
| 139 | | }; |
| 115 | pool.mutex.lock(); |
| 140 | 116 | |
| 141 | | pool.run_queue.prepend(&closure.runnable.node); |
| 117 | const gpa = pool.allocator; |
| 118 | const closure = gpa.create(Closure) catch { |
| 142 | 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 | }; |
| 143 | 140 | } |
| 144 | 141 | |
| 145 | | // Notify waiting threads outside the lock to try and keep the critical section small. |
| 142 | pool.mutex.unlock(); |
| 146 | 143 | pool.cond.signal(); |
| 147 | 144 | } |
| 148 | 145 | |
| ... | ... | @@ -179,31 +176,39 @@ pub fn spawnWgId(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, ar |
| 179 | 176 | } |
| 180 | 177 | }; |
| 181 | 178 | |
| 182 | | { |
| 183 | | pool.mutex.lock(); |
| 184 | | |
| 185 | | const closure = pool.allocator.create(Closure) catch { |
| 186 | | const id: ?usize = pool.ids.getIndex(std.Thread.getCurrentId()); |
| 187 | | pool.mutex.unlock(); |
| 188 | | @call(.auto, func, .{id.?} ++ args); |
| 189 | | wait_group.finish(); |
| 190 | | return; |
| 191 | | }; |
| 192 | | closure.* = .{ |
| 193 | | .arguments = args, |
| 194 | | .pool = pool, |
| 195 | | .wait_group = wait_group, |
| 196 | | }; |
| 179 | pool.mutex.lock(); |
| 197 | 180 | |
| 198 | | pool.run_queue.prepend(&closure.runnable.node); |
| 181 | const gpa = pool.allocator; |
| 182 | const closure = gpa.create(Closure) catch { |
| 183 | const id: ?usize = pool.ids.getIndex(std.Thread.getCurrentId()); |
| 199 | 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 | }; |
| 200 | 205 | } |
| 201 | 206 | |
| 202 | | // Notify waiting threads outside the lock to try and keep the critical section small. |
| 207 | pool.mutex.unlock(); |
| 203 | 208 | pool.cond.signal(); |
| 204 | 209 | } |
| 205 | 210 | |
| 206 | | pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { |
| 211 | pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void { |
| 207 | 212 | if (builtin.single_threaded) { |
| 208 | 213 | @call(.auto, func, args); |
| 209 | 214 | return; |
| ... | ... | @@ -222,20 +227,32 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { |
| 222 | 227 | } |
| 223 | 228 | }; |
| 224 | 229 | |
| 225 | | { |
| 226 | | pool.mutex.lock(); |
| 227 | | defer pool.mutex.unlock(); |
| 230 | pool.mutex.lock(); |
| 228 | 231 | |
| 229 | | const closure = try pool.allocator.create(Closure); |
| 230 | | closure.* = .{ |
| 231 | | .arguments = args, |
| 232 | | .pool = pool, |
| 233 | | }; |
| 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); |
| 234 | 244 | |
| 235 | | pool.run_queue.prepend(&closure.runnable.node); |
| 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 | }; |
| 236 | 253 | } |
| 237 | 254 | |
| 238 | | // Notify waiting threads outside the lock to try and keep the critical section small. |
| 255 | pool.mutex.unlock(); |
| 239 | 256 | pool.cond.signal(); |
| 240 | 257 | } |
| 241 | 258 | |
| ... | ... | @@ -254,7 +271,7 @@ test spawn { |
| 254 | 271 | .allocator = std.testing.allocator, |
| 255 | 272 | }); |
| 256 | 273 | defer pool.deinit(); |
| 257 | | try pool.spawn(TestFn.checkRun, .{&completed}); |
| 274 | pool.spawn(TestFn.checkRun, .{&completed}); |
| 258 | 275 | } |
| 259 | 276 | |
| 260 | 277 | try std.testing.expectEqual(true, completed); |
| ... | ... | @@ -306,15 +323,17 @@ pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void { |
| 306 | 323 | } |
| 307 | 324 | |
| 308 | 325 | pub fn getIdCount(pool: *Pool) usize { |
| 309 | | return @intCast(1 + pool.threads.len); |
| 326 | return @intCast(1 + pool.threads.items.len); |
| 310 | 327 | } |
| 311 | 328 | |
| 312 | | pub fn io(pool: *Pool) std.Io { |
| 329 | pub fn io(pool: *Pool) Io { |
| 313 | 330 | return .{ |
| 314 | 331 | .userdata = pool, |
| 315 | 332 | .vtable = &.{ |
| 316 | 333 | .@"async" = @"async", |
| 317 | 334 | .@"await" = @"await", |
| 335 | .cancel = cancel, |
| 336 | .cancelRequested = cancelRequested, |
| 318 | 337 | .createFile = createFile, |
| 319 | 338 | .openFile = openFile, |
| 320 | 339 | .closeFile = closeFile, |
| ... | ... | @@ -326,15 +345,17 @@ pub fn io(pool: *Pool) std.Io { |
| 326 | 345 | |
| 327 | 346 | const AsyncClosure = struct { |
| 328 | 347 | func: *const fn (context: *anyopaque, result: *anyopaque) void, |
| 329 | | run_node: std.Thread.Pool.RunQueue.Node = .{ .data = .{ .runFn = runFn } }, |
| 348 | runnable: Runnable = .{ .runFn = runFn }, |
| 330 | 349 | reset_event: std.Thread.ResetEvent, |
| 350 | cancel_flag: bool, |
| 331 | 351 | context_offset: usize, |
| 332 | 352 | result_offset: usize, |
| 333 | 353 | |
| 334 | 354 | fn runFn(runnable: *std.Thread.Pool.Runnable, _: ?usize) void { |
| 335 | | const run_node: *std.Thread.Pool.RunQueue.Node = @fieldParentPtr("data", runnable); |
| 336 | | const closure: *AsyncClosure = @alignCast(@fieldParentPtr("run_node", run_node)); |
| 355 | const closure: *AsyncClosure = @alignCast(@fieldParentPtr("runnable", runnable)); |
| 356 | current_closure = closure; |
| 337 | 357 | closure.func(closure.contextPointer(), closure.resultPointer()); |
| 358 | current_closure = null; |
| 338 | 359 | closure.reset_event.set(); |
| 339 | 360 | } |
| 340 | 361 | |
| ... | ... | @@ -359,16 +380,23 @@ const AsyncClosure = struct { |
| 359 | 380 | const base: [*]u8 = @ptrCast(closure); |
| 360 | 381 | return base + closure.context_offset; |
| 361 | 382 | } |
| 383 | |
| 384 | fn waitAndFree(closure: *AsyncClosure, gpa: Allocator, result: []u8) void { |
| 385 | closure.reset_event.wait(); |
| 386 | const base: [*]align(@alignOf(AsyncClosure)) u8 = @ptrCast(closure); |
| 387 | @memcpy(result, closure.resultPointer()[0..result.len]); |
| 388 | gpa.free(base[0 .. closure.result_offset + result.len]); |
| 389 | } |
| 362 | 390 | }; |
| 363 | 391 | |
| 364 | | pub fn @"async"( |
| 392 | fn @"async"( |
| 365 | 393 | userdata: ?*anyopaque, |
| 366 | 394 | result: []u8, |
| 367 | 395 | result_alignment: std.mem.Alignment, |
| 368 | 396 | context: []const u8, |
| 369 | 397 | context_alignment: std.mem.Alignment, |
| 370 | 398 | start: *const fn (context: *const anyopaque, result: *anyopaque) void, |
| 371 | | ) ?*std.Io.AnyFuture { |
| 399 | ) ?*Io.AnyFuture { |
| 372 | 400 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 373 | 401 | pool.mutex.lock(); |
| 374 | 402 | |
| ... | ... | @@ -386,46 +414,87 @@ pub fn @"async"( |
| 386 | 414 | .context_offset = context_offset, |
| 387 | 415 | .result_offset = result_offset, |
| 388 | 416 | .reset_event = .{}, |
| 417 | .cancel_flag = false, |
| 389 | 418 | }; |
| 390 | 419 | @memcpy(closure.contextPointer()[0..context.len], context); |
| 391 | | pool.run_queue.prepend(&closure.run_node); |
| 392 | | pool.mutex.unlock(); |
| 420 | pool.run_queue.prepend(&closure.runnable.node); |
| 421 | |
| 422 | if (pool.threads.items.len < pool.threads.capacity) { |
| 423 | pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{ |
| 424 | .stack_size = pool.stack_size, |
| 425 | .allocator = gpa, |
| 426 | }, worker, .{pool}) catch t: { |
| 427 | pool.threads.items.len -= 1; |
| 428 | break :t undefined; |
| 429 | }; |
| 430 | } |
| 393 | 431 | |
| 432 | pool.mutex.unlock(); |
| 394 | 433 | pool.cond.signal(); |
| 395 | 434 | |
| 396 | 435 | return @ptrCast(closure); |
| 397 | 436 | } |
| 398 | 437 | |
| 399 | | pub fn @"await"(userdata: ?*anyopaque, any_future: *std.Io.AnyFuture, result: []u8) void { |
| 400 | | const thread_pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 438 | fn @"await"(userdata: ?*anyopaque, any_future: *Io.AnyFuture, result: []u8) void { |
| 439 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 401 | 440 | const closure: *AsyncClosure = @ptrCast(@alignCast(any_future)); |
| 402 | | closure.reset_event.wait(); |
| 403 | | const base: [*]align(@alignOf(AsyncClosure)) u8 = @ptrCast(closure); |
| 404 | | @memcpy(result, closure.resultPointer()[0..result.len]); |
| 405 | | thread_pool.allocator.free(base[0 .. closure.result_offset + result.len]); |
| 441 | closure.waitAndFree(pool.allocator, result); |
| 442 | } |
| 443 | |
| 444 | fn cancel(userdata: ?*anyopaque, any_future: *Io.AnyFuture, result: []u8) void { |
| 445 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 446 | const closure: *AsyncClosure = @ptrCast(@alignCast(any_future)); |
| 447 | @atomicStore(bool, &closure.cancel_flag, true, .seq_cst); |
| 448 | closure.waitAndFree(pool.allocator, result); |
| 449 | } |
| 450 | |
| 451 | fn cancelRequested(userdata: ?*anyopaque) bool { |
| 452 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 453 | _ = pool; |
| 454 | const closure = current_closure orelse return false; |
| 455 | return @atomicLoad(bool, &closure.cancel_flag, .unordered); |
| 456 | } |
| 457 | |
| 458 | fn checkCancel(pool: *Pool) error{AsyncCancel}!void { |
| 459 | if (cancelRequested(pool)) return error.AsyncCancel; |
| 406 | 460 | } |
| 407 | 461 | |
| 408 | | pub fn createFile(userdata: ?*anyopaque, dir: std.fs.Dir, sub_path: []const u8, flags: std.fs.File.CreateFlags) std.fs.File.OpenError!std.fs.File { |
| 409 | | _ = userdata; |
| 462 | pub fn createFile( |
| 463 | userdata: ?*anyopaque, |
| 464 | dir: std.fs.Dir, |
| 465 | sub_path: []const u8, |
| 466 | flags: std.fs.File.CreateFlags, |
| 467 | ) Io.FileOpenError!std.fs.File { |
| 468 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 469 | try pool.checkCancel(); |
| 410 | 470 | return dir.createFile(sub_path, flags); |
| 411 | 471 | } |
| 412 | 472 | |
| 413 | | pub fn openFile(userdata: ?*anyopaque, dir: std.fs.Dir, sub_path: []const u8, flags: std.fs.File.OpenFlags) std.fs.File.OpenError!std.fs.File { |
| 414 | | _ = userdata; |
| 473 | pub fn openFile( |
| 474 | userdata: ?*anyopaque, |
| 475 | dir: std.fs.Dir, |
| 476 | sub_path: []const u8, |
| 477 | flags: std.fs.File.OpenFlags, |
| 478 | ) Io.FileOpenError!std.fs.File { |
| 479 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 480 | try pool.checkCancel(); |
| 415 | 481 | return dir.openFile(sub_path, flags); |
| 416 | 482 | } |
| 417 | 483 | |
| 418 | 484 | pub fn closeFile(userdata: ?*anyopaque, file: std.fs.File) void { |
| 419 | | _ = userdata; |
| 485 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 486 | _ = pool; |
| 420 | 487 | return file.close(); |
| 421 | 488 | } |
| 422 | 489 | |
| 423 | | pub fn read(userdata: ?*anyopaque, file: std.fs.File, buffer: []u8) std.fs.File.ReadError!usize { |
| 424 | | _ = userdata; |
| 490 | pub fn read(userdata: ?*anyopaque, file: std.fs.File, buffer: []u8) Io.FileReadError!usize { |
| 491 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 492 | try pool.checkCancel(); |
| 425 | 493 | return file.read(buffer); |
| 426 | 494 | } |
| 427 | 495 | |
| 428 | | pub fn write(userdata: ?*anyopaque, file: std.fs.File, buffer: []const u8) std.fs.File.WriteError!usize { |
| 429 | | _ = userdata; |
| 496 | pub fn write(userdata: ?*anyopaque, file: std.fs.File, buffer: []const u8) Io.FileWriteError!usize { |
| 497 | const pool: *std.Thread.Pool = @alignCast(@ptrCast(userdata)); |
| 498 | try pool.checkCancel(); |
| 430 | 499 | return file.write(buffer); |
| 431 | 500 | } |