From 360df0acea1cef715789129a32367dfbd2d73470 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Wed, 2 Sep 2026 19:00:27 +0200 Subject: [PATCH] feat: extend db with KeyedIndex and keyIndexLess --- src/db.zig | 196 +++++++++++++++++++++++++++++++++++++---------------- 1 file changed, 136 insertions(+), 60 deletions(-) diff --git a/src/db.zig b/src/db.zig index 84fc18d..e07f39f 100644 --- a/src/db.zig +++ b/src/db.zig @@ -4,7 +4,10 @@ const memtable = @import("memtable.zig"); const sstable = @import("sstable.zig"); const cache = @import("cache.zig"); const statsmod = @import("stats.zig"); -const Wal = @import("wal.zig").Wal; +const heap = @import("heap.zig"); +const version = @import("version.zig"); +const walmod = @import("wal.zig"); +const Wal = walmod.Wal; const Allocator = std.mem.Allocator; const Kind = encode.Kind; @@ -13,6 +16,8 @@ const Memtable = memtable.Memtable; const Reader = sstable.Reader; const Builder = sstable.Builder; const BlockCache = cache.BlockCache; +const VersionSet = version.VersionSet; +const VersionEdit = version.VersionEdit; pub const WriteBatch = @import("batch.zig").WriteBatch; pub const Stats = statsmod.Stats; @@ -22,7 +27,6 @@ fn order(a: []const u8, b: []const u8) std.math.Order { } const walName = "wal.log"; -const manifestName = "MANIFEST"; pub const Options = struct { flushThreshold: usize = 4 * 1024 * 1024, @@ -46,6 +50,8 @@ pub const Db = struct { wal: Wal, sstables: std.ArrayList(*Reader), blockCache: BlockCache, + versionSet: VersionSet, + versionState: std.AutoHashMap(u64, usize), nextFileId: u64, nextSeq: u64, flushThreshold: usize, @@ -91,10 +97,11 @@ pub const Db = struct { if (id > maxId) maxId = id; } - var levelMap = try loadManifest(allocator, dir); - defer levelMap.deinit(); + var opened = try VersionSet.open(allocator, dir); + errdefer opened.set.close(); + errdefer opened.state.deinit(); for (sstables.items) |reader| { - reader.level = levelMap.get(reader.fileId) orelse 0; + reader.level = opened.state.get(reader.fileId) orelse 0; } var table = try Memtable.init(allocator); @@ -104,6 +111,20 @@ pub const Db = struct { const walMaxSeq = try wal.replay(&table); if (walMaxSeq > maxSeq) maxSeq = walMaxSeq; + // The directory scan alone can regress: a leveled compaction that + // fully garbage-collects a run of tombstones produces zero output + // files (see runMerge), so no fresh fileId is minted and the + // deleted inputs' seq/fileId watermark vanishes from disk. The + // manifest's persisted counters (stamped on every flush/compact, + // see VersionEdit.nextFileId/nextSeq) are the authoritative floor + // for exactly that case; take whichever recovery source is higher + // so this is purely additive safety, never a regression versus the + // scan-only behaviour this project already relied on. + const scanNextFileId = maxId + 1; + const scanNextSeq = maxSeq + 1; + const nextFileId = @max(scanNextFileId, opened.set.counters.nextFileId orelse 0); + const nextSeq = @max(scanNextSeq, opened.set.counters.nextSeq orelse 0); + const dirPathCopy = try allocator.dupe(u8, dirPath); return .{ @@ -114,8 +135,10 @@ pub const Db = struct { .wal = wal, .sstables = sstables, .blockCache = BlockCache.init(allocator, options.blockCacheBytes), - .nextFileId = maxId + 1, - .nextSeq = maxSeq + 1, + .versionSet = opened.set, + .versionState = opened.state, + .nextFileId = nextFileId, + .nextSeq = nextSeq, .flushThreshold = options.flushThreshold, .l0Trigger = options.l0CompactionTrigger, .l1BaseBytes = options.l1BaseBytes, @@ -132,6 +155,8 @@ pub const Db = struct { } self.sstables.deinit(); self.blockCache.deinit(); + self.versionSet.close(); + self.versionState.deinit(); self.memtable.deinit(); self.dir.close(); self.allocator.free(self.dirPath); @@ -234,7 +259,14 @@ pub const Db = struct { self.nextFileId += 1; self.memtable.clear(); try self.wal.reset(); - try self.writeManifest(); + _ = try walmod.bumpGeneration(self.dir); + + var edit = VersionEdit.init(self.allocator); + defer edit.deinit(); + try edit.addFile(fileId, 0); + edit.setNextFileId(self.nextFileId); + edit.setNextSeq(self.nextSeq); + try self.versionSet.logAndApply(&edit, &self.versionState); std.posix.fsync(self.dir.fd) catch {}; } @@ -254,7 +286,9 @@ pub const Db = struct { { var builder = try Builder.init(self.allocator, self.dir, name, expected); defer builder.deinit(); - while (pickWinner(cursors)) |winner| { + var merger = try SstableMerger.init(self.allocator, cursors); + defer merger.deinit(); + while (try merger.next()) |winner| { if (winner.kind == .value) try builder.add(winner); } try builder.finish(); @@ -265,7 +299,10 @@ pub const Db = struct { reader.* = try Reader.open(self.allocator, self.dir, name, fileId); errdefer reader.deinit(); + var edit = VersionEdit.init(self.allocator); + defer edit.deinit(); for (self.sstables.items) |old| { + try edit.deleteFile(old.fileId); var oldNameBuf: [64]u8 = undefined; const oldName = sstName(&oldNameBuf, old.fileId) catch unreachable; old.deinit(); @@ -276,7 +313,10 @@ pub const Db = struct { reader.level = 0; try self.sstables.append(reader); self.nextFileId += 1; - try self.writeManifest(); + try edit.addFile(fileId, 0); + edit.setNextFileId(self.nextFileId); + edit.setNextSeq(self.nextSeq); + try self.versionSet.logAndApply(&edit, &self.versionState); std.posix.fsync(self.dir.fd) catch {}; } @@ -365,7 +405,9 @@ pub const Db = struct { var outputName: [64]u8 = undefined; errdefer if (builder) |*b| b.deinit(); - while (pickWinner(cursors)) |winner| { + var merger = try SstableMerger.init(self.allocator, cursors); + defer merger.deinit(); + while (try merger.next()) |winner| { if (dropTombstones and winner.kind == .tombstone) continue; if (builder == null) { outputId = self.nextFileId; @@ -385,7 +427,10 @@ pub const Db = struct { try self.finishOutput(done, outputId, targetLevel, &outputs); } + var edit = VersionEdit.init(self.allocator); + defer edit.deinit(); for (inputs) |old| { + try edit.deleteFile(old.fileId); self.removeReader(old); var oldNameBuf: [64]u8 = undefined; const oldName = sstName(&oldNameBuf, old.fileId) catch unreachable; @@ -393,10 +438,21 @@ pub const Db = struct { self.allocator.destroy(old); self.dir.deleteFile(oldName) catch {}; } - for (outputs.items) |produced| try self.sstables.append(produced); + for (outputs.items) |produced| { + try edit.addFile(produced.fileId, produced.level); + try self.sstables.append(produced); + } outputs.deinit(); - try self.writeManifest(); + // Stamped even when this round produced zero output files (every + // input was a dropped tombstone): that is exactly the case where no + // fresh fileId gets minted and a directory scan on the next restart + // could otherwise regress below the seq/fileId already consumed by + // the now-deleted inputs (see the comment on VersionEdit.nextFileId + // and Db.open's counter merge). + edit.setNextFileId(self.nextFileId); + edit.setNextSeq(self.nextSeq); + try self.versionSet.logAndApply(&edit, &self.versionState); std.posix.fsync(self.dir.fd) catch {}; } @@ -465,32 +521,6 @@ pub const Db = struct { return false; } - fn writeManifest(self: *Db) !void { - var buf = std.ArrayList(u8).init(self.allocator); - defer buf.deinit(); - const bodyWriter = encode.ByteWriter.init(&buf); - try bodyWriter.putU32(@intCast(self.sstables.items.len)); - for (self.sstables.items) |reader| { - try bodyWriter.putU64(reader.fileId); - try bodyWriter.putU32(@intCast(reader.level)); - } - - var frame = std.ArrayList(u8).init(self.allocator); - defer frame.deinit(); - const frameWriter = encode.ByteWriter.init(&frame); - try frameWriter.putU32(encode.crc(buf.items)); - try frameWriter.slice(buf.items); - - const tmpName = "MANIFEST.tmp"; - { - const file = try self.dir.createFile(tmpName, .{ .truncate = true }); - defer file.close(); - try file.writeAll(frame.items); - try file.sync(); - } - try self.dir.rename(tmpName, manifestName); - } - pub fn levelOf(self: *Db, fileId: u64) ?usize { for (self.sstables.items) |reader| { if (reader.fileId == fileId) return reader.level; @@ -548,6 +578,12 @@ pub const Db = struct { cursors[0] = .{ .mem = self.memtable.cursor(options.start) }; for (self.sstables.items, 0..) |reader, i| cursors[i + 1] = .{ .sst = reader.cursor(options.start) }; + var mergeHeap = AscKeyHeap.init(allocator, {}); + errdefer mergeHeap.deinit(); + for (cursors, 0..) |*c, i| { + if (c.valid()) try mergeHeap.push(.{ .key = c.record().key, .idx = i }); + } + var startCopy: ?[]u8 = null; errdefer if (startCopy) |s| allocator.free(s); if (options.start) |bound| startCopy = try allocator.dupe(u8, bound); @@ -558,6 +594,7 @@ pub const Db = struct { .resultAllocator = allocator, .arena = std.heap.ArenaAllocator.init(allocator), .cursors = cursors, + .heap = mergeHeap, .start = startCopy, .end = endCopy, .startInclusive = options.startInclusive, @@ -596,6 +633,12 @@ pub const Db = struct { built += 1; } + var mergeHeap = DescKeyHeap.init(allocator, {}); + errdefer mergeHeap.deinit(); + for (cursors, 0..) |*c, i| { + if (c.valid()) try mergeHeap.push(.{ .key = c.record().key, .idx = i }); + } + var startCopy: ?[]u8 = null; if (options.start) |bound| startCopy = try allocator.dupe(u8, bound); @@ -603,6 +646,7 @@ pub const Db = struct { .resultAllocator = allocator, .arena = std.heap.ArenaAllocator.init(allocator), .cursors = cursors, + .heap = mergeHeap, .start = startCopy, .startInclusive = options.startInclusive, .maxSeq = maxSeq, @@ -678,30 +722,68 @@ const SourceCursor = union(enum) { } }; -fn pickWinner(cursors: []sstable.Cursor) ?Record { - var minKey: ?[]const u8 = null; - for (cursors) |*c| { - if (c.valid()) { - const key = c.record().key; - if (minKey == null or order(key, minKey.?) == .lt) minKey = key; +/// One live position in a k-way merge: the key it currently points at +/// (a snapshot valid for the underlying source's whole lifetime, since +/// sstable/memtable cursors never overwrite already-decoded key bytes) +/// plus the index of the source cursor it belongs to. +const KeyedIndex = struct { key: []const u8, idx: usize }; + +fn keyIndexLess(_: void, a: KeyedIndex, b: KeyedIndex) bool { + return order(a.key, b.key) == .lt; +} +fn keyIndexGreater(_: void, a: KeyedIndex, b: KeyedIndex) bool { + return order(a.key, b.key) == .gt; +} + +/// Ascending k-way merge heap: root always holds the smallest live key. +const AscKeyHeap = heap.MinHeap(KeyedIndex, void, keyIndexLess); +/// Descending k-way merge heap: root always holds the largest live key. +const DescKeyHeap = heap.MinHeap(KeyedIndex, void, keyIndexGreater); + +/// Merges N sorted sstable cursors (compaction inputs) in ascending key +/// order using a binary heap: O(log N) per advance instead of the O(N) +/// full rescan a naive linear-scan merge needs for every emitted record. +/// Ties (same key present in several inputs) resolve to the entry with +/// the highest seq, matching last-writer-wins semantics. +const SstableMerger = struct { + cursors: []sstable.Cursor, + heap: AscKeyHeap, + + fn init(allocator: Allocator, cursors: []sstable.Cursor) !SstableMerger { + var h = AscKeyHeap.init(allocator, {}); + errdefer h.deinit(); + for (cursors, 0..) |*c, i| { + if (c.valid()) try h.push(.{ .key = c.record().key, .idx = i }); } + return .{ .cursors = cursors, .heap = h }; } - if (minKey == null) return null; - var winner: ?Record = null; - for (cursors) |*c| { - if (c.valid() and order(c.record().key, minKey.?) == .eq) { + + fn deinit(self: *SstableMerger) void { + self.heap.deinit(); + } + + fn next(self: *SstableMerger) !?Record { + const top = self.heap.peek() orelse return null; + const minKey = top.key; + var winner: ?Record = null; + while (self.heap.peek()) |entry| { + if (order(entry.key, minKey) != .eq) break; + const popped = self.heap.pop().?; + var c = &self.cursors[popped.idx]; const rec = c.record(); if (winner == null or rec.seq > winner.?.seq) winner = rec; c.advance(); + if (c.valid()) try self.heap.push(.{ .key = c.record().key, .idx = popped.idx }); } + return winner.?; } - return winner.?; -} +}; pub const RangeIterator = struct { resultAllocator: Allocator, arena: std.heap.ArenaAllocator, cursors: []SourceCursor, + heap: AscKeyHeap, start: ?[]u8, end: ?[]u8, startInclusive: bool, @@ -710,16 +792,10 @@ pub const RangeIterator = struct { pub fn next(self: *RangeIterator) !?Entry { while (true) { - var minKey: ?[]const u8 = null; - for (self.cursors) |*c| { - if (c.valid()) { - const key = c.record().key; - if (minKey == null or order(key, minKey.?) == .lt) minKey = key; - } - } - if (minKey == null) return null; + const top = self.heap.peek() orelse return null; + const minKey = top.key; if (self.end) |bound| { - const past = if (self.endInclusive) order(minKey.?, bound) == .gt else order(minKey.?, bound) != .lt; + const past = if (self.endInclusive) order(minKey, bound) == .gt else order(minKey, bound) != .lt; if (past) return null; }