zzstoatzz.io
forked from
karlseguin.tngl.sh/http.zig
An HTTP/1.1 server for zig
1const std = @import("std");
2const metrics = @import("metrics.zig");
3const blockingMode = @import("httpz.zig").blockingMode;
4
5const Io = std.Io;
6const Mutex = Io.Mutex;
7const Allocator = std.mem.Allocator;
8
9pub const Buffer = struct {
10 data: []u8,
11 type: Type,
12
13 const Type = enum {
14 arena,
15 static,
16 pooled,
17 dynamic,
18 };
19};
20
21pub const Pool = struct {
22 io: Io,
23 mutex: Mutex,
24 available: usize,
25 buffers: []Buffer,
26 allocator: Allocator,
27 buffer_size: usize,
28
29 pub fn init(io: Io, allocator: Allocator, count: usize, buffer_size: usize) !Pool {
30 const buffers = try allocator.alloc(Buffer, count);
31 errdefer allocator.free(buffers);
32
33 var initialized: usize = 0;
34 errdefer {
35 for (0..initialized) |i| {
36 allocator.free(buffers[i].data);
37 }
38 }
39
40 for (0..count) |i| {
41 buffers[i] = .{
42 .type = .pooled,
43 .data = try allocator.alloc(u8, buffer_size),
44 };
45 initialized += 1;
46 }
47
48 return .{
49 .io = io,
50 .mutex = .init,
51 .buffers = buffers,
52 .available = count,
53 .allocator = allocator,
54 .buffer_size = buffer_size,
55 };
56 }
57
58 pub fn deinit(self: *Pool) void {
59 const allocator = self.allocator;
60 for (self.buffers) |buf| {
61 allocator.free(buf.data);
62 }
63 allocator.free(self.buffers);
64 }
65
66 pub fn static(self: Pool, size: usize) !Buffer {
67 return .{
68 .type = .static,
69 .data = try self.allocator.alloc(u8, size),
70 };
71 }
72
73 pub fn alloc(self: *Pool, size: usize) !Buffer {
74 return self.allocType(self.allocator, .dynamic, size);
75 }
76
77 pub fn arenaAlloc(self: *Pool, arena: Allocator, size: usize) !Buffer {
78 return self.allocType(arena, .arena, size);
79 }
80
81 // Returns a pooled buffer, or null if the pool is empty. Never falls back
82 // to allocating. Useful when the caller wants to use the pool as a hot path
83 // and handle the empty case with a different strategy.
84 pub fn tryAlloc(self: *Pool) ?Buffer {
85 self.mutex.lockUncancelable(self.io);
86 defer self.mutex.unlock(self.io);
87
88 const available = self.available;
89 if (available == 0) {
90 return null;
91 }
92 const index = available - 1;
93 const buffer = self.buffers[index];
94 self.available = index;
95 return buffer;
96 }
97
98 fn allocType(self: *Pool, allocator: Allocator, buffer_type: Buffer.Type, size: usize) !Buffer {
99 if (size > self.buffer_size) {
100 metrics.allocBufferLarge(size);
101 return .{
102 .type = buffer_type,
103 .data = try allocator.alloc(u8, size),
104 };
105 }
106
107 self.mutex.lockUncancelable(self.io);
108 const available = self.available;
109 if (available == 0) {
110 self.mutex.unlock(self.io);
111 metrics.allocBufferEmpty(size);
112 return .{
113 .type = buffer_type,
114 .data = try allocator.alloc(u8, size),
115 };
116 }
117 defer self.mutex.unlock(self.io);
118
119 const index = available - 1;
120 const buffer = self.buffers[index];
121 self.available = index;
122 return buffer;
123 }
124
125 pub fn free(self: *Pool, buffer: Buffer) void {
126 switch (buffer.type) {
127 .arena => {},
128 .pooled => self.release(buffer),
129 .static => self.allocator.free(buffer.data),
130 .dynamic => self.allocator.free(buffer.data),
131 }
132 }
133
134 pub fn release(self: *Pool, buffer: Buffer) void {
135 switch (buffer.type) {
136 .static, .arena => {},
137 .dynamic => self.allocator.free(buffer.data),
138 .pooled => {
139 self.mutex.lockUncancelable(self.io);
140 defer self.mutex.unlock(self.io);
141 const available = self.available;
142 self.buffers[available] = buffer;
143 self.available = available + 1;
144 },
145 }
146 }
147};
148
149const t = @import("t.zig");
150
151test "BufferPool race" {
152 if (comptime blockingMode()) return error.SkipZigTest;
153
154 var pool = try Pool.init(t.io, t.allocator, 1, 64);
155 defer pool.deinit();
156
157 const Ctx = struct {
158 fn worker(p: *Pool) void {
159 var i: usize = 0;
160 while (i < 500_000) : (i += 1) {
161 const buf = p.tryAlloc() orelse continue;
162 std.atomic.spinLoopHint();
163 p.release(buf);
164 }
165 }
166 };
167
168 var threads: [4]std.Thread = undefined;
169 for (&threads) |*thr| thr.* = try std.Thread.spawn(.{}, Ctx.worker, .{&pool});
170 for (&threads) |*thr| thr.join();
171}
172
173test "BufferPool" {
174 var pool = try Pool.init(t.io, t.allocator, 2, 10);
175 defer pool.deinit();
176
177 {
178 // bigger than our buffers in pool
179 const buffer = try pool.alloc(11);
180 defer pool.release(buffer);
181 try t.expectEqual(.dynamic, buffer.type);
182 try t.expectEqual(11, buffer.data.len);
183 }
184
185 {
186 // smaller than our buffers in pool
187 const buf1 = try pool.alloc(4);
188 try t.expectEqual(.pooled, buf1.type);
189 try t.expectEqual(10, buf1.data.len);
190
191 const buf2 = try pool.alloc(5);
192 try t.expectEqual(.pooled, buf2.type);
193 try t.expectEqual(10, buf2.data.len);
194
195 try t.expectEqual(false, &buf1.data[0] == &buf2.data[0]);
196
197 // no more buffers in the pool, creates a dynamic buffer
198 const buf3 = try pool.alloc(6);
199 try t.expectEqual(.dynamic, buf3.type);
200 try t.expectEqual(6, buf3.data.len);
201
202 pool.release(buf1);
203
204 // now has items!
205 const buf4 = try pool.alloc(6);
206 try t.expectEqual(.pooled, buf4.type);
207 try t.expectEqual(10, buf4.data.len);
208
209 pool.release(buf2);
210 pool.release(buf3);
211 pool.release(buf4);
212 }
213}