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,
99run_queue: RunQueue,
1010end_flag: bool,
1111allocator: std.mem.Allocator,
12threads: []std.Thread,
12workers: []Worker,
1313job_server_options: Options.JobServer,
1414job_server: ?*JobServer,
1515
......@@ -20,6 +20,20 @@ const Runnable = struct {
2020
2121const 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
2337pub const Options = struct {
2438 /// Not required to be thread-safe; protected by the pool's mutex.
2539 allocator: std.mem.Allocator,
......@@ -59,7 +73,7 @@ pub fn init(pool: *Pool, options: Options) !void {
5973 .run_queue = .{},
6074 .end_flag = false,
6175 .allocator = allocator,
62 .threads = &.{},
76 .workers = &.{},
6377 .job_server_options = options.job_server,
6478 .job_server = null,
6579 };
......@@ -71,19 +85,26 @@ pub fn init(pool: *Pool, options: Options) !void {
7185 assert(thread_count > 0);
7286
7387 // 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);
7589 var spawned: usize = 0;
7690 errdefer pool.join(spawned);
7791
78 for (pool.threads) |*thread| {
79 thread.* = try std.Thread.spawn(.{}, worker, .{pool});
92 const temporary_connection_memory: std.posix.fd_t = -1;
93
94 for (pool.workers) |*worker| {
95 worker.* = .{
96 .connection = &temporary_connection_memory,
97 .thread = try std.Thread.spawn(.{}, workerRun, .{ pool, spawned }),
98 };
8099 spawned += 1;
81100 }
82101
83102 switch (options.job_server) {
84103 .abstain, .connect => {},
85104 .host => |addr| {
86 var server = try addr.listen(.{});
105 var server = try addr.listen(.{
106 .kernel_backlog = std.math.maxInt(u31),
107 });
87108 errdefer server.deinit();
88109
89110 const pollfds = try allocator.alloc(std.posix.pollfd, thread_count + 1);
......@@ -104,7 +125,7 @@ pub fn init(pool: *Pool, options: Options) !void {
104125}
105126
106127pub fn deinit(pool: *Pool) void {
107 pool.join(pool.threads.len);
128 pool.join(pool.workers.len);
108129 pool.* = undefined;
109130}
110131
......@@ -118,6 +139,12 @@ fn join(pool: *Pool, spawned: usize) void {
118139
119140 // Ensure future worker threads exit the dequeue loop.
120141 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 }
121148 }
122149
123150 // Wake up any sleeping threads (this can be done outside the mutex) then
......@@ -132,10 +159,10 @@ fn join(pool: *Pool, spawned: usize) void {
132159 job_server.thread.join();
133160 }
134161
135 for (pool.threads[0..spawned]) |thread|
136 thread.join();
162 for (pool.workers[0..spawned]) |worker|
163 worker.thread.join();
137164
138 pool.allocator.free(pool.threads);
165 pool.allocator.free(pool.workers);
139166}
140167
141168pub const JobServer = struct {
......@@ -310,36 +337,6 @@ pub fn spawn(pool: *Pool, comptime func: anytype, args: anytype) void {
310337 pool.cond.signal();
311338}
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
343340fn releaseThreadToken(fd_ptr: *std.posix.fd_t) void {
344341 const fd = fd_ptr.*;
345342 if (fd >= 0) {
......@@ -348,29 +345,65 @@ fn releaseThreadToken(fd_ptr: *std.posix.fd_t) void {
348345 }
349346}
350347
351fn worker(pool: *Pool) void {
348fn workerRun(pool: *Pool, worker_index: usize) void {
352349 var connection: std.posix.fd_t = -1;
353 defer releaseThreadToken(&connection);
354350
355351 pool.mutex.lock();
356 defer pool.mutex.unlock();
357
358 while (true) {
359 const work_available = pool.run_queue.first != null;
360 if (work_available) {
361 pool.mutex.unlock();
362 defer pool.mutex.lock();
363 acquireThreadToken(pool.job_server_options, &connection);
352 pool.workers[worker_index].connection = &connection;
353
354 while (!pool.end_flag) {
355 if (connection == -1 and pool.run_queue.first != null) token: {
356 switch (pool.job_server_options) {
357 .abstain => {},
358 .connect, .host => |addr| {
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 }
364394 }
365395 while (pool.run_queue.popFirst()) |run_node| {
366396 pool.mutex.unlock();
367 defer pool.mutex.lock();
368397 run_node.data.runFn(&run_node.data);
398 pool.mutex.lock();
399 if (pool.end_flag) break;
369400 }
370 if (pool.end_flag) return;
371401 releaseThreadToken(&connection);
372402 pool.cond.wait(&pool.mutex);
373403 }
404
405 releaseThreadToken(&connection);
406 pool.mutex.unlock();
374407}
375408
376409pub fn waitAndWork(pool: *Pool, wait_group: *WaitGroup) void {