lib/memtrace/src/census/census.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const pretty_json = @import("pretty").json;
3 const memtrace = @import("../root.zig");
4 const model = @import("model.zig");
5 const storage_mod = @import("storage.zig");
6
7 const Entry = model.Entry;
8 const SummaryOptions = model.SummaryOptions;
9 const Storage = storage_mod.Storage;
10
11 const RecordingExhaustion = error{
12 CategoryCapacityExceeded,
13 LabelBytesCapacityExceeded,
14 JoinedLabelCapacityExceeded,
15 CensusCounterOverflow,
16 };
17
18 pub const Exhaustion = storage_mod.Exhaustion || RecordingExhaustion;
19
20 pub const Status = struct {
21 categories: usize,
22 category_capacity: usize,
23 label_bytes_used: usize,
24 label_bytes_capacity: usize,
25 joined_label_bytes_capacity: usize,
26 storage_bytes: usize,
27 };
28
29 const Lookup = union(enum) {
30 existing: *model.Slot,
31 vacant: *model.Slot,
32 };
33
34 pub const Census = struct {
35 storage: *Storage,
36 slots: []model.Slot,
37 summary: []Entry,
38 labels: []u8,
39 joined_label: []u8,
40 category_count: usize = 0,
41 label_bytes_used: usize = 0,
42
43 pub fn init(storage: *Storage) storage_mod.Exhaustion!Census {
44 const regions = try storage.acquire();
45 return .{
46 .storage = storage,
47 .slots = regions.slots,
48 .summary = regions.summary,
49 .labels = regions.labels,
50 .joined_label = regions.joined_label,
51 };
52 }
53
54 pub fn deinit(self: *Census) void {
55 const storage = self.storage;
56 self.* = undefined;
57 storage.release();
58 }
59
60 pub fn record(self: *Census, label: []const u8, bytes: usize) Exhaustion!void {
61 try self.recordCount(label, 1, bytes);
62 }
63
64 pub fn recordPrefixed(
65 self: *Census,
66 prefix: []const u8,
67 label: []const u8,
68 bytes: usize,
69 ) Exhaustion!void {
70 try self.recordCountPrefixed(prefix, label, 1, bytes);
71 }
72
73 pub fn recordCount(
74 self: *Census,
75 label: []const u8,
76 count: usize,
77 bytes: usize,
78 ) Exhaustion!void {
79 const hash = std.hash.Wyhash.hash(0, label);
80 switch (self.locate(label, hash)) {
81 .existing => |slot| {
82 const next_items = std.math.add(usize, slot.entry.items, count) catch
83 return error.CensusCounterOverflow;
84 const next_bytes = std.math.add(usize, slot.entry.bytes, bytes) catch
85 return error.CensusCounterOverflow;
86 slot.entry.items = next_items;
87 slot.entry.bytes = next_bytes;
88 },
89 .vacant => |slot| {
90 if (self.category_count >= self.summary.len) return error.CategoryCapacityExceeded;
91 if (label.len > self.labels.len - self.label_bytes_used) {
92 return error.LabelBytesCapacityExceeded;
93 }
94 const start = self.label_bytes_used;
95 const end = start + label.len;
96 const owned = self.labels[start..end];
97 @memcpy(owned, label);
98 slot.* = .{
99 .entry = .{
100 .label = owned,
101 .items = count,
102 .bytes = bytes,
103 },
104 .hash = hash,
105 .occupied = true,
106 };
107 self.label_bytes_used = end;
108 self.category_count += 1;
109 },
110 }
111 }
112
113 pub fn recordCountPrefixed(
114 self: *Census,
115 prefix: []const u8,
116 label: []const u8,
117 count: usize,
118 bytes: usize,
119 ) Exhaustion!void {
120 if (prefix.len == 0) return try self.recordCount(label, count, bytes);
121 const joined_len = std.math.add(usize, prefix.len, label.len) catch
122 return error.JoinedLabelCapacityExceeded;
123 if (joined_len > self.joined_label.len) return error.JoinedLabelCapacityExceeded;
124 @memcpy(self.joined_label[0..prefix.len], prefix);
125 @memcpy(self.joined_label[prefix.len..joined_len], label);
126 try self.recordCount(self.joined_label[0..joined_len], count, bytes);
127 }
128
129 pub fn writeSummary(
130 self: *Census,
131 writer: *std.Io.Writer,
132 options: SummaryOptions,
133 ) !void {
134 var item_count: usize = 0;
135 for (self.slots) |slot| {
136 if (!slot.occupied or slot.entry.bytes < options.min_bytes) continue;
137 std.debug.assert(item_count < self.summary.len);
138 self.summary[item_count] = slot.entry;
139 item_count += 1;
140 }
141 const items = self.summary[0..item_count];
142 std.mem.sort(Entry, items, {}, entryGreaterThan);
143
144 try writer.print("memtrace census categories={d}\n", .{items.len});
145 const limit = @min(options.top, items.len);
146 for (items[0..limit]) |entry| {
147 try writer.print(
148 "{s} items={d} bytes={d}\n",
149 .{ entry.label, entry.items, entry.bytes },
150 );
151 }
152 }
153
154 pub fn writeJsonl(self: *const Census, writer: *std.Io.Writer) !void {
155 var seq: u64 = 0;
156 for (self.slots) |slot| {
157 if (!slot.occupied) continue;
158 seq = std.math.add(u64, seq, 1) catch return error.CensusCounterOverflow;
159 var stream = pretty_json.Writer.init(writer, .minified);
160 const object = try stream.object();
161 try object.field("v", memtrace.event.format_version);
162 try object.field("seq", seq);
163 try object.field("kind", "census");
164 try object.field("label", slot.entry.label);
165 try object.field("items", slot.entry.items);
166 try object.field("bytes", slot.entry.bytes);
167 try object.endLine();
168 }
169 }
170
171 pub fn status(self: *const Census) Status {
172 return .{
173 .categories = self.category_count,
174 .category_capacity = self.summary.len,
175 .label_bytes_used = self.label_bytes_used,
176 .label_bytes_capacity = self.labels.len,
177 .joined_label_bytes_capacity = self.joined_label.len,
178 .storage_bytes = self.storage.capacity.storage_bytes,
179 };
180 }
181
182 fn locate(self: *Census, label: []const u8, hash: u64) Lookup {
183 std.debug.assert(self.slots.len >= 2);
184 std.debug.assert(std.math.isPowerOfTwo(self.slots.len));
185 const mask = self.slots.len - 1;
186 const start = @as(usize, @truncate(hash)) & mask;
187 for (0..self.slots.len) |probe| {
188 const slot = &self.slots[(start +% probe) & mask];
189 if (!slot.occupied) return .{ .vacant = slot };
190 if (slot.hash == hash and std.mem.eql(u8, slot.entry.label, label)) {
191 return .{ .existing = slot };
192 }
193 }
194 unreachable;
195 }
196 };
197
198 fn entryGreaterThan(_: void, left: Entry, right: Entry) bool {
199 if (left.bytes != right.bytes) return left.bytes > right.bytes;
200 if (left.items != right.items) return left.items > right.items;
201 return std.mem.lessThan(u8, left.label, right.label);
202 }