karlseguin.tngl.sh
A poolable string builder (aka string buffer) for Zig
4.3 kB
152 lines
1const std = @import("std");
2const builtin = @import("builtin");
3
4const Buffer = @import("buffer.zig").Buffer;
5
6const Mutex = std.Thread.Mutex;
7const Allocator = std.mem.Allocator;
8
9pub const Pool = struct {
10 mutex: Mutex,
11 available: usize,
12 allocator: Allocator,
13 buffer_size: usize,
14 buffers: []*Buffer,
15
16 pub fn init(allocator: Allocator, pool_size: u16, buffer_size: usize) !Pool {
17 const buffers = try allocator.alloc(*Buffer, pool_size);
18
19 for (0..pool_size) |i| {
20 const sb = try allocator.create(Buffer);
21 sb.* = try Buffer.init(allocator, buffer_size);
22 buffers[i] = sb;
23 }
24
25 return .{ .mutex = .{}, .buffers = buffers, .allocator = allocator, .available = pool_size, .buffer_size = buffer_size };
26 }
27
28 pub fn deinit(self: *Pool) void {
29 const allocator = self.allocator;
30 for (self.buffers) |sb| {
31 sb.deinit();
32 allocator.destroy(sb);
33 }
34 allocator.free(self.buffers);
35 }
36
37 pub fn acquire(self: *Pool) !*Buffer {
38 return self.acquireWithAllocator(self.allocator);
39 }
40
41 pub fn acquireWithAllocator(self: *Pool, dyn_allocator: Allocator) !*Buffer {
42 const buffers = self.buffers;
43
44 self.mutex.lock();
45 const available = self.available;
46 if (available == 0) {
47 // dont hold the lock over factory
48 self.mutex.unlock();
49
50 const allocator = self.allocator;
51 const sb = try allocator.create(Buffer);
52 sb.* = try Buffer.init(allocator, self.buffer_size);
53 sb._da = dyn_allocator;
54 return sb;
55 }
56 const index = available - 1;
57 const sb = buffers[index];
58 self.available = index;
59 self.mutex.unlock();
60 sb._da = dyn_allocator;
61 return sb;
62 }
63
64 pub fn release(self: *Pool, sb: *Buffer) void {
65 sb.reset();
66 self.mutex.lock();
67
68 var buffers = self.buffers;
69 const available = self.available;
70 if (available == buffers.len) {
71 self.mutex.unlock();
72 const allocator = self.allocator;
73 sb.deinit();
74 allocator.destroy(sb);
75 return;
76 }
77 buffers[available] = sb;
78 self.available = available + 1;
79 self.mutex.unlock();
80 }
81};
82
83const t = @import("t.zig");
84test "pool: acquire and release" {
85 var p = try Pool.init(t.allocator, 2, 100);
86 defer p.deinit();
87
88 const sb1a = p.acquire() catch unreachable;
89 const sb2a = p.acquire() catch unreachable;
90 const sb3a = p.acquire() catch unreachable; // this should be dynamically generated
91
92 try t.expectEqual(false, sb1a == sb2a);
93 try t.expectEqual(false, sb2a == sb3a);
94
95 p.release(sb1a);
96
97 const sb1b = p.acquire() catch unreachable;
98 try t.expectEqual(true, sb1a == sb1b);
99
100 p.release(sb3a);
101 p.release(sb2a);
102 p.release(sb1b);
103}
104
105test "pool: dynamic allocator" {
106 var p = try Pool.init(t.allocator, 2, 5);
107 defer p.deinit();
108
109 var arena = std.heap.ArenaAllocator.init(t.allocator);
110 defer arena.deinit();
111
112 var sb = p.acquireWithAllocator(arena.allocator()) catch unreachable;
113 try sb.write("hello world how's it going?");
114 try sb.write("he");
115 try sb.write("hello world");
116 try sb.write("are you doing well? I hope so, I don't love how this is being implemented, but I think the feature is worthwhile");
117 p.release(sb);
118}
119
120test "pool: threadsafety" {
121 var p = try Pool.init(t.allocator, 3, 20);
122 defer p.deinit();
123
124 // initialize this to 0 since we're asserting that it's 0
125 for (p.buffers) |sb| {
126 sb.buf[0] = 0;
127 }
128
129 const t1 = try std.Thread.spawn(.{}, testPool, .{&p});
130 const t2 = try std.Thread.spawn(.{}, testPool, .{&p});
131 const t3 = try std.Thread.spawn(.{}, testPool, .{&p});
132
133 t1.join();
134 t2.join();
135 t3.join();
136}
137
138fn testPool(p: *Pool) void {
139 var r = t.getRandom();
140 const random = r.random();
141
142 for (0..5000) |_| {
143 var sb = p.acquire() catch unreachable;
144 // no other thread should have set this to 255
145 std.debug.assert(sb.buf[0] == 0);
146
147 sb.buf[0] = 255;
148 std.Thread.sleep(random.uintAtMost(u32, 100000));
149 sb.buf[0] = 0;
150 p.release(sb);
151 }
152}