| ... | @@ -6,331 +6,88 @@ const WaitGroup = std.Thread.WaitGroup; | ... | @@ -6,331 +6,88 @@ const WaitGroup = std.Thread.WaitGroup; |
| 6 | const Io = std.Io; | 6 | const Io = std.Io; |
| 7 | const Pool = @This(); | 7 | const Pool = @This(); |
| 8 | | 8 | |
| 9 | /// Must be a thread-safe allocator. | 9 | /// Thread-safe. |
| 10 | allocator: std.mem.Allocator, | 10 | allocator: Allocator, |
| 11 | mutex: std.Thread.Mutex = .{}, | 11 | mutex: std.Thread.Mutex = .{}, |
| 12 | cond: std.Thread.Condition = .{}, | 12 | cond: std.Thread.Condition = .{}, |
| 13 | run_queue: std.SinglyLinkedList = .{}, | 13 | run_queue: std.SinglyLinkedList = .{}, |
| 14 | is_running: bool = true, | 14 | join_requested: bool = false, |
| 15 | threads: std.ArrayListUnmanaged(std.Thread), | 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 | stack_size: usize, | 16 | stack_size: usize, |
| | 17 | cpu_count: std.Thread.CpuCountError!usize, |
| | 18 | parallel_count: usize, |
| 23 | | 19 | |
| 24 | threadlocal var current_closure: ?*AsyncClosure = null; | 20 | threadlocal var current_closure: ?*AsyncClosure = null; |
| 25 | | 21 | |
| 26 | pub const Runnable = struct { | 22 | pub const Runnable = struct { |
| 27 | runFn: RunProto, | 23 | start: Start, |
| 28 | node: std.SinglyLinkedList.Node = .{}, | 24 | node: std.SinglyLinkedList.Node = .{}, |
| 29 | }; | 25 | is_parallel: bool, |
| 30 | | | |
| 31 | pub const RunProto = *const fn (*Runnable, id: ?usize) void; | | |
| 32 | | 26 | |
| 33 | pub const Options = struct { | 27 | pub const Start = *const fn (*Runnable) void; |
| 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 | }; | 28 | }; |
| 39 | | 29 | |
| 40 | pub fn init(pool: *Pool, options: Options) !void { | 30 | pub const InitError = std.Thread.CpuCountError || Allocator.Error; |
| 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 | | 31 | |
| 46 | pool.* = .{ | 32 | pub fn init(gpa: Allocator) Pool { |
| | 33 | var pool: Pool = .{ |
| 47 | .allocator = gpa, | 34 | .allocator = gpa, |
| 48 | .threads = .initBuffer(threads), | 35 | .threads = .empty, |
| 49 | .ids = .{}, | 36 | .stack_size = std.Thread.SpawnConfig.default_stack_size, |
| 50 | .stack_size = options.stack_size, | 37 | .cpu_count = std.Thread.getCpuCount(), |
| | 38 | .parallel_count = 0, |
| 51 | }; | 39 | }; |
| 52 | | 40 | if (pool.cpu_count) |n| { |
| 53 | if (builtin.single_threaded) return; | 41 | pool.threads.ensureTotalCapacityPrecise(gpa, n - 1) catch {}; |
| 54 | | 42 | } else |_| {} |
| 55 | if (options.track_ids) { | 43 | return pool; |
| 56 | try pool.ids.ensureTotalCapacity(gpa, 1 + thread_count); | | |
| 57 | pool.ids.putAssumeCapacityNoClobber(std.Thread.getCurrentId(), {}); | | |
| 58 | } | | |
| 59 | } | 44 | } |
| 60 | | 45 | |
| 61 | pub fn deinit(pool: *Pool) void { | 46 | pub fn deinit(pool: *Pool) void { |
| 62 | const gpa = pool.allocator; | 47 | const gpa = pool.allocator; |
| 63 | pool.join(); | 48 | pool.join(); |
| 64 | pool.threads.deinit(gpa); | 49 | pool.threads.deinit(gpa); |
| 65 | pool.ids.deinit(gpa); | | |
| 66 | pool.* = undefined; | 50 | pool.* = undefined; |
| 67 | } | 51 | } |
| 68 | | 52 | |
| 69 | fn join(pool: *Pool) void { | 53 | fn join(pool: *Pool) void { |
| 70 | if (builtin.single_threaded) return; | 54 | if (builtin.single_threaded) return; |
| 71 | | | |
| 72 | { | 55 | { |
| 73 | pool.mutex.lock(); | 56 | pool.mutex.lock(); |
| 74 | defer pool.mutex.unlock(); | 57 | defer pool.mutex.unlock(); |
| 75 | | 58 | pool.join_requested = true; |
| 76 | // ensure future worker threads exit the dequeue loop | | |
| 77 | pool.is_running = false; | | |
| 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 | pool.cond.broadcast(); | 60 | pool.cond.broadcast(); |
| 83 | for (pool.threads.items) |thread| thread.join(); | 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 | fn worker(pool: *Pool) void { | 64 | fn worker(pool: *Pool) void { |
| 281 | pool.mutex.lock(); | 65 | pool.mutex.lock(); |
| 282 | defer pool.mutex.unlock(); | 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 | while (true) { | 68 | while (true) { |
| 288 | while (pool.run_queue.popFirst()) |run_node| { | 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 | pool.mutex.unlock(); | 70 | pool.mutex.unlock(); |
| 314 | const runnable: *Runnable = @fieldParentPtr("node", run_node); | 71 | const runnable: *Runnable = @fieldParentPtr("node", run_node); |
| 315 | runnable.runFn(runnable, id); | 72 | runnable.start(runnable); |
| 316 | continue; | 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 | | 79 | if (pool.join_requested) break; |
| 319 | pool.mutex.unlock(); | 80 | pool.cond.wait(&pool.mutex); |
| 320 | wait_group.wait(); | | |
| 321 | return; | | |
| 322 | } | 81 | } |
| 323 | } | 82 | } |
| 324 | | 83 | |
| 325 | pub fn getIdCount(pool: *Pool) usize { | | |
| 326 | return @intCast(1 + pool.threads.items.len); | | |
| 327 | } | | |
| 328 | | | |
| 329 | pub fn io(pool: *Pool) Io { | 84 | pub fn io(pool: *Pool) Io { |
| 330 | return .{ | 85 | return .{ |
| 331 | .userdata = pool, | 86 | .userdata = pool, |
| 332 | .vtable = &.{ | 87 | .vtable = &.{ |
| 333 | .async = async, | 88 | .async = async, |
| | 89 | .asyncConcurrent = asyncParallel, |
| | 90 | .asyncParallel = asyncParallel, |
| 334 | .await = await, | 91 | .await = await, |
| 335 | .asyncDetached = asyncDetached, | 92 | .asyncDetached = asyncDetached, |
| 336 | .cancel = cancel, | 93 | .cancel = cancel, |
| ... | @@ -357,7 +114,7 @@ pub fn io(pool: *Pool) Io { | ... | @@ -357,7 +114,7 @@ pub fn io(pool: *Pool) Io { |
| 357 | | 114 | |
| 358 | const AsyncClosure = struct { | 115 | const AsyncClosure = struct { |
| 359 | func: *const fn (context: *anyopaque, result: *anyopaque) void, | 116 | func: *const fn (context: *anyopaque, result: *anyopaque) void, |
| 360 | runnable: Runnable = .{ .runFn = runFn }, | 117 | runnable: Runnable, |
| 361 | reset_event: std.Thread.ResetEvent, | 118 | reset_event: std.Thread.ResetEvent, |
| 362 | select_condition: ?*std.Thread.ResetEvent, | 119 | select_condition: ?*std.Thread.ResetEvent, |
| 363 | cancel_tid: std.Thread.Id, | 120 | cancel_tid: std.Thread.Id, |
| ... | @@ -375,7 +132,7 @@ const AsyncClosure = struct { | ... | @@ -375,7 +132,7 @@ const AsyncClosure = struct { |
| 375 | else => @compileError("unsupported std.Thread.Id: " ++ @typeName(std.Thread.Id)), | 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 | const closure: *AsyncClosure = @alignCast(@fieldParentPtr("runnable", runnable)); | 136 | const closure: *AsyncClosure = @alignCast(@fieldParentPtr("runnable", runnable)); |
| 380 | const tid = std.Thread.getCurrentId(); | 137 | const tid = std.Thread.getCurrentId(); |
| 381 | if (@cmpxchgStrong( | 138 | if (@cmpxchgStrong( |
| ... | @@ -387,6 +144,7 @@ const AsyncClosure = struct { | ... | @@ -387,6 +144,7 @@ const AsyncClosure = struct { |
| 387 | .acquire, | 144 | .acquire, |
| 388 | )) |cancel_tid| { | 145 | )) |cancel_tid| { |
| 389 | assert(cancel_tid == canceling_tid); | 146 | assert(cancel_tid == canceling_tid); |
| | 147 | closure.reset_event.set(); |
| 390 | return; | 148 | return; |
| 391 | } | 149 | } |
| 392 | current_closure = closure; | 150 | current_closure = closure; |
| ... | @@ -438,9 +196,13 @@ const AsyncClosure = struct { | ... | @@ -438,9 +196,13 @@ const AsyncClosure = struct { |
| 438 | | 196 | |
| 439 | fn waitAndFree(closure: *AsyncClosure, gpa: Allocator, result: []u8) void { | 197 | fn waitAndFree(closure: *AsyncClosure, gpa: Allocator, result: []u8) void { |
| 440 | closure.reset_event.wait(); | 198 | closure.reset_event.wait(); |
| 441 | const base: [*]align(@alignOf(AsyncClosure)) u8 = @ptrCast(closure); | | |
| 442 | @memcpy(result, closure.resultPointer()[0..result.len]); | 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,18 +214,26 @@ fn async( |
| 452 | context_alignment: std.mem.Alignment, | 214 | context_alignment: std.mem.Alignment, |
| 453 | start: *const fn (context: *const anyopaque, result: *anyopaque) void, | 215 | start: *const fn (context: *const anyopaque, result: *anyopaque) void, |
| 454 | ) ?*Io.AnyFuture { | 216 | ) ?*Io.AnyFuture { |
| | 217 | if (builtin.single_threaded) { |
| | 218 | start(context.ptr, result.ptr); |
| | 219 | return null; |
| | 220 | } |
| 455 | const pool: *Pool = @alignCast(@ptrCast(userdata)); | 221 | const pool: *Pool = @alignCast(@ptrCast(userdata)); |
| 456 | pool.mutex.lock(); | 222 | const cpu_count = pool.cpu_count catch { |
| 457 | | 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 | const gpa = pool.allocator; | 228 | const gpa = pool.allocator; |
| 459 | const context_offset = context_alignment.forward(@sizeOf(AsyncClosure)); | 229 | const context_offset = context_alignment.forward(@sizeOf(AsyncClosure)); |
| 460 | const result_offset = result_alignment.forward(context_offset + context.len); | 230 | const result_offset = result_alignment.forward(context_offset + context.len); |
| 461 | const n = result_offset + result.len; | 231 | const n = result_offset + result.len; |
| 462 | const closure: *AsyncClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(AsyncClosure), n) catch { | 232 | const closure: *AsyncClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(AsyncClosure), n) catch { |
| 463 | pool.mutex.unlock(); | | |
| 464 | start(context.ptr, result.ptr); | 233 | start(context.ptr, result.ptr); |
| 465 | return null; | 234 | return null; |
| 466 | })); | 235 | })); |
| | 236 | |
| 467 | closure.* = .{ | 237 | closure.* = .{ |
| 468 | .func = start, | 238 | .func = start, |
| 469 | .context_offset = context_offset, | 239 | .context_offset = context_offset, |
| ... | @@ -471,37 +241,124 @@ fn async( | ... | @@ -471,37 +241,124 @@ fn async( |
| 471 | .reset_event = .{}, | 241 | .reset_event = .{}, |
| 472 | .cancel_tid = 0, | 242 | .cancel_tid = 0, |
| 473 | .select_condition = null, | 243 | .select_condition = null, |
| | 244 | .runnable = .{ |
| | 245 | .start = AsyncClosure.start, |
| | 246 | .is_parallel = false, |
| | 247 | }, |
| 474 | }; | 248 | }; |
| | 249 | |
| 475 | @memcpy(closure.contextPointer()[0..context.len], context); | 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 | pool.run_queue.prepend(&closure.runnable.node); | 263 | pool.run_queue.prepend(&closure.runnable.node); |
| 477 | | 264 | |
| 478 | if (pool.threads.items.len < pool.threads.capacity) { | 265 | if (pool.threads.items.len < thread_capacity) { |
| 479 | pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{ | 266 | const thread = std.Thread.spawn(.{ .stack_size = pool.stack_size }, worker, .{pool}) catch { |
| 480 | .stack_size = pool.stack_size, | 267 | if (pool.threads.items.len == 0) { |
| 481 | .allocator = gpa, | 268 | assert(pool.run_queue.popFirst() == &closure.runnable.node); |
| 482 | }, worker, .{pool}) catch t: { | 269 | pool.mutex.unlock(); |
| 483 | pool.threads.items.len -= 1; | 270 | closure.free(gpa, result.len); |
| 484 | break :t undefined; | 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 | pool.mutex.unlock(); | 282 | pool.mutex.unlock(); |
| 489 | pool.cond.signal(); | 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 | return @ptrCast(closure); | 344 | return @ptrCast(closure); |
| 492 | } | 345 | } |
| 493 | | 346 | |
| 494 | const DetachedClosure = struct { | 347 | const DetachedClosure = struct { |
| 495 | pool: *Pool, | 348 | pool: *Pool, |
| 496 | func: *const fn (context: *anyopaque) void, | 349 | func: *const fn (context: *anyopaque) void, |
| 497 | runnable: Runnable = .{ .runFn = runFn }, | 350 | runnable: Runnable, |
| 498 | context_alignment: std.mem.Alignment, | 351 | context_alignment: std.mem.Alignment, |
| 499 | context_len: usize, | 352 | context_len: usize, |
| 500 | | 353 | |
| 501 | fn runFn(runnable: *Pool.Runnable, _: ?usize) void { | 354 | fn start(runnable: *Runnable) void { |
| 502 | const closure: *DetachedClosure = @alignCast(@fieldParentPtr("runnable", runnable)); | 355 | const closure: *DetachedClosure = @alignCast(@fieldParentPtr("runnable", runnable)); |
| 503 | closure.func(closure.contextPointer()); | 356 | closure.func(closure.contextPointer()); |
| 504 | const gpa = closure.pool.allocator; | 357 | const gpa = closure.pool.allocator; |
| | 358 | free(closure, gpa); |
| | 359 | } |
| | 360 | |
| | 361 | fn free(closure: *DetachedClosure, gpa: Allocator) void { |
| 505 | const base: [*]align(@alignOf(DetachedClosure)) u8 = @ptrCast(closure); | 362 | const base: [*]align(@alignOf(DetachedClosure)) u8 = @ptrCast(closure); |
| 506 | gpa.free(base[0..contextEnd(closure.context_alignment, closure.context_len)]); | 363 | gpa.free(base[0..contextEnd(closure.context_alignment, closure.context_len)]); |
| 507 | } | 364 | } |
| ... | @@ -526,33 +383,46 @@ fn asyncDetached( | ... | @@ -526,33 +383,46 @@ fn asyncDetached( |
| 526 | context_alignment: std.mem.Alignment, | 383 | context_alignment: std.mem.Alignment, |
| 527 | start: *const fn (context: *const anyopaque) void, | 384 | start: *const fn (context: *const anyopaque) void, |
| 528 | ) void { | 385 | ) void { |
| | 386 | if (builtin.single_threaded) return start(context.ptr); |
| 529 | const pool: *Pool = @alignCast(@ptrCast(userdata)); | 387 | const pool: *Pool = @alignCast(@ptrCast(userdata)); |
| 530 | pool.mutex.lock(); | 388 | const cpu_count = pool.cpu_count catch 1; |
| 531 | | | |
| 532 | const gpa = pool.allocator; | 389 | const gpa = pool.allocator; |
| 533 | const n = DetachedClosure.contextEnd(context_alignment, context.len); | 390 | const n = DetachedClosure.contextEnd(context_alignment, context.len); |
| 534 | const closure: *DetachedClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(DetachedClosure), n) catch { | 391 | const closure: *DetachedClosure = @alignCast(@ptrCast(gpa.alignedAlloc(u8, .of(DetachedClosure), n) catch { |
| 535 | pool.mutex.unlock(); | 392 | return start(context.ptr); |
| 536 | start(context.ptr); | | |
| 537 | return; | | |
| 538 | })); | 393 | })); |
| 539 | closure.* = .{ | 394 | closure.* = .{ |
| 540 | .pool = pool, | 395 | .pool = pool, |
| 541 | .func = start, | 396 | .func = start, |
| 542 | .context_alignment = context_alignment, | 397 | .context_alignment = context_alignment, |
| 543 | .context_len = context.len, | 398 | .context_len = context.len, |
| | 399 | .runnable = .{ |
| | 400 | .start = DetachedClosure.start, |
| | 401 | .is_parallel = false, |
| | 402 | }, |
| 544 | }; | 403 | }; |
| 545 | @memcpy(closure.contextPointer()[0..context.len], context); | 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 | pool.run_queue.prepend(&closure.runnable.node); | 416 | pool.run_queue.prepend(&closure.runnable.node); |
| 547 | | 417 | |
| 548 | if (pool.threads.items.len < pool.threads.capacity) { | 418 | if (pool.threads.items.len < thread_capacity) { |
| 549 | pool.threads.addOneAssumeCapacity().* = std.Thread.spawn(.{ | 419 | const thread = std.Thread.spawn(.{ .stack_size = pool.stack_size }, worker, .{pool}) catch { |
| 550 | .stack_size = pool.stack_size, | 420 | assert(pool.run_queue.popFirst() == &closure.runnable.node); |
| 551 | .allocator = gpa, | 421 | pool.mutex.unlock(); |
| 552 | }, worker, .{pool}) catch t: { | 422 | closure.free(gpa); |
| 553 | pool.threads.items.len -= 1; | 423 | return start(context.ptr); |
| 554 | break :t undefined; | | |
| 555 | }; | 424 | }; |
| | 425 | pool.threads.appendAssumeCapacity(thread); |
| 556 | } | 426 | } |
| 557 | | 427 | |
| 558 | pool.mutex.unlock(); | 428 | pool.mutex.unlock(); |
Andrew Kelley <andrew@ziglang.org> 2025-07-17 20:26:07-07:00