dualfroz/cobalt
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: extend db with KeyedIndex and keyIndexLess

Browse Source
This commit is contained in:
dualfroz committed 2026-09-02 19:00:27 +02:00
1 parent 4e87469ab3
commit 360df0acea
1 file changed
+134 -58
Download Patch File Download Diff File Expand all files Collapse all files
src/db.zig
+134 -58
View File
@@ -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;
}
if (minKey == null) return null;
/// 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 };
}
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;
for (cursors) |*c| {
if (c.valid() and order(c.record().key, minKey.?) == .eq) {
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.?;
}
};
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;
}