harden durable PKBFI scans

.gitignore

@@ -9,4 +9,5 @@ tools/
 .dbb
 test_nov.js
 bin/
+/pizzakv_*
 .DS_Store

engine.zig

@@ -18,6 +18,16 @@ pub const Value = struct {
     lsn: u64,
 };
 
+pub const CompareCheck = struct {
+    key: []const u8,
+    expected_lsn: u64,
+};
+
+pub const CompareBatchResult = struct {
+    lsn: u64,
+    committed: bool,
+};
+
 pub const ScanEntry = struct {
     key: []u8,
     value: ?[]u8,
@@ -70,9 +80,11 @@ const Root = struct {
 const PendingWrite = struct {
     operations: []const Operation,
     metadata: []const u8,
+    checks: []const CompareCheck = &.{},
     next: ?*PendingWrite = null,
     completion: *WriteCompletion,
     lsn: u64 = 0,
+    committed: bool = true,
     frame_position: usize = 0,
     changed: bool = false,
     prepared_position: usize = 0,
@@ -586,6 +598,24 @@ pub const Engine = struct {
         return pending.lsn;
     }
 
+    pub fn compareBatchWrite(self: *Engine, checks: []const CompareCheck, operations: []const Operation, metadata: []const u8) !CompareBatchResult {
+        if (operations.len == 0 or operations.len > pkvdb.max_operations or metadata.len > pkvdb.max_transaction_size) return error.InvalidLength;
+        if (checks.len > pkvdb.max_operations) return error.InvalidLength;
+        for (operations) |operation| {
+            if (operation.key.len > pkvdb.max_key_size or operation.value.len > pkvdb.max_value_size) return error.InvalidLength;
+            if (operation.opcode == .delete and operation.value.len != 0) return error.InvalidLength;
+        }
+        for (checks) |check| {
+            if (check.key.len > pkvdb.max_key_size) return error.InvalidLength;
+        }
+        const frame_length = try transactionLength(operations, metadata);
+        var completion = WriteCompletion{ .remaining = 1 };
+        var pending = PendingWrite{ .operations = operations, .metadata = metadata, .checks = checks, .bytes = frame_length, .completion = &completion };
+        try self.enqueueAndWait(&.{&pending}, &completion);
+        if (completion.failure) |failure| return failure;
+        return .{ .lsn = pending.lsn, .committed = pending.committed };
+    }
+
     pub fn putMany(self: *Engine, operations: []const Operation, lsns: []u64) !void {
         if (operations.len == 0 or operations.len != lsns.len or operations.len > max_group_transactions) return error.InvalidLength;
         const pending = try self.allocator.alloc(PendingWrite, operations.len);
@@ -634,8 +664,9 @@ pub const Engine = struct {
             _ = self.queue_condition.timedWait(&self.queue_mutex, group_wait_ns) catch {};
             var count: usize = 0;
             var bytes: usize = pkvdb.group_header_size;
+            var conditional: bool = false;
             while (self.queue_head) |pending| {
-                if (count != 0 and (count == group.len or bytes + pending.bytes > max_group_bytes)) break;
+                if (count != 0 and (count == group.len or bytes + pending.bytes > max_group_bytes or conditional or pending.checks.len != 0)) break;
                 self.queue_head = pending.next;
                 if (self.queue_head == null) self.queue_tail = null;
                 pending.next = null;
@@ -643,6 +674,7 @@ pub const Engine = struct {
                 count += 1;
                 bytes += pending.bytes;
                 self.queued_bytes -= pending.bytes;
+                conditional = pending.checks.len != 0;
             }
             self.queue_condition.broadcast();
             self.queue_mutex.unlock();
@@ -675,6 +707,18 @@ pub const Engine = struct {
         }
     }
 
+    fn checkConditions(self: *Engine, checks: []const CompareCheck) !bool {
+        for (checks) |check| {
+            const record = try self.directory.get(self.file, check.key);
+            if (check.expected_lsn == 0) {
+                if (record != null) return false;
+            } else if (record == null or record.?.lsn != check.expected_lsn) {
+                return false;
+            }
+        }
+        return true;
+    }
+
     fn processGroup(self: *Engine, group: []*PendingWrite, payload_length: usize) !void {
         var operation_count: usize = 0;
         for (group) |pending| operation_count = try std.math.add(usize, operation_count, pending.operations.len);
@@ -683,6 +727,22 @@ pub const Engine = struct {
         self.ordered_gate.lock();
         defer self.ordered_gate.unlock();
         self.lock.lock();
+        if (group[0].checks.len != 0) {
+            if (group.len != 1) {
+                self.lock.unlock();
+                return error.InvalidConditionalGroup;
+            }
+            const conditions_match = self.checkConditions(group[0].checks) catch |failure| {
+                self.lock.unlock();
+                return failure;
+            };
+            if (!conditions_match) {
+                group[0].committed = false;
+                group[0].lsn = 0;
+                self.lock.unlock();
+                return;
+            }
+        }
         self.directory.ensureAdditional(self.file, operation_count) catch |failure| {
             self.lock.unlock();
             return failure;
@@ -1107,6 +1167,32 @@ test "atomic batch checkpoint tail and ordered scan" {
     try std.testing.expectEqualStrings("p/3", batch.entries[0].key);
 }
 
+test "scan supports full-size PKBFI pages and accounts for entry framing" {
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    var path_buffer: [std.fs.max_path_bytes]u8 = undefined;
+    const path = try testPath(&tmp, "large-scan.pkvdb", &path_buffer);
+    defer std.testing.allocator.free(path);
+    var engine = try Engine.open(std.testing.allocator, path);
+    defer engine.close();
+
+    const value = try std.testing.allocator.alloc(u8, 2 * 1024 * 1024);
+    defer std.testing.allocator.free(value);
+    @memset(value, 'x');
+    _ = try engine.put("large/key", value);
+
+    const entry_size = 16 + "large/key".len + value.len;
+    var batch = try engine.scan(std.testing.allocator, "large/", "", 1, true, @intCast(entry_size));
+    defer batch.deinit(std.testing.allocator);
+    try std.testing.expectEqual(@as(usize, 1), batch.entries.len);
+    try std.testing.expectEqual(value.len, batch.entries[0].value.?.len);
+
+    try std.testing.expectError(
+        error.ScanEntryTooLarge,
+        engine.scan(std.testing.allocator, "large/", "", 1, true, @intCast(entry_size - 1)),
+    );
+}
+
 fn testPath(tmp: *std.testing.TmpDir, name: []const u8, buffer: *[std.fs.max_path_bytes]u8) ![]const u8 {
     const directory = try tmp.dir.realpath(".", buffer);
     return std.fs.path.join(std.testing.allocator, &.{ directory, name });
@@ -1419,3 +1505,93 @@ test "concurrent delete reports one removal" {
     try std.testing.expect(!failed.load(.acquire));
     try std.testing.expect(results[0] != results[1]);
 }
+
+test "compare batch write absent match and stale conflict" {
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    var path_buffer: [std.fs.max_path_bytes]u8 = undefined;
+    const path = try testPath(&tmp, "compare.pkvdb", &path_buffer);
+    defer std.testing.allocator.free(path);
+    var engine = try Engine.open(std.testing.allocator, path);
+
+    const absent_checks = [_]CompareCheck{.{ .key = "k", .expected_lsn = 0 }};
+    var result = try engine.compareBatchWrite(&absent_checks, &.{.{ .opcode = .put, .key = "k", .value = "one" }}, "meta");
+    try std.testing.expect(result.committed);
+    try std.testing.expectEqual(@as(u64, 1), result.lsn);
+
+    const match_checks = [_]CompareCheck{.{ .key = "k", .expected_lsn = result.lsn }};
+    result = try engine.compareBatchWrite(&match_checks, &.{.{ .opcode = .put, .key = "k", .value = "two" }}, "");
+    try std.testing.expect(result.committed);
+    try std.testing.expectEqual(@as(u64, 2), result.lsn);
+
+    const stale_checks = [_]CompareCheck{.{ .key = "k", .expected_lsn = 1 }};
+    result = try engine.compareBatchWrite(&stale_checks, &.{.{ .opcode = .put, .key = "k", .value = "three" }}, "");
+    try std.testing.expect(!result.committed);
+    try std.testing.expectEqual(@as(u64, 0), result.lsn);
+
+    var value = (try engine.get(std.testing.allocator, "k")).?;
+    try std.testing.expectEqualStrings("two", value.bytes);
+    try std.testing.expectEqual(@as(u64, 2), value.lsn);
+    std.testing.allocator.free(value.bytes);
+
+    engine.close();
+    engine = try Engine.open(std.testing.allocator, path);
+    defer engine.close();
+    value = (try engine.get(std.testing.allocator, "k")).?;
+    defer std.testing.allocator.free(value.bytes);
+    try std.testing.expectEqualStrings("two", value.bytes);
+    try std.testing.expectEqual(@as(u64, 2), value.lsn);
+}
+
+test "compare absent check rejects existing key" {
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    var path_buffer: [std.fs.max_path_bytes]u8 = undefined;
+    const path = try testPath(&tmp, "compare-absent.pkvdb", &path_buffer);
+    defer std.testing.allocator.free(path);
+    var engine = try Engine.open(std.testing.allocator, path);
+    defer engine.close();
+    _ = try engine.put("k", "seed");
+    const checks = [_]CompareCheck{.{ .key = "k", .expected_lsn = 0 }};
+    const result = try engine.compareBatchWrite(&checks, &.{.{ .opcode = .put, .key = "k", .value = "x" }}, "");
+    try std.testing.expect(!result.committed);
+    try std.testing.expectEqual(@as(u64, 0), result.lsn);
+    const value = (try engine.get(std.testing.allocator, "k")).?;
+    defer std.testing.allocator.free(value.bytes);
+    try std.testing.expectEqualStrings("seed", value.bytes);
+}
+
+test "concurrent compare batch writes expecting same lsn commit exactly one" {
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    var path_buffer: [std.fs.max_path_bytes]u8 = undefined;
+    const path = try testPath(&tmp, "compare-race.pkvdb", &path_buffer);
+    defer std.testing.allocator.free(path);
+    var engine = try Engine.open(std.testing.allocator, path);
+    defer engine.close();
+    _ = try engine.put("key", "seed");
+    var results: [2]CompareBatchResult = undefined;
+    var failed = std.atomic.Value(bool).init(false);
+    const Worker = struct {
+        fn run(target: *Engine, result: *CompareBatchResult, failure: *std.atomic.Value(bool)) void {
+            const checks = [_]CompareCheck{.{ .key = "key", .expected_lsn = 1 }};
+            const operations = [_]Operation{.{ .opcode = .put, .key = "key", .value = "winner" }};
+            result.* = target.compareBatchWrite(&checks, &operations, "") catch {
+                failure.store(true, .release);
+                return;
+            };
+        }
+    };
+    const first = try std.Thread.spawn(.{}, Worker.run, .{ &engine, &results[0], &failed });
+    const second = try std.Thread.spawn(.{}, Worker.run, .{ &engine, &results[1], &failed });
+    first.join();
+    second.join();
+    try std.testing.expect(!failed.load(.acquire));
+    const committed = @as(u64, @intFromBool(results[0].committed)) + @as(u64, @intFromBool(results[1].committed));
+    try std.testing.expectEqual(@as(u64, 1), committed);
+    const winner = if (results[0].committed) results[0] else results[1];
+    try std.testing.expectEqual(@as(u64, 2), winner.lsn);
+    const value = (try engine.get(std.testing.allocator, "key")).?;
+    defer std.testing.allocator.free(value.bytes);
+    try std.testing.expectEqualStrings("winner", value.bytes);
+}

pkbfi.zig

@@ -17,6 +17,7 @@ pub const Opcode = enum(u16) {
     scan_open = 9,
     scan_next = 10,
     scan_close = 11,
+    compare_batch_write = 12,
 };
 
 pub const Frame = struct {
@@ -106,6 +107,7 @@ pub const Session = struct {
             },
             .multi_get => try self.multiGet(engine, frame.payload, body),
             .batch_write => try self.batchWrite(engine, frame.payload, body),
+            .compare_batch_write => try self.compareBatchWrite(engine, frame.payload, body),
             .scan_open => try self.scanOpen(frame.payload, body),
             .scan_next => try self.scanNext(engine, frame.payload, body),
             .scan_close => try self.scanClose(frame.payload, body),
@@ -172,6 +174,52 @@ pub const Session = struct {
         try appendInt(u64, body, self.allocator, try engine.batchWrite(operations, metadata));
     }
 
+    fn compareBatchWrite(self: *Session, engine: *engine_mod.Engine, payload: []const u8, body: *std.ArrayListUnmanaged(u8)) !void {
+        if (payload.len < 16) return error.InvalidPayload;
+        const check_count = readInt(u32, payload, 0);
+        const op_count = readInt(u32, payload, 4);
+        const metadata_length = readInt(u32, payload, 8);
+        if (op_count == 0 or op_count > pkvdb.max_operations or check_count > pkvdb.max_operations) return error.InvalidPayload;
+        const checks = try self.allocator.alloc(engine_mod.CompareCheck, check_count);
+        defer self.allocator.free(checks);
+        var position: usize = 16;
+        for (checks) |*check| {
+            if (position > payload.len or payload.len - position < 16) return error.InvalidPayload;
+            const key_length = readInt(u32, payload, position);
+            position += 4;
+            position += 4;
+            const expected_lsn = readInt(u64, payload, position);
+            position += 8;
+            const end = try std.math.add(usize, position, key_length);
+            if (end > payload.len or key_length > pkvdb.max_key_size) return error.InvalidPayload;
+            check.* = .{ .key = payload[position..end], .expected_lsn = expected_lsn };
+            position = end;
+        }
+        if (position > payload.len or payload.len - position < metadata_length) return error.InvalidPayload;
+        const metadata_end = position + metadata_length;
+        const metadata = payload[position..metadata_end];
+        const operations = try self.allocator.alloc(engine_mod.Operation, op_count);
+        defer self.allocator.free(operations);
+        position = metadata_end;
+        for (operations) |*operation| {
+            if (position > payload.len or payload.len - position < 12) return error.InvalidPayload;
+            const opcode: pkvdb.Opcode = std.meta.intToEnum(pkvdb.Opcode, payload[position]) catch return error.InvalidPayload;
+            const key_length = readInt(u32, payload, position + 4);
+            const value_length = readInt(u32, payload, position + 8);
+            position += 12;
+            const key_end = try std.math.add(usize, position, key_length);
+            const value_end = try std.math.add(usize, key_end, value_length);
+            if (value_end > payload.len or key_length > pkvdb.max_key_size or value_length > pkvdb.max_value_size or (opcode == .delete and value_length != 0)) return error.InvalidPayload;
+            operation.* = .{ .opcode = opcode, .key = payload[position..key_end], .value = payload[key_end..value_end] };
+            position = value_end;
+        }
+        if (position != payload.len) return error.InvalidPayload;
+        const result = try engine.compareBatchWrite(checks, operations, metadata);
+        try body.append(self.allocator, @intFromBool(result.committed));
+        try body.appendNTimes(self.allocator, 0, 7);
+        try appendInt(u64, body, self.allocator, result.lsn);
+    }
+
     fn scanOpen(self: *Session, payload: []const u8, body: *std.ArrayListUnmanaged(u8)) !void {
         if (payload.len < 12 or self.scans.items.len >= 64) return error.InvalidPayload;
         const include_values = payload[0] != 0;
@@ -340,3 +388,48 @@ test "PKBFI binary point batch and streaming scan" {
     const next_frame = try parse(next_response);
     try std.testing.expectEqual(@as(u32, 1), readInt(u32, next_frame.payload, 6));
 }
+
+test "PKBFI compare batch write commit and conflict" {
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    var path_buffer: [std.fs.max_path_bytes]u8 = undefined;
+    const directory = try tmp.dir.realpath(".", &path_buffer);
+    const path = try std.fmt.allocPrint(std.testing.allocator, "{s}/pkbfi-compare.pkvdb", .{directory});
+    defer std.testing.allocator.free(path);
+    var engine = try engine_mod.Engine.open(std.testing.allocator, path);
+    defer engine.close();
+    var session = Session.init(std.testing.allocator);
+    defer session.deinit();
+
+    var payload = [_]u8{0} ** 47;
+    writeInt(u32, &payload, 0, 1);
+    writeInt(u32, &payload, 4, 1);
+    writeInt(u32, &payload, 8, 0);
+    var position: usize = 16;
+    writeInt(u32, &payload, position, 1);
+    writeInt(u64, &payload, position + 8, 0);
+    position += 16;
+    payload[position] = 'k';
+    position += 1;
+    payload[position] = @intFromEnum(pkvdb.Opcode.put);
+    writeInt(u32, &payload, position + 4, 1);
+    writeInt(u32, &payload, position + 8, 1);
+    position += 12;
+    payload[position] = 'k';
+    position += 1;
+    payload[position] = 'v';
+
+    const response = try session.execute(&engine, .{ .opcode = .compare_batch_write, .flags = 0, .request_id = 1, .payload = &payload, .consumed = 0 });
+    defer std.testing.allocator.free(response);
+    const frame = try parse(response);
+    try std.testing.expectEqual(@as(u16, 0), readInt(u16, frame.payload, 0));
+    try std.testing.expectEqual(@as(u8, 1), frame.payload[2]);
+    try std.testing.expectEqual(@as(u64, 1), readInt(u64, frame.payload, 10));
+
+    const conflict_response = try session.execute(&engine, .{ .opcode = .compare_batch_write, .flags = 0, .request_id = 2, .payload = &payload, .consumed = 0 });
+    defer std.testing.allocator.free(conflict_response);
+    const conflict_frame = try parse(conflict_response);
+    try std.testing.expectEqual(@as(u16, 0), readInt(u16, conflict_frame.payload, 0));
+    try std.testing.expectEqual(@as(u8, 0), conflict_frame.payload[2]);
+    try std.testing.expectEqual(@as(u64, 0), readInt(u64, conflict_frame.payload, 10));
+}