| ... | ... | @@ -6,331 +6,88 @@ const WaitGroup = std.Thread.WaitGroup; |
| 6 | 6 | const Io = std.Io; |
| 7 | 7 | const Pool = @This(); |
| 8 | 8 | |
| 9 | | /// Must be a thread-safe allocator. |
| 10 | | allocator: std.mem.Allocator, |
| 9 | /// Thread-safe. |
| 10 | allocator: Allocator, |
| 11 | 11 | mutex: std.Thread.Mutex = .{}, |
| 12 | 12 | cond: std.Thread.Condition = .{}, |
| 13 | 13 | run_queue: std.SinglyLinkedList = .{}, |
| 14 | | is_running: bool = true, |
| 14 | join_requested: bool = false, |
| 15 | 15 | threads: std.ArrayListUnmanaged(std.Thread), |
| 16 | | ids: 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), |
| 22 | 16 | stack_size: usize, |
| 17 | cpu_count: std.Thread.CpuCountError!usize, |
| 18 | parallel_count: usize, |
| 23 | 19 | |
| 24 | 20 | threadlocal var current_closure: ?*AsyncClosure = null; |
| 25 | 21 | |
| 26 | 22 | pub const Runnable = struct { |
| 27 | | runFn: RunProto, |
| 23 | start: Start, |
| 28 | 24 | node: std.SinglyLinkedList.Node = .{}, |
| 29 | | }; |
| 30 | | |
| 31 | | pub const RunProto = *const fn (*Runnable, id: ?usize) void; |
| 25 | is_parallel: bool, |
| 32 | 26 | |
| 33 | | pub 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, |
| 27 | pub const Start = *const fn (*Runnable) void; |
| 38 | 28 | }; |
| 39 | 29 | |
| 40 | | pub 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); |
| 30 | pub const InitError = std.Thread.CpuCountError || Allocator.Error; |
| 45 | 31 | |
| 46 | | pool.* = .{ |
| 32 | pub fn init(gpa: Allocator) Pool { |
| 33 | var pool: Pool = .{ |
| 47 | 34 | .allocator = gpa, |
| 48 | | .threads = .initBuffer(threads), |
| 49 | | .ids = .{}, |
| 50 | | .stack_size = options.stack_size, |
| 35 | .threads = .empty, |
| 36 | .stack_size = std.Thread.SpawnConfig.default_stack_size, |
| 37 | .cpu_count = std.Thread.getCpuCount(), |
| 38 | .parallel_count = 0, |
| 51 | 39 | }; |
| 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 | | } |
| 40 | if (pool.cpu_count) |n| { |
| 41 | pool.threads.ensureTotalCapacityPrecise(gpa, n - 1) catch {}; |
| 42 | } else |_| {} |
| 43 | return pool; |
| 59 | 44 | } |
| 60 | 45 | |
| 61 | 46 | pub fn deinit(pool: *Pool) void { |
| 62 | 47 | const gpa = pool.allocator; |
| 63 | 48 | pool.join(); |
| 64 | 49 | pool.threads.deinit(gpa); |
| 65 | | pool.ids.deinit(gpa); |
| 66 | 50 | pool.* = undefined; |
| 67 | 51 | } |
| 68 | 52 | |
| 69 | 53 | fn join(pool: *Pool) void { |
| 70 | 54 | if (builtin.single_threaded) return; |
| 71 | | |
| 72 | 55 | { |
| 73 | 56 | pool.mutex.lock(); |
| 74 | 57 | defer pool.mutex.unlock(); |
| 75 | | |
| 76 | | // ensure future worker threads exit the dequeue loop |
| 77 | | pool.is_running = false; |
| 58 | pool.join_requested = true; |
| 78 | 59 | } |
| 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 | 60 | pool.cond.broadcast(); |
| 83 | 61 | for (pool.threads.items) |thread| thread.join(); |
| 84 | 62 | } |
| 85 | 63 | |
| 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. |
| 91 | | pub 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. |
| 155 | | pub 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 | | |
| 211 | | pub 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 | | |
| 259 | | test 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 | | |
| 280 | 64 | fn worker(pool: *Pool) void { |
| 281 | 65 | pool.mutex.lock(); |
| 282 | 66 | defer pool.mutex.unlock(); |
| 283 | 67 | |
| 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 | 68 | while (true) { |
| 288 | 69 | 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 | | |
| 306 | | pub 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 | 70 | pool.mutex.unlock(); |
| 314 | 71 | const runnable: *Runnable = @fieldParentPtr("node", run_node); |
| 315 | | runnable.runFn(runnable, id); |
| 316 | | continue; |
| 72 | runnable.start(runnable); |
| 73 | pool.mutex.lock(); |
| 74 | if (runnable.is_parallel) { |
| 75 | // TODO also pop thread and join sometimes |
| 76 | pool.parallel_count -= 1; |
| 77 | } |
| 317 | 78 | } |
| 318 | | |
| 319 | | pool.mutex.unlock(); |
| 320 | | wait_group.wait(); |
| 321 | | return; |
| 79 | if (pool.join_requested) break; |
| 80 | pool.cond.wait(&pool.mutex); |
| 322 | 81 | } |
| 323 | 82 | } |
| 324 | 83 | |
| 325 | | pub fn getIdCount(pool: *Pool) usize { |
| 326 | | return @intCast(1 + pool.threads.items.len); |
| 327 | | } |
| 328 | | |
| 329 | 84 | pub fn io(pool: *Pool) Io { |
| 330 | 85 | return .{ |
| 331 | 86 | .userdata = pool, |
| 332 | 87 | .vtable = &.{ |
| 333 | 88 | .async = async, |
| 89 | .asyncConcurrent = asyncParallel, |
| 90 | .asyncParallel = asyncParallel, |
| 334 | 91 | .await = await, |
| 335 | 92 | .asyncDetached = asyncDetached, |
| 336 | 93 | .cancel = cancel, |
| ... | ... | @@ -357,7 +114,7 @@ pub fn io(pool: *Pool) Io { |
| 357 | 114 | |
| 358 | 115 | const AsyncClosure = struct { |
| 359 | 116 | func: *const fn (context: *anyopaque, result: *anyopaque) void, |
| 360 | | runnable: Runnable = .{ .runFn = runFn }, |
| 117 | runnable: Runnable, |
| 361 | 118 | reset_event: std.Thread.ResetEvent, |
| 362 | 119 | select_condition: ?*std.Thread.ResetEvent, |
| 363 | 120 | cancel_tid: std.Thread.Id, |
| ... | ... | @@ -375,7 +132,7 @@ const AsyncClosure = struct { |
| 375 | 132 | else => @compileError("unsupported std.Thread.Id: " ++ @typeName(std.Thread.Id)), |
| 376 | 133 | }; |
| 377 | 134 | |
| 378 | | fn runFn(runnable: *Pool.Runnable, _: ?usize) void { |
| 135 | fn start(runnable: *Runnable) void { |
| 379 | 136 | const closure: *AsyncClosure = @alignCast(@fieldParentPtr("runnable", runnable)); |
| 380 | 137 | const tid = std.Thread.getCurrentId(); |
| 381 | 138 | if (@cmpxchgStrong( |
| ... | ... | @@ -387,6 +144,7 @@ const AsyncClosure = struct { |
| 387 | 144 | .acquire, |
| 388 | 145 | )) |cancel_tid| { |
| 389 | 146 | assert(cancel_tid == canceling_tid); |
| 147 | closure.reset_event.set(); |
| 390 | 148 | return; |
| 391 | 149 | } |
| 392 | 150 | current_closure = closure; |
| ... | ... | @@ -438,9 +196,13 @@ const AsyncClosure = struct { |
| 438 | 196 | |
| 439 | 197 | fn waitAndFree(closure: *AsyncClosure, gpa: Allocator, result: []u8) void { |
| 440 | 198 | closure.reset_event.wait(); |
| 441 | | const base: [*]align(@alignOf(AsyncClosure)) u8 = @ptrCast(closure); |
| 442 | 199 | @memcpy(result, closure.resultPointer()[0..result.len]); |
| 443 | | gpa.free(base[0 .. closure.result_offset + result.len]); |
| 200 | free(closure, gpa, result.len); |
| 201 | } |
| 202 | |
| 203 | fn free(closure: *AsyncClosure, gpa: Allocator, result_len: usize) void { |
| 204 | const base: [*]align(@alignOf(AsyncClosure)) u8 = @ptrCast(closure); |
| 205 | gpa.free(base[0 .. closure.result_offset + result_len]); |
| 444 | 206 | } |
| 445 | 207 | }; |
| 446 | 208 | |
| ... | ... | @@ -452,18 +214,26 @@ fn async( |
| 452 | 214 | context_alignment: std.mem.Alignment, |
| 453 | 215 | start: *const fn (context: *const anyopaque, result: *anyopaque) void, |
| 454 | 216 | ) ?*Io.AnyFuture { |
| 217 | if (builtin.single_threaded) { |
| 218 | start(context.ptr, result.ptr); |
| 219 | return null; |
| 220 | } |
| 455 | 221 | const pool: *Pool = @alignCast(@ptrCast(userdata)); |
| 456 | | pool.mutex.lock(); |
| 457 | | |
| 222 | const cpu_count = pool.cpu_count catch { |
| 223 | return asyncParallel(userdata, result.len, result_alignment, context, context_alignment, start) orelse { |
| 224 | start(context.ptr, result.ptr); |
| 225 | return null; |
| 226 | }; |
| 227 | }; |
| 458 | 228 | const gpa = pool.allocator; |
| 459 | 229 | const context_offset = context_alignment.forward(@sizeOf(AsyncClosure)); |
| 460 | 230 | const result_offset = result_alignment.forward(context_offset + context.len); |
| 461 | 231 | const n = result_offset + result.len; |
| 462 | 232 | const closure: *AsyncClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(AsyncClosure), n) catch { |
| 463 | | pool.mutex.unlock(); |
| 464 | 233 | start(context.ptr, result.ptr); |
| 465 | 234 | return null; |
| 466 | 235 | })); |
| 236 | |
| 467 | 237 | closure.* = .{ |
| 468 | 238 | .func = start, |
| 469 | 239 | .context_offset = context_offset, |
| ... | ... | @@ -471,37 +241,124 @@ fn async( |
| 471 | 241 | .reset_event = .{}, |
| 472 | 242 | .cancel_tid = 0, |
| 473 | 243 | .select_condition = null, |
| 244 | .runnable = .{ |
| 245 | .start = AsyncClosure.start, |
| 246 | .is_parallel = false, |
| 247 | }, |
| 474 | 248 | }; |
| 249 | |
| 475 | 250 | @memcpy(closure.contextPointer()[0..context.len], context); |
| 251 | |
| 252 | pool.mutex.lock(); |
| 253 | |
| 254 | const thread_capacity = cpu_count - 1 + pool.parallel_count; |
| 255 | |
| 256 | pool.threads.ensureTotalCapacityPrecise(gpa, thread_capacity) catch { |
| 257 | pool.mutex.unlock(); |
| 258 | closure.free(gpa, result.len); |
| 259 | start(context.ptr, result.ptr); |
| 260 | return null; |
| 261 | }; |
| 262 | |
| 476 | 263 | pool.run_queue.prepend(&closure.runnable.node); |
| 477 | 264 | |
| 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; |
| 265 | if (pool.threads.items.len < thread_capacity) { |
| 266 | const thread = std.Thread.spawn(.{ .stack_size = pool.stack_size }, worker, .{pool}) catch { |
| 267 | if (pool.threads.items.len == 0) { |
| 268 | assert(pool.run_queue.popFirst() == &closure.runnable.node); |
| 269 | pool.mutex.unlock(); |
| 270 | closure.free(gpa, result.len); |
| 271 | start(context.ptr, result.ptr); |
| 272 | return null; |
| 273 | } |
| 274 | // Rely on other workers to do it. |
| 275 | pool.mutex.unlock(); |
| 276 | pool.cond.signal(); |
| 277 | return @ptrCast(closure); |
| 485 | 278 | }; |
| 279 | pool.threads.appendAssumeCapacity(thread); |
| 486 | 280 | } |
| 487 | 281 | |
| 488 | 282 | pool.mutex.unlock(); |
| 489 | 283 | pool.cond.signal(); |
| 284 | return @ptrCast(closure); |
| 285 | } |
| 286 | |
| 287 | fn asyncParallel( |
| 288 | userdata: ?*anyopaque, |
| 289 | result_len: usize, |
| 290 | result_alignment: std.mem.Alignment, |
| 291 | context: []const u8, |
| 292 | context_alignment: std.mem.Alignment, |
| 293 | start: *const fn (context: *const anyopaque, result: *anyopaque) void, |
| 294 | ) ?*Io.AnyFuture { |
| 295 | if (builtin.single_threaded) return null; |
| 296 | |
| 297 | const pool: *Pool = @alignCast(@ptrCast(userdata)); |
| 298 | const cpu_count = pool.cpu_count catch 1; |
| 299 | const gpa = pool.allocator; |
| 300 | const context_offset = context_alignment.forward(@sizeOf(AsyncClosure)); |
| 301 | const result_offset = result_alignment.forward(context_offset + context.len); |
| 302 | const n = result_offset + result_len; |
| 303 | const closure: *AsyncClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(AsyncClosure), n) catch return null)); |
| 304 | |
| 305 | closure.* = .{ |
| 306 | .func = start, |
| 307 | .context_offset = context_offset, |
| 308 | .result_offset = result_offset, |
| 309 | .reset_event = .{}, |
| 310 | .cancel_tid = 0, |
| 311 | .select_condition = null, |
| 312 | .runnable = .{ |
| 313 | .start = AsyncClosure.start, |
| 314 | .is_parallel = true, |
| 315 | }, |
| 316 | }; |
| 317 | @memcpy(closure.contextPointer()[0..context.len], context); |
| 318 | |
| 319 | pool.mutex.lock(); |
| 320 | |
| 321 | pool.parallel_count += 1; |
| 322 | const thread_capacity = cpu_count - 1 + pool.parallel_count; |
| 323 | |
| 324 | pool.threads.ensureTotalCapacity(gpa, thread_capacity) catch { |
| 325 | pool.mutex.unlock(); |
| 326 | closure.free(gpa, result_len); |
| 327 | return null; |
| 328 | }; |
| 329 | |
| 330 | pool.run_queue.prepend(&closure.runnable.node); |
| 331 | |
| 332 | if (pool.threads.items.len < thread_capacity) { |
| 333 | const thread = std.Thread.spawn(.{ .stack_size = pool.stack_size }, worker, .{pool}) catch { |
| 334 | assert(pool.run_queue.popFirst() == &closure.runnable.node); |
| 335 | pool.mutex.unlock(); |
| 336 | closure.free(gpa, result_len); |
| 337 | return null; |
| 338 | }; |
| 339 | pool.threads.appendAssumeCapacity(thread); |
| 340 | } |
| 490 | 341 | |
| 342 | pool.mutex.unlock(); |
| 343 | pool.cond.signal(); |
| 491 | 344 | return @ptrCast(closure); |
| 492 | 345 | } |
| 493 | 346 | |
| 494 | 347 | const DetachedClosure = struct { |
| 495 | 348 | pool: *Pool, |
| 496 | 349 | func: *const fn (context: *anyopaque) void, |
| 497 | | runnable: Runnable = .{ .runFn = runFn }, |
| 350 | runnable: Runnable, |
| 498 | 351 | context_alignment: std.mem.Alignment, |
| 499 | 352 | context_len: usize, |
| 500 | 353 | |
| 501 | | fn runFn(runnable: *Pool.Runnable, _: ?usize) void { |
| 354 | fn start(runnable: *Runnable) void { |
| 502 | 355 | const closure: *DetachedClosure = @alignCast(@fieldParentPtr("runnable", runnable)); |
| 503 | 356 | closure.func(closure.contextPointer()); |
| 504 | 357 | const gpa = closure.pool.allocator; |
| 358 | free(closure, gpa); |
| 359 | } |
| 360 | |
| 361 | fn free(closure: *DetachedClosure, gpa: Allocator) void { |
| 505 | 362 | const base: [*]align(@alignOf(DetachedClosure)) u8 = @ptrCast(closure); |
| 506 | 363 | gpa.free(base[0..contextEnd(closure.context_alignment, closure.context_len)]); |
| 507 | 364 | } |
| ... | ... | @@ -526,33 +383,46 @@ fn asyncDetached( |
| 526 | 383 | context_alignment: std.mem.Alignment, |
| 527 | 384 | start: *const fn (context: *const anyopaque) void, |
| 528 | 385 | ) void { |
| 386 | if (builtin.single_threaded) return start(context.ptr); |
| 529 | 387 | const pool: *Pool = @alignCast(@ptrCast(userdata)); |
| 530 | | pool.mutex.lock(); |
| 531 | | |
| 388 | const cpu_count = pool.cpu_count catch 1; |
| 532 | 389 | const gpa = pool.allocator; |
| 533 | 390 | const n = DetachedClosure.contextEnd(context_alignment, context.len); |
| 534 | 391 | const closure: *DetachedClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(DetachedClosure), n) catch { |
| 535 | | pool.mutex.unlock(); |
| 536 | | start(context.ptr); |
| 537 | | return; |
| 392 | return start(context.ptr); |
| 538 | 393 | })); |
| 539 | 394 | closure.* = .{ |
| 540 | 395 | .pool = pool, |
| 541 | 396 | .func = start, |
| 542 | 397 | .context_alignment = context_alignment, |
| 543 | 398 | .context_len = context.len, |
| 399 | .runnable = .{ |
| 400 | .start = DetachedClosure.start, |
| 401 | .is_parallel = false, |
| 402 | }, |
| 544 | 403 | }; |
| 545 | 404 | @memcpy(closure.contextPointer()[0..context.len], context); |
| 405 | |
| 406 | pool.mutex.lock(); |
| 407 | |
| 408 | const thread_capacity = cpu_count - 1 + pool.parallel_count; |
| 409 | |
| 410 | pool.threads.ensureTotalCapacityPrecise(gpa, thread_capacity) catch { |
| 411 | pool.mutex.unlock(); |
| 412 | closure.free(gpa); |
| 413 | return start(context.ptr); |
| 414 | }; |
| 415 | |
| 546 | 416 | pool.run_queue.prepend(&closure.runnable.node); |
| 547 | 417 | |
| 548 | | if (pool.threads.items.len < pool.threads.capacity) { |
| 549 | | pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{ |
| 550 | | .stack_size = pool.stack_size, |
| 551 | | .allocator = gpa, |
| 552 | | }, worker, .{pool}) catch t: { |
| 553 | | pool.threads.items.len -= 1; |
| 554 | | break :t undefined; |
| 418 | if (pool.threads.items.len < thread_capacity) { |
| 419 | const thread = std.Thread.spawn(.{ .stack_size = pool.stack_size }, worker, .{pool}) catch { |
| 420 | assert(pool.run_queue.popFirst() == &closure.runnable.node); |
| 421 | pool.mutex.unlock(); |
| 422 | closure.free(gpa); |
| 423 | return start(context.ptr); |
| 555 | 424 | }; |
| 425 | pool.threads.appendAssumeCapacity(thread); |
| 556 | 426 | } |
| 557 | 427 | |
| 558 | 428 | pool.mutex.unlock(); |
Andrew Kelley <andrew@ziglang.org> 2025-07-17 20:26:07-07:00