Something went wrong. Try again.
zig utility library forked from karlseguin.tngl.sh/zul
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159const std = @import("std");
const Io = std.Io;const Thread = std.Thread;const Allocator = std.mem.Allocator;
pub const Opts = struct { count: u32 = 1, backlog: u32 = 128,};
pub fn ThreadPool(comptime F: anytype) type { const Args = std.meta.ArgsTuple(@TypeOf(F)); return struct { io: Io, stop: bool, push: usize, pull: usize, pending: usize, queue: []Args, threads: []Thread, mutex: Io.Mutex, sem: Io.Semaphore, cond: Io.Condition, queue_end: usize,
const Self = @This();
pub fn init(io: Io, allocator: Allocator, opts: Opts) !*Self { const queue = try allocator.alloc(Args, opts.backlog); errdefer allocator.free(queue);
const threads = try allocator.alloc(Thread, opts.count); errdefer allocator.free(threads);
const thread_pool = try allocator.create(Self); errdefer allocator.destroy(thread_pool);
thread_pool.* = .{ .io = io, .pull = 0, .push = 0, .pending = 0, .cond = .init, .mutex = .init, .stop = false, .threads = threads, .queue = queue, .queue_end = queue.len - 1, .sem = .{ .permits = queue.len }, };
var started: usize = 0; errdefer { thread_pool.stop = true; thread_pool.cond.broadcast(io); for (0..started) |i| { threads[i].join(); } }
for (0..threads.len) |i| { threads[i] = try Thread.spawn(.{}, Self.worker, .{thread_pool}); started += 1; }
return thread_pool; }
pub fn deinit(self: *Self, allocator: Allocator) void { self.mutex.lockUncancelable(self.io); self.stop = true; self.mutex.unlock(self.io);
self.cond.broadcast(self.io); for (self.threads) |thrd| { thrd.join(); } allocator.free(self.threads); allocator.free(self.queue);
allocator.destroy(self); }
pub fn empty(self: *Self) bool { self.mutex.lockUncancelable(self.io); defer self.mutex.unlock(self.io); return self.pull == self.push; }
pub fn spawn(self: *Self, args: Args) !void { const io = self.io; self.sem.waitUncancelable(io); self.mutex.lockUncancelable(io); const push = self.push; self.queue[push] = args; self.push = if (push == self.queue_end) 0 else push + 1; self.pending += 1; self.mutex.unlock(io); self.cond.signal(io); }
fn worker(self: *Self) void { const io = self.io; while (true) { self.mutex.lockUncancelable(io); while (self.pending == 0) { if (self.stop) { self.mutex.unlock(io); return; } self.cond.waitUncancelable(io, &self.mutex); } const pull = self.pull; const args = self.queue[pull]; self.pull = if (pull == self.queue_end) 0 else pull + 1; self.pending -= 1; self.mutex.unlock(io); self.sem.post(io); @call(.auto, F, args); } } };}
const t = @import("zul.zig").testing;test "ThreadPool: small fuzz" { testSum = 0; // global defined near the end of this file var tp = try ThreadPool(testIncr).init(t.io, t.allocator, .{ .count = 3, .backlog = 3 });
for (0..50_000) |_| { try tp.spawn(.{1}); } while (tp.empty() == false) { try Io.sleep(t.io, .fromMilliseconds(1), .real); } tp.deinit(t.allocator); try t.expectEqual(50_000, testSum);}
test "ThreadPool: large fuzz" { testSum = 0; // global defined near the end of this file var tp = try ThreadPool(testIncr).init(t.io, t.allocator, .{ .count = 50, .backlog = 1000 });
for (0..50_000) |_| { try tp.spawn(.{1}); } while (tp.empty() == false) { try Io.sleep(t.io, .fromMilliseconds(1), .real); } tp.deinit(t.allocator); try t.expectEqual(50_000, testSum);}
var testSum: u64 = 0;fn testIncr(c: u64) void { _ = @atomicRmw(u64, &testSum, .Add, c, .monotonic);}