lib/trace/src/store/reader/stream.zig
daab053ee43316e1809a84551d573ddd1e5bf3d2
1 const std = @import("std");
2 const sql = @import("sql");
3 const sys = @import("sys");
4 const event = @import("../../root.zig").event;
5 const format = @import("../format/root.zig");
6 const block = format.block;
7 const codec = format.codec;
8 const database = format.database;
9 const capacity_mod = @import("capacity.zig");
10 const parse = @import("parse.zig");
11 const storage_mod = @import("storage.zig");
12
13 pub const Limits = capacity_mod.Limits;
14 pub const Capacity = capacity_mod.Capacity;
15 pub const DeriveError = capacity_mod.DeriveError;
16 pub const Storage = storage_mod.Storage;
17 pub const StorageStatus = storage_mod.Status;
18 pub const StorageExhaustion = storage_mod.Exhaustion;
19
20 pub const Reader = struct {
21 storage: ?*Storage = null,
22 root_path: []u8 = &.{},
23 header: format.Header = undefined,
24 block_buffer: []u8 = &.{},
25 block_view: ?block.View = null,
26 sql_allocator: std.heap.FixedBufferAllocator = undefined,
27 database_workspace: sql.FileDatabase.Workspace = undefined,
28 database_workspace_open: bool = false,
29 file: sql.FileDatabase = undefined,
30 file_open: bool = false,
31 events: sql.Tree = undefined,
32 chunk_event_count: u64 = 0,
33 chunk_byte_count: u64 = 0,
34 event_hasher: std.hash.Wyhash = format.newHasher(),
35 loaded_chunk_count: u64 = 0,
36 loaded_event_count: u64 = 0,
37 loaded_event_bytes: u64 = 0,
38 loaded_block_count: u64 = 0,
39 loaded_block_bytes: u64 = 0,
40 current: ?event.Event = null,
41 initialized: bool = false,
42 verified: bool = false,
43
44 pub fn open(self: *Reader, storage: *Storage, root_path: []const u8) !void {
45 if (self.initialized) return error.TraceReaderAlreadyOpen;
46 const regions = try storage.acquire(.{ .root_path_bytes = root_path.len });
47 var acquired = true;
48 errdefer if (acquired) storage.release();
49 std.mem.copyForwards(u8, regions.root_path, root_path);
50 const header_value = try readHeader(
51 regions.root_path,
52 regions.input,
53 regions.path,
54 regions.json_stack,
55 storage.capacity.json_nesting,
56 );
57 try storage.admit(header_value.limits);
58 if (!header_value.status_closed) return error.TraceNotClosed;
59 const block_bytes = block.capacity(header_value.limits.max_event_bytes) catch
60 return error.InvalidTraceManifest;
61 std.debug.assert(block_bytes <= regions.input.len);
62 self.* = .{
63 .storage = storage,
64 .root_path = regions.root_path,
65 .header = header_value,
66 .block_buffer = regions.input[0..block_bytes],
67 .sql_allocator = std.heap.FixedBufferAllocator.init(regions.sql),
68 .initialized = true,
69 };
70 acquired = false;
71 errdefer self.deinit();
72 self.database_workspace = try sql.FileDatabase.Workspace.allocate(
73 self.sql_allocator.allocator(),
74 database.workspaceLimits(storage.capacity.database),
75 );
76 self.database_workspace_open = true;
77 self.file = try database.openReader(
78 self.sql_allocator.allocator(),
79 &self.database_workspace,
80 self.root_path,
81 storage.capacity.database,
82 );
83 self.file_open = true;
84 self.events = try database.events(&self.file);
85 const identity = try self.events.identity();
86 const expected_key_bytes = std.math.mul(u64, self.header.block_count, database.key_bytes) catch
87 return error.InvalidTraceManifest;
88 if (identity.entries != self.header.block_count or
89 identity.key_bytes != expected_key_bytes or
90 identity.value_bytes != self.header.block_bytes)
91 {
92 return error.InvalidTraceManifest;
93 }
94 }
95
96 pub fn deinit(self: *Reader) void {
97 if (!self.initialized) return;
98 if (self.file_open) {
99 self.file.deinit();
100 self.file_open = false;
101 }
102 if (self.database_workspace_open) {
103 self.database_workspace.deallocate(self.sql_allocator.allocator());
104 self.database_workspace_open = false;
105 }
106 const storage = self.storage.?;
107 self.* = .{};
108 storage.release();
109 }
110
111 pub fn source(self: *Reader) event.Source {
112 return .{
113 .context = self,
114 .peekFn = peekFromSource,
115 .advanceFn = advanceFromSource,
116 .countFn = countFromSource,
117 };
118 }
119
120 pub fn peek(self: *Reader) !?*const event.Event {
121 if (!self.initialized) return error.TraceReaderNotOpen;
122 if (self.current) |*item| return item;
123 if (self.verified) return null;
124 const maybe_encoded = try self.peekEncoded();
125 if (self.loaded_event_count == self.header.event_count) {
126 if (maybe_encoded != null) return error.InvalidTraceManifest;
127 try self.finishTrace();
128 return null;
129 }
130 const encoded = maybe_encoded orelse return error.InvalidTraceManifest;
131 if (encoded.len == 0 or encoded.len > self.header.limits.max_event_bytes) {
132 return error.InvalidTraceEvent;
133 }
134 const prospective_chunk_bytes = std.math.add(u64, self.chunk_byte_count, encoded.len) catch
135 return error.InvalidTraceManifest;
136 var prior_chunk_event_count = self.chunk_event_count;
137 var prior_chunk_byte_count = self.chunk_byte_count;
138 var next_loaded_chunk_count = self.loaded_chunk_count;
139 if (self.chunk_event_count != 0 and prospective_chunk_bytes > self.header.limits.max_chunk_bytes) {
140 next_loaded_chunk_count = std.math.add(u64, self.loaded_chunk_count, 1) catch
141 return error.InvalidTraceManifest;
142 prior_chunk_event_count = 0;
143 prior_chunk_byte_count = 0;
144 }
145 const next_chunk_event_count = std.math.add(u64, prior_chunk_event_count, 1) catch
146 return error.InvalidTraceManifest;
147 const next_chunk_byte_count = std.math.add(u64, prior_chunk_byte_count, encoded.len) catch
148 return error.InvalidTraceManifest;
149 const next_event_count = std.math.add(u64, self.loaded_event_count, 1) catch
150 return error.InvalidTraceManifest;
151 const next_event_bytes = std.math.add(u64, self.loaded_event_bytes, encoded.len) catch
152 return error.InvalidTraceManifest;
153 if (next_chunk_byte_count > self.header.limits.max_chunk_bytes or
154 next_event_count > self.header.event_count or
155 next_event_bytes > self.header.event_bytes)
156 {
157 return error.InvalidTraceManifest;
158 }
159 const parsed = codec.decode(encoded) catch return error.InvalidTraceEvent;
160 var key_buffer: [database.key_bytes]u8 = undefined;
161 const key = database.sequenceKey(&key_buffer, self.loaded_event_count);
162 self.advanceEncoded();
163 database.updateChecksum(&self.event_hasher, key, encoded);
164 self.loaded_chunk_count = next_loaded_chunk_count;
165 self.chunk_event_count = next_chunk_event_count;
166 self.chunk_byte_count = next_chunk_byte_count;
167 self.loaded_event_count = next_event_count;
168 self.loaded_event_bytes = next_event_bytes;
169 self.current = parsed;
170 return &self.current.?;
171 }
172
173 fn peekEncoded(self: *Reader) !?[]u8 {
174 while (true) {
175 if (self.block_view) |*view| {
176 if (view.peek()) |encoded| return encoded;
177 self.block_view = null;
178 }
179 if (self.loaded_block_count == self.header.block_count) return null;
180 try self.loadBlock();
181 }
182 }
183
184 fn advanceEncoded(self: *Reader) void {
185 self.block_view.?.advance();
186 }
187
188 fn loadBlock(self: *Reader) !void {
189 var key_buffer: [database.key_bytes]u8 = undefined;
190 const key = database.sequenceKey(&key_buffer, self.loaded_block_count);
191 const encoded = self.events.getInto(key, self.block_buffer) catch |err| switch (err) {
192 error.OutputTooSmall => return error.InvalidTraceEvent,
193 else => |read_err| return read_err,
194 } orelse return error.InvalidTraceManifest;
195 const view = block.View.init(encoded) catch return error.InvalidTraceEvent;
196 const next_block_count = std.math.add(u64, self.loaded_block_count, 1) catch
197 return error.InvalidTraceManifest;
198 const next_block_bytes = std.math.add(u64, self.loaded_block_bytes, encoded.len) catch
199 return error.InvalidTraceManifest;
200 if (next_block_count > self.header.block_count or next_block_bytes > self.header.block_bytes) {
201 return error.InvalidTraceManifest;
202 }
203 self.block_view = view;
204 self.loaded_block_count = next_block_count;
205 self.loaded_block_bytes = next_block_bytes;
206 }
207
208 pub fn advance(self: *Reader) void {
209 std.debug.assert(self.current != null);
210 self.current = null;
211 }
212
213 pub fn eventCount(self: *const Reader) u64 {
214 return self.header.event_count;
215 }
216
217 pub fn verify(self: *Reader) !VerifyResult {
218 while (try self.peek()) |_| self.advance();
219 return .{
220 .chunk_count = self.loaded_chunk_count,
221 .event_count = self.loaded_event_count,
222 .event_bytes = self.loaded_event_bytes,
223 .checksum = self.event_hasher.final(),
224 };
225 }
226
227 fn finishTrace(self: *Reader) !void {
228 if (self.chunk_event_count != 0) {
229 self.loaded_chunk_count = std.math.add(u64, self.loaded_chunk_count, 1) catch
230 return error.InvalidTraceManifest;
231 self.chunk_event_count = 0;
232 self.chunk_byte_count = 0;
233 }
234 if (self.loaded_chunk_count != self.header.chunk_count or
235 self.loaded_event_count != self.header.event_count or
236 self.loaded_event_bytes != self.header.event_bytes or
237 self.loaded_block_count != self.header.block_count or
238 self.loaded_block_bytes != self.header.block_bytes)
239 {
240 return error.InvalidTraceManifest;
241 }
242 if (self.event_hasher.final() != self.header.event_checksum) return error.TraceChecksumMismatch;
243 self.verified = true;
244 }
245
246 fn peekFromSource(context: *anyopaque) !?*const event.Event {
247 const self: *Reader = @ptrCast(@alignCast(context));
248 return try self.peek();
249 }
250
251 fn advanceFromSource(context: *anyopaque) void {
252 const self: *Reader = @ptrCast(@alignCast(context));
253 self.advance();
254 }
255
256 fn countFromSource(context: *anyopaque) u64 {
257 const self: *Reader = @ptrCast(@alignCast(context));
258 return self.eventCount();
259 }
260 };
261
262 fn testReaderForBlock(
263 block_bytes: []const u8,
264 event_bytes: usize,
265 max_event_bytes: usize,
266 ) !Reader {
267 return .{
268 .header = .{
269 .format_version = event.trace_format_version,
270 .status_closed = true,
271 .limits = .{
272 .max_event_bytes = max_event_bytes,
273 .max_chunk_bytes = 512,
274 },
275 .chunk_count = 1,
276 .event_count = 1,
277 .event_bytes = event_bytes,
278 .event_checksum = 0,
279 .block_count = 1,
280 .block_bytes = block_bytes.len,
281 },
282 .block_view = try block.View.init(@constCast(block_bytes)),
283 .loaded_block_count = 1,
284 .loaded_block_bytes = block_bytes.len,
285 .initialized = true,
286 };
287 }
288
289 fn expectPeekFailureRetryable(reader: *Reader, expected: anyerror) !void {
290 const initial_offset = reader.block_view.?.offset;
291 const initial_remaining = reader.block_view.?.remaining;
292
293 try std.testing.expectError(expected, reader.peek());
294 try std.testing.expectError(expected, reader.peek());
295 try std.testing.expectEqual(initial_offset, reader.block_view.?.offset);
296 try std.testing.expectEqual(initial_remaining, reader.block_view.?.remaining);
297 try std.testing.expectEqual(@as(u64, 0), reader.loaded_event_count);
298 try std.testing.expectEqual(@as(?event.Event, null), reader.current);
299 }
300
301 test "reader decode failure leaves the encoded event pending for retry" {
302 var bytes: [512]u8 = undefined;
303 var builder = block.Builder.init(&bytes);
304 const item = event.Event.user(.{ .thread_id = 1, .seq = 1 }, "label", "payload");
305 const encoded = try builder.append(item, try codec.encodedSize(item));
306 const block_bytes = builder.finish();
307 encoded[0] = 0;
308
309 var reader = try testReaderForBlock(block_bytes, encoded.len, 512);
310 try expectPeekFailureRetryable(&reader, error.InvalidTraceEvent);
311 }
312
313 test "reader limit failure leaves the encoded event pending for retry" {
314 var bytes: [512]u8 = undefined;
315 var builder = block.Builder.init(&bytes);
316 const item = event.Event.user(.{ .thread_id = 1, .seq = 1 }, "label", "payload");
317 const encoded = try builder.append(item, try codec.encodedSize(item));
318 const block_bytes = builder.finish();
319
320 var reader = try testReaderForBlock(block_bytes, encoded.len, encoded.len - 1);
321 try expectPeekFailureRetryable(&reader, error.InvalidTraceEvent);
322 }
323
324 pub const VerifyResult = struct {
325 chunk_count: u64,
326 event_count: u64,
327 event_bytes: u64,
328 checksum: u64,
329 };
330
331 fn readHeader(
332 root_path: []const u8,
333 input: []u8,
334 path_buffer: []u8,
335 json_stack: []u8,
336 json_nesting: usize,
337 ) !format.Header {
338 const path = joinedPath(root_path, format.manifest_path, path_buffer);
339 var file = try sys.fs.cwd().openFile(sys.fs.debugIo(), path, .{});
340 defer file.close(sys.fs.debugIo());
341 const size_u64 = (try file.stat(sys.fs.debugIo())).size;
342 if (size_u64 == 0 or size_u64 > format.max_manifest_bytes) return error.InvalidTraceManifest;
343 const size = std.math.cast(usize, size_u64) orelse return error.InvalidTraceManifest;
344 std.debug.assert(size <= input.len);
345 const read = try file.readPositionalAll(sys.fs.debugIo(), input[0..size], 0);
346 if (read != size) return error.InvalidTraceManifest;
347 return parse.header(input[0..size], json_stack, json_nesting) catch |err| switch (err) {
348 error.JsonNestingCapacityExceeded => return err,
349 error.UnsupportedTraceVersion => return err,
350 error.InvalidDocument => return error.InvalidTraceManifest,
351 };
352 }
353
354 fn joinedPath(root_path: []const u8, relative: []const u8, buffer: []u8) []const u8 {
355 std.debug.assert(root_path.len != 0);
356 std.debug.assert(relative.len <= format.max_relative_path_bytes);
357 var length = root_path.len;
358 @memcpy(buffer[0..length], root_path);
359 if (!sys.path.isSeparator(buffer[length - 1])) {
360 buffer[length] = sys.path.separator;
361 length += 1;
362 }
363 @memcpy(buffer[length..][0..relative.len], relative);
364 length += relative.len;
365 std.debug.assert(length <= buffer.len);
366 return buffer[0..length];
367 }