Something went wrong. Try again.
websocket forked from zzstoatzz.io/websocket.zig
Something went wrong. Try again.
7.3 kB · 233 lines
Zig
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234const std = @import("std");const c = std.c;
const Io = std.Io;const Thread = std.Thread;const Allocator = std.mem.Allocator;
// nanosleep via libc — no Io-based uncancelable sleep existsfn sleep(ns: u64) void { const secs = ns / std.time.ns_per_s; const nsecs = ns % std.time.ns_per_s; var ts: c.timespec = .{ .sec = @intCast(secs), .nsec = @intCast(nsecs), }; _ = c.nanosleep(&ts, null);}
pub const Opts = struct { count: u16, backlog: u32, buffer_size: usize,};
pub fn ThreadPool(comptime F: anytype) type { // When the worker thread calls F, it'll inject its static buffer. // So F would be: handle(server: *Server, conn: *Conn, buf: []u8) // and FullArgs would be our 3 args.... const FullArgs = std.meta.ArgsTuple(@TypeOf(F)); const full_field_types = std.meta.fieldTypes(FullArgs); const ARG_COUNT = full_field_types.len - 1;
// Args will be FullArgs[0..len-1], so in the above example, args would be // (*Server, *Conn) // Args is what we expect the caller to pass to spawn. The worker thread // will convert an Args into FullArgs by injecting its static buffer as // the final argument.
// TODO: We could verify that the last argument to FullArgs is, in fact, a // []u8. But this ThreadPool is private and being used for 2 specific cases // that we control.
// 0.16: @Type was removed, use @Tuple instead var field_types: [ARG_COUNT]type = undefined; inline for (full_field_types[0..ARG_COUNT], 0..) |field_type, index| { field_types[index] = field_type; } const Args = @Tuple(&field_types);
return struct { stopped: bool, push: usize, pull: usize, pending: usize, queue: []Args, threads: []Thread, mutex: Io.Mutex, pull_cond: Io.Condition, push_cond: Io.Condition, queue_end: usize, allocator: Allocator, io: Io,
const Self = @This();
pub fn init(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);
const io = std.Options.debug_io;
thread_pool.* = .{ .pull = 0, .push = 0, .pending = 0, .io = io, .mutex = .init, .stopped = false, .queue = queue, .pull_cond = .init, .push_cond = .init, .threads = threads, .allocator = allocator, .queue_end = queue.len - 1, };
var started: usize = 0; errdefer { thread_pool.stopped = true; thread_pool.pull_cond.broadcast(io); for (0..started) |i| { threads[i].join(); } }
for (0..threads.len) |i| { // This becomes owned by the thread, it'll free it as it ends const buffer = try allocator.alloc(u8, opts.buffer_size); errdefer allocator.free(buffer);
threads[i] = try Thread.spawn(.{}, Self.worker, .{ thread_pool, buffer }); started += 1; }
return thread_pool; }
pub fn deinit(self: *Self) void { const allocator = self.allocator; self.stop(); allocator.free(self.threads); allocator.free(self.queue);
allocator.destroy(self); }
pub fn stop(self: *Self) void { const io = self.io; { self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); if (self.stopped == true) { return; } self.stopped = true; }
self.pull_cond.broadcast(io); for (self.threads) |thrd| { thrd.join(); } }
pub fn empty(self: *Self) bool { const io = self.io; self.mutex.lockUncancelable(io); defer self.mutex.unlock(io); return self.pull == self.push; }
pub fn spawn(self: *Self, args: Args) void { const queue = self.queue; const len = queue.len; const io = self.io;
self.mutex.lockUncancelable(io); while (self.pending == len) { self.push_cond.waitUncancelable(io, &self.mutex); }
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.pull_cond.signal(io); }
fn worker(self: *Self, buffer: []u8) void { defer self.allocator.free(buffer); const io = self.io;
while (true) { self.mutex.lockUncancelable(io); while (self.pending == 0) { if (self.stopped) { self.mutex.unlock(io); return; } self.pull_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.push_cond.signal(io);
// convert Args to FullArgs, i.e. inject buffer as the last argument var full_args: FullArgs = undefined; full_args[ARG_COUNT] = buffer; inline for (0..ARG_COUNT) |i| { full_args[i] = args[i]; } @call(.auto, F, full_args); } } };}
const t = @import("../t.zig");test "ThreadPool: small fuzz" { testSum = 0; // global defined near the end of this file var tp = try ThreadPool(testIncr).init(t.allocator, .{ .count = 3, .backlog = 3, .buffer_size = 512 });
for (0..50_000) |_| { tp.spawn(.{1}); } while (tp.empty() == false) { sleep(std.time.ns_per_ms); } tp.deinit(); 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.allocator, .{ .count = 50, .backlog = 1000, .buffer_size = 512 });
for (0..50_000) |_| { tp.spawn(.{1}); } while (tp.empty() == false) { sleep(std.time.ns_per_ms); } tp.deinit(); try t.expectEqual(50_000, testSum);}
var testSum: u64 = 0;fn testIncr(val: u64, buf: []u8) void { std.debug.assert(buf.len == 512); _ = @atomicRmw(u64, &testSum, .Add, val, .monotonic); // let the threadpool queue get backed up sleep(std.time.ns_per_us * 100);}