diff --git a/src/sstable.zig b/src/sstable.zig index ac4fa92..45f6c0d 100644 --- a/src/sstable.zig +++ b/src/sstable.zig @@ -612,6 +612,113 @@ const Decoded = struct { next: usize, }; +const ExpandedBlock = struct { + bytes: []u8, + // Flat-block offsets (into `bytes`) of every restart record, ascending. + // Consumed by `Reader.seekInBlock` for binary-search point lookups. + restarts: []u32, +}; + +/// Reconstructs a physical block (prefix-compressed records plus its +/// restart-offset trailer, as written by `Builder`) back into the flat +/// every-record-carries-its-full-key layout that `decodeRecord` and every +/// cursor in this file expect. Expansion happens once, right after +/// decompression in `Reader.open`, so `dataBuf`/the block cache keep +/// storing fully materialized blocks and the rest of the reader never has +/// to know keys were delta-encoded on disk - in particular this preserves +/// the invariant `db.zig`'s k-way merge relies on, that a cursor's current +/// key slice stays valid until that same cursor is advanced again. +/// +/// Also validates the restart trailer against the record stream: every +/// declared restart offset must land exactly on a record whose `shared` +/// field is 0, and every record with `shared` 0 at that exact interval +/// must have a matching restart entry. This catches a truncated or +/// shifted restart table that a whole-block CRC would not (the CRC only +/// proves the compressed bytes were not corrupted in transit/on disk, not +/// that this structural relationship inside them still holds). +/// +/// Also re-derives, in the *flattened* block's own coordinates, the offset +/// of each restart record (`ExpandedBlock.restarts`) so the read side can +/// still binary-search to a restart-sized (<= restartInterval records) +/// linear-scan window even though the on-disk shared/unshared deltas are +/// gone post-expansion. +fn expandBlock(allocator: Allocator, physical: []const u8) !ExpandedBlock { + if (physical.len < 4) return error.CorruptSstable; + const numRestarts = readU32(physical, physical.len - 4); + const restartTableBytes = @as(usize, numRestarts) * 4; + if (restartTableBytes + 4 > physical.len) return error.CorruptSstable; + const recordsEnd = physical.len - 4 - restartTableBytes; + const restartTable = physical[recordsEnd .. recordsEnd + restartTableBytes]; + + var out = std.ArrayList(u8).init(allocator); + errdefer out.deinit(); + + var flatRestarts = std.ArrayList(u32).init(allocator); + errdefer flatRestarts.deinit(); + + var prevKey = std.ArrayList(u8).init(allocator); + defer prevKey.deinit(); + + const perRecordFixed = 1 + 8 + 4 + 4 + 4; + var pos: usize = 0; + var restartIdx: usize = 0; + var recordIdx: usize = 0; + while (pos < recordsEnd) { + if (pos + perRecordFixed > recordsEnd) return error.CorruptSstable; + const kind = physical[pos]; + const seq = readU64(physical, pos + 1); + const shared = readU32(physical, pos + 9); + const unshared = readU32(physical, pos + 13); + const valLen = readU32(physical, pos + 17); + const keyStart = pos + perRecordFixed; + if (shared > prevKey.items.len) return error.CorruptSstable; + if (keyStart + @as(usize, unshared) + @as(usize, valLen) > recordsEnd) return error.CorruptSstable; + + // A genuine block (as `Builder` writes it) marks exactly every + // `restartInterval`-th record (by position within the block) as a + // restart, in order, with shared == 0. Checking both directions - + // a restart-interval position must be declared, and a declared + // offset must fall on a restart-interval position - catches a + // restart table that is present but shifted, truncated, or + // fabricated, not just one that is structurally impossible. + const shouldBeRestart = recordIdx % restartInterval == 0; + const isDeclaredRestart = restartIdx < numRestarts and readU32(restartTable, restartIdx * 4) == pos; + if (shouldBeRestart != isDeclaredRestart) return error.CorruptSstable; + if (shouldBeRestart) { + if (shared != 0) return error.CorruptSstable; + restartIdx += 1; + try flatRestarts.append(@intCast(out.items.len)); + } + + const unsharedBytes = physical[keyStart .. keyStart + unshared]; + const valStart = keyStart + unshared; + const value = physical[valStart .. valStart + valLen]; + + const fullKey = try allocator.alloc(u8, shared + unshared); + defer allocator.free(fullKey); + @memcpy(fullKey[0..shared], prevKey.items[0..shared]); + @memcpy(fullKey[shared..], unsharedBytes); + + const writer = encode.ByteWriter.init(&out); + try writer.byte(kind); + try writer.putU64(seq); + try writer.putU32(@intCast(fullKey.len)); + try writer.putU32(valLen); + try writer.slice(fullKey); + try writer.slice(value); + + prevKey.clearRetainingCapacity(); + try prevKey.appendSlice(fullKey); + + pos = valStart + valLen; + recordIdx += 1; + } + if (pos != recordsEnd) return error.CorruptSstable; + if (restartIdx != numRestarts) return error.CorruptSstable; + + return .{ .bytes = try out.toOwnedSlice(), .restarts = try flatRestarts.toOwnedSlice() }; +} + fn decodeRecord(block: []const u8, pos: usize) Decoded { const kind: Kind = @enumFromInt(block[pos]); const seq = readU64(block, pos + 1); @@ -697,6 +804,73 @@ test "sstable cursor seek" { try std.testing.expect(!beyond.valid()); } +test "restart index lets point lookups skip straight past earlier records in a block" { + const allocator = std.testing.allocator; + var tmp = std.testing.tmpDir(.{}); + defer tmp.cleanup(); + + try buildFixture(allocator, tmp.dir, "test.sst", 500); + var reader = try Reader.open(allocator, tmp.dir, "test.sst", 1); + defer reader.deinit(); + + // Every block must have carried its restart offsets through expansion: + // the first is always the block's own first record, and multi-restart + // blocks (any block with more than restartInterval records) must have + // strictly increasing keys at each restart, matching overall key order. + var sawMultiRestartBlock = false; + var blockIdx: usize = 0; + while (blockIdx < reader.index.len) : (blockIdx += 1) { + const entry = reader.index[blockIdx]; + try std.testing.expect(entry.restarts.len >= 1); + try std.testing.expectEqual(@as(u32, 0), entry.restarts[0]); + if (entry.restarts.len > 1) sawMultiRestartBlock = true; + + const block = reader.blockBytes(blockIdx); + var prevKey: []const u8 = &.{}; + for (entry.restarts) |r| { + const rec = decodeRecord(block, r); + if (prevKey.len > 0) try std.testing.expect(order(prevKey, rec.record.key) == .lt); + prevKey = rec.record.key; + } + } + // 500 records at ~87/block (blockTarget 4096 / ~47 bytes/record) and + // restartInterval 16 guarantees at least one block with several + // restarts - if this ever goes false the fixture stopped exercising + // the multi-restart path the rest of this test relies on. + try std.testing.expect(sawMultiRestartBlock); + + // A probe exactly equal to a later restart's key must land on that + // restart's own offset (the binary search must not undershoot to an + // earlier restart when there is an exact match). + const entry0 = reader.index[0]; + try std.testing.expect(entry0.restarts.len > 2); + const block0 = reader.blockBytes(0); + const secondRestartOffset = entry0.restarts[1]; + const secondRestartKey = decodeRecord(block0, secondRestartOffset).record.key; + try std.testing.expectEqual(secondRestartOffset, Reader.seekInBlock(block0, entry0.restarts, secondRestartKey)); + + // A probe strictly between the first and second restart's keys must + // stay on the first restart (must not overshoot to the second). + const firstRestartKey = decodeRecord(block0, entry0.restarts[0]).record.key; + var midBuf: [16]u8 = undefined; + @memcpy(midBuf[0..firstRestartKey.len], firstRestartKey); + midBuf[firstRestartKey.len - 1] += 1; // one key past first restart's key + try std.testing.expect(order(midBuf[0..firstRestartKey.len], secondRestartKey) == .lt); + try std.testing.expectEqual(entry0.restarts[0], Reader.seekInBlock(block0, entry0.restarts, midBuf[0..firstRestartKey.len])); + + // A probe past every restart's key must resolve to the last restart. + const lastRestartKey = decodeRecord(block0, entry0.restarts[entry0.restarts.len - 1]).record.key; + var pastBuf: [16]u8 = undefined; + @memcpy(pastBuf[0..lastRestartKey.len], lastRestartKey); + pastBuf[lastRestartKey.len - 1] += 1; + try std.testing.expectEqual(entry0.restarts[entry0.restarts.len - 1], Reader.seekInBlock(block0, entry0.restarts, pastBuf[0..lastRestartKey.len])); + + // End to end: lookups still resolve correctly for keys sitting right at + // and around restart boundaries, now that the scan starts mid-block. + try std.testing.expectEqualStrings("payload-value", reader.find(firstRestartKey).?.value); + try std.testing.expectEqualStrings("payload-value", reader.find(secondRestartKey).?.value); +} + test "sstable reverse cursor descends fully" { const allocator = std.testing.allocator; var tmp = std.testing.tmpDir(.{});