aboutsummaryrefslogtreecommitdiff
path: root/src/ThreadPool.zig
diff options
context:
space:
mode:
Diffstat (limited to 'src/ThreadPool.zig')
-rw-r--r--src/ThreadPool.zig116
1 files changed, 116 insertions, 0 deletions
diff --git a/src/ThreadPool.zig b/src/ThreadPool.zig
new file mode 100644
index 0000000000..a5f59a30e7
--- /dev/null
+++ b/src/ThreadPool.zig
@@ -0,0 +1,116 @@
+const std = @import("std");
+const ThreadPool = @This();
+
+lock: std.Mutex = .{},
+is_running: bool = true,
+allocator: *std.mem.Allocator,
+running: usize = 0,
+threads: []*std.Thread,
+run_queue: RunQueue = .{},
+idle_queue: IdleQueue = .{},
+
+const IdleQueue = std.SinglyLinkedList(std.AutoResetEvent);
+const RunQueue = std.SinglyLinkedList(Runnable);
+const Runnable = struct {
+ runFn: fn (*Runnable) void,
+};
+
+pub fn init(self: *ThreadPool, allocator: *std.mem.Allocator) !void {
+ self.* = .{
+ .allocator = allocator,
+ .threads = &[_]*std.Thread{},
+ };
+
+ errdefer self.deinit();
+
+ var num_threads = std.Thread.cpuCount() catch 1;
+ if (num_threads > 0)
+ self.threads = try allocator.alloc(*std.Thread, num_threads);
+
+ while (num_threads > 0) : (num_threads -= 1) {
+ const thread = try std.Thread.spawn(self, runWorker);
+ self.threads[self.running] = thread;
+ self.running += 1;
+ }
+}
+
+pub fn deinit(self: *ThreadPool) void {
+ self.shutdown();
+
+ std.debug.assert(!self.is_running);
+ for (self.threads[0..self.running]) |thread|
+ thread.wait();
+
+ defer self.threads = &[_]*std.Thread{};
+ if (self.running > 0)
+ self.allocator.free(self.threads);
+}
+
+pub fn shutdown(self: *ThreadPool) void {
+ const held = self.lock.acquire();
+
+ if (!self.is_running)
+ return held.release();
+
+ var idle_queue = self.idle_queue;
+ self.idle_queue = .{};
+ self.is_running = false;
+ held.release();
+
+ while (idle_queue.popFirst()) |idle_node|
+ idle_node.data.set();
+}
+
+pub fn spawn(self: *ThreadPool, comptime func: anytype, args: anytype) !void {
+ const Args = @TypeOf(args);
+ const Closure = struct {
+ arguments: Args,
+ pool: *ThreadPool,
+ run_node: RunQueue.Node = .{ .data = .{ .runFn = runFn } },
+
+ fn runFn(runnable: *Runnable) void {
+ const run_node = @fieldParentPtr(RunQueue.Node, "data", runnable);
+ const closure = @fieldParentPtr(@This(), "run_node", run_node);
+ const result = @call(.{}, func, closure.arguments);
+ closure.pool.allocator.destroy(closure);
+ }
+ };
+
+ const closure = try self.allocator.create(Closure);
+ errdefer self.allocator.destroy(closure);
+ closure.* = .{
+ .arguments = args,
+ .pool = self,
+ };
+
+ const held = self.lock.acquire();
+ self.run_queue.prepend(&closure.run_node);
+
+ const idle_node = self.idle_queue.popFirst();
+ held.release();
+
+ if (idle_node) |node|
+ node.data.set();
+}
+
+fn runWorker(self: *ThreadPool) void {
+ while (true) {
+ const held = self.lock.acquire();
+
+ if (self.run_queue.popFirst()) |run_node| {
+ held.release();
+ (run_node.data.runFn)(&run_node.data);
+ continue;
+ }
+
+ if (!self.is_running) {
+ held.release();
+ return;
+ }
+
+ var idle_node = IdleQueue.Node{ .data = .{} };
+ self.idle_queue.prepend(&idle_node);
+ held.release();
+ idle_node.data.wait();
+ }
+}