diff --git a/build.zig b/build.zig index ed3424f..ff059f2 100644 --- a/build.zig +++ b/build.zig @@ -26,6 +26,7 @@ pub fn build(b: *std.Build) void { .{ .name = "websocket", .module = websocket.module("websocket") }, .{ .name = "pg", .module = pg.module("pg") }, .{ .name = "rocksdb", .module = rocksdb.module("bindings") }, + .{ .name = "rocksdb_c", .module = rocksdb.module("rocksdb") }, }; // relay executable diff --git a/src/collection_index.zig b/src/collection_index.zig index 5398d59..e5dc7a5 100644 --- a/src/collection_index.zig +++ b/src/collection_index.zig @@ -11,71 +11,156 @@ const std = @import("std"); const rocksdb = @import("rocksdb"); +const rdb = @import("rocksdb_c"); const Allocator = std.mem.Allocator; const log = std.log.scoped(.collection_index); const separator = '\x00'; +/// shared block cache size — bounds total read memory (indexes + data + filters) +/// across all column families. with cache_index_and_filter_blocks=true, this is +/// the hard cap on RocksDB's read-path memory. 256 MB is conservative for the +/// collection index workload (mostly sequential writes + prefix scans). +const block_cache_bytes: usize = 256 * 1024 * 1024; + pub const CollectionIndex = struct { - db: rocksdb.DB, - rbc: rocksdb.ColumnFamilyHandle, - cbr: rocksdb.ColumnFamilyHandle, + db: *rdb.rocksdb_t, + rbc: *rdb.rocksdb_column_family_handle_t, + cbr: *rdb.rocksdb_column_family_handle_t, + default_cf: *rdb.rocksdb_column_family_handle_t, + cache: *rdb.rocksdb_cache_t, allocator: Allocator, pub fn open(allocator: Allocator, data_dir: []const u8) !CollectionIndex { - var err_str: ?rocksdb.Data = null; - - const db, const families = rocksdb.DB.open( - allocator, - data_dir, - .{ - .create_if_missing = true, - .create_missing_column_families = true, - }, - &.{ - .{ .name = "default" }, - .{ .name = "rbc" }, - .{ .name = "cbr" }, - }, - false, - &err_str, - ) catch { - if (err_str) |e| { - log.err("rocksdb open failed: {s}", .{e.data}); - e.deinit(); + // shared LRU block cache — bounds total read memory across all CFs + const cache = rdb.rocksdb_cache_create_lru(block_cache_bytes) orelse return error.RocksDBOpen; + errdefer rdb.rocksdb_cache_destroy(cache); + + // block-based table options — force indexes/filters into the bounded cache + const bbo = rdb.rocksdb_block_based_options_create() orelse return error.RocksDBOpen; + defer rdb.rocksdb_block_based_options_destroy(bbo); + rdb.rocksdb_block_based_options_set_block_cache(bbo, cache); + rdb.rocksdb_block_based_options_set_block_size(bbo, 16 * 1024); // 16 KB (reduces index size vs 4 KB default) + rdb.rocksdb_block_based_options_set_cache_index_and_filter_blocks(bbo, 1); + rdb.rocksdb_block_based_options_set_pin_l0_filter_and_index_blocks_in_cache(bbo, 1); + + // per-CF options: 32 MB write buffer, max 2 concurrent memtables + const cf_opts = rdb.rocksdb_options_create() orelse return error.RocksDBOpen; + defer rdb.rocksdb_options_destroy(cf_opts); + rdb.rocksdb_options_set_block_based_table_factory(cf_opts, bbo); + rdb.rocksdb_options_set_write_buffer_size(cf_opts, 32 * 1024 * 1024); + rdb.rocksdb_options_set_max_write_buffer_number(cf_opts, 2); + + // DB-level options + const db_opts = rdb.rocksdb_options_create() orelse return error.RocksDBOpen; + defer rdb.rocksdb_options_destroy(db_opts); + rdb.rocksdb_options_set_create_if_missing(db_opts, 1); + rdb.rocksdb_options_set_create_missing_column_families(db_opts, 1); + + // null-terminate path for C API + var path_buf: [4096]u8 = @splat(0); + if (data_dir.len >= path_buf.len) return error.RocksDBOpen; + @memcpy(path_buf[0..data_dir.len], data_dir); + + const cf_names: [3][*c]const u8 = .{ "default", "rbc", "cbr" }; + const cf_options: [3]?*const rdb.rocksdb_options_t = .{ cf_opts, cf_opts, cf_opts }; + var cf_handles: [3]?*rdb.rocksdb_column_family_handle_t = .{ null, null, null }; + var err: ?[*:0]u8 = null; + + const db = rdb.rocksdb_open_column_families( + db_opts, + &path_buf, + 3, + &cf_names, + @ptrCast(&cf_options), + @ptrCast(&cf_handles), + @ptrCast(&err), + ) orelse { + if (err) |e| { + log.err("rocksdb open failed: {s}", .{std.mem.span(e)}); + rdb.rocksdb_free(@ptrCast(e)); } return error.RocksDBOpen; }; - defer allocator.free(families); - - // find column family handles by name - var rbc: ?rocksdb.ColumnFamilyHandle = null; - var cbr: ?rocksdb.ColumnFamilyHandle = null; - for (families) |cf| { - if (std.mem.eql(u8, cf.name, "rbc")) rbc = cf.handle; - if (std.mem.eql(u8, cf.name, "cbr")) cbr = cf.handle; - } - log.info("collection index opened at {s}", .{data_dir}); + log.info("collection index opened at {s} (block cache: {d} MB, indexes in cache)", .{ + data_dir, + block_cache_bytes / (1024 * 1024), + }); return .{ .db = db, - .rbc = rbc orelse return error.MissingColumnFamily, - .cbr = cbr orelse return error.MissingColumnFamily, + .default_cf = cf_handles[0] orelse return error.MissingColumnFamily, + .rbc = cf_handles[1] orelse return error.MissingColumnFamily, + .cbr = cf_handles[2] orelse return error.MissingColumnFamily, + .cache = cache, .allocator = allocator, }; } pub fn deinit(self: *CollectionIndex) void { - self.db.deinit(); + rdb.rocksdb_column_family_handle_destroy(self.default_cf); + rdb.rocksdb_column_family_handle_destroy(self.rbc); + rdb.rocksdb_column_family_handle_destroy(self.cbr); + rdb.rocksdb_close(self.db); + rdb.rocksdb_cache_destroy(self.cache); + } + + // --- raw C API helpers (replacing rocksdb-zig wrapper methods) --- + + fn dbWrite(self: *CollectionIndex, batch: rocksdb.WriteBatch) !void { + const opts = rdb.rocksdb_writeoptions_create() orelse return error.WriteFailed; + defer rdb.rocksdb_writeoptions_destroy(opts); + var err: ?[*:0]u8 = null; + rdb.rocksdb_write(self.db, opts, @ptrCast(batch.inner), @ptrCast(&err)); + if (err) |e| { + defer rdb.rocksdb_free(@ptrCast(e)); + return error.WriteFailed; + } + } + + fn dbIterator(self: *CollectionIndex, cf: *rdb.rocksdb_column_family_handle_t, direction: rocksdb.IteratorDirection, start: ?[]const u8) rocksdb.Iterator { + const opts = rdb.rocksdb_readoptions_create(); + defer if (opts) |o| rdb.rocksdb_readoptions_destroy(o); + const it: *rdb.rocksdb_iterator_t = rdb.rocksdb_create_iterator_cf(self.db, opts, cf) orelse unreachable; + if (start) |s| { + switch (direction) { + .forward => rdb.rocksdb_iter_seek(it, s.ptr, s.len), + .reverse => rdb.rocksdb_iter_seek_for_prev(it, s.ptr, s.len), + } + } else { + switch (direction) { + .forward => rdb.rocksdb_iter_seek_to_first(it), + .reverse => rdb.rocksdb_iter_seek_to_last(it), + } + } + return .{ + .raw = .{ .inner = @ptrCast(it) }, + .direction = direction, + .done = false, + }; + } + + fn dbGet(self: *CollectionIndex, cf: *rdb.rocksdb_column_family_handle_t, key: []const u8) ?rocksdb.Data { + var val_len: usize = 0; + const opts = rdb.rocksdb_readoptions_create(); + defer if (opts) |o| rdb.rocksdb_readoptions_destroy(o); + var err: ?[*:0]u8 = null; + const val = rdb.rocksdb_get_cf(self.db, opts, cf, key.ptr, key.len, &val_len, @ptrCast(&err)); + if (err) |e| { + rdb.rocksdb_free(@ptrCast(e)); + return null; + } + if (val) |v| { + return .{ .data = v[0..val_len], .free = @ptrCast(&rdb.rocksdb_free) }; + } + return null; } /// add a (collection, did) entry to both indexes. /// idempotent — overwrites are no-ops for empty values. pub fn addCollection(self: *CollectionIndex, did: []const u8, collection: []const u8) !void { - var err_str: ?rocksdb.Data = null; - const rbc_key = makeKey(self.allocator, collection, did) catch return; defer self.allocator.free(rbc_key); const cbr_key = makeKey(self.allocator, did, collection) catch return; @@ -87,19 +172,11 @@ pub const CollectionIndex = struct { batch.put(self.rbc, rbc_key, ""); batch.put(self.cbr, cbr_key, ""); - self.db.write(batch, &err_str) catch { - if (err_str) |e| { - log.warn("addCollection write failed: {s}", .{e.data}); - e.deinit(); - } - return error.WriteFailed; - }; + self.dbWrite(batch) catch return error.WriteFailed; } /// remove a (collection, did) entry from both indexes. pub fn removeCollection(self: *CollectionIndex, did: []const u8, collection: []const u8) !void { - var err_str: ?rocksdb.Data = null; - const rbc_key = makeKey(self.allocator, collection, did) catch return; defer self.allocator.free(rbc_key); const cbr_key = makeKey(self.allocator, did, collection) catch return; @@ -111,13 +188,7 @@ pub const CollectionIndex = struct { batch.delete(self.rbc, rbc_key); batch.delete(self.cbr, cbr_key); - self.db.write(batch, &err_str) catch { - if (err_str) |e| { - log.warn("removeCollection write failed: {s}", .{e.data}); - e.deinit(); - } - return error.WriteFailed; - }; + self.dbWrite(batch) catch return error.WriteFailed; } /// remove all collection entries for a DID (e.g. account tombstone). @@ -136,7 +207,7 @@ pub const CollectionIndex = struct { defer batch.deinit(); // scan cbr for all collections belonging to this DID - var iter = self.db.iterator(self.cbr, .forward, prefix); + var iter = self.dbIterator(self.cbr, .forward, prefix); defer iter.deinit(); while (try iter.next(&err_str)) |entry| { @@ -160,13 +231,7 @@ pub const CollectionIndex = struct { batch.delete(self.cbr, cbr_key); } - self.db.write(batch, &err_str) catch { - if (err_str) |e| { - log.warn("removeAll write failed: {s}", .{e.data}); - e.deinit(); - } - return error.WriteFailed; - }; + self.dbWrite(batch) catch return error.WriteFailed; } /// process commit ops to update collection index. @@ -233,7 +298,7 @@ pub const CollectionIndex = struct { const prefix = seek_key[0 .. collection.len + 1]; // collection\0 - var iter = self.db.iterator(self.rbc, .forward, seek_key); + var iter = self.dbIterator(self.rbc, .forward, seek_key); defer iter.deinit(); var count: usize = 0; @@ -280,7 +345,7 @@ pub const CollectionIndex = struct { defer seen.deinit(allocator); // full scan of RBC — keys are collection\0did - var iter = self.db.iterator(self.rbc, .forward, ""); + var iter = self.dbIterator(self.rbc, .forward, ""); defer iter.deinit(); while (try iter.next(&err_str)) |entry| { @@ -317,7 +382,7 @@ pub const CollectionIndex = struct { prefix_buf[did.len] = separator; const prefix = prefix_buf[0 .. did.len + 1]; - var iter = self.db.iterator(self.cbr, .forward, prefix); + var iter = self.dbIterator(self.cbr, .forward, prefix); defer iter.deinit(); if (iter.next(&err_str) catch null) |entry| { @@ -388,10 +453,9 @@ test "collection index: basic put and get" { try ci.addCollection("did:plc:alice", "app.bsky.feed.post"); // verify via direct rocksdb get - var err_str: ?rocksdb.Data = null; const rbc_key = try makeKey(allocator, "app.bsky.feed.post", "did:plc:alice"); defer allocator.free(rbc_key); - const val = try ci.db.get(ci.rbc, rbc_key, &err_str); + const val = ci.dbGet(ci.rbc, rbc_key); try testing.expect(val != null); if (val) |v| v.deinit(); }