calabro.io
forked from
zzstoatzz.io/websocket.zig
websocket
11 kB
384 lines
1const std = @import("std");
2
3const Allocator = std.mem.Allocator;
4
5pub const Buffer = struct {
6 data: []u8,
7 type: Type,
8
9 const Type = enum {
10 static,
11 pooled,
12 dynamic,
13 };
14};
15
16pub const Writer = struct {
17 buf: []u8,
18 pos: usize = 0,
19 pooled: bool,
20 provider: *Provider,
21 interface: std.Io.Writer,
22
23 pub fn init(buf: []u8, pooled: bool, provider: *Provider) Writer {
24 return .{
25 .buf = buf,
26 .pooled = pooled,
27 .provider = provider,
28 .interface = .{
29 .buffer = &.{},
30 .vtable = &.{
31 .drain = drain,
32 .flush = std.Io.Writer.defaultFlush,
33 .rebase = std.Io.Writer.failingRebase,
34 },
35 },
36 };
37 }
38
39 pub fn deinit(self: *Writer) void {
40 if (self.pooled) {
41 self.provider.pool.release(self.buf);
42 } else {
43 self.provider.allocator.free(self.buf);
44 }
45 }
46
47 pub fn drain(io_w: *std.Io.Writer, data: []const []const u8, splat: usize) error{WriteFailed}!usize {
48 const self: *Writer = @alignCast(@fieldParentPtr("interface", io_w));
49 var consumed: usize = 0;
50 for (data[0 .. data.len - 1]) |bytes| {
51 self.writeAll(bytes) catch return error.WriteFailed;
52 consumed += bytes.len;
53 }
54 const pattern = data[data.len - 1];
55 for (0..splat) |_| {
56 self.writeAll(pattern) catch return error.WriteFailed;
57 consumed += pattern.len;
58 }
59 return consumed;
60 }
61
62 pub fn writeAll(self: *Writer, data: []const u8) !void {
63 const pos = self.pos;
64 const total_len = pos + data.len;
65 if (total_len > self.provider.max_buffer_size) {
66 return error.TooLarge;
67 }
68 try self.ensureTotalCapacity(total_len);
69
70 @memcpy(self.buf[pos..total_len], data);
71 self.pos = total_len;
72 }
73
74 fn ensureTotalCapacity(self: *Writer, required_capacity: usize) !void {
75 const buf = self.buf;
76 if (required_capacity <= buf.len) {
77 return;
78 }
79
80 // from std.ArrayList
81 var new_capacity = buf.len;
82 while (true) {
83 new_capacity +|= new_capacity / 2 + 8;
84 if (new_capacity >= required_capacity) break;
85 }
86
87 const allocator = self.provider.allocator;
88 if (self.pooled or !allocator.resize(buf, new_capacity)) {
89 const new_buffer = try allocator.alloc(u8, new_capacity);
90 @memcpy(new_buffer[0..buf.len], buf);
91
92 if (self.pooled) {
93 self.provider.pool.release(buf);
94 } else {
95 allocator.free(buf);
96 }
97
98 self.buf = new_buffer;
99 self.pooled = false;
100 } else {
101 const new_buffer = buf.ptr[0..new_capacity];
102 self.buf = new_buffer;
103 }
104 }
105};
106
107pub const Config = struct {
108 count: u16 = 1,
109 size: usize = 65536,
110 max: usize = 65536,
111};
112
113// Manages all buffer access and types. It's where code goes to ask
114// for and release buffers.
115pub const Provider = struct {
116 pool: Pool,
117 allocator: Allocator,
118
119 max_buffer_size: usize,
120
121 // If this is 0, pool is undefined. We need this field here anyways.
122 pool_buffer_size: usize,
123
124 pub fn init(allocator: Allocator, config: Config) !Provider {
125 const size = config.size;
126 const count = config.count;
127
128 if (count == 0 or size == 0) {
129
130 // Large buffering can be disabled, in which case any large buffers will
131 // be dynamically allocated using the allocator (assuming the requested
132 // size is less than the max_message_size)
133 return .{
134 // this is safe to do, because we set size = 0, so we'll
135 // never try to access the pool
136 .pool = undefined,
137 .pool_buffer_size = 0,
138 .allocator = allocator,
139 .max_buffer_size = config.max,
140 };
141 }
142
143 return .{
144 .allocator = allocator,
145 .pool_buffer_size = size,
146 .max_buffer_size = config.max,
147 .pool = try Pool.init(allocator, count, size),
148 };
149 }
150
151 pub fn deinit(self: *Provider) void {
152 if (self.pool_buffer_size > 0) {
153 // else, pool is undefined
154 self.pool.deinit();
155 }
156 }
157
158 pub fn alloc(self: *Provider, size: usize) !Buffer {
159 if (size > self.max_buffer_size) {
160 return error.TooLarge;
161 }
162
163 // remember: if self.pool_buffer_size == 0, then self.pool is undefined.
164 if (size <= self.pool_buffer_size) {
165 if (self.pool.acquire()) |buffer| {
166 // See the Reader struct comment to see why this is necessary
167 var copy = buffer;
168 copy.len = size;
169 return .{ .type = .pooled, .data = copy };
170 }
171 }
172
173 return .{
174 .type = .dynamic,
175 .data = try self.allocator.alloc(u8, size),
176 };
177 }
178
179 pub fn grow(self: *Provider, buffer: Buffer, current_size: usize, new_size: usize) !Buffer {
180 if (new_size > self.max_buffer_size) {
181 return error.TooLarge;
182 }
183
184 if (buffer.type == .dynamic) {
185 var copy = buffer;
186 copy.data = try self.allocator.realloc(buffer.data, new_size);
187 return copy;
188 }
189
190 defer self.release(buffer);
191
192 const new_buffer = try self.alloc(new_size);
193 @memcpy(new_buffer.data[0..current_size], buffer.data[0..current_size]);
194 return new_buffer;
195 }
196
197 pub fn free(self: *Provider, buffer: Buffer) void {
198 switch (buffer.type) {
199 .pooled => {
200 // this resize is necessary because on alloc, we potentially shrink data
201 var copy = buffer.data;
202 copy.len = self.pool_buffer_size;
203 self.pool.release(copy);
204 },
205 .static => self.allocator.free(buffer.data),
206 .dynamic => self.allocator.free(buffer.data),
207 }
208 }
209
210 pub fn release(self: *Provider, buffer: Buffer) void {
211 switch (buffer.type) {
212 .static => {},
213 .pooled => {
214 // this resize is necessary because on alloc, we potentially shrink data
215 var copy = buffer.data;
216 copy.len = self.pool_buffer_size;
217 self.pool.release(copy);
218 },
219 .dynamic => self.allocator.free(buffer.data),
220 }
221 }
222};
223
224pub const Pool = struct {
225 buffer_size: usize,
226 available: usize,
227 buffers: [][]u8,
228 allocator: Allocator,
229 mutex: std.Io.Mutex,
230 io: std.Io,
231
232 pub fn init(allocator: Allocator, count: usize, buffer_size: usize) !Pool {
233 const buffers = try allocator.alloc([]u8, count);
234
235 for (0..count) |i| {
236 buffers[i] = try allocator.alloc(u8, buffer_size);
237 }
238
239 return .{
240 .mutex = .init,
241 .io = std.Options.debug_io,
242 .buffers = buffers,
243 .available = count,
244 .allocator = allocator,
245 .buffer_size = buffer_size,
246 };
247 }
248
249 pub fn deinit(self: *Pool) void {
250 const allocator = self.allocator;
251 for (self.buffers) |buf| {
252 allocator.free(buf);
253 }
254 allocator.free(self.buffers);
255 }
256
257 pub fn acquire(self: *Pool) ?[]u8 {
258 const buffers = self.buffers;
259
260 self.mutex.lockUncancelable(self.io);
261 defer self.mutex.unlock(self.io);
262 const available = self.available;
263 if (available == 0) {
264 return null;
265 }
266 const index = available - 1;
267 const buffer = buffers[index];
268 self.available = index;
269 return buffer;
270 }
271
272 pub fn acquireOrCreate(self: *Pool) ![]u8 {
273 return self.acquire() orelse self.allocator.alloc(u8, self.buffer_size);
274 }
275
276 pub fn release(self: *Pool, buffer: []u8) void {
277 var buffers = self.buffers;
278
279 self.mutex.lockUncancelable(self.io);
280 const available = self.available;
281 if (available == buffers.len) {
282 self.mutex.unlock(self.io);
283 self.allocator.free(buffer);
284 return;
285 }
286 buffers[available] = buffer;
287 self.available = available + 1;
288 self.mutex.unlock(self.io);
289 }
290};
291
292const t = @import("t.zig");
293test "buffer: no pool" {
294 var p = try Provider.init(t.allocator, .{ .count = 0, .size = 0, .max = 100 });
295
296 const buffer = try p.alloc(100);
297 defer p.free(buffer);
298 try t.expectEqual(.dynamic, buffer.type);
299 try t.expectEqual(100, buffer.data.len);
300}
301
302test "buffer: pool" {
303 var p = try Provider.init(t.allocator, .{ .count = 2, .size = 10, .max = 15 });
304 defer p.deinit();
305
306 {
307 // bigger than allowed
308 try t.expectError(error.TooLarge, p.alloc(16));
309 }
310
311 {
312 // bigger than our buffers in pool
313 const buffer = try p.alloc(15);
314 defer p.free(buffer);
315 try t.expectEqual(.dynamic, buffer.type);
316 try t.expectEqual(15, buffer.data.len);
317 }
318
319 {
320 // smaller than our buffers in pool
321 const buf1 = try p.alloc(4);
322 try t.expectEqual(.pooled, buf1.type);
323 try t.expectEqual(4, buf1.data.len);
324
325 const buf2 = try p.alloc(5);
326 try t.expectEqual(.pooled, buf2.type);
327 try t.expectEqual(5, buf2.data.len);
328 try t.expectEqual(true, buf1.data.ptr != buf2.data.ptr);
329
330 // no more buffers in the pool, creats a dynamic buffer
331 const buf3 = try p.alloc(6);
332 try t.expectEqual(.dynamic, buf3.type);
333 try t.expectEqual(6, buf3.data.len);
334
335 p.release(buf1);
336
337 const buf4 = try p.alloc(7);
338 try t.expectEqual(.pooled, buf4.type);
339 try t.expectEqual(7, buf4.data.len);
340 try t.expectEqual(true, buf1.data.ptr == buf4.data.ptr);
341
342 p.release(buf2);
343 p.release(buf3);
344 }
345}
346
347test "buffer: grow" {
348 var p = try Provider.init(t.allocator, .{ .count = 1, .size = 10, .max = 30 });
349 defer p.deinit();
350
351 {
352 // grow a dynamic buffer
353 var buf1 = try p.alloc(15);
354 @memcpy(buf1.data[0..5], "hello");
355 const buf2 = try p.grow(buf1, 5, 20);
356 defer p.free(buf2);
357 try t.expectEqual(20, buf2.data.len);
358 try t.expectString("hello", buf2.data[0..5]);
359 }
360
361 {
362 // grow a static buffer
363 var buf1 = Buffer{ .type = .static, .data = try t.allocator.alloc(u8, 15) };
364 defer t.allocator.free(buf1.data);
365 @memcpy(buf1.data[0..6], "hello2");
366
367 const buf2 = try p.grow(buf1, 6, 21);
368 defer p.free(buf2);
369 try t.expectEqual(21, buf2.data.len);
370 try t.expectString("hello2", buf2.data[0..6]);
371 }
372
373 {
374 // grow a pooled buffer
375 var buf1 = try p.alloc(8);
376
377 @memcpy(buf1.data[0..7], "hello2a");
378 const buf2 = try p.grow(buf1, 7, 14);
379 defer p.free(buf2);
380 try t.expectEqual(14, buf2.data.len);
381 try t.expectString("hello2a", buf2.data[0..7]);
382 try t.expectEqual(1, p.pool.available);
383 }
384}