commands: add concurrent insert/find stress test; fix lock defer scoping in dispatch

The original dispatch restructure placed the unlock defers inside switch
prongs, where Zig runs them when the prong block exits — immediately after
acquisition. All commands therefore ran with no lock at all, which the new
stress test caught deterministically (two writers inside Engine.insert at
once). Lock acquisition now happens in the prongs and the command body runs
via dispatch_impl, so the defers (still prong-scoped) release the lock only
after the command finishes.

src/commands.zig

@@ -59,17 +59,25 @@ fn command_kind(name: []const u8) CommandKind {
 
 pub fn dispatch(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
     switch (command_kind(msg.command_name())) {
-        .none => {},
+        .none => return dispatch_impl(ctx, msg, reply),
         .read => {
             try ctx.engine.lock_read();
+            // Defers are block-scoped: this one is registered in the prong
+            // block, so it runs when the prong exits — after dispatch_impl
+            // returns. The shared lock is thus held for the whole command.
             defer ctx.engine.unlock_read();
+            return dispatch_impl(ctx, msg, reply);
         },
         .write => {
             try ctx.engine.lock();
             defer ctx.engine.unlock();
+            return dispatch_impl(ctx, msg, reply);
         },
     }
+}
 
+/// Runs the command with the engine lock held (or lock-free for `.none`).
+fn dispatch_impl(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void {
     const name = msg.command_name();
     if (std.mem.eql(u8, name, "hello")) return cmd_hello(ctx, reply);
     if (std.mem.eql(u8, name, "isMaster") or std.mem.eql(u8, name, "ismaster")) return cmd_is_master(ctx, reply);
@@ -1292,3 +1300,101 @@ fn parse_fake_msg(name: []const u8, value: bson.Value, extra: []const bson.Pair)
     std.mem.writeInt(i32, msg.items[12..16], wire.op_code_msg, .little);
     return wire.Message.parse(testing.allocator, msg.items);
 }
+
+test "concurrent insert/find commands on a threaded Io" {
+    // Exercises dispatch's lock classification end-to-end: writer fibers run
+    // `insert` under the exclusive lock, reader fibers run `count`/`find`
+    // under the shared lock. Every committed insert must be visible once all
+    // writers finish, and no reader may observe more docs than can exist.
+    var threaded: std.Io.Threaded = .init(testing.allocator, .{});
+    defer threaded.deinit();
+    const io = threaded.io();
+    var gen = bson.ObjectIdGen.init(io);
+
+    var tmp = std.testing.tmpDir(.{});
+    defer tmp.cleanup();
+    const path = try std.fmt.allocPrint(testing.allocator, ".zig-cache/tmp/{s}/conc.log", .{tmp.sub_path});
+    defer testing.allocator.free(path);
+    var engine = try db.Engine.open(testing.allocator, io, path);
+    defer engine.deinit();
+    const boot_ts = std.Io.Timestamp.now(io, .real);
+
+    const writers = 4;
+    const readers = 4;
+    const per_writer: i32 = 150;
+    const total: i32 = writers * per_writer;
+    var next_id = std.atomic.Value(i32).init(1);
+    var remaining = std.atomic.Value(usize).init(@intCast(total));
+
+    const Worker = struct {
+        fn writer(iow: std.Io, eng: *db.Engine, id_counter: *std.atomic.Value(i32), pending: *std.atomic.Value(usize), fiber_id: u32, total_writes: i32, start_ts: std.Io.Timestamp) error{Canceled}!void {
+            var wgen = bson.ObjectIdGen.init(iow);
+            var ctx = Context{
+                .gpa = testing.allocator,
+                .io = iow,
+                .oid_gen = &wgen,
+                .connection_id = fiber_id,
+                .client_desc = "test",
+                .engine = eng,
+                .server_start = start_ts,
+            };
+            while (true) {
+                const id = id_counter.fetchAdd(1, .monotonic);
+                if (id > total_writes) return;
+                const docs = [_]bson.Value{.{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = id } }} }};
+                var msg = parse_fake_msg("insert", .{ .string = "users" }, &.{
+                    .{ .key = "documents", .value = .{ .array = &docs } },
+                }) catch return error.Canceled;
+                defer msg.deinit();
+                var reply = wire.Reply.init(testing.allocator);
+                defer reply.deinit();
+                dispatch(&ctx, &msg, &reply) catch return error.Canceled;
+                _ = pending.fetchSub(1, .monotonic);
+            }
+        }
+
+        fn reader(iow: std.Io, eng: *db.Engine, pending: *std.atomic.Value(usize), fiber_id: u32, total_writes: i32, start_ts: std.Io.Timestamp) error{Canceled}!void {
+            var rgen = bson.ObjectIdGen.init(iow);
+            var ctx = Context{
+                .gpa = testing.allocator,
+                .io = iow,
+                .oid_gen = &rgen,
+                .connection_id = fiber_id,
+                .client_desc = "test",
+                .engine = eng,
+                .server_start = start_ts,
+            };
+            while (pending.load(.acquire) > 0) {
+                var msg = parse_fake_msg("count", .{ .string = "users" }, &.{}) catch return error.Canceled;
+                defer msg.deinit();
+                var reply = wire.Reply.init(testing.allocator);
+                defer reply.deinit();
+                dispatch(&ctx, &msg, &reply) catch return error.Canceled;
+                const n = bson.get_pair(reply.pairs.items, "n") orelse return error.Canceled;
+                if (n.int32 > total_writes) return error.Canceled;
+            }
+        }
+    };
+
+    var group: std.Io.Group = .init;
+    defer group.cancel(io);
+    for (0..readers) |i| group.async(io, Worker.reader, .{ io, &engine, &remaining, @intCast(i + 1), total, boot_ts });
+    for (0..writers) |i| group.async(io, Worker.writer, .{ io, &engine, &next_id, &remaining, @intCast(i + 1), total, boot_ts });
+    try group.await(io);
+
+    var ctx = Context{
+        .gpa = testing.allocator,
+        .io = io,
+        .oid_gen = &gen,
+        .connection_id = 0,
+        .client_desc = "test",
+        .engine = &engine,
+        .server_start = boot_ts,
+    };
+    var msg = try parse_fake_msg("count", .{ .string = "users" }, &.{});
+    defer msg.deinit();
+    var reply = wire.Reply.init(testing.allocator);
+    defer reply.deinit();
+    try dispatch(&ctx, &msg, &reply);
+    try testing.expectEqual(total, bson.get_pair(reply.pairs.items, "n").?.int32);
+}