zzstoatzz.io
websocket
7.3 kB
233 lines
1const std = @import("std");
2const c = std.c;
3
4const Io = std.Io;
5const Thread = std.Thread;
6const Allocator = std.mem.Allocator;
7
8// nanosleep via libc — no Io-based uncancelable sleep exists
9fn sleep(ns: u64) void {
10 const secs = ns / std.time.ns_per_s;
11 const nsecs = ns % std.time.ns_per_s;
12 var ts: c.timespec = .{
13 .sec = @intCast(secs),
14 .nsec = @intCast(nsecs),
15 };
16 _ = c.nanosleep(&ts, null);
17}
18
19pub const Opts = struct {
20 count: u16,
21 backlog: u32,
22 buffer_size: usize,
23};
24
25pub fn ThreadPool(comptime F: anytype) type {
26 // When the worker thread calls F, it'll inject its static buffer.
27 // So F would be: handle(server: *Server, conn: *Conn, buf: []u8)
28 // and FullArgs would be our 3 args....
29 const FullArgs = std.meta.ArgsTuple(@TypeOf(F));
30 const full_fields = std.meta.fields(FullArgs);
31 const ARG_COUNT = full_fields.len - 1;
32
33 // Args will be FullArgs[0..len-1], so in the above example, args would be
34 // (*Server, *Conn)
35 // Args is what we expect the caller to pass to spawn. The worker thread
36 // will convert an Args into FullArgs by injecting its static buffer as
37 // the final argument.
38
39 // TODO: We could verify that the last argument to FullArgs is, in fact, a
40 // []u8. But this ThreadPool is private and being used for 2 specific cases
41 // that we control.
42
43 // 0.16: @Type was removed, use @Tuple instead
44 var field_types: [ARG_COUNT]type = undefined;
45 inline for (full_fields[0..ARG_COUNT], 0..) |field, index| {
46 field_types[index] = field.type;
47 }
48 const Args = @Tuple(&field_types);
49
50 return struct {
51 stopped: bool,
52 push: usize,
53 pull: usize,
54 pending: usize,
55 queue: []Args,
56 threads: []Thread,
57 mutex: Io.Mutex,
58 pull_cond: Io.Condition,
59 push_cond: Io.Condition,
60 queue_end: usize,
61 allocator: Allocator,
62 io: Io,
63
64 const Self = @This();
65
66 pub fn init(allocator: Allocator, opts: Opts) !*Self {
67 const queue = try allocator.alloc(Args, opts.backlog);
68 errdefer allocator.free(queue);
69
70 const threads = try allocator.alloc(Thread, opts.count);
71 errdefer allocator.free(threads);
72
73 const thread_pool = try allocator.create(Self);
74 errdefer allocator.destroy(thread_pool);
75
76 const io = std.Options.debug_io;
77
78 thread_pool.* = .{
79 .pull = 0,
80 .push = 0,
81 .pending = 0,
82 .io = io,
83 .mutex = .init,
84 .stopped = false,
85 .queue = queue,
86 .pull_cond = .init,
87 .push_cond = .init,
88 .threads = threads,
89 .allocator = allocator,
90 .queue_end = queue.len - 1,
91 };
92
93 var started: usize = 0;
94 errdefer {
95 thread_pool.stopped = true;
96 thread_pool.pull_cond.broadcast(io);
97 for (0..started) |i| {
98 threads[i].join();
99 }
100 }
101
102 for (0..threads.len) |i| {
103 // This becomes owned by the thread, it'll free it as it ends
104 const buffer = try allocator.alloc(u8, opts.buffer_size);
105 errdefer allocator.free(buffer);
106
107 threads[i] = try Thread.spawn(.{}, Self.worker, .{ thread_pool, buffer });
108 started += 1;
109 }
110
111 return thread_pool;
112 }
113
114 pub fn deinit(self: *Self) void {
115 const allocator = self.allocator;
116 self.stop();
117 allocator.free(self.threads);
118 allocator.free(self.queue);
119
120 allocator.destroy(self);
121 }
122
123 pub fn stop(self: *Self) void {
124 const io = self.io;
125 {
126 self.mutex.lockUncancelable(io);
127 defer self.mutex.unlock(io);
128 if (self.stopped == true) {
129 return;
130 }
131 self.stopped = true;
132 }
133
134 self.pull_cond.broadcast(io);
135 for (self.threads) |thrd| {
136 thrd.join();
137 }
138 }
139
140 pub fn empty(self: *Self) bool {
141 const io = self.io;
142 self.mutex.lockUncancelable(io);
143 defer self.mutex.unlock(io);
144 return self.pull == self.push;
145 }
146
147 pub fn spawn(self: *Self, args: Args) void {
148 const queue = self.queue;
149 const len = queue.len;
150 const io = self.io;
151
152 self.mutex.lockUncancelable(io);
153 while (self.pending == len) {
154 self.push_cond.waitUncancelable(io, &self.mutex);
155 }
156
157 const push = self.push;
158 self.queue[push] = args;
159 self.push = if (push == self.queue_end) 0 else push + 1;
160 self.pending += 1;
161 self.mutex.unlock(io);
162
163 self.pull_cond.signal(io);
164 }
165
166 fn worker(self: *Self, buffer: []u8) void {
167 defer self.allocator.free(buffer);
168 const io = self.io;
169
170 while (true) {
171 self.mutex.lockUncancelable(io);
172 while (self.pending == 0) {
173 if (self.stopped) {
174 self.mutex.unlock(io);
175 return;
176 }
177 self.pull_cond.waitUncancelable(io, &self.mutex);
178 }
179 const pull = self.pull;
180 const args = self.queue[pull];
181 self.pull = if (pull == self.queue_end) 0 else pull + 1;
182 self.pending -= 1;
183 self.mutex.unlock(io);
184 self.push_cond.signal(io);
185
186 // convert Args to FullArgs, i.e. inject buffer as the last argument
187 var full_args: FullArgs = undefined;
188 full_args[ARG_COUNT] = buffer;
189 inline for (0..ARG_COUNT) |i| {
190 full_args[i] = args[i];
191 }
192 @call(.auto, F, full_args);
193 }
194 }
195 };
196}
197
198const t = @import("../t.zig");
199test "ThreadPool: small fuzz" {
200 testSum = 0; // global defined near the end of this file
201 var tp = try ThreadPool(testIncr).init(t.allocator, .{ .count = 3, .backlog = 3, .buffer_size = 512 });
202
203 for (0..50_000) |_| {
204 tp.spawn(.{1});
205 }
206 while (tp.empty() == false) {
207 sleep(std.time.ns_per_ms);
208 }
209 tp.deinit();
210 try t.expectEqual(50_000, testSum);
211}
212
213test "ThreadPool: large fuzz" {
214 testSum = 0; // global defined near the end of this file
215 var tp = try ThreadPool(testIncr).init(t.allocator, .{ .count = 50, .backlog = 1000, .buffer_size = 512 });
216
217 for (0..50_000) |_| {
218 tp.spawn(.{1});
219 }
220 while (tp.empty() == false) {
221 sleep(std.time.ns_per_ms);
222 }
223 tp.deinit();
224 try t.expectEqual(50_000, testSum);
225}
226
227var testSum: u64 = 0;
228fn testIncr(val: u64, buf: []u8) void {
229 std.debug.assert(buf.len == 512);
230 _ = @atomicRmw(u64, &testSum, .Add, val, .monotonic);
231 // let the threadpool queue get backed up
232 sleep(std.time.ns_per_us * 100);
233}