| ... | @@ -2,13 +2,18 @@ const std = @import("std"); | ... | @@ -2,13 +2,18 @@ const std = @import("std"); |
| 2 | const builtin = @import("builtin"); | 2 | const builtin = @import("builtin"); |
| 3 | const Pool = @This(); | 3 | const Pool = @This(); |
| 4 | const WaitGroup = @import("WaitGroup.zig"); | 4 | const WaitGroup = @import("WaitGroup.zig"); |
| | 5 | const assert = std.debug.assert; |
| 5 | | 6 | |
| 6 | mutex: std.Thread.Mutex = .{}, | 7 | mutex: std.Thread.Mutex, |
| 7 | cond: std.Thread.Condition = .{}, | 8 | cond: std.Thread.Condition, |
| 8 | run_queue: RunQueue = .{}, | 9 | run_queue: RunQueue, |
| 9 | is_running: bool = true, | 10 | run_queue_len: usize, |
| | 11 | end_flag: bool, |
| 10 | allocator: std.mem.Allocator, | 12 | allocator: std.mem.Allocator, |
| 11 | threads: []std.Thread, | 13 | threads_buffer: []std.Thread, |
| | 14 | threads_len: usize, |
| | 15 | job_server_options: Options.JobServer, |
| | 16 | job_server: ?*JobServer, |
| 12 | | 17 | |
| 13 | const RunQueue = std.SinglyLinkedList(Runnable); | 18 | const RunQueue = std.SinglyLinkedList(Runnable); |
| 14 | const Runnable = struct { | 19 | const Runnable = struct { |
| ... | @@ -18,63 +23,187 @@ const Runnable = struct { | ... | @@ -18,63 +23,187 @@ const Runnable = struct { |
| 18 | const RunProto = *const fn (*Runnable) void; | 23 | const RunProto = *const fn (*Runnable) void; |
| 19 | | 24 | |
| 20 | pub const Options = struct { | 25 | pub const Options = struct { |
| 21 | allocator: std.mem.Allocator, | 26 | /// Max number of threads to be actively working at the same time. |
| | 27 | /// |
| | 28 | /// `null` means to use the logical core count, leaving the main thread to |
| | 29 | /// fill in the last slot. |
| | 30 | /// |
| | 31 | /// `0` is an illegal value. |
| 22 | n_jobs: ?u32 = null, | 32 | n_jobs: ?u32 = null, |
| 23 | }; | | |
| 24 | | 33 | |
| 25 | pub fn init(pool: *Pool, options: Options) !void { | 34 | /// For coordinating amongst an entire process tree. |
| 26 | const allocator = options.allocator; | 35 | job_server: Options.JobServer = .abstain, |
| | 36 | |
| | 37 | pub const JobServer = union(enum) { |
| | 38 | /// The thread pool neither hosts a jobserver nor connects to an existing one. |
| | 39 | abstain, |
| | 40 | /// The thread pool uses the Jobserver2 protocol to coordinate a global |
| | 41 | /// thread pool across the entire process tree, avoiding cache |
| | 42 | /// thrashing. |
| | 43 | connect: std.net.Address, |
| | 44 | /// The thread pool assumes the role of the root process and spawns a |
| | 45 | /// dedicated thread for hosting the Jobserver2 protocol. |
| | 46 | /// |
| | 47 | /// Suggested to use a UNIX domain socket. |
| | 48 | host: std.net.Address, |
| | 49 | }; |
| | 50 | }; |
| 27 | | 51 | |
| 28 | pool.* = .{ | 52 | /// After initializing the thread pool and spawning work, the main thread must |
| | 53 | /// call `waitAndWork`. |
| | 54 | pub fn init( |
| | 55 | /// Not required to be thread-safe; protected by the pool's mutex. |
| | 56 | allocator: std.mem.Allocator, |
| | 57 | options: Options, |
| | 58 | ) !Pool { |
| | 59 | var pool: Pool = .{ |
| | 60 | .mutex = .{}, |
| | 61 | .cond = .{}, |
| | 62 | .run_queue = .{}, |
| | 63 | .run_queue_len = 0, |
| | 64 | .end_flag = false, |
| 29 | .allocator = allocator, | 65 | .allocator = allocator, |
| 30 | .threads = &[_]std.Thread{}, | 66 | .threads_buffer = &.{}, |
| | 67 | .threads_len = 0, |
| | 68 | .job_server_options = options.job_server, |
| | 69 | .job_server = null, |
| 31 | }; | 70 | }; |
| 32 | | 71 | |
| 33 | if (builtin.single_threaded) { | 72 | if (builtin.single_threaded) |
| 34 | return; | 73 | return; |
| 35 | } | | |
| 36 | | 74 | |
| 37 | const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1); | 75 | const thread_count = options.n_jobs orelse @max(1, std.Thread.getCpuCount() catch 1); |
| | 76 | assert(thread_count > 0); |
| | 77 | |
| | 78 | pool.threads_buffer = try allocator.alloc(std.Thread, thread_count); |
| | 79 | errdefer allocator.free(pool.threads_buffer); |
| | 80 | |
| | 81 | switch (options.job_server) { |
| | 82 | .abstain, .connect => {}, |
| | 83 | .host => |addr| { |
| | 84 | var server = try addr.listen(.{}); |
| | 85 | errdefer server.deinit(); |
| 38 | | 86 | |
| 39 | // kill and join any threads we spawned and free memory on error. | 87 | const pollfds = try allocator.alloc(std.posix.pollfd, thread_count); |
| 40 | pool.threads = try allocator.alloc(std.Thread, thread_count); | 88 | errdefer allocator.free(pollfds); |
| 41 | var spawned: usize = 0; | | |
| 42 | errdefer pool.join(spawned); | | |
| 43 | | 89 | |
| 44 | for (pool.threads) |*thread| { | 90 | const job_server = try allocator.create(JobServer); |
| 45 | thread.* = try std.Thread.spawn(.{}, worker, .{pool}); | 91 | errdefer allocator.destroy(job_server); |
| 46 | spawned += 1; | 92 | |
| | 93 | job_server.* = .{ |
| | 94 | .server = server, |
| | 95 | .pollfds = pollfds, |
| | 96 | .thread = try std.Thread.spawn(.{}, JobServer.run, .{job_server}), |
| | 97 | }; |
| | 98 | |
| | 99 | pool.job_server = job_server; |
| | 100 | }, |
| 47 | } | 101 | } |
| 48 | } | | |
| 49 | | 102 | |
| 50 | pub fn deinit(pool: *Pool) void { | 103 | return pool; |
| 51 | pool.join(pool.threads.len); // kill and join all threads. | | |
| 52 | pool.* = undefined; | | |
| 53 | } | 104 | } |
| 54 | | 105 | |
| 55 | fn join(pool: *Pool, spawned: usize) void { | 106 | pub fn deinit(pool: *Pool) void { |
| 56 | if (builtin.single_threaded) { | 107 | if (builtin.single_threaded) |
| 57 | return; | 108 | return; |
| 58 | } | | |
| 59 | | 109 | |
| 60 | { | 110 | { |
| 61 | pool.mutex.lock(); | 111 | pool.mutex.lock(); |
| 62 | defer pool.mutex.unlock(); | 112 | defer pool.mutex.unlock(); |
| 63 | | 113 | |
| 64 | // ensure future worker threads exit the dequeue loop | 114 | // Ensure future worker threads exit the dequeue loop. |
| 65 | pool.is_running = false; | 115 | pool.end_flag = true; |
| 66 | } | 116 | } |
| 67 | | 117 | |
| 68 | // wake up any sleeping threads (this can be done outside the mutex) | 118 | // Wake up any sleeping threads (this can be done outside the mutex) then |
| 69 | // then wait for all the threads we know are spawned to complete. | 119 | // wait for all the threads we know are spawned to complete. |
| 70 | pool.cond.broadcast(); | 120 | pool.cond.broadcast(); |
| 71 | for (pool.threads[0..spawned]) |thread| { | 121 | |
| 72 | thread.join(); | 122 | if (pool.job_server) |job_server| { |
| | 123 | // Interrupt the jobserver thread from accepting connections. |
| | 124 | // Since the server fd is also in the poll set, this handles both |
| | 125 | // places where control flow could be blocked. |
| | 126 | std.posix.shutdown(job_server.server.stream.handle, .both) catch {}; |
| | 127 | job_server.thread.join(); |
| 73 | } | 128 | } |
| 74 | | 129 | |
| 75 | pool.allocator.free(pool.threads); | 130 | // Since we set end_flag with the mutex locked, no more threads could have |
| | 131 | // been created. |
| | 132 | const threads = pool.threads_buffer[0..pool.threads_len]; |
| | 133 | |
| | 134 | for (threads) |thread| |
| | 135 | thread.join(); |
| | 136 | |
| | 137 | pool.allocator.free(pool.threads_buffer); |
| | 138 | pool.* = undefined; |
| 76 | } | 139 | } |
| 77 | | 140 | |
| | 141 | pub const JobServer = struct { |
| | 142 | server: std.net.Server, |
| | 143 | /// Has length n_jobs + 1. The first entry contains the server socket |
| | 144 | /// itself, so that calling shutdown() in the other thread will both cause |
| | 145 | /// the accept to return error.SocketNotListening and cause the poll() to |
| | 146 | /// return. |
| | 147 | pollfds: []std.posix.pollfd, |
| | 148 | thread: std.Thread, |
| | 149 | |
| | 150 | pub fn run(js: *JobServer) void { |
| | 151 | @memset(js.pollfds, .{ |
| | 152 | .fd = -1, |
| | 153 | // Only interested in errors and hangups. |
| | 154 | .events = 0, |
| | 155 | .revents = 0, |
| | 156 | }); |
| | 157 | |
| | 158 | js.pollfds[0].fd = js.server.stream.handle; |
| | 159 | |
| | 160 | main_loop: while (true) { |
| | 161 | for (js.pollfds[1..]) |*pollfd| { |
| | 162 | const err_event = (pollfd.revents & std.posix.POLL.ERR) != 0; |
| | 163 | const hup_event = (pollfd.revents & std.posix.POLL.HUP) != 0; |
| | 164 | if (err_event or hup_event) { |
| | 165 | std.posix.close(pollfd.fd); |
| | 166 | pollfd.fd = -1; |
| | 167 | pollfd.revents = 0; |
| | 168 | } |
| | 169 | |
| | 170 | if (pollfd.fd >= 0) continue; |
| | 171 | |
| | 172 | const connection = js.server.accept() catch |err| switch (err) { |
| | 173 | error.SocketNotListening => break :main_loop, // Indicates a shutdown request. |
| | 174 | else => |e| { |
| | 175 | std.log.debug("job server accept failure: {s}", .{@errorName(e)}); |
| | 176 | continue; |
| | 177 | }, |
| | 178 | }; |
| | 179 | _ = std.posix.send(connection.stream.handle, &.{0}, std.posix.MSG.NOSIGNAL) catch { |
| | 180 | connection.stream.close(); |
| | 181 | continue; |
| | 182 | }; |
| | 183 | pollfd.fd = connection.stream.handle; |
| | 184 | } |
| | 185 | |
| | 186 | _ = std.posix.poll(js.pollfds, -1) catch continue; |
| | 187 | } |
| | 188 | |
| | 189 | // Closes the active connections as well as the server itself. |
| | 190 | for (js.pollfds) |pollfd| { |
| | 191 | if (pollfd.fd >= 0) { |
| | 192 | std.posix.close(pollfd.fd); |
| | 193 | } |
| | 194 | } |
| | 195 | |
| | 196 | // Delete the UNIX domain socket. |
| | 197 | switch (js.server.listen_address.any.family) { |
| | 198 | std.posix.AF.UNIX => { |
| | 199 | const path = std.mem.sliceTo(&js.server.listen_address.un.path, 0); |
| | 200 | std.fs.cwd().deleteFile(path) catch {}; |
| | 201 | }, |
| | 202 | else => {}, |
| | 203 | } |
| | 204 | } |
| | 205 | }; |
| | 206 | |
| 78 | /// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and | 207 | /// Runs `func` in the thread pool, calling `WaitGroup.start` beforehand, and |
| 79 | /// `WaitGroup.finish` after it returns. | 208 | /// `WaitGroup.finish` after it returns. |
| 80 | /// | 209 | /// |
| ... | @@ -127,6 +256,22 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args | ... | @@ -127,6 +256,22 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args |
| 127 | }; | 256 | }; |
| 128 | | 257 | |
| 129 | pool.run_queue.prepend(&closure.run_node); | 258 | pool.run_queue.prepend(&closure.run_node); |
| | 259 | pool.run_queue_len += 1; |
| | 260 | |
| | 261 | // If there was already any queued work, spawn a new thread if we are |
| | 262 | // under the max. |
| | 263 | if (pool.run_queue_len > 1 and pool.threads_len < pool.threads_buffer.len) { |
| | 264 | if (std.Thread.spawn(.{}, worker, .{pool})) |new_thread| { |
| | 265 | pool.threads_buffer[pool.threads_len] = new_thread; |
| | 266 | pool.threads_len += 1; |
| | 267 | } else |_| if (pool.threads_len == 0) { |
| | 268 | pool.mutex.unlock(); |
| | 269 | @call(.auto, func, args); |
| | 270 | wait_group.finish(); |
| | 271 | return; |
| | 272 | } |
| | 273 | } |
| | 274 | |
| 130 | pool.mutex.unlock(); | 275 | pool.mutex.unlock(); |
| 131 | } | 276 | } |
| 132 | | 277 | |
| ... | @@ -134,7 +279,7 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args | ... | @@ -134,7 +279,7 @@ pub fn spawnWg(pool: *Pool, wait_group: *WaitGroup, comptime func: anytype, args |
| 134 | pool.cond.signal(); | 279 | pool.cond.signal(); |
| 135 | } | 280 | } |
| 136 | | 281 | |
| 137 | pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { | 282 | pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void { |
| 138 | if (builtin.single_threaded) { | 283 | if (builtin.single_threaded) { |
| 139 | @call(.auto, func, args); | 284 | @call(.auto, func, args); |
| 140 | return; | 285 | return; |
| ... | @@ -162,15 +307,34 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { | ... | @@ -162,15 +307,34 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { |
| 162 | | 307 | |
| 163 | { | 308 | { |
| 164 | pool.mutex.lock(); | 309 | pool.mutex.lock(); |
| 165 | defer pool.mutex.unlock(); | | |
| 166 | | 310 | |
| 167 | const closure = try pool.allocator.create(Closure); | 311 | const closure = pool.allocator.create(Closure) catch { |
| | 312 | pool.mutex.unlock(); |
| | 313 | @call(.auto, func, args); |
| | 314 | return; |
| | 315 | }; |
| 168 | closure.* = .{ | 316 | closure.* = .{ |
| 169 | .arguments = args, | 317 | .arguments = args, |
| 170 | .pool = pool, | 318 | .pool = pool, |
| 171 | }; | 319 | }; |
| 172 | | 320 | |
| 173 | pool.run_queue.prepend(&closure.run_node); | 321 | pool.run_queue.prepend(&closure.run_node); |
| | 322 | pool.run_queue_len += 1; |
| | 323 | |
| | 324 | // If there was already any queued work, spawn a new thread if we are |
| | 325 | // under the max. |
| | 326 | if (pool.run_queue_len > 1 and pool.threads_len < pool.threads_buffer.len) { |
| | 327 | if (std.Thread.spawn(.{}, worker, .{pool})) |new_thread| { |
| | 328 | pool.threads_buffer[pool.threads_len] = new_thread; |
| | 329 | pool.threads_len += 1; |
| | 330 | } else |_| if (pool.threads_len == 0) { |
| | 331 | pool.mutex.unlock(); |
| | 332 | @call(.auto, func, args); |
| | 333 | return; |
| | 334 | } |
| | 335 | } |
| | 336 | |
| | 337 | pool.mutex.unlock(); |
| 174 | } | 338 | } |
| 175 | | 339 | |
| 176 | // Notify waiting threads outside the lock to try and keep the critical section small. | 340 | // Notify waiting threads outside the lock to try and keep the critical section small. |
| ... | @@ -178,25 +342,45 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { | ... | @@ -178,25 +342,45 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) !void { |
| 178 | } | 342 | } |
| 179 | | 343 | |
| 180 | fn worker(pool: *Pool) void { | 344 | fn worker(pool: *Pool) void { |
| | 345 | var trash_buf: [1]u8 = undefined; |
| | 346 | var connection: ?std.net.Stream = null; |
| | 347 | defer if (connection) |stream| stream.close(); |
| | 348 | |
| 181 | pool.mutex.lock(); | 349 | pool.mutex.lock(); |
| 182 | defer pool.mutex.unlock(); | 350 | defer pool.mutex.unlock(); |
| 183 | | 351 | |
| 184 | while (true) { | 352 | while (true) { |
| 185 | while (pool.run_queue.popFirst()) |run_node| { | 353 | while (pool.run_queue.popFirst()) |run_node| { |
| 186 | // Temporarily unlock the mutex in order to execute the run_node | 354 | pool.run_queue_len -= 1; |
| | 355 | |
| | 356 | // Temporarily unlock the mutex in order to execute the run_node. |
| 187 | pool.mutex.unlock(); | 357 | pool.mutex.unlock(); |
| 188 | defer pool.mutex.lock(); | 358 | defer pool.mutex.lock(); |
| 189 | | 359 | |
| | 360 | if (connection == null) switch (pool.job_server_options) { |
| | 361 | .abstain => {}, |
| | 362 | .connect, .host => |addr| { |
| | 363 | if (std.net.tcpConnectToAddress(addr)) |stream| { |
| | 364 | connection = stream; |
| | 365 | _ = stream.readAll(&trash_buf) catch 1; |
| | 366 | } else |_| {} |
| | 367 | }, |
| | 368 | }; |
| | 369 | |
| 190 | const runFn = run_node.data.runFn; | 370 | const runFn = run_node.data.runFn; |
| 191 | runFn(&run_node.data); | 371 | runFn(&run_node.data); |
| 192 | } | 372 | } |
| 193 | | 373 | |
| 194 | // Stop executing instead of waiting if the thread pool is no longer running. | 374 | // Stop executing instead of waiting if the thread pool is no longer running. |
| 195 | if (pool.is_running) { | 375 | if (pool.end_flag) |
| 196 | pool.cond.wait(&pool.mutex); | | |
| 197 | } else { | | |
| 198 | break; | 376 | break; |
| | 377 | |
| | 378 | if (connection) |stream| { |
| | 379 | stream.close(); |
| | 380 | connection = null; |
| 199 | } | 381 | } |
| | 382 | |
| | 383 | pool.cond.wait(&pool.mutex); |
| 200 | } | 384 | } |
| 201 | } | 385 | } |
| 202 | | 386 | |
| ... | @@ -207,6 +391,7 @@ pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void { | ... | @@ -207,6 +391,7 @@ pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void { |
| 207 | defer pool.mutex.unlock(); | 391 | defer pool.mutex.unlock(); |
| 208 | break :blk pool.run_queue.popFirst(); | 392 | break :blk pool.run_queue.popFirst(); |
| 209 | }) |run_node| { | 393 | }) |run_node| { |
| | 394 | pool.run_queue_len -= 1; |
| 210 | run_node.data.runFn(&run_node.data); | 395 | run_node.data.runFn(&run_node.data); |
| 211 | continue; | 396 | continue; |
| 212 | } | 397 | } |