authorgravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-07-01 15:32:46-07:00
committergravatar for andrew@ziglang.orgAndrew Kelley <andrew@ziglang.org> 2024-07-01 15:32:46-07:00
loga239e65b54a5030310ce721593a94a6289b625fe
treed26385cfbc6716ce51a28b23237ddcd13f63ecaf
parent86e68dbf53b40480e6b57b0f184308fd83801b1f

std.Thread.Pool: call shutdown on workers blocking on connect or read

and increase kernel backlog to the maximum number. Without increasing the kernel backlog to the maximum number, I observed connect() to block indefinitely, even when another thread calls shutdown() on that socket file descriptor.

1 files changed, 85 insertions(+), 52 deletions(-)

lib/std/Thread/Pool.zig+85-52
...@@ -9,7 +9,7 @@ cond: std.Thread.Condition,...@@ -9,7 +9,7 @@ cond: std.Thread.Condition,
9run_queue: RunQueue,9run_queue: RunQueue,
10end_flag: bool,10end_flag: bool,
11allocator: std.mem.Allocator,11allocator: std.mem.Allocator,
12threads: []std.Thread,12workers: []Worker,
13job_server_options: Options.JobServer,13job_server_options: Options.JobServer,
14job_server: ?*JobServer,14job_server: ?*JobServer,
1515
...@@ -20,6 +20,20 @@ const Runnable = struct {...@@ -20,6 +20,20 @@ const Runnable = struct {
2020
21const RunProto = *const fn (*Runnable) void;21const RunProto = *const fn (*Runnable) void;
2222
23pub const Worker = struct {
24 thread: std.Thread,
25 /// This data is shared with the thread pool for deinitialization purposes:
26 /// calling shutdown on this file descriptor will wake up any workers
27 /// blocking on acquiring a thread token.
28 ///
29 /// Pointers to file descriptors are used so that the actual loads and
30 /// stores of file descriptors take place in thread-local data, avoiding
31 /// false sharing.
32 ///
33 /// Protected by the pool mutex.
34 connection: *const std.posix.fd_t,
35};
36
23pub const Options = struct {37pub const Options = struct {
24 /// Not required to be thread-safe; protected by the pool's mutex.38 /// Not required to be thread-safe; protected by the pool's mutex.
25 allocator: std.mem.Allocator,39 allocator: std.mem.Allocator,
...@@ -59,7 +73,7 @@ pub fn init(pool: *Pool, options: Options) !void {...@@ -59,7 +73,7 @@ pub fn init(pool: *Pool, options: Options) !void {
59 .run_queue = .{},73 .run_queue = .{},
60 .end_flag = false,74 .end_flag = false,
61 .allocator = allocator,75 .allocator = allocator,
62 .threads = &.{},76 .workers = &.{},
63 .job_server_options = options.job_server,77 .job_server_options = options.job_server,
64 .job_server = null,78 .job_server = null,
65 };79 };
...@@ -71,19 +85,26 @@ pub fn init(pool: *Pool, options: Options) !void {...@@ -71,19 +85,26 @@ pub fn init(pool: *Pool, options: Options) !void {
71 assert(thread_count > 0);85 assert(thread_count > 0);
7286
73 // Kill and join any threads we spawned and free memory on error.87 // Kill and join any threads we spawned and free memory on error.
74 pool.threads = try allocator.alloc(std.Thread, thread_count);88 pool.workers = try allocator.alloc(Worker, thread_count);
75 var spawned: usize = 0;89 var spawned: usize = 0;
76 errdefer pool.join(spawned);90 errdefer pool.join(spawned);
7791
78 for (pool.threads) |*thread| {92 const temporary_connection_memory: std.posix.fd_t = -1;
79 thread.* = try std.Thread.spawn(.{}, worker, .{pool});93
94 for (pool.workers) |*worker| {
95 worker.* = .{
96 .connection = &temporary_connection_memory,
97 .thread = try std.Thread.spawn(.{}, workerRun, .{ pool, spawned }),
98 };
80 spawned += 1;99 spawned += 1;
81 }100 }
82101
83 switch (options.job_server) {102 switch (options.job_server) {
84 .abstain, .connect => {},103 .abstain, .connect => {},
85 .host => |addr| {104 .host => |addr| {
86 var server = try addr.listen(.{});105 var server = try addr.listen(.{
106 .kernel_backlog = std.math.maxInt(u31),
107 });
87 errdefer server.deinit();108 errdefer server.deinit();
88109
89 const pollfds = try allocator.alloc(std.posix.pollfd, thread_count + 1);110 const pollfds = try allocator.alloc(std.posix.pollfd, thread_count + 1);
...@@ -104,7 +125,7 @@ pub fn init(pool: *Pool, options: Options) !void {...@@ -104,7 +125,7 @@ pub fn init(pool: *Pool, options: Options) !void {
104}125}
105126
106pub fn deinit(pool: *Pool) void {127pub fn deinit(pool: *Pool) void {
107 pool.join(pool.threads.len);128 pool.join(pool.workers.len);
108 pool.* = undefined;129 pool.* = undefined;
109}130}
110131
...@@ -118,6 +139,12 @@ fn join(pool: *Pool, spawned: usize) void {...@@ -118,6 +139,12 @@ fn join(pool: *Pool, spawned: usize) void {
118139
119 // Ensure future worker threads exit the dequeue loop.140 // Ensure future worker threads exit the dequeue loop.
120 pool.end_flag = true;141 pool.end_flag = true;
142
143 // Wake up any workers blocking on connect or read.
144 for (pool.workers[0..spawned]) |worker| {
145 const fd = worker.connection.*;
146 if (fd >= 0) std.posix.shutdown(fd, .both) catch {};
147 }
121 }148 }
122149
123 // Wake up any sleeping threads (this can be done outside the mutex) then150 // Wake up any sleeping threads (this can be done outside the mutex) then
...@@ -132,10 +159,10 @@ fn join(pool: *Pool, spawned: usize) void {...@@ -132,10 +159,10 @@ fn join(pool: *Pool, spawned: usize) void {
132 job_server.thread.join();159 job_server.thread.join();
133 }160 }
134161
135 for (pool.threads[0..spawned]) |thread|162 for (pool.workers[0..spawned]) |worker|
136 thread.join();163 worker.thread.join();
137164
138 pool.allocator.free(pool.threads);165 pool.allocator.free(pool.workers);
139}166}
140167
141pub const JobServer = struct {168pub const JobServer = struct {
...@@ -310,36 +337,6 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void {...@@ -310,36 +337,6 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void {
310 pool.cond.signal();337 pool.cond.signal();
311}338}
312339
313fn acquireThreadToken(job_server_options: Options.JobServer, fd_ptr: *std.posix.fd_t) void {
314 if (fd_ptr.* >= 0) return;
315
316 switch (job_server_options) {
317 .abstain => {},
318 .connect, .host => |addr| {
319 const sockfd = std.posix.socket(
320 std.posix.AF.UNIX,
321 std.posix.SOCK.STREAM | std.posix.SOCK.CLOEXEC,
322 0,
323 ) catch |err| {
324 std.log.debug("failed to make socket: {s}", .{@errorName(err)});
325 return;
326 };
327 fd_ptr.* = sockfd;
328
329 std.posix.connect(sockfd, &addr.any, addr.getOsSockLen()) catch |err| {
330 std.log.debug("failed to connect: {s}", .{@errorName(err)});
331 return;
332 };
333
334 var trash_buf: [1]u8 = undefined;
335 _ = std.posix.read(sockfd, &trash_buf) catch |err| {
336 std.log.debug("failed to read: {s}", .{@errorName(err)});
337 return;
338 };
339 },
340 }
341}
342
343fn releaseThreadToken(fd_ptr: *std.posix.fd_t) void {340fn releaseThreadToken(fd_ptr: *std.posix.fd_t) void {
344 const fd = fd_ptr.*;341 const fd = fd_ptr.*;
345 if (fd >= 0) {342 if (fd >= 0) {
...@@ -348,29 +345,65 @@ fn releaseThreadToken(fd_ptr: *std.posix.fd_t) void {...@@ -348,29 +345,65 @@ fn releaseThreadToken(fd_ptr: *std.posix.fd_t) void {
348 }345 }
349}346}
350347
351fn worker(pool: *Pool) void {348fn workerRun(pool: *Pool, worker_index: usize) void {
352 var connection: std.posix.fd_t = -1;349 var connection: std.posix.fd_t = -1;
353 defer releaseThreadToken(&connection);
354350
355 pool.mutex.lock();351 pool.mutex.lock();
356 defer pool.mutex.unlock();352 pool.workers[worker_index].connection = &connection;
357353
358 while (true) {354 while (!pool.end_flag) {
359 const work_available = pool.run_queue.first != null;355 if (connection == -1 and pool.run_queue.first != null) token: {
360 if (work_available) {356 switch (pool.job_server_options) {
361 pool.mutex.unlock();357 .abstain => {},
362 defer pool.mutex.lock();358 .connect, .host => |addr| {
363 acquireThreadToken(pool.job_server_options, &connection);359 pool.mutex.unlock();
360 const sockfd = std.posix.socket(
361 std.posix.AF.UNIX,
362 std.posix.SOCK.STREAM | std.posix.SOCK.CLOEXEC,
363 0,
364 ) catch |err| {
365 std.log.debug("failed to make socket: {s}", .{@errorName(err)});
366 pool.mutex.lock();
367 if (pool.end_flag) break;
368 break :token;
369 };
370 pool.mutex.lock();
371 connection = sockfd;
372 if (pool.end_flag) break;
373 pool.mutex.unlock();
374
375 std.posix.connect(sockfd, &addr.any, addr.getOsSockLen()) catch |err| {
376 std.log.debug("failed to connect: {s}", .{@errorName(err)});
377 pool.mutex.lock();
378 if (pool.end_flag) break;
379 break :token;
380 };
381
382 var trash_buf: [1]u8 = undefined;
383 _ = std.posix.read(sockfd, &trash_buf) catch |err| {
384 std.log.debug("failed to read: {s}", .{@errorName(err)});
385 pool.mutex.lock();
386 if (pool.end_flag) break;
387 break :token;
388 };
389
390 pool.mutex.lock();
391 if (pool.end_flag) break;
392 },
393 }
364 }394 }
365 while (pool.run_queue.popFirst()) |run_node| {395 while (pool.run_queue.popFirst()) |run_node| {
366 pool.mutex.unlock();396 pool.mutex.unlock();
367 defer pool.mutex.lock();
368 run_node.data.runFn(&run_node.data);397 run_node.data.runFn(&run_node.data);
398 pool.mutex.lock();
399 if (pool.end_flag) break;
369 }400 }
370 if (pool.end_flag) return;
371 releaseThreadToken(&connection);401 releaseThreadToken(&connection);
372 pool.cond.wait(&pool.mutex);402 pool.cond.wait(&pool.mutex);
373 }403 }
404
405 releaseThreadToken(&connection);
406 pool.mutex.unlock();
374}407}
375408
376pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void {409pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void {