diff --git a/src/commands.zig b/src/commands.zig index 292aab8..5d3f6ea 100644 --- a/src/commands.zig +++ b/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); +}