Something went wrong. Try again.
atproto relay in zig zlay.waow.tech
relay zig atproto
Something went wrong. Try again.
Zig
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368const std = @import("std");const index = @import("internal/collection_index/index.zig");const zio = @import("zio");const backend = @import("backend_config");const Io = std.Io;const collection = "app.bsky.feed.post";const Pair = struct { did: []const u8 };const Batch = struct { const Request = struct { item: Pair, done: Io.Event = .unset, outcome: anyerror!void = {}, queued_at: i96 = 0, completed_at: i96 = 0, };};fn did(buffer: []u8, n: usize) []const u8 { return std.fmt.bufPrint(buffer, "did:plc:{d:0>24}", .{n}) catch unreachable;}fn now(io: Io) i96 { return Io.Clock.awake.now(io).nanoseconds;}fn waitUntil(io: Io, at: i96) void { (Io.Clock.Timestamp{ .raw = .{ .nanoseconds = at }, .clock = .awake }).wait(io) catch {};}fn percentile(samples: []i64, p: usize) i64 { if (samples.len == 0) return 0; std.mem.sort(i64, samples, {}, std.sort.asc(i64)); return samples[@min(samples.len - 1, samples.len * p / 100)];}const Ingress = struct { ci: *index.CollectionIndex, io: Io, mutex: Io.Mutex = .init, changed: Io.Condition = .init, space: Io.Condition = .init, blocked: usize = 0, queue: [4096]*Batch.Request = undefined, head: usize = 0, count: usize = 0, high_water: usize = 0, closing: bool = false, threads: [16]std.Thread = undefined, started: usize = 0, fn start(self: *Ingress) !void { errdefer self.close(); for (&self.threads) |*thread| { thread.* = try std.Thread.spawn(.{}, run, .{self}); self.started += 1; } } fn submit(self: *Ingress, request: *Batch.Request) void { self.mutex.lockUncancelable(self.io); defer self.mutex.unlock(self.io); if (self.count == self.queue.len) self.blocked += 1; while (self.count == self.queue.len) self.space.waitUncancelable(self.io, &self.mutex); self.queue[(self.head + self.count) % self.queue.len] = request; self.count += 1; self.high_water = @max(self.high_water, self.count); self.changed.signal(self.io); } fn close(self: *Ingress) void { self.mutex.lockUncancelable(self.io); self.closing = true; self.changed.broadcast(self.io); self.mutex.unlock(self.io); for (self.threads[0..self.started]) |thread| thread.join(); self.started = 0; } fn run(self: *Ingress) void { while (true) { self.mutex.lockUncancelable(self.io); while (self.count == 0 and !self.closing) self.changed.waitUncancelable(self.io, &self.mutex); if (self.count == 0) { self.mutex.unlock(self.io); return; } const request = self.queue[self.head]; self.head = (self.head + 1) % self.queue.len; self.count -= 1; self.space.signal(self.io); self.mutex.unlock(self.io); request.queued_at = now(self.io); request.outcome = self.ci.addCollection(request.item.did, collection); request.completed_at = now(self.io); request.done.set(self.io); } }};const Producer = struct { ingress: *Ingress, io: Io, requests: []Batch.Request, accepted: []bool, schedule_lag: []i64, start: i96, rate: usize, fn run(self: *Producer) void { for (self.requests, 0..) |*request, i| { const scheduled = self.start + @divTrunc(@as(i96, @intCast(i)) * 1_000_000_000, @as(i96, @intCast(self.rate))); waitUntil(self.io, scheduled); self.schedule_lag[i] = @intCast(now(self.io) - scheduled); self.ingress.submit(request); self.accepted[i] = true; } }};const Reader = struct { ci: *index.CollectionIndex, io: Io, start: i96, rate: usize, samples: []i64, seed_count: usize, outcome: anyerror!void = {}, first_dispatch_lag: i96 = 0, fn run(self: *Reader) void { self.outcome = self.read(); } fn read(self: *Reader) !void { var bytes: [4096]u8 = undefined; var spans: [100]index.Span = undefined; var expected: [64]u8 = undefined; for (self.samples, 0..) |*sample, i| { const scheduled = self.start + @divTrunc(@as(i96, @intCast(i)) * 1_000_000_000, @as(i96, @intCast(self.rate))); waitUntil(self.io, scheduled); if (i == 0) self.first_dispatch_lag = now(self.io) - scheduled; const offset = i * 997 % (self.seed_count - 100); var cursor: [64]u8 = undefined; const page = try self.ci.listReposByCollection(collection, 100, did(&cursor, offset), &bytes, &spans); if (page.count != 100) return error.ShortPage; for (0..100) |j| if (!std.mem.eql(u8, page.getDid(j), did(&expected, offset + j + 1))) return error.WrongPage; sample.* = @intCast(now(self.io) - scheduled); } }};
const ProfileControl = struct { io: Io, control: Io.File, acknowledgement: Io.File,
fn open(init: std.process.Init, io: Io) !?ProfileControl { const control_path = init.environ_map.get("ZKV_PERF_CONTROL") orelse { if (init.environ_map.get("ZKV_PERF_ACK") != null) return error.IncompleteProfileControl; return null; }; const ack_path = init.environ_map.get("ZKV_PERF_ACK") orelse return error.IncompleteProfileControl; const control = try Io.Dir.cwd().openFile(io, control_path, .{ .mode = .read_write }); errdefer control.close(io); const acknowledgement = try Io.Dir.cwd().openFile(io, ack_path, .{ .mode = .read_write }); return .{ .io = io, .control = control, .acknowledgement = acknowledgement }; }
fn deinit(self: ProfileControl) void { self.control.close(self.io); self.acknowledgement.close(self.io); }
fn command(self: ProfileControl, bytes: []const u8) !void { try self.control.writeStreamingAll(self.io, bytes); var buffer: [16]u8 = undefined; var reader = self.acknowledgement.readerStreaming(self.io, &buffer); const ack = try reader.interface.takeDelimiterExclusive('\n'); if (!std.mem.eql(u8, ack, "ack")) return error.InvalidProfileAcknowledgement; }};
fn storageState(ci: *index.CollectionIndex) [4]u64 { var tables: u64 = 0; for (0..ci.db.version.tables.len) |level| tables += ci.db.version.levelBytes(level); return .{ @intCast(ci.db.memtableBytes()), ci.db.unflushedSegments(), tables, @intCast(ci.db.version.tables[0].len) };}
pub fn main(init: std.process.Init) !void { const all_args = try init.minimal.args.toSlice(init.arena.allocator()); const allocator_flag = all_args.len > 1 and std.mem.startsWith(u8, all_args[all_args.len - 1], "--allocator="); const allocator_mode = if (allocator_flag) all_args[all_args.len - 1]["--allocator=".len..] else "c"; const engine_allocator = if (std.mem.eql(u8, allocator_mode, "c")) std.heap.c_allocator else if (std.mem.eql(u8, allocator_mode, "process")) init.gpa else return error.InvalidAllocator; const args = all_args[0 .. all_args.len - @intFromBool(allocator_flag)]; if (args.len < 8 or args.len > 13) return error.ExpectedPathRateSecondsBatchNovelPercentReadRateSync; const rate = try std.fmt.parseInt(usize, args[2], 10); const seconds = try std.fmt.parseInt(usize, args[3], 10); const batch_size = try std.fmt.parseInt(usize, args[4], 10); const requested_delay_us = if (args.len >= 12) try std.fmt.parseInt(usize, args[11], 10) else 1000; if (requested_delay_us > 10_000) return error.InvalidDelay; const delay_us = if (batch_size == 1) 0 else requested_delay_us; const novelty = try std.fmt.parseInt(usize, args[5], 10); const queue_limit = if (args.len >= 9) try std.fmt.parseInt(usize, args[8], 10) else 16; const seed_count = if (args.len >= 10) try std.fmt.parseInt(usize, args[9], 10) else 4096; if (seed_count < 128 or seed_count % 64 != 0 or seed_count > 10_000_000) return error.InvalidSeedCount; const read_rate = try std.fmt.parseInt(usize, args[6], 10); if (rate == 0 or rate > 100_000 or seconds == 0 or seconds > 120 or novelty > 100 or read_rate > 1000) return error.InvalidLoad; var threaded_io = Io.Threaded.init(engine_allocator, .{}); defer threaded_io.deinit(); const io = threaded_io.io(); const profile = try ProfileControl.open(init, io); defer if (profile) |control| control.deinit(); const executors = if (args.len >= 11) try std.fmt.parseInt(u8, args[10], 10) else 4; if (executors == 0 or executors > 64 or queue_limit < 16 or batch_size == 0 or batch_size > 64 or rate * seconds > 2_000_000) return error.InvalidLoad; var runtime: ?*zio.Runtime = null; defer if (runtime) |rt| rt.deinit(); const read_io = if (comptime backend.use_zio) blk: { runtime = try zio.Runtime.init(engine_allocator, .{ .executors = .exact(executors) }); break :blk runtime.?.io(); } else io; var ci = try index.CollectionIndex.openWithReaderIo(engine_allocator, io, read_io, args[1]); var index_open = true; defer if (index_open) ci.deinit(); const prepared = std.mem.endsWith(u8, args[7], "-prepared"); const mode = if (prepared) args[7][0 .. args[7].len - "-prepared".len] else args[7]; const worker_reads = std.mem.eql(u8, mode, "owner-reader"); const bounded_gc = worker_reads or std.mem.eql(u8, mode, "owner-bounded"); const background_writes = bounded_gc or std.mem.eql(u8, mode, "owner-async"); const after_ack = background_writes or std.mem.eql(u8, mode, "owner-stepped"); const stepped = after_ack or std.mem.eql(u8, mode, "stepped"); if (!stepped and !std.mem.eql(u8, mode, "default")) return error.InvalidSyncMode; for (0..seed_count / 64) |chunk| { var buffers: [64][64]u8 = undefined; var ids: [64][]const u8 = undefined; for (&buffers, &ids, 0..) |*buf, *id, i| id.* = did(buf, chunk * 64 + i); try ci.addCollectionPage(&ids, collection); } if (prepared) { try ci.db.flush(); while (try ci.db.compact()) {} try ci.db.collect(); } const seed_state = storageState(&ci); const max_segments = if (args.len >= 13) try std.fmt.parseInt(u64, args[12], 10) else 256; if (max_segments == 0 or max_segments > 4096) return error.InvalidSegmentLimit; ci.maintenance_policy.max_segments = max_segments; ci.maintenance_policy.step_bytes = if (stepped) 64 << 10 else 0; ci.maintenance_policy.log_gc_names = if (bounded_gc) 8 else 0; ci.maintenance_detail = .{}; ci.maintenance_runs = 0; ci.maintenance_ns = 0; ci.maintenance_max_ns = 0; if (worker_reads) try ci.enablePageWorker(16); { var warm_bytes: [4096]u8 = undefined; var warm_spans: [100]index.Span = undefined; const page = try ci.listReposByCollection(collection, 100, null, &warm_bytes, &warm_spans); if (page.count != 100) return error.WarmupFailed; } const n = rate * seconds; const allocator = init.arena.allocator(); const requests = try allocator.alloc(Batch.Request, n); const buffers = try allocator.alloc([64]u8, n); const accepted = try allocator.alloc(bool, n); @memset(accepted, false); for (requests, buffers, 0..) |*request, *buf, i| request.* = .{ .item = .{ .did = did(buf, if (i * 37 % 100 < novelty) seed_count + i else i % seed_count) } }; try ci.enableBatching(.{ .max_batch = batch_size, .queue_limit = queue_limit, .delay_ns = @as(i96, @intCast(delay_us)) * 1000, .maintenance_after_ack = after_ack, .background_writes = background_writes }); var ingress: Ingress = .{ .ci = &ci, .io = io }; try ingress.start(); var ingress_closed = false; defer if (!ingress_closed) ingress.close(); const read_samples = try allocator.alloc(i64, read_rate * seconds); if (profile) |control| try control.command("enable\n"); const start = now(io) + 100_000_000; var reader: Reader = .{ .ci = &ci, .io = read_io, .start = start, .rate = read_rate, .samples = read_samples, .seed_count = seed_count }; var reader_future = try read_io.concurrent(Reader.run, .{&reader}); var reader_joined = false; defer if (!reader_joined) reader_future.cancel(read_io); const schedule_lag = try allocator.alloc(i64, n); var producer: Producer = .{ .ingress = &ingress, .io = io, .requests = requests, .accepted = accepted, .schedule_lag = schedule_lag, .start = start, .rate = rate }; const producer_thread = try std.Thread.spawn(.{}, Producer.run, .{&producer}); var producer_joined = false; defer if (!producer_joined) producer_thread.join(); reader_future.await(read_io); reader_joined = true; producer_thread.join(); producer_joined = true; waitUntil(io, start + @as(i96, @intCast(seconds)) * 1_000_000_000); ingress.close(); ingress_closed = true; ci.write_owner.?.close(); if (ci.page_owner) |owner| owner.close(); if (ci.failed.load(.acquire)) { ci.async_store.cancelPut(); return error.IndexUnavailable; } if (ci.async_store.pending != null or ci.log_gc_pending or ci.log_collector.names != null) return error.UndrainedMaintenance; const batches = ci.write_owner.?.batches; const queue_high_water = ci.write_owner.?.high_water; const detail = ci.maintenance_detail; const page_worker_high_water = if (ci.page_owner) |owner| owner.high_water else 0; const page_worker_completed = if (ci.page_owner) |owner| owner.batches else 0; if (ci.page_owner) |owner| if (owner.outstanding != 0) return error.UndrainedReads; const log_gc_names = ci.maintenance_policy.log_gc_names; const maintenance_runs = ci.maintenance_runs; const maintenance_ns = ci.maintenance_ns; const maintenance_max_ns = ci.maintenance_max_ns; const background_puts = ci.maintenance_store.writes; const background_put_ns = ci.maintenance_store.elapsed_ns; const end = now(io); if (profile) |control| try control.command("disable\n"); const end_state = storageState(&ci); const maintenance_drain_started = now(io); try ci.db.flush(); while (try ci.db.compact()) {} try ci.db.collect(); const maintenance_drain_ns = now(io) - maintenance_drain_started; try reader.outcome; var failures: usize = 0; var acknowledgements: std.ArrayList(i64) = .empty; var queue_time: std.ArrayList(i64) = .empty; for (requests, accepted, 0..) |*request, ok, i| { if (!ok) continue; if (!request.done.isSet()) return error.Undrained; if (request.outcome) |_| {} else |_| { failures += 1; continue; } const scheduled = start + @divTrunc(@as(i96, @intCast(i)) * 1_000_000_000, @as(i96, @intCast(rate))); try acknowledgements.append(allocator, @intCast(request.completed_at - scheduled)); try queue_time.append(allocator, @intCast(request.completed_at - request.queued_at)); } const scan_p99 = percentile(read_samples, 99); const ack_p50 = percentile(acknowledgements.items, 50); const ack_p99 = percentile(acknowledgements.items, 99); const queue_p99 = percentile(queue_time.items, 99); const schedule_p99 = percentile(schedule_lag, 99); ci.deinit(); index_open = false; const reopen_started = now(io); ci = try index.CollectionIndex.openWithReaderIo(engine_allocator, io, read_io, args[1]); const reopen_ns = now(io) - reopen_started; index_open = true; var expected: usize = 0; var output: [4096]u8 = undefined; var spans: [100]index.Span = undefined; var cursor_buffer: [64]u8 = undefined; var cursor: ?[]const u8 = null; var expected_buffer: [64]u8 = undefined; var verified: usize = 0; while (true) { const page = try ci.listReposByCollection(collection, 100, cursor, &output, &spans); for (0..page.count) |j| { while (expected >= seed_count and expected < seed_count + n and (!accepted[expected - seed_count] or (expected - seed_count) * 37 % 100 >= novelty)) expected += 1; if (expected >= seed_count + n) return error.UnexpectedPair; const want = did(&expected_buffer, expected); if (!std.mem.eql(u8, page.getDid(j), want) or !ci.hasCollections(want)) return error.WrongReopenedPair; expected += 1; verified += 1; } const last = page.last_did orelse break; @memcpy(cursor_buffer[0..last.len], last); cursor = cursor_buffer[0..last.len]; } var want_count: usize = seed_count; for (accepted, 0..) |ok, i| if (ok and i * 37 % 100 < novelty) { want_count += 1; }; if (verified != want_count or failures != 0) return error.IncompleteLoad; var out = Io.File.stdout().writer(io, &.{}); try out.interface.writeAll("engine_allocator,rate,seconds,batch,batch_delay_us,novel_percent,read_rate,queue_limit,seed_count,submitted,blocked_submissions,failed,batches,queue_high_water,ingress_high_water,ack_p50_ns,ack_p99_ns,queue_p99_ns,schedule_p99_ns,scan_p99_ns,scan_first_dispatch_lag_ns,drain_ns,adapter_maintenance_runs,adapter_maintenance_ns,adapter_maintenance_max_ns,verified_pairs,max_unflushed_segments,changed_commits,flushed_memtable_bytes,flushed_segments,flush_ns,compact_ns,collect_ns,compactions,ack_p999_ns,ack_max_ns,reopen_ns,compaction_steps,maintenance_drain_ns,seed_zkv_memtable_bytes,seed_zkv_unflushed_segments,seed_sst_bytes,seed_l0_tables,end_zkv_memtable_bytes,end_zkv_unflushed_segments,end_sst_bytes,end_l0_tables,maintenance_after_ack,background_writes,background_puts,background_put_ns,maintenance_pending_polls,log_gc_names,collect_max_ns,log_gc_steps,log_gc_max_ns,page_worker,page_worker_high_water,page_worker_completed\n"); try out.interface.print("{s},", .{allocator_mode}); try out.interface.print("{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d}", .{ rate, seconds, batch_size, delay_us, novelty, read_rate, queue_limit, seed_count, n, ingress.blocked, failures, batches, queue_high_water, ingress.high_water, ack_p50, ack_p99, queue_p99, schedule_p99, scan_p99, reader.first_dispatch_lag, end - start - @as(i96, @intCast(seconds)) * 1_000_000_000, maintenance_runs, maintenance_ns, maintenance_max_ns, verified }); try out.interface.print(",{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d},{d}", .{ max_segments, detail.changed_commits, detail.flushed_bytes, detail.flushed_segments, detail.flush_ns, detail.compact_ns, detail.collect_ns, detail.compactions, acknowledgements.items[@min(acknowledgements.items.len - 1, acknowledgements.items.len * 999 / 1000)], acknowledgements.items[acknowledgements.items.len - 1], reopen_ns, detail.compaction_steps, maintenance_drain_ns }); for (seed_state) |value| try out.interface.print(",{d}", .{value}); for (end_state) |value| try out.interface.print(",{d}", .{value}); try out.interface.print(",{d},{d},{d},{d},{d}", .{ @intFromBool(after_ack), @intFromBool(background_writes), background_puts, background_put_ns, detail.pending_polls }); try out.interface.print(",{d},{d},{d},{d}", .{ log_gc_names, detail.collect_max_ns, detail.log_gc_steps, detail.log_gc_max_ns }); try out.interface.print(",{d},{d},{d}\n", .{ @intFromBool(worker_reads), page_worker_high_water, page_worker_completed });}