diff --git a/Dockerfile.runtime b/Dockerfile.runtime index 1c4865b..61a21be 100644 --- a/Dockerfile.runtime +++ b/Dockerfile.runtime @@ -1,6 +1,7 @@ FROM debian:bookworm-slim RUN apt-get update && apt-get install -y --no-install-recommends ca-certificates && rm -rf /var/lib/apt/lists/* COPY zig-out/bin/zlay /usr/local/bin/zlay +COPY zig-out/bin/collectiondir /usr/local/bin/collectiondir RUN mkdir -p /data/events /data/collection-index ENV RELAY_DATA_DIR=/data/events EXPOSE 3000 3001 diff --git a/build.zig b/build.zig index 44c6f03..ce7a4e6 100644 --- a/build.zig +++ b/build.zig @@ -64,6 +64,27 @@ pub fn build(b: *std.Build) void { const run_step = b.step("run", "run the relay"); run_step.dependOn(&run_relay.step); + // collectiondir executable (shadow collection directory service) + const collectiondir_mod = b.createModule(.{ + .root_source_file = b.path("src/collectiondir.zig"), + .target = target, + .optimize = optimize, + .imports = imports, + }); + collectiondir_mod.addImport("build_options", build_options.createModule()); + const collectiondir = b.addExecutable(.{ + .name = "collectiondir", + .root_module = collectiondir_mod, + }); + collectiondir.linkLibC(); + collectiondir.linkLibCpp(); + b.installArtifact(collectiondir); + + const run_collectiondir = b.addRunArtifact(collectiondir); + if (b.args) |args| run_collectiondir.addArgs(args); + const run_collectiondir_step = b.step("run-collectiondir", "run the collection directory service"); + run_collectiondir_step.dependOn(&run_collectiondir.step); + // tests const test_step = b.step("test", "run unit tests"); const test_files = .{ diff --git a/src/collection_index.zig b/src/collection_index.zig index 5c69be7..86aa98d 100644 --- a/src/collection_index.zig +++ b/src/collection_index.zig @@ -19,6 +19,7 @@ const separator = '\x00'; pub const CollectionIndex = struct { db: rocksdb.DB, + default: rocksdb.ColumnFamilyHandle, rbc: rocksdb.ColumnFamilyHandle, cbr: rocksdb.ColumnFamilyHandle, allocator: Allocator, @@ -50,9 +51,11 @@ pub const CollectionIndex = struct { defer allocator.free(families); // find column family handles by name + var default: ?rocksdb.ColumnFamilyHandle = null; var rbc: ?rocksdb.ColumnFamilyHandle = null; var cbr: ?rocksdb.ColumnFamilyHandle = null; for (families) |cf| { + if (std.mem.eql(u8, cf.name, "default")) default = cf.handle; if (std.mem.eql(u8, cf.name, "rbc")) rbc = cf.handle; if (std.mem.eql(u8, cf.name, "cbr")) cbr = cf.handle; } @@ -61,12 +64,57 @@ pub const CollectionIndex = struct { return .{ .db = db, + .default = default orelse return error.MissingColumnFamily, .rbc = rbc orelse return error.MissingColumnFamily, .cbr = cbr orelse return error.MissingColumnFamily, .allocator = allocator, }; } + // --- metadata helpers (default CF) --- + + /// read a metadata value from the default column family. + /// caller must call .deinit() on the returned Data to free rocksdb memory. + pub fn getMeta(self: *CollectionIndex, key: []const u8) ?rocksdb.Data { + var err_str: ?rocksdb.Data = null; + const val = self.db.get(self.default, key, &err_str) catch { + if (err_str) |e| { + log.debug("getMeta({s}) failed: {s}", .{ key, e.data }); + e.deinit(); + } + return null; + }; + return val; + } + + /// write a metadata value to the default column family + pub fn putMeta(self: *CollectionIndex, key: []const u8, value: []const u8) void { + var err_str: ?rocksdb.Data = null; + var batch = rocksdb.WriteBatch.init(); + defer batch.deinit(); + batch.put(self.default, key, value); + self.db.write(batch, &err_str) catch { + if (err_str) |e| { + log.warn("putMeta({s}) failed: {s}", .{ key, e.data }); + e.deinit(); + } + }; + } + + /// delete a metadata key from the default column family + pub fn deleteMeta(self: *CollectionIndex, key: []const u8) void { + var err_str: ?rocksdb.Data = null; + var batch = rocksdb.WriteBatch.init(); + defer batch.deinit(); + batch.delete(self.default, key); + self.db.write(batch, &err_str) catch { + if (err_str) |e| { + log.warn("deleteMeta({s}) failed: {s}", .{ key, e.data }); + e.deinit(); + } + }; + } + pub fn deinit(self: *CollectionIndex) void { self.db.deinit(); } diff --git a/src/collectiondir.zig b/src/collectiondir.zig new file mode 100644 index 0000000..cfe6750 --- /dev/null +++ b/src/collectiondir.zig @@ -0,0 +1,407 @@ +//! collectiondir — standalone collection directory service +//! +//! subscribes to the relay's firehose and maintains an independent RocksDB +//! collection index. no postgres — all state (cursor, backfill progress, +//! collection index) lives in RocksDB. +//! +//! reuses collection_index.zig (unchanged) and resync.zig (unchanged, stage 2). +//! backfill uses collectiondir_backfill.zig (RocksDB-backed progress). +//! +//! LISTEN_PORT (default 2510): HTTP API +//! GET /xrpc/com.atproto.sync.listReposByCollection +//! GET /_health — readiness (subscriber connected + cursor fresh) +//! GET /_healthz — liveness +//! GET /metrics +//! GET /admin/backfill-collections — backfill status +//! POST /admin/backfill-collections — trigger backfill +//! +//! env: +//! RELAY_URL — relay firehose host (e.g. "zlay" for in-cluster) +//! RELAY_PORT — relay firehose port (default 3000) +//! RELAY_TLS — "true" for wss:// (default "false") +//! COLLECTION_INDEX_DIR — RocksDB data directory +//! LISTEN_PORT — HTTP listen port (default 2510) +//! RELAY_ADMIN_PASSWORD — admin auth token + +const std = @import("std"); +const http = std.http; +const collection_index_mod = @import("collection_index.zig"); +const backfill_mod = @import("collectiondir_backfill.zig"); +const collectiondir_sub = @import("collectiondir_subscriber.zig"); + +const log = std.log.scoped(.collectiondir); + +/// 8 MiB stacks — matches relay's default_stack_size for ReleaseSafe +pub const default_stack_size = 8 * 1024 * 1024; + +var shutdown_flag: std.atomic.Value(bool) = .{ .raw = false }; + +pub fn main() !void { + const allocator = std.heap.c_allocator; + + // parse config from env + const listen_port = parseEnvInt(u16, "LISTEN_PORT", 2510); + const relay_url = std.posix.getenv("RELAY_URL") orelse "zlay"; + const relay_port = parseEnvInt(u16, "RELAY_PORT", 3000); + const use_tls = if (std.posix.getenv("RELAY_TLS")) |v| std.mem.eql(u8, v, "true") else false; + const ci_dir = std.posix.getenv("COLLECTION_INDEX_DIR") orelse "data/collection-index"; + + // install signal handlers + installSignalHandlers(); + + // init RocksDB collection index (all state lives here) + var ci = collection_index_mod.CollectionIndex.open(allocator, ci_dir) catch |err| { + log.err("failed to open collection index at {s}: {s}", .{ ci_dir, @errorName(err) }); + return err; + }; + defer ci.deinit(); + + // init backfiller (RocksDB-backed progress) + var backfiller = backfill_mod.Backfiller.init(allocator, &ci); + + // init firehose subscriber (cursor in RocksDB) + var subscriber = collectiondir_sub.Subscriber.init( + allocator, + relay_url, + relay_port, + use_tls, + &ci, + &shutdown_flag, + ); + + // start subscriber thread + const sub_thread = try std.Thread.spawn(.{ .stack_size = default_stack_size }, collectiondir_sub.Subscriber.run, .{&subscriber}); + + // start HTTP server + log.info("collectiondir listening on :{d}", .{listen_port}); + log.info("relay upstream: {s}:{d} (tls={s})", .{ relay_url, relay_port, if (use_tls) "true" else "false" }); + log.info("collection index: {s}", .{ci_dir}); + + var http_ctx = HttpContext{ + .collection_index = &ci, + .backfiller = &backfiller, + .subscriber = &subscriber, + }; + + const address = std.net.Address.initIp4(.{ 0, 0, 0, 0 }, listen_port); + var server = address.listen(.{ .reuse_address = true }) catch |err| { + log.err("failed to listen on :{d}: {s}", .{ listen_port, @errorName(err) }); + return err; + }; + + while (!shutdown_flag.load(.acquire)) { + const conn = server.accept() catch |err| { + if (shutdown_flag.load(.acquire)) break; + log.debug("accept error: {s}", .{@errorName(err)}); + continue; + }; + const timeout = std.posix.timeval{ .sec = 5, .usec = 0 }; + std.posix.setsockopt(conn.stream.handle, std.posix.SOL.SOCKET, std.posix.SO.RCVTIMEO, std.mem.asBytes(&timeout)) catch {}; + handleConnection(conn.stream, &http_ctx); + } + + log.info("shutdown signal received, stopping...", .{}); + server.stream.close(); + sub_thread.join(); + log.info("collectiondir stopped cleanly", .{}); +} + +const HttpContext = struct { + collection_index: *collection_index_mod.CollectionIndex, + backfiller: *backfill_mod.Backfiller, + subscriber: *collectiondir_sub.Subscriber, +}; + +fn handleConnection(stream: std.net.Stream, ctx: *HttpContext) void { + defer stream.close(); + + var recv_buf: [4096]u8 = undefined; + var send_buf: [4096]u8 = undefined; + var connection_reader = stream.reader(&recv_buf); + var connection_writer = stream.writer(&send_buf); + var server = http.Server.init(connection_reader.interface(), &connection_writer.interface); + + var request = server.receiveHead() catch return; + const target = request.head.target; + const method = request.head.method; + + const qmark = std.mem.indexOfScalar(u8, target, '?'); + const path = target[0..(qmark orelse target.len)]; + const query = if (qmark) |q| target[q + 1 ..] else ""; + + if (method == .GET) { + if (std.mem.eql(u8, path, "/_healthz")) { + respond(&request, .ok, "application/json", "{\"status\":\"ok\"}"); + } else if (std.mem.eql(u8, path, "/_health")) { + handleReadiness(&request, ctx); + } else if (std.mem.eql(u8, path, "/metrics")) { + var metrics_buf: [8192]u8 = undefined; + const body = ctx.subscriber.formatMetrics(&metrics_buf); + respond(&request, .ok, "text/plain; version=0.0.4; charset=utf-8", body); + } else if (std.mem.eql(u8, path, "/xrpc/com.atproto.sync.listReposByCollection")) { + handleListReposByCollection(&request, query, ctx.collection_index); + } else if (std.mem.eql(u8, path, "/admin/backfill-collections")) { + if (!checkAdmin(&request)) return; + handleBackfillStatus(&request, ctx.backfiller); + } else { + respond(&request, .not_found, "text/plain", "not found"); + } + } else if (method == .POST) { + if (std.mem.eql(u8, path, "/admin/backfill-collections")) { + if (!checkAdmin(&request)) return; + handleBackfillTrigger(&request, query, ctx.backfiller); + } else { + respond(&request, .not_found, "text/plain", "not found"); + } + } else { + respond(&request, .method_not_allowed, "text/plain", "method not allowed"); + } +} + +/// readiness: subscriber connected + cursor advancing +fn handleReadiness(request: *http.Server.Request, ctx: *HttpContext) void { + const stats = &ctx.subscriber.stats; + const connected = stats.connected.load(.monotonic); + const last_event_time = stats.last_event_time.load(.monotonic); + const now = std.time.timestamp(); + + // subscriber must be connected and have received an event within the last 60s + // (allows for startup grace — last_event_time == 0 means never received) + const stale_threshold = 60; + const cursor_fresh = connected and (last_event_time == 0 or (now - last_event_time) < stale_threshold); + + if (cursor_fresh) { + var buf: [256]u8 = undefined; + const body = std.fmt.bufPrint(&buf, "{{\"status\":\"ok\",\"connected\":{},\"last_seq\":{d}}}", .{ + connected, stats.last_seq.load(.monotonic), + }) catch "{\"status\":\"ok\"}"; + respond(request, .ok, "application/json", body); + } else { + var buf: [256]u8 = undefined; + const body = std.fmt.bufPrint(&buf, "{{\"status\":\"unhealthy\",\"connected\":{},\"cursor_fresh\":{}}}", .{ + connected, cursor_fresh, + }) catch "{\"status\":\"unhealthy\"}"; + respond(request, .service_unavailable, "application/json", body); + } +} + +fn handleListReposByCollection( + request: *http.Server.Request, + query: []const u8, + ci: *collection_index_mod.CollectionIndex, +) void { + const collection = queryParam(query, "collection") orelse { + respond(request, .bad_request, "application/json", "{\"error\":\"BadRequest\",\"message\":\"collection parameter required\"}"); + return; + }; + + if (collection.len == 0 or !std.mem.containsAtLeast(u8, collection, 1, ".")) { + respond(request, .bad_request, "application/json", "{\"error\":\"BadRequest\",\"message\":\"invalid collection NSID\"}"); + return; + } + + const limit_str = queryParam(query, "limit") orelse "500"; + const limit = std.fmt.parseInt(usize, limit_str, 10) catch { + respond(request, .bad_request, "application/json", "{\"error\":\"BadRequest\",\"message\":\"invalid limit\"}"); + return; + }; + if (limit < 1 or limit > 2000) { + respond(request, .bad_request, "application/json", "{\"error\":\"BadRequest\",\"message\":\"limit must be 1..2000\"}"); + return; + } + + var cursor_buf: [256]u8 = undefined; + const cursor_did = queryParamDecoded(query, "cursor", &cursor_buf); + + var did_buf: [65536]u8 = undefined; + const ci_result = ci.listReposByCollection(collection, limit, cursor_did, &did_buf) catch { + respond(request, .internal_server_error, "application/json", "{\"error\":\"InternalError\",\"message\":\"index scan failed\"}"); + return; + }; + + var buf: [65536]u8 = undefined; + var fbs = std.io.fixedBufferStream(&buf); + const w = fbs.writer(); + + w.writeAll("{\"repos\":[") catch return; + for (0..ci_result.count) |i| { + if (i > 0) w.writeByte(',') catch return; + w.writeAll("{\"did\":\"") catch return; + w.writeAll(ci_result.getDid(i)) catch return; + w.writeAll("\"}") catch return; + } + w.writeByte(']') catch return; + + if (ci_result.last_did) |last| { + if (ci_result.count >= limit) { + w.writeAll(",\"cursor\":\"") catch return; + w.writeAll(last) catch return; + w.writeAll("\"") catch return; + } + } + + w.writeByte('}') catch return; + respond(request, .ok, "application/json", fbs.getWritten()); +} + +fn handleBackfillStatus(request: *http.Server.Request, backfiller: *backfill_mod.Backfiller) void { + const body = backfiller.getStatus(backfiller.allocator) catch { + respond(request, .internal_server_error, "application/json", "{\"error\":\"failed to query backfill status\"}"); + return; + }; + defer backfiller.allocator.free(body); + respond(request, .ok, "application/json", body); +} + +fn handleBackfillTrigger(request: *http.Server.Request, query: []const u8, backfiller: *backfill_mod.Backfiller) void { + const source = queryParam(query, "source") orelse "bsky.network"; + backfiller.start(source) catch |err| { + switch (err) { + error.AlreadyRunning => { + respond(request, .conflict, "application/json", "{\"error\":\"backfill already in progress\"}"); + }, + else => { + respond(request, .internal_server_error, "application/json", "{\"error\":\"failed to start backfill\"}"); + }, + } + return; + }; + var buf: [256]u8 = undefined; + const resp_body = std.fmt.bufPrint(&buf, "{{\"status\":\"started\",\"source\":\"{s}\"}}", .{source}) catch { + respond(request, .ok, "application/json", "{\"status\":\"started\"}"); + return; + }; + respond(request, .ok, "application/json", resp_body); +} + +// --- auth --- + +fn checkAdmin(request: *http.Server.Request) bool { + const admin_pw = std.posix.getenv("RELAY_ADMIN_PASSWORD") orelse { + respond(request, .forbidden, "application/json", "{\"error\":\"admin endpoint not configured\"}"); + return false; + }; + + var it = request.iterateHeaders(); + while (it.next()) |header| { + if (std.ascii.eqlIgnoreCase(header.name, "authorization")) { + const bearer_prefix = "Bearer "; + if (!std.mem.startsWith(u8, header.value, bearer_prefix)) { + respond(request, .unauthorized, "application/json", "{\"error\":\"invalid authorization scheme\"}"); + return false; + } + const token = header.value[bearer_prefix.len..]; + if (!std.mem.eql(u8, token, admin_pw)) { + respond(request, .unauthorized, "application/json", "{\"error\":\"invalid token\"}"); + return false; + } + return true; + } + } + + respond(request, .unauthorized, "application/json", "{\"error\":\"missing authorization header\"}"); + return false; +} + +// --- response helpers --- + +fn respond(request: *http.Server.Request, status: http.Status, content_type: []const u8, body: []const u8) void { + request.respond(body, .{ + .status = status, + .keep_alive = false, + .extra_headers = &.{ + .{ .name = "content-type", .value = content_type }, + .{ .name = "server", .value = "collectiondir (zlay)" }, + }, + }) catch {}; +} + +// --- query string helpers --- + +fn queryParam(query: []const u8, name: []const u8) ?[]const u8 { + if (query.len == 0) return null; + var iter = std.mem.splitScalar(u8, query, '&'); + while (iter.next()) |pair| { + const eq = std.mem.indexOfScalar(u8, pair, '=') orelse continue; + if (std.mem.eql(u8, pair[0..eq], name)) { + return pair[eq + 1 ..]; + } + } + return null; +} + +fn queryParamDecoded(query: []const u8, name: []const u8, buf: []u8) ?[]const u8 { + const raw = queryParam(query, name) orelse return null; + var i: usize = 0; + var out: usize = 0; + while (i < raw.len) { + if (raw[i] == '%' and i + 2 < raw.len) { + const hi = hexVal(raw[i + 1]) orelse { + if (out >= buf.len) return null; + buf[out] = raw[i]; + out += 1; + i += 1; + continue; + }; + const lo = hexVal(raw[i + 2]) orelse { + if (out >= buf.len) return null; + buf[out] = raw[i]; + out += 1; + i += 1; + continue; + }; + if (out >= buf.len) return null; + buf[out] = (@as(u8, hi) << 4) | @as(u8, lo); + out += 1; + i += 3; + } else if (raw[i] == '+') { + if (out >= buf.len) return null; + buf[out] = ' '; + out += 1; + i += 1; + } else { + if (out >= buf.len) return null; + buf[out] = raw[i]; + out += 1; + i += 1; + } + } + return buf[0..out]; +} + +fn hexVal(c: u8) ?u4 { + return switch (c) { + '0'...'9' => @intCast(c - '0'), + 'a'...'f' => @intCast(c - 'a' + 10), + 'A'...'F' => @intCast(c - 'A' + 10), + else => null, + }; +} + +// --- helpers --- + +fn signalHandler(_: c_int) callconv(.c) void { + shutdown_flag.store(true, .release); +} + +fn installSignalHandlers() void { + const act: std.posix.Sigaction = .{ + .handler = .{ .handler = signalHandler }, + .mask = std.posix.sigemptyset(), + .flags = 0, + }; + std.posix.sigaction(std.posix.SIG.INT, &act, null); + std.posix.sigaction(std.posix.SIG.TERM, &act, null); + + const ignore_act: std.posix.Sigaction = .{ + .handler = .{ .handler = std.posix.SIG.IGN }, + .mask = std.posix.sigemptyset(), + .flags = 0, + }; + std.posix.sigaction(std.posix.SIG.PIPE, &ignore_act, null); +} + +fn parseEnvInt(comptime T: type, key: []const u8, default: T) T { + const val = std.posix.getenv(key) orelse return default; + return std.fmt.parseInt(T, val, 10) catch default; +} diff --git a/src/collectiondir_backfill.zig b/src/collectiondir_backfill.zig new file mode 100644 index 0000000..7f6cde4 --- /dev/null +++ b/src/collectiondir_backfill.zig @@ -0,0 +1,423 @@ +//! collectiondir backfill — discovers collections and imports DIDs from a source relay +//! +//! same logic as backfill.zig but stores progress in RocksDB default CF instead +//! of postgres. no external database dependency. +//! +//! progress keys: "bf:{collection}\0{source}" → "{cursor}\0{count}\0{0|1}" +//! (cursor string, imported count, completed flag) + +const std = @import("std"); +const http = std.http; +const collection_index_mod = @import("collection_index.zig"); + +const Allocator = std.mem.Allocator; +const log = std.log.scoped(.backfill); + +pub const Backfiller = struct { + allocator: Allocator, + collection_index: *collection_index_mod.CollectionIndex, + running: std.atomic.Value(bool), + thread: ?std.Thread, + source: []const u8, + + pub fn init( + allocator: Allocator, + collection_index: *collection_index_mod.CollectionIndex, + ) Backfiller { + return .{ + .allocator = allocator, + .collection_index = collection_index, + .running = .{ .raw = false }, + .thread = null, + .source = "", + }; + } + + pub fn isRunning(self: *Backfiller) bool { + return self.running.load(.acquire); + } + + pub fn start(self: *Backfiller, source: []const u8) !void { + if (self.running.cmpxchgStrong(false, true, .acq_rel, .acquire) != null) { + return error.AlreadyRunning; + } + errdefer self.running.store(false, .release); + + self.source = try self.allocator.dupe(u8, source); + self.thread = try std.Thread.spawn(.{ .stack_size = 8 * 1024 * 1024 }, run, .{self}); + } + + fn run(self: *Backfiller) void { + defer { + self.allocator.free(self.source); + self.source = ""; + self.thread = null; + self.running.store(false, .release); + } + + const collections = self.discoverCollections() catch |err| { + log.err("collection discovery failed: {s}", .{@errorName(err)}); + return; + }; + defer { + for (collections) |c| self.allocator.free(c); + self.allocator.free(collections); + } + + log.info("discovered {d} collections to backfill from {s}", .{ collections.len, self.source }); + + for (collections) |collection| { + self.backfillCollection(collection) catch |err| { + log.warn("backfill failed for {s}: {s}", .{ collection, @errorName(err) }); + }; + } + + log.info("backfill complete", .{}); + } + + fn discoverCollections(self: *Backfiller) ![][]const u8 { + var seen: std.StringHashMapUnmanaged(void) = .{}; + defer seen.deinit(self.allocator); + + // source 1: lexicon garden + const garden = self.fetchLexiconGarden() catch |err| blk: { + log.warn("lexicon garden fetch failed: {s}", .{@errorName(err)}); + break :blk &[_][]const u8{}; + }; + defer { + for (garden) |c| self.allocator.free(c); + self.allocator.free(garden); + } + for (garden) |c| { + if (!seen.contains(c)) { + const duped = try self.allocator.dupe(u8, c); + errdefer self.allocator.free(duped); + try seen.put(self.allocator, duped, {}); + } + } + + // source 2: observed collections from RBC scan + const observed = self.collection_index.listKnownCollections(self.allocator) catch |err| blk: { + log.warn("RBC scan failed: {s}", .{@errorName(err)}); + break :blk &[_][]const u8{}; + }; + defer { + for (observed) |c| self.allocator.free(c); + self.allocator.free(observed); + } + for (observed) |c| { + if (!seen.contains(c)) { + const duped = try self.allocator.dupe(u8, c); + errdefer self.allocator.free(duped); + try seen.put(self.allocator, duped, {}); + } + } + + const result = try self.allocator.alloc([]const u8, seen.count()); + var i: usize = 0; + var key_iter = seen.keyIterator(); + while (key_iter.next()) |key| { + result[i] = key.*; + i += 1; + } + return result; + } + + fn fetchLexiconGarden(self: *Backfiller) ![][]const u8 { + var client: http.Client = .{ .allocator = self.allocator }; + defer client.deinit(); + + var aw: std.Io.Writer.Allocating = .init(self.allocator); + defer aw.deinit(); + + const result = client.fetch(.{ + .location = .{ .url = "https://lexicon.garden/llms.txt" }, + .response_writer = &aw.writer, + .method = .GET, + }) catch return error.FetchFailed; + + if (result.status != .ok) return error.FetchFailed; + + const body = aw.written(); + + var nsids: std.ArrayListUnmanaged([]const u8) = .{}; + defer nsids.deinit(self.allocator); + + var lines = std.mem.splitScalar(u8, body, '\n'); + while (lines.next()) |line| { + const backtick_start = std.mem.indexOf(u8, line, "- [`") orelse continue; + const nsid_start = backtick_start + 4; + const rest = line[nsid_start..]; + const backtick_end = std.mem.indexOf(u8, rest, "`](") orelse continue; + const nsid = rest[0..backtick_end]; + if (!std.mem.containsAtLeast(u8, nsid, 1, ".")) continue; + + const duped = try self.allocator.dupe(u8, nsid); + errdefer self.allocator.free(duped); + try nsids.append(self.allocator, duped); + } + + log.info("lexicon garden: found {d} NSIDs", .{nsids.items.len}); + return try nsids.toOwnedSlice(self.allocator); + } + + // --- progress storage in RocksDB default CF --- + // key: "bf:{collection}\0{source}" + // value: "{cursor}\0{count}\0{0|1}" (cursor, imported count, completed) + + fn progressKey(self: *const Backfiller, collection: []const u8, buf: []u8) ?[]const u8 { + const prefix = "bf:"; + const needed = prefix.len + collection.len + 1 + self.source.len; + if (needed > buf.len) return null; + @memcpy(buf[0..prefix.len], prefix); + @memcpy(buf[prefix.len..][0..collection.len], collection); + buf[prefix.len + collection.len] = 0; + @memcpy(buf[prefix.len + collection.len + 1 ..][0..self.source.len], self.source); + return buf[0..needed]; + } + + const Progress = struct { + cursor: []const u8, + imported: i64, + completed: bool, + /// rocksdb-owned memory — caller must call deinit() when done + _data: @import("rocksdb").Data, + + pub fn deinit(self: Progress) void { + self._data.deinit(); + } + }; + + fn loadProgress(self: *Backfiller, collection: []const u8) ?Progress { + var key_buf: [512]u8 = undefined; + const key = self.progressKey(collection, &key_buf) orelse return null; + const data = self.collection_index.getMeta(key) orelse return null; + const val = data.data; + + // parse: "{cursor}\0{count}\0{0|1}" + const sep1 = std.mem.indexOfScalar(u8, val, 0) orelse { + data.deinit(); + return null; + }; + const rest = val[sep1 + 1 ..]; + const sep2 = std.mem.indexOfScalar(u8, rest, 0) orelse { + data.deinit(); + return null; + }; + + return .{ + .cursor = val[0..sep1], + .imported = std.fmt.parseInt(i64, rest[0..sep2], 10) catch 0, + .completed = rest.len > sep2 + 1 and rest[sep2 + 1] == '1', + ._data = data, + }; + } + + fn saveProgress(self: *Backfiller, collection: []const u8, cursor: []const u8, imported: i64, completed: bool) void { + var key_buf: [512]u8 = undefined; + const key = self.progressKey(collection, &key_buf) orelse return; + + var val_buf: [512]u8 = undefined; + var pos: usize = 0; + if (cursor.len > 0) { + @memcpy(val_buf[0..cursor.len], cursor); + pos = cursor.len; + } + val_buf[pos] = 0; + pos += 1; + const count_str = std.fmt.bufPrint(val_buf[pos..], "{d}", .{imported}) catch return; + pos += count_str.len; + val_buf[pos] = 0; + pos += 1; + val_buf[pos] = if (completed) '1' else '0'; + pos += 1; + + self.collection_index.putMeta(key, val_buf[0..pos]); + } + + fn backfillCollection(self: *Backfiller, collection: []const u8) !void { + // check saved progress + var cursor: ?[]const u8 = null; + defer if (cursor) |c| self.allocator.free(c); + var imported: i64 = 0; + + if (self.loadProgress(collection)) |progress| { + defer progress.deinit(); + if (progress.completed) return; // already done + imported = progress.imported; + if (progress.cursor.len > 0) { + cursor = try self.allocator.dupe(u8, progress.cursor); + log.info("{s}: resuming from cursor (imported {d} so far)", .{ collection, imported }); + } + } + + var client: http.Client = .{ .allocator = self.allocator }; + defer client.deinit(); + + var page_count: usize = 0; + while (true) { + const fetch_result = self.fetchPage(&client, collection, cursor) catch |err| { + log.warn("{s}: fetch page failed: {s}", .{ collection, @errorName(err) }); + break; + }; + defer { + for (fetch_result.dids) |d| self.allocator.free(d); + self.allocator.free(fetch_result.dids); + if (fetch_result.next_cursor) |nc| self.allocator.free(nc); + } + + for (fetch_result.dids) |did| { + self.collection_index.addCollection(did, collection) catch {}; + imported += 1; + } + + page_count += 1; + + const new_cursor = fetch_result.next_cursor orelse ""; + self.saveProgress(collection, new_cursor, imported, false); + + if (fetch_result.next_cursor) |nc| { + if (cursor) |old| self.allocator.free(old); + cursor = self.allocator.dupe(u8, nc) catch break; + std.posix.nanosleep(0, 100 * std.time.ns_per_ms); + } else { + self.saveProgress(collection, "", imported, true); + log.info("{s}: complete ({d} DIDs, {d} pages)", .{ collection, imported, page_count }); + break; + } + } + } + + fn fetchPage(self: *Backfiller, client: *http.Client, collection: []const u8, cursor: ?[]const u8) !FetchResult { + var url_buf: [1024]u8 = undefined; + const url = if (cursor) |c| + std.fmt.bufPrint(&url_buf, "https://{s}/xrpc/com.atproto.sync.listReposByCollection?collection={s}&limit=1000&cursor={s}", .{ self.source, collection, c }) catch return error.UrlTooLong + else + std.fmt.bufPrint(&url_buf, "https://{s}/xrpc/com.atproto.sync.listReposByCollection?collection={s}&limit=1000", .{ self.source, collection }) catch return error.UrlTooLong; + + var aw: std.Io.Writer.Allocating = .init(self.allocator); + defer aw.deinit(); + + const result = client.fetch(.{ + .location = .{ .url = url }, + .response_writer = &aw.writer, + .method = .GET, + }) catch return error.FetchFailed; + + if (result.status != .ok) return error.FetchFailed; + + const body = aw.written(); + + const parsed = std.json.parseFromSlice(ListReposResponse, self.allocator, body, .{ .ignore_unknown_fields = true }) catch return error.ParseFailed; + defer parsed.deinit(); + + const repos = parsed.value.repos orelse return .{ + .dids = try self.allocator.alloc([]const u8, 0), + .next_cursor = null, + }; + + var dids: std.ArrayListUnmanaged([]const u8) = .{}; + defer dids.deinit(self.allocator); + + for (repos) |repo| { + const duped = self.allocator.dupe(u8, repo.did) catch continue; + dids.append(self.allocator, duped) catch { + self.allocator.free(duped); + continue; + }; + } + + const next_cursor = if (parsed.value.cursor) |c| + self.allocator.dupe(u8, c) catch null + else + null; + + return .{ + .dids = dids.toOwnedSlice(self.allocator) catch return error.OutOfMemory, + .next_cursor = next_cursor, + }; + } + + const FetchResult = struct { + dids: [][]const u8, + next_cursor: ?[]const u8, + }; + + const ListReposResponse = struct { + repos: ?[]const RepoEntry = null, + cursor: ?[]const u8 = null, + }; + + const RepoEntry = struct { + did: []const u8, + }; + + /// return status summary for the admin endpoint + pub fn getStatus(self: *Backfiller, allocator: Allocator) ![]u8 { + var list: std.ArrayListUnmanaged(u8) = .{}; + defer list.deinit(allocator); + const w = list.writer(allocator); + + var total: usize = 0; + var completed: usize = 0; + var total_imported: i64 = 0; + + // scan all bf: keys in the default CF + const ci = self.collection_index; + var err_str: ?@import("rocksdb").Data = null; + var iter = ci.db.iterator(ci.default, .forward, "bf:"); + defer iter.deinit(); + + var first_detail = true; + var details: std.ArrayListUnmanaged(u8) = .{}; + defer details.deinit(allocator); + const dw = details.writer(allocator); + + while (iter.next(&err_str) catch null) |entry| { + const key_data = entry[0].data; + if (!std.mem.startsWith(u8, key_data, "bf:")) break; + + const val_data = entry[1].data; + + // parse key: "bf:{collection}\0{source}" + const key_rest = key_data[3..]; // skip "bf:" + const sep = std.mem.indexOfScalar(u8, key_rest, 0) orelse continue; + const collection = key_rest[0..sep]; + const source = key_rest[sep + 1 ..]; + + // parse value: "{cursor}\0{count}\0{0|1}" + const val_sep1 = std.mem.indexOfScalar(u8, val_data, 0) orelse continue; + const val_rest = val_data[val_sep1 + 1 ..]; + const val_sep2 = std.mem.indexOfScalar(u8, val_rest, 0) orelse continue; + const count = std.fmt.parseInt(i64, val_rest[0..val_sep2], 10) catch 0; + const is_completed = val_rest.len > val_sep2 + 1 and val_rest[val_sep2 + 1] == '1'; + + total += 1; + if (is_completed) completed += 1; + total_imported += count; + + if (!first_detail) dw.writeByte(',') catch {}; + first_detail = false; + + std.fmt.format(dw, "{{\"collection\":\"{s}\",\"source\":\"{s}\",\"imported\":{d},\"completed\":{}}}", .{ + collection, + source, + count, + is_completed, + }) catch {}; + } + + std.fmt.format(w, "{{\"running\":{},\"total\":{d},\"completed\":{d},\"in_progress\":{d},\"total_imported\":{d},\"collections\":[", .{ + self.isRunning(), + total, + completed, + total - completed, + total_imported, + }) catch return error.FormatError; + + w.writeAll(details.items) catch {}; + w.writeAll("]}") catch {}; + + return try list.toOwnedSlice(allocator); + } +}; diff --git a/src/collectiondir_subscriber.zig b/src/collectiondir_subscriber.zig new file mode 100644 index 0000000..72bfcba --- /dev/null +++ b/src/collectiondir_subscriber.zig @@ -0,0 +1,314 @@ +//! collectiondir firehose subscriber — consumes relay's subscribeRepos stream +//! +//! simpler than the relay's subscriber.zig — no validation, no host management, +//! no rate limiting. just CBOR frame decode → collection index updates. +//! +//! on #commit: extract ops, call trackCommitOps +//! on #account: if inactive, removeAll from index +//! on #sync: count only (resync needs PDS hostname — requires DID doc resolution, stage 2) +//! +//! cursor persisted to RocksDB default CF (key "cursor"), not postgres. + +const std = @import("std"); +const websocket = @import("websocket"); +const zat = @import("zat"); +const collection_index_mod = @import("collection_index.zig"); + +const Allocator = std.mem.Allocator; +const log = std.log.scoped(.collectiondir); + +const cursor_flush_interval_sec = 4; + +pub const Stats = struct { + commits: std.atomic.Value(u64) = .{ .raw = 0 }, + accounts: std.atomic.Value(u64) = .{ .raw = 0 }, + syncs: std.atomic.Value(u64) = .{ .raw = 0 }, + identities: std.atomic.Value(u64) = .{ .raw = 0 }, + decode_errors: std.atomic.Value(u64) = .{ .raw = 0 }, + ops_tracked: std.atomic.Value(u64) = .{ .raw = 0 }, + accounts_removed: std.atomic.Value(u64) = .{ .raw = 0 }, + reconnects: std.atomic.Value(u64) = .{ .raw = 0 }, + last_seq: std.atomic.Value(u64) = .{ .raw = 0 }, + connected: std.atomic.Value(bool) = .{ .raw = false }, + last_event_time: std.atomic.Value(i64) = .{ .raw = 0 }, +}; + +pub const Subscriber = struct { + allocator: Allocator, + relay_url: []const u8, + relay_port: u16, + use_tls: bool, + collection_index: *collection_index_mod.CollectionIndex, + shutdown: *std.atomic.Value(bool), + stats: Stats = .{}, + last_seq: ?u64 = null, + last_cursor_flush: i64 = 0, + + pub fn init( + allocator: Allocator, + relay_url: []const u8, + relay_port: u16, + use_tls: bool, + collection_index: *collection_index_mod.CollectionIndex, + shutdown: *std.atomic.Value(bool), + ) Subscriber { + return .{ + .allocator = allocator, + .relay_url = relay_url, + .relay_port = relay_port, + .use_tls = use_tls, + .collection_index = collection_index, + .shutdown = shutdown, + }; + } + + /// load saved cursor from RocksDB default CF + pub fn loadCursor(self: *Subscriber) void { + const val = self.collection_index.getMeta("cursor") orelse return; + defer val.deinit(); + const seq = std.fmt.parseInt(u64, val.data, 10) catch return; + if (seq > 0) { + self.last_seq = seq; + log.info("resuming from cursor {d}", .{seq}); + } + } + + /// save cursor to RocksDB default CF + fn flushCursor(self: *Subscriber) void { + const seq = self.last_seq orelse return; + var buf: [20]u8 = undefined; + const val = std.fmt.bufPrint(&buf, "{d}", .{seq}) catch return; + self.collection_index.putMeta("cursor", val); + } + + /// run the subscriber loop with reconnect + backoff + pub fn run(self: *Subscriber) void { + self.loadCursor(); + + var backoff: u64 = 1; + const max_backoff: u64 = 60; + + while (!self.shutdown.load(.acquire)) { + log.info("connecting to {s}:{d}...", .{ self.relay_url, self.relay_port }); + + if (self.connectAndRead()) { + // clean disconnect (EOF) — reset backoff + backoff = 1; + } else |err| { + if (self.shutdown.load(.acquire)) return; + log.err("connection error: {s}, reconnecting in {d}s...", .{ @errorName(err), backoff }); + _ = self.stats.reconnects.fetchAdd(1, .monotonic); + } + + if (self.shutdown.load(.acquire)) return; + + // flush cursor before backoff + self.flushCursor(); + + // backoff sleep in 1s chunks + var remaining: u64 = backoff; + while (remaining > 0 and !self.shutdown.load(.acquire)) { + std.posix.nanosleep(1, 0); + remaining -= 1; + } + backoff = @min(backoff * 2, max_backoff); + } + + // final cursor flush + self.flushCursor(); + } + + fn connectAndRead(self: *Subscriber) !void { + var path_buf: [256]u8 = undefined; + var w: std.Io.Writer = .fixed(&path_buf); + + try w.writeAll("/xrpc/com.atproto.sync.subscribeRepos"); + if (self.last_seq) |cursor| { + try w.print("?cursor={d}", .{cursor}); + } + const path = w.buffered(); + + var client = try websocket.Client.init(self.allocator, .{ + .host = self.relay_url, + .port = self.relay_port, + .tls = self.use_tls, + .max_size = 5 * 1024 * 1024, + }); + defer client.deinit(); + + var host_header_buf: [256]u8 = undefined; + const host_header = std.fmt.bufPrint( + &host_header_buf, + "Host: {s}\r\n", + .{self.relay_url}, + ) catch self.relay_url; + + try client.handshake(path, .{ .headers = host_header }); + log.info("connected to {s}", .{self.relay_url}); + self.stats.connected.store(true, .release); + defer self.stats.connected.store(false, .release); + + var handler = FrameHandler{ .sub = self }; + try client.readLoop(&handler); + } + + pub fn formatMetrics(self: *Subscriber, buf: []u8) []const u8 { + var fbs = std.io.fixedBufferStream(buf); + const w = fbs.writer(); + + const commits = self.stats.commits.load(.monotonic); + const accounts = self.stats.accounts.load(.monotonic); + const syncs = self.stats.syncs.load(.monotonic); + const identities = self.stats.identities.load(.monotonic); + const decode_errors = self.stats.decode_errors.load(.monotonic); + const ops_tracked = self.stats.ops_tracked.load(.monotonic); + const accounts_removed = self.stats.accounts_removed.load(.monotonic); + const reconnects = self.stats.reconnects.load(.monotonic); + const last_seq = self.stats.last_seq.load(.monotonic); + const connected: u64 = if (self.stats.connected.load(.monotonic)) 1 else 0; + const last_event_time = self.stats.last_event_time.load(.monotonic); + + std.fmt.format(w, + \\# HELP collectiondir_events_total events processed by type + \\# TYPE collectiondir_events_total counter + \\collectiondir_events_total{{type="commit"}} {d} + \\collectiondir_events_total{{type="account"}} {d} + \\collectiondir_events_total{{type="sync"}} {d} + \\collectiondir_events_total{{type="identity"}} {d} + \\# HELP collectiondir_decode_errors_total CBOR decode failures + \\# TYPE collectiondir_decode_errors_total counter + \\collectiondir_decode_errors_total {d} + \\# HELP collectiondir_ops_tracked_total collection index ops processed + \\# TYPE collectiondir_ops_tracked_total counter + \\collectiondir_ops_tracked_total {d} + \\# HELP collectiondir_accounts_removed_total accounts removed from index (inactive) + \\# TYPE collectiondir_accounts_removed_total counter + \\collectiondir_accounts_removed_total {d} + \\# HELP collectiondir_reconnects_total WebSocket reconnect count + \\# TYPE collectiondir_reconnects_total counter + \\collectiondir_reconnects_total {d} + \\# HELP collectiondir_last_seq last processed firehose sequence number + \\# TYPE collectiondir_last_seq gauge + \\collectiondir_last_seq {d} + \\# HELP collectiondir_connected whether firehose subscriber is connected + \\# TYPE collectiondir_connected gauge + \\collectiondir_connected {d} + \\# HELP collectiondir_last_event_time_seconds unix timestamp of last processed event + \\# TYPE collectiondir_last_event_time_seconds gauge + \\collectiondir_last_event_time_seconds {d} + \\ + , .{ + commits, + accounts, + syncs, + identities, + decode_errors, + ops_tracked, + accounts_removed, + reconnects, + last_seq, + connected, + last_event_time, + }) catch {}; + + return fbs.getWritten(); + } +}; + +const FrameHandler = struct { + sub: *Subscriber, + + pub fn serverMessage(self: *FrameHandler, data: []const u8) !void { + const sub = self.sub; + + var arena = std.heap.ArenaAllocator.init(sub.allocator); + defer arena.deinit(); + const alloc = arena.allocator(); + + // decode header + const header_result = zat.cbor.decode(alloc, data) catch |err| { + log.debug("frame header decode failed: {s}", .{@errorName(err)}); + _ = sub.stats.decode_errors.fetchAdd(1, .monotonic); + return; + }; + const header = header_result.value; + const payload_data = data[header_result.consumed..]; + + const op = header.getInt("op") orelse return; + if (op != 1) return; // skip error frames + + const frame_type = header.getString("t") orelse return; + const payload = zat.cbor.decodeAll(alloc, payload_data) catch |err| { + log.debug("frame payload decode failed: {s} (type={s})", .{ @errorName(err), frame_type }); + _ = sub.stats.decode_errors.fetchAdd(1, .monotonic); + return; + }; + + const upstream_seq = payload.getUint("seq"); + + // route by frame type + const is_commit = std.mem.eql(u8, frame_type, "#commit"); + const is_account = std.mem.eql(u8, frame_type, "#account"); + const is_sync = std.mem.eql(u8, frame_type, "#sync"); + const is_identity = std.mem.eql(u8, frame_type, "#identity"); + + if (is_commit) { + _ = sub.stats.commits.fetchAdd(1, .monotonic); + + const did = payload.getString("repo") orelse { + if (upstream_seq) |s| sub.last_seq = s; + return; + }; + + // track collection ops — the core purpose of this service + if (payload.get("ops")) |ops| { + sub.collection_index.trackCommitOps(did, ops); + + // count individual ops for metrics + switch (ops) { + .array => |items| { + _ = sub.stats.ops_tracked.fetchAdd(items.len, .monotonic); + }, + else => {}, + } + } + } else if (is_account) { + _ = sub.stats.accounts.fetchAdd(1, .monotonic); + + const did = payload.getString("did") orelse { + if (upstream_seq) |s| sub.last_seq = s; + return; + }; + + const is_active = payload.getBool("active") orelse false; + if (!is_active) { + sub.collection_index.removeAll(did) catch |err| { + log.debug("removeAll failed for {s}: {s}", .{ did, @errorName(err) }); + }; + _ = sub.stats.accounts_removed.fetchAdd(1, .monotonic); + } + } else if (is_sync) { + _ = sub.stats.syncs.fetchAdd(1, .monotonic); + // stage 1: count only — resync requires PDS hostname via DID doc resolution (stage 2) + } else if (is_identity) { + _ = sub.stats.identities.fetchAdd(1, .monotonic); + } else { + if (upstream_seq) |s| sub.last_seq = s; + return; + } + + // advance cursor + if (upstream_seq) |s| { + sub.last_seq = s; + sub.stats.last_seq.store(s, .monotonic); + sub.stats.last_event_time.store(std.time.timestamp(), .monotonic); + } + + // periodic cursor flush + const now = std.time.timestamp(); + if (now - sub.last_cursor_flush >= cursor_flush_interval_sec) { + sub.flushCursor(); + sub.last_cursor_flush = now; + } + } +};