feat: extend db with KeyedIndex and keyIndexLess

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);
+        var opened = try VersionSet.open(allocator, dir);
-        defer levelMap.deinit();
+        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,
+            .versionSet = opened.set,
-            .nextSeq = maxSeq + 1,
+            .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 {
+/// One live position in a k-way merge: the key it currently points at
-    var minKey: ?[]const u8 = null;
+/// (a snapshot valid for the underlying source's whole lifetime, since
-    for (cursors) |*c| {
+/// sstable/memtable cursors never overwrite already-decoded key bytes)
-        if (c.valid()) {
+/// plus the index of the source cursor it belongs to.
-            const key = c.record().key;
+const KeyedIndex = struct { key: []const u8, idx: usize };
-            if (minKey == null or order(key, minKey.?) == .lt) minKey = key;
+
+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;
+    fn deinit(self: *SstableMerger) void {
-    for (cursors) |*c| {
+        self.heap.deinit();
-        if (c.valid() and order(c.record().key, minKey.?) == .eq) {
+    }
+
+    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;
+            const top = self.heap.peek() orelse return null;
-            for (self.cursors) |*c| {
+            const minKey = top.key;
-                if (c.valid()) {
-                    const key = c.record().key;
-                    if (minKey == null or order(key, minKey.?) == .lt) minKey = key;
-                }
-            }
-            if (minKey == null) return null;
             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;
             }