zzstoatzz.io
forked from
karlseguin.tngl.sh/http.zig
An HTTP/1.1 server for zig
1const std = @import("std");
2
3const Io = std.Io;
4const Thread = std.Thread;
5const Allocator = std.mem.Allocator;
6
7pub const Opts = struct {
8 count: u32,
9 backlog: u32,
10 buffer_size: usize,
11};
12
13pub fn ThreadPool(comptime F: anytype) type {
14 const BATCH_SIZE = 16;
15
16 // When the worker thread calls F, it'll inject its static buffer.
17 // So F would be: handle(server: *Server, conn: *Conn, buf: []u8)
18 // and FullArgs would be our 3 args....
19 const FullArgs = std.meta.ArgsTuple(@TypeOf(F));
20 const Args = SpawnArgs(FullArgs);
21
22 return struct {
23 stopped: bool,
24 threads: []Thread,
25 worker_index: usize,
26 workers: []Worker(F),
27 arena: std.heap.ArenaAllocator,
28
29 // we queue jobs here before batching them to a worker. We do this
30 // to minimze the amount of locking we need to do.
31 batch: [BATCH_SIZE]Args,
32 batch_size: usize,
33
34 const Self = @This();
35
36 // we expect allocator to be an Arena
37 pub fn init(io: Io, allocator: Allocator, opts: Opts) !Self {
38 var arena = std.heap.ArenaAllocator.init(allocator);
39 errdefer arena.deinit();
40
41 const aa = arena.allocator();
42
43 const threads = try aa.alloc(Thread, opts.count);
44 const workers = try aa.alloc(Worker(F), opts.count);
45
46 var started: usize = 0;
47 errdefer for (0..started) |i| {
48 workers[i].stop();
49 threads[i].join();
50 };
51
52 for (0..workers.len) |i| {
53 workers[i] = try Worker(F).init(io, aa, &workers[@mod(i + i, workers.len)], opts);
54 }
55 for (0..workers.len) |i| {
56 threads[i] = try Thread.spawn(.{}, Worker(F).run, .{&workers[i]});
57 started += 1;
58 }
59
60 return .{
61 .arena = arena,
62 .worker_index = 0,
63 .stopped = false,
64 .workers = workers,
65 .threads = threads,
66 .batch = undefined,
67 .batch_size = 0,
68 };
69 }
70
71 pub fn deinit(self: *Self) void {
72 self.arena.deinit();
73 }
74
75 pub fn stop(self: *Self) void {
76 if (@atomicRmw(bool, &self.stopped, .Xchg, true, .monotonic) == true) {
77 return;
78 }
79
80 for (self.workers, self.threads) |*worker, *thread| {
81 worker.stop();
82 thread.join();
83 }
84 }
85
86 pub fn spawn(self: *Self, args: Args) void {
87 var i = self.batch_size;
88 self.batch[i] = args;
89 i += 1;
90
91 if (i == BATCH_SIZE) {
92 self.flush(i);
93 i = 0;
94 }
95 self.batch_size = i;
96 }
97
98 pub fn spawnOne(self: *Self, args: Args) void {
99 const worker_index = self.worker_index +% 1;
100 self.worker_index = worker_index;
101 const workers = self.workers;
102 workers[@mod(worker_index, workers.len)].spawn(&.{args});
103 }
104
105 pub fn flush(self: *Self, batch_size: usize) void {
106 self.batch_size = 0;
107
108 const worker_index = self.worker_index +% 1;
109 self.worker_index = worker_index;
110 const workers = self.workers;
111 workers[@mod(worker_index, workers.len)].spawn(self.batch[0..batch_size]);
112 }
113
114 pub fn empty(self: *Self) bool {
115 for (self.workers) |*w| {
116 if (w.empty() == false) {
117 return false;
118 }
119 }
120 return true;
121 }
122 };
123}
124
125fn Worker(comptime F: anytype) type {
126 // When the worker thread calls F, it'll inject its static buffer.
127 // So F would be: handle(server: *Server, conn: *Conn, buf: []u8)
128 // and FullArgs would be our 3 args....
129 const FullArgs = std.meta.ArgsTuple(@TypeOf(F));
130 const Args = SpawnArgs(FullArgs);
131
132 return struct {
133 io: Io,
134
135 // position in queue to read from
136 tail: usize,
137
138 // position in the queue to write to
139 head: usize,
140
141 // pending jobs
142 queue: []Args,
143
144 buffer: []u8,
145
146 stopped: bool,
147 mutex: Io.Mutex,
148 read_cond: Io.Condition,
149 write_cond: Io.Condition,
150 peer: *Worker(F),
151
152 const Self = @This();
153
154 // we expect allocator to be an Arena
155 pub fn init(io: Io, allocator: Allocator, peer: *Worker(F), opts: Opts) !Self {
156 const queue = try allocator.alloc(Args, if (opts.backlog == 0 or opts.backlog == 1) 2 else opts.backlog);
157 const buffer = try allocator.alloc(u8, opts.buffer_size);
158
159 return .{
160 .io = io,
161 .tail = 0,
162 .head = 0,
163 .peer = peer,
164 .mutex = .init,
165 .stopped = false,
166 .queue = queue,
167 .read_cond = .init,
168 .write_cond = .init,
169 .buffer = buffer,
170 };
171 }
172
173 pub fn stop(self: *Self) void {
174 const io = self.io;
175 {
176 // allow stop to be called as part of server.stop()
177 // but also in server.deinit(), or in both.
178 self.mutex.lockUncancelable(io);
179 defer self.mutex.unlock(io);
180 if (self.stopped) {
181 return;
182 }
183 self.stopped = true;
184 }
185 self.read_cond.broadcast(io);
186 }
187
188 pub fn empty(self: *Self) bool {
189 const io = self.io;
190 self.mutex.lockUncancelable(io);
191 defer self.mutex.unlock(io);
192 return self.head == self.tail;
193 }
194
195 pub fn spawn(self: *Self, args: []const Args) void {
196 const io = self.io;
197 var pending = args;
198 var capacity: usize = 0;
199
200 const queue = self.queue;
201 const queue_end = queue.len - 1;
202
203 while (true) {
204 self.mutex.lockUncancelable(io);
205 var head = self.head;
206 var tail = self.tail;
207 while (true) {
208 capacity = if (head < tail) tail - head - 1 else queue_end - head + tail;
209 if (capacity > 0) {
210 break;
211 }
212 self.write_cond.waitUncancelable(io, &self.mutex);
213 head = self.head;
214 tail = self.tail;
215 }
216
217 const ready = if (capacity >= pending.len) pending else pending[0..capacity];
218 for (ready) |a| {
219 queue[head] = a;
220 head = if (head == queue_end) 0 else head + 1;
221 }
222 self.head = head;
223 self.mutex.unlock(io);
224 self.read_cond.signal(io);
225 if (ready.len == pending.len) {
226 break;
227 }
228 pending = pending[ready.len..];
229 }
230 }
231
232 // Having a re-usable buffer per thread is the most efficient way
233 // we can do any dynamic allocations. We'll pair this later with
234 // a FallbackAllocator. The main issue is that some data must outlive
235 // the worker thread (in nonblocking mode), but this isn't something
236 // we need to worry about here. As far as this worker thread is
237 // concerned, it has a chunk of memory (buffer) which it'll pass
238 // to the callback function to do with as it wants.
239 fn run(self: *Self) void {
240 const buffer = self.buffer;
241 while (true) {
242 const args = self.getNext(true) orelse return;
243
244 // convert Args to FullArgs, i.e. inject buffer as the last argument
245 var full_args: FullArgs = undefined;
246 const ARG_COUNT = std.meta.fields(FullArgs).len - 1;
247 full_args[ARG_COUNT] = buffer;
248 inline for (0..ARG_COUNT) |i| {
249 full_args[i] = args[i];
250 }
251 @call(.auto, F, full_args);
252 }
253 }
254
255 fn getNext(self: *Self, block: bool) ?Args {
256 const io = self.io;
257 const queue = self.queue;
258 const queue_end = queue.len - 1;
259
260 self.mutex.lockUncancelable(io);
261 while (self.tail == self.head) {
262 if (block == false or self.stopped) {
263 self.mutex.unlock(io);
264 return null;
265 }
266
267 self.mutex.unlock(io);
268 if (self.peer.getNext(false)) |args| {
269 return args;
270 }
271 self.mutex.lockUncancelable(io);
272 if (self.tail == self.head) {
273 self.read_cond.waitUncancelable(io, &self.mutex);
274 } else {
275 break;
276 }
277 }
278
279 const tail = self.tail;
280 const args = queue[tail];
281 self.tail = if (tail == queue_end) 0 else tail + 1;
282 self.mutex.unlock(io);
283 self.write_cond.signal(io);
284 return args;
285 }
286 };
287}
288
289fn SpawnArgs(FullArgs: anytype) type {
290 const full_fields = std.meta.fields(FullArgs);
291 const ARG_COUNT = full_fields.len - 1;
292
293 // Args will be FullArgs[0..len-1], so in the above example, args would be
294 // (*Server, *Conn)
295 // Args is what we expect the caller to pass to spawn. The worker thread
296 // will convert an Args into FullArgs by injecting its static buffer as
297 // the final argument.
298
299 // TODO: We could verify that the last argument to FullArgs is, in fact, a
300 // []u8. But this ThreadPool is private and being used for 2 specific cases
301 // that we control.
302
303 var field_types: [ARG_COUNT]type = undefined;
304 inline for (full_fields[0..ARG_COUNT], 0..) |field, i| {
305 field_types[i] = field.type;
306 }
307 return @Tuple(&field_types);
308}
309
310const t = @import("t.zig");
311test "ThreadPool: batch add" {
312 defer t.reset();
313
314 const counts = [_]u32{ 1, 2, 3, 4, 5, 6 };
315 const backlogs = [_]u32{ 1, 2, 3, 4, 5, 6 };
316 for (counts) |count| {
317 for (backlogs) |backlog| {
318 testSum = 0; // global defined near the end of this file
319 testCount = 0; // global defined near the end of this file
320 testC1 = 0;
321 testC2 = 0;
322 testC3 = 0;
323 testC4 = 0;
324 testC5 = 0;
325 testC6 = 0;
326 var tp = try ThreadPool(testIncr).init(t.io, t.arena.allocator(), .{ .count = count, .backlog = backlog, .buffer_size = 512 });
327 defer tp.deinit();
328
329 for (0..1_000) |_| {
330 tp.spawn(.{1});
331 tp.spawn(.{2});
332 tp.spawn(.{3});
333 tp.spawn(.{4});
334 }
335 while (tp.empty() == false) {
336 try t.io.sleep(.fromMilliseconds(1), .awake);
337 }
338 tp.stop();
339 try t.expectEqual(10_000, testSum);
340 try t.expectEqual(4_000, testCount);
341
342 try t.expectEqual(1000, testC1);
343 try t.expectEqual(1000, testC2);
344 try t.expectEqual(1000, testC3);
345 try t.expectEqual(1000, testC4);
346 try t.expectEqual(0, testC5);
347 try t.expectEqual(0, testC6);
348 }
349 }
350}
351
352test "ThreadPool: small fuzz" {
353 defer t.reset();
354
355 testSum = 0; // global defined near the end of this file
356 testCount = 0; // global defined near the end of this file
357 testC1 = 0;
358 testC2 = 0;
359 testC3 = 0;
360 testC4 = 0;
361 testC5 = 0;
362 testC6 = 0;
363 var tp = try ThreadPool(testIncr).init(t.io, t.arena.allocator(), .{ .count = 3, .backlog = 3, .buffer_size = 512 });
364 defer tp.deinit();
365
366 for (0..10_000) |_| {
367 tp.spawn(.{1});
368 tp.spawn(.{2});
369 tp.spawn(.{3});
370 }
371 while (tp.empty() == false) {
372 try t.io.sleep(.fromMilliseconds(1), .awake);
373 }
374 tp.stop();
375 try t.expectEqual(60_000, testSum);
376 try t.expectEqual(30_000, testCount);
377 try t.expectEqual(10_000, testC1);
378 try t.expectEqual(10_000, testC2);
379 try t.expectEqual(10_000, testC3);
380 try t.expectEqual(0, testC4);
381 try t.expectEqual(0, testC5);
382 try t.expectEqual(0, testC6);
383}
384
385test "ThreadPool: large fuzz" {
386 defer t.reset();
387
388 testSum = 0; // global defined near the end of this file
389 testCount = 0; // global defined near the end of this file
390 testC1 = 0;
391 testC2 = 0;
392 testC3 = 0;
393 testC4 = 0;
394 testC5 = 0;
395 testC6 = 0;
396 var tp = try ThreadPool(testIncr).init(t.io, t.arena.allocator(), .{ .count = 50, .backlog = 1000, .buffer_size = 512 });
397 defer tp.deinit();
398
399 for (0..10_000) |_| {
400 tp.spawn(.{1});
401 tp.spawn(.{2});
402 tp.spawn(.{3});
403 tp.spawn(.{4});
404 tp.spawn(.{5});
405 tp.spawn(.{6});
406 }
407 while (tp.empty() == false) {
408 try t.io.sleep(.fromMilliseconds(1), .awake);
409 }
410 tp.stop();
411 try t.expectEqual(210_000, testSum);
412 try t.expectEqual(60_000, testCount);
413 try t.expectEqual(10_000, testC1);
414 try t.expectEqual(10_000, testC2);
415 try t.expectEqual(10_000, testC3);
416 try t.expectEqual(10_000, testC4);
417 try t.expectEqual(10_000, testC5);
418 try t.expectEqual(10_000, testC6);
419}
420
421var testSum: u64 = 0;
422var testCount: u64 = 0;
423var testC1: u64 = 0;
424var testC2: u64 = 0;
425var testC3: u64 = 0;
426var testC4: u64 = 0;
427var testC5: u64 = 0;
428var testC6: u64 = 0;
429fn testIncr(c: u64, buf: []u8) void {
430 std.debug.assert(buf.len == 512);
431 _ = @atomicRmw(u64, &testSum, .Add, c, .monotonic);
432 _ = @atomicRmw(u64, &testCount, .Add, 1, .monotonic);
433 switch (c) {
434 1 => _ = @atomicRmw(u64, &testC1, .Add, 1, .monotonic),
435 2 => _ = @atomicRmw(u64, &testC2, .Add, 1, .monotonic),
436 3 => _ = @atomicRmw(u64, &testC3, .Add, 1, .monotonic),
437 4 => _ = @atomicRmw(u64, &testC4, .Add, 1, .monotonic),
438 5 => _ = @atomicRmw(u64, &testC5, .Add, 1, .monotonic),
439 6 => _ = @atomicRmw(u64, &testC6, .Add, 1, .monotonic),
440 else => unreachable,
441 }
442 // let the threadpool queue get backed up
443 t.io.sleep(.fromMicroseconds(20), .awake) catch unreachable;
444}