lib/bench/src/allocation.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2
3 const timing = @import("timing.zig");
4
5 const Allocator = std.mem.Allocator;
6
7 /// Records the allocation traffic and the outstanding-byte high water of one
8 /// counted allocator. A counted allocator holds one of these and updates it on
9 /// every request that passes through it.
10 pub const Counts = struct {
11 /// Counts the allocation requests the backing allocator satisfied.
12 alloc_count: u64 = 0,
13 /// Counts the free requests.
14 free_count: u64 = 0,
15 /// Counts the bytes handed out by allocations and by resizes that grew.
16 alloc_bytes: u64 = 0,
17 /// Bytes handed out and not yet freed. The total carries across a reset.
18 live_bytes: u64 = 0,
19 /// The largest live-byte total reached since the counted allocator started
20 /// or since its last reset.
21 peak_live_bytes: u64 = 0,
22 };
23
24 pub const CountingAllocator = struct {
25 backing: Allocator,
26 counts: Counts = .{},
27 mutex: std.atomic.Mutex = .unlocked,
28
29 pub fn init(backing: Allocator) CountingAllocator {
30 return .{ .backing = backing };
31 }
32
33 pub fn allocator(self: *CountingAllocator) Allocator {
34 return .{ .ptr = self, .vtable = &vtable };
35 }
36
37 /// Zeroes the traffic counters and sets the high water back to the live
38 /// total. The live-byte total survives the reset, because memory handed out
39 /// earlier is still outstanding. The high water restarts at the live total,
40 /// so the next window reports the peak that window itself reached.
41 pub fn reset(self: *CountingAllocator) void {
42 self.lock();
43 defer self.mutex.unlock();
44 const live = self.counts.live_bytes;
45 self.counts = .{ .live_bytes = live, .peak_live_bytes = live };
46 }
47
48 const vtable: Allocator.VTable = .{
49 .alloc = rawAlloc,
50 .resize = rawResize,
51 .remap = rawRemap,
52 .free = rawFree,
53 };
54
55 fn rawAlloc(ctx: *anyopaque, len: usize, alignment: std.mem.Alignment, ret_addr: usize) ?[*]u8 {
56 const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
57 const ptr = self.backing.rawAlloc(len, alignment, ret_addr) orelse return null;
58 self.lock();
59 defer self.mutex.unlock();
60 self.counts.alloc_count +|= 1;
61 self.counts.alloc_bytes +|= timing.saturatingU64(len);
62 self.growLive(timing.saturatingU64(len));
63 return ptr;
64 }
65
66 fn rawResize(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) bool {
67 const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
68 const resized = self.backing.rawResize(memory, alignment, new_len, ret_addr);
69 if (resized) self.recordLength(memory.len, new_len);
70 return resized;
71 }
72
73 fn rawRemap(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, new_len: usize, ret_addr: usize) ?[*]u8 {
74 const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
75 const ptr = self.backing.rawRemap(memory, alignment, new_len, ret_addr) orelse return null;
76 self.recordLength(memory.len, new_len);
77 return ptr;
78 }
79
80 fn rawFree(ctx: *anyopaque, memory: []u8, alignment: std.mem.Alignment, ret_addr: usize) void {
81 const self: *CountingAllocator = @ptrCast(@alignCast(ctx));
82 self.backing.rawFree(memory, alignment, ret_addr);
83 self.lock();
84 defer self.mutex.unlock();
85 self.counts.free_count +|= 1;
86 self.counts.live_bytes -|= timing.saturatingU64(memory.len);
87 }
88
89 fn recordLength(self: *CountingAllocator, old_len: usize, new_len: usize) void {
90 if (new_len > old_len) {
91 const grew = timing.saturatingU64(new_len - old_len);
92 self.lock();
93 defer self.mutex.unlock();
94 self.counts.alloc_bytes +|= grew;
95 self.counts.live_bytes +|= grew;
96 self.counts.peak_live_bytes = @max(
97 self.counts.peak_live_bytes,
98 self.counts.live_bytes,
99 );
100 return;
101 }
102 if (new_len == old_len) return;
103 self.lock();
104 defer self.mutex.unlock();
105 self.counts.live_bytes -|= timing.saturatingU64(old_len - new_len);
106 }
107
108 fn growLive(self: *CountingAllocator, bytes: u64) void {
109 self.counts.live_bytes +|= bytes;
110 self.counts.peak_live_bytes = @max(
111 self.counts.peak_live_bytes,
112 self.counts.live_bytes,
113 );
114 }
115
116 fn lock(self: *CountingAllocator) void {
117 while (!self.mutex.tryLock()) std.atomic.spinLoopHint();
118 }
119 };
120
121 test "allocation counters retain concurrent updates" {
122 const worker_count = 4;
123 const allocations_per_worker = 256;
124 const allocation_size = 64;
125 const Worker = struct {
126 fn run(allocator: Allocator) void {
127 for (0..allocations_per_worker) |_| {
128 const memory = allocator.alloc(u8, allocation_size) catch @panic("allocation failed");
129 allocator.free(memory);
130 }
131 }
132 };
133 var counting = CountingAllocator.init(std.testing.allocator);
134 const allocator = counting.allocator();
135 var threads: [worker_count]std.Thread = undefined;
136 for (&threads) |*thread| {
137 thread.* = try std.Thread.spawn(.{}, Worker.run, .{allocator});
138 }
139 for (threads) |thread| thread.join();
140 const expected = worker_count * allocations_per_worker;
141 try std.testing.expectEqual(@as(u64, expected), counting.counts.alloc_count);
142 try std.testing.expectEqual(@as(u64, expected), counting.counts.free_count);
143 try std.testing.expectEqual(
144 @as(u64, expected * allocation_size),
145 counting.counts.alloc_bytes,
146 );
147 try std.testing.expectEqual(@as(u64, 0), counting.counts.live_bytes);
148 }
149
150 test "allocation counters report a window high water" {
151 const block_bytes = 4096;
152 var counting = CountingAllocator.init(std.testing.allocator);
153 const allocator = counting.allocator();
154 const retained = try allocator.alloc(u8, block_bytes);
155 defer allocator.free(retained);
156 counting.reset();
157 try std.testing.expectEqual(@as(u64, block_bytes), counting.counts.live_bytes);
158 try std.testing.expectEqual(@as(u64, block_bytes), counting.counts.peak_live_bytes);
159
160 const transient = try allocator.alloc(u8, block_bytes * 3);
161 allocator.free(transient);
162 try std.testing.expectEqual(@as(u64, block_bytes), counting.counts.live_bytes);
163 try std.testing.expectEqual(@as(u64, block_bytes * 4), counting.counts.peak_live_bytes);
164 try std.testing.expectEqual(@as(u64, block_bytes * 3), counting.counts.alloc_bytes);
165 }