From a38ddc2f50e196e2c533868955b43276d713d0e2 Mon Sep 17 00:00:00 2001 From: Aleksey Shakhmatov Date: Sun, 2 Aug 2026 12:21:22 +0300 Subject: [PATCH] =?UTF-8?q?index:=20secondary=20index=20core=20=E2=80=94?= =?UTF-8?q?=20entries,=20search,=20planner,=20=5Fid=20fast=20path?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds src/index.zig with the full secondary-index machinery: entry generation mirroring field_matches (array value + elements), BSON-order sorted entries with binary search, compound prefix and range lookups, unique/sparse options, the query planner (longest equality/$in run + optional range, $in cartesian cap, sparse/null bail), and the _id_ fast path guarded against serialization-ambiguous values (numbers, strings, symbols, codes, opaque payloads). query.collect_values is now pub so entry generation can mirror it exactly. storage.zig gains record_type_index_create/drop; lib.zig exports index. --- src/bson.zig | 123 ++--- src/commands.zig | 859 +++++++++++++-------------------- src/db.zig | 193 ++++---- src/index.zig | 1190 ++++++++++++++++++++++++++++++++++++++++++++++ src/lib.zig | 2 + src/query.zig | 57 +-- src/server.zig | 19 +- src/storage.zig | 32 +- src/update.zig | 45 +- src/wire.zig | 86 +++- 10 files changed, 1824 insertions(+), 782 deletions(-) create mode 100644 src/index.zig diff --git a/src/bson.zig b/src/bson.zig index 2dd3828..14cedac 100644 --- a/src/bson.zig +++ b/src/bson.zig @@ -128,14 +128,17 @@ pub const Document = struct { }; pub fn get_pair(pairs: []const Pair, key: []const u8) ?Value { - for (pairs) |p| { - if (std.mem.eql(u8, p.key, key)) return p.value; - } - return null; + const i = get_pair_index(pairs, key) orelse return null; + return pairs[i].value; } -pub fn parse_doc(allocator: std.mem.Allocator, bytes: []const u8) !Document { - return Document.parse(allocator, bytes); +/// Position of `key` in `pairs` — for callers that need to replace or remove +/// the pair in place rather than just read it. +pub fn get_pair_index(pairs: []const Pair, key: []const u8) ?usize { + for (pairs, 0..) |p, i| { + if (std.mem.eql(u8, p.key, key)) return i; + } + return null; } // --------------------------------------------------------------------------- @@ -145,10 +148,6 @@ pub fn parse_doc(allocator: std.mem.Allocator, bytes: []const u8) !Document { const Parser = struct { arena: *std.heap.ArenaAllocator, bytes: []const u8, - - fn fail() error{InvalidBson} { - return error.InvalidBson; - } }; const ParseError = error{ InvalidBson, OutOfMemory }; @@ -162,14 +161,21 @@ fn ensure_available(bytes: []const u8, idx: usize, n: usize) error{InvalidBson}! if (bytes.len -| idx < n) return error.InvalidBson; } +/// Validate the length-prefixed frame starting at `start` — documents and +/// arrays share it — and return the offset just past its terminating byte. +fn doc_extent(p: Parser, start: usize) ParseError!usize { + try ensure_available(p.bytes, start, 4); + const total: u32 = std.mem.readInt(u32, p.bytes[start..][0..4], .little); + if (total < 5) return error.InvalidBson; + if (p.bytes.len - start < total) return error.InvalidBson; + const end = start + total; + if (p.bytes[end - 1] != 0x00) return error.InvalidBson; + return end; +} + fn parse_doc_inner(p: Parser, idx: *usize) ParseError![]const Pair { const start = idx.*; - ensure_available(p.bytes, start, 4) catch return Parser.fail(); - const total: u32 = std.mem.readInt(u32, p.bytes[start..][0..4], .little); - if (total < 5) return Parser.fail(); - if (p.bytes.len - start < total) return Parser.fail(); - const end = start + total; - if (p.bytes[end - 1] != 0x00) return Parser.fail(); + const end = try doc_extent(p, start); const gpa = p.arena.allocator(); var pairs: std.ArrayListUnmanaged(Pair) = .empty; @@ -180,13 +186,13 @@ fn parse_doc_inner(p: Parser, idx: *usize) ParseError![]const Pair { const value = try parse_element(p, idx); try pairs.append(gpa, value); } - if (idx.* != end - 1) return Parser.fail(); + if (idx.* != end - 1) return error.InvalidBson; idx.* = end; return pairs.toOwnedSlice(gpa); } fn parse_element(p: Parser, idx: *usize) ParseError!Pair { - ensure_available(p.bytes, idx.*, 1) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 1); const tag = p.bytes[idx.*]; idx.* += 1; const key = try parse_cstring(p, idx); @@ -197,7 +203,7 @@ fn parse_element(p: Parser, idx: *usize) ParseError!Pair { fn parse_cstring(p: Parser, idx: *usize) ParseError![]const u8 { const start = idx.*; while (idx.* < p.bytes.len and p.bytes[idx.*] != 0) idx.* += 1; - if (idx.* >= p.bytes.len) return Parser.fail(); + if (idx.* >= p.bytes.len) return error.InvalidBson; idx.* += 1; // Strings are copied into the arena so documents are self-contained and // outlive the input buffer (wire messages and log records are transient). @@ -205,11 +211,11 @@ fn parse_cstring(p: Parser, idx: *usize) ParseError![]const u8 { } fn parse_string(p: Parser, idx: *usize) ParseError![]const u8 { - ensure_available(p.bytes, idx.*, 4) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 4); const len: u32 = std.mem.readInt(u32, p.bytes[idx.*..][0..4], .little); - if (len == 0 or len > p.bytes.len - (idx.* + 4)) return Parser.fail(); + if (len == 0 or len > p.bytes.len - (idx.* + 4)) return error.InvalidBson; const str = p.bytes[idx.* + 4 .. idx.* + 4 + len]; - if (str[len - 1] != 0) return Parser.fail(); + if (str[len - 1] != 0) return error.InvalidBson; idx.* += 4 + len; return p.arena.allocator().dupe(u8, str[0 .. len - 1]); } @@ -217,7 +223,7 @@ fn parse_string(p: Parser, idx: *usize) ParseError![]const u8 { fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value { return switch (tag) { 0x01 => blk: { - ensure_available(p.bytes, idx.*, 8) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 8); const v: f64 = @bitCast(std.mem.readInt(u64, p.bytes[idx.*..][0..8], .little)); idx.* += 8; break :blk .{ .double = v }; @@ -226,30 +232,30 @@ fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value { 0x03 => .{ .doc = try parse_doc_inner(p, idx) }, 0x04 => .{ .array = try parse_array(p, idx) }, 0x05 => blk: { - ensure_available(p.bytes, idx.*, 5) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 5); const len: u32 = std.mem.readInt(u32, p.bytes[idx.*..][0..4], .little); const subtype = p.bytes[idx.* + 4]; - ensure_available(p.bytes, idx.* + 5, len) catch return Parser.fail(); + try ensure_available(p.bytes, idx.* + 5, len); const data = p.bytes[idx.* + 5 .. idx.* + 5 + len]; idx.* += 5 + len; break :blk .{ .binary = .{ .subtype = subtype, .data = try p.arena.allocator().dupe(u8, data) } }; }, 0x06 => .{ .opaque_val = .{ .kind = 0x06, .data = &.{} } }, 0x07 => blk: { - ensure_available(p.bytes, idx.*, 12) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 12); const oid: ObjectId = p.bytes[idx.*..][0..12].*; idx.* += 12; break :blk .{ .object_id = oid }; }, 0x08 => blk: { - ensure_available(p.bytes, idx.*, 1) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 1); const v = p.bytes[idx.*]; - if (v > 1) return Parser.fail(); + if (v > 1) return error.InvalidBson; idx.* += 1; break :blk .{ .bool = v == 1 }; }, 0x09 => blk: { - ensure_available(p.bytes, idx.*, 8) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 8); const v: i64 = std.mem.readInt(i64, p.bytes[idx.*..][0..8], .little); idx.* += 8; break :blk .{ .datetime = v }; @@ -262,7 +268,7 @@ fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value { 0x0C => blk: { const start = idx.*; _ = try parse_string(p, idx); - ensure_available(p.bytes, idx.*, 12) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 12); idx.* += 12; // Like all other types, the payload is copied into the arena so // documents stay valid after the input buffer is reused. @@ -271,49 +277,44 @@ fn parse_value(p: Parser, tag: u8, idx: *usize) ParseError!Value { 0x0D => .{ .code = try parse_string(p, idx) }, 0x0E => .{ .symbol = try parse_string(p, idx) }, 0x0F => blk: { - ensure_available(p.bytes, idx.*, 4) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 4); const total: u32 = std.mem.readInt(u32, p.bytes[idx.*..][0..4], .little); - if (total < 4 or total > p.bytes.len - idx.*) return Parser.fail(); + if (total < 4 or total > p.bytes.len - idx.*) return error.InvalidBson; const data = p.bytes[idx.* .. idx.* + total]; idx.* += total; break :blk .{ .opaque_val = .{ .kind = 0x0F, .data = try p.arena.allocator().dupe(u8, data) } }; }, 0x10 => blk: { - ensure_available(p.bytes, idx.*, 4) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 4); const v: i32 = std.mem.readInt(i32, p.bytes[idx.*..][0..4], .little); idx.* += 4; break :blk .{ .int32 = v }; }, 0x11 => blk: { - ensure_available(p.bytes, idx.*, 8) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 8); const v: u64 = std.mem.readInt(u64, p.bytes[idx.*..][0..8], .little); idx.* += 8; break :blk .{ .timestamp = v }; }, 0x12 => blk: { - ensure_available(p.bytes, idx.*, 8) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 8); const v: i64 = std.mem.readInt(i64, p.bytes[idx.*..][0..8], .little); idx.* += 8; break :blk .{ .int64 = v }; }, 0x13 => blk: { - ensure_available(p.bytes, idx.*, 16) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 16); const v: [16]u8 = p.bytes[idx.*..][0..16].*; idx.* += 16; break :blk .{ .decimal128 = v }; }, - else => return Parser.fail(), + else => return error.InvalidBson, }; } fn parse_array(p: Parser, idx: *usize) ParseError![]const Value { const start = idx.*; - ensure_available(p.bytes, start, 4) catch return Parser.fail(); - const total: u32 = std.mem.readInt(u32, p.bytes[start..][0..4], .little); - if (total < 5) return Parser.fail(); - if (p.bytes.len - start < total) return Parser.fail(); - const end = start + total; - if (p.bytes[end - 1] != 0x00) return Parser.fail(); + const end = try doc_extent(p, start); const gpa = p.arena.allocator(); var values: std.ArrayListUnmanaged(Value) = .empty; @@ -321,13 +322,13 @@ fn parse_array(p: Parser, idx: *usize) ParseError![]const Value { idx.* = start + 4; while (idx.* < end - 1) { - ensure_available(p.bytes, idx.*, 1) catch return Parser.fail(); + try ensure_available(p.bytes, idx.*, 1); const tag = p.bytes[idx.*]; idx.* += 1; _ = try parse_cstring(p, idx); try values.append(gpa, try parse_value(p, tag, idx)); } - if (idx.* != end - 1) return Parser.fail(); + if (idx.* != end - 1) return error.InvalidBson; idx.* = end; return values.toOwnedSlice(gpa); } @@ -417,22 +418,31 @@ pub fn write_element(pair: Pair, gpa: std.mem.Allocator, out: *std.ArrayListUnma try write_value(pair.value, gpa, out); } -/// Write a length-prefixed document. Length is patched in after the body. -pub fn write_doc(pairs: []const Pair, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) SerializeError!void { +/// Reserve the 4-byte length prefix of a document or array frame. Returns the +/// offset to hand back to `end_frame`. +fn begin_frame(gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) SerializeError!usize { const len_pos = out.items.len; - var zero: [4]u8 = [4]u8{ 0, 0, 0, 0 }; - try out.appendSlice(gpa, &zero); - for (pairs) |p| try write_element(p, gpa, out); + try out.appendSlice(gpa, &[4]u8{ 0, 0, 0, 0 }); + return len_pos; +} + +/// Terminate the frame opened at `len_pos` and patch in its total length. +fn end_frame(gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8), len_pos: usize) SerializeError!void { try out.append(gpa, 0); const total = out.items.len - len_pos; if (total > std.math.maxInt(u32)) return error.BsonTooLarge; std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little); } -fn write_array(items: []const Value, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) !void { - const len_pos = out.items.len; - var zero: [4]u8 = [4]u8{ 0, 0, 0, 0 }; - try out.appendSlice(gpa, &zero); +/// Write a length-prefixed document. Length is patched in after the body. +pub fn write_doc(pairs: []const Pair, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) SerializeError!void { + const len_pos = try begin_frame(gpa, out); + for (pairs) |p| try write_element(p, gpa, out); + try end_frame(gpa, out, len_pos); +} + +fn write_array(items: []const Value, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) SerializeError!void { + const len_pos = try begin_frame(gpa, out); var buf: [16]u8 = undefined; for (items, 0..) |item, i| { try out.append(gpa, item.type_tag()); @@ -440,10 +450,7 @@ fn write_array(items: []const Value, gpa: std.mem.Allocator, out: *std.ArrayList try write_cstring(key, gpa, out); try write_value(item, gpa, out); } - try out.append(gpa, 0); - const total = out.items.len - len_pos; - if (total > std.math.maxInt(u32)) return error.BsonTooLarge; - std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little); + try end_frame(gpa, out, len_pos); } /// Serialize a single value with its type byte (no key) — used for `_id` diff --git a/src/commands.zig b/src/commands.zig index 5d3f6ea..90a6f42 100644 --- a/src/commands.zig +++ b/src/commands.zig @@ -38,82 +38,74 @@ pub const ErrorCode = enum(i32) { /// `.none` commands must not touch the engine at all. const CommandKind = enum { none, read, write }; -fn command_kind(name: []const u8) CommandKind { +const Command = struct { + name: []const u8, + kind: CommandKind, + handler: *const fn (*Context, *wire.Message, *wire.Reply) anyerror!void, +}; + +/// The one place a command exists: name, lock class, and handler declared +/// together so a new command cannot be given a handler but no lock. +const command_table = [_]Command{ + // Handshake, admin info, no-ops: never touch the engine. + .{ .name = "hello", .kind = .none, .handler = cmd_hello }, + .{ .name = "isMaster", .kind = .none, .handler = cmd_is_master }, + .{ .name = "ismaster", .kind = .none, .handler = cmd_is_master }, + .{ .name = "ping", .kind = .none, .handler = cmd_ping }, + .{ .name = "buildInfo", .kind = .none, .handler = cmd_build_info }, + .{ .name = "getParameter", .kind = .none, .handler = cmd_get_parameter }, + .{ .name = "whatsmyuri", .kind = .none, .handler = cmd_whatsmyuri }, + .{ .name = "hostInfo", .kind = .none, .handler = cmd_host_info }, + .{ .name = "getCmdLineOpts", .kind = .none, .handler = cmd_get_cmd_line_opts }, + .{ .name = "serverStatus", .kind = .none, .handler = cmd_server_status }, + .{ .name = "endSessions", .kind = .none, .handler = cmd_end_sessions }, + .{ .name = "connectionStatus", .kind = .none, .handler = cmd_connection_status }, + .{ .name = "getMore", .kind = .none, .handler = cmd_get_more }, + .{ .name = "killCursors", .kind = .none, .handler = cmd_kill_cursors }, // Read-only: scan the engine without mutating it. - if (std.mem.eql(u8, name, "find") or - std.mem.eql(u8, name, "count") or - std.mem.eql(u8, name, "aggregate") or - std.mem.eql(u8, name, "listDatabases") or - std.mem.eql(u8, name, "listCollections")) return .read; + .{ .name = "find", .kind = .read, .handler = cmd_find }, + .{ .name = "count", .kind = .read, .handler = cmd_count }, + .{ .name = "aggregate", .kind = .read, .handler = cmd_aggregate }, + .{ .name = "listDatabases", .kind = .read, .handler = cmd_list_databases }, + .{ .name = "listCollections", .kind = .read, .handler = cmd_list_collections }, // Writes: exclusive, totally ordered. - if (std.mem.eql(u8, name, "create") or - std.mem.eql(u8, name, "drop") or - std.mem.eql(u8, name, "dropDatabase") or - std.mem.eql(u8, name, "insert") or - std.mem.eql(u8, name, "update") or - std.mem.eql(u8, name, "delete") or - std.mem.eql(u8, name, "findAndModify")) return .write; - // Everything else (handshake, admin info, no-ops) needs no lock. - return .none; -} + .{ .name = "create", .kind = .write, .handler = cmd_create }, + .{ .name = "drop", .kind = .write, .handler = cmd_drop }, + .{ .name = "dropDatabase", .kind = .write, .handler = cmd_drop_database }, + .{ .name = "insert", .kind = .write, .handler = cmd_insert }, + .{ .name = "update", .kind = .write, .handler = cmd_update }, + .{ .name = "delete", .kind = .write, .handler = cmd_delete }, + .{ .name = "findAndModify", .kind = .write, .handler = cmd_find_and_modify }, +}; pub fn dispatch(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { - switch (command_kind(msg.command_name())) { - .none => return dispatch_impl(ctx, msg, reply), + const name = msg.command_name(); + const cmd = for (&command_table) |*c| { + if (std.mem.eql(u8, c.name, name)) break c; + } else { + var buf: [256]u8 = undefined; + const errmsg = try std.fmt.bufPrint(&buf, "no such command: '{s}'", .{name}); + return reply.put_error(@intFromEnum(ErrorCode.command_not_found), "CommandNotFound", errmsg); + }; + + switch (cmd.kind) { + .none => return cmd.handler(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 + // block, so it runs when the prong exits — after the handler // returns. The shared lock is thus held for the whole command. defer ctx.engine.unlock_read(); - return dispatch_impl(ctx, msg, reply); + return cmd.handler(ctx, msg, reply); }, .write => { try ctx.engine.lock(); defer ctx.engine.unlock(); - return dispatch_impl(ctx, msg, reply); + return cmd.handler(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); - if (std.mem.eql(u8, name, "ping")) return cmd_ping(reply); - if (std.mem.eql(u8, name, "buildInfo")) return cmd_build_info(reply); - if (std.mem.eql(u8, name, "getParameter")) return cmd_get_parameter(ctx, msg, reply); - if (std.mem.eql(u8, name, "whatsmyuri")) return cmd_whatsmyuri(ctx, reply); - if (std.mem.eql(u8, name, "hostInfo")) return cmd_host_info(ctx, reply); - if (std.mem.eql(u8, name, "getCmdLineOpts")) return cmd_get_cmd_line_opts(reply); - if (std.mem.eql(u8, name, "serverStatus")) return cmd_server_status(ctx, reply); - if (std.mem.eql(u8, name, "endSessions")) return cmd_end_sessions(reply); - if (std.mem.eql(u8, name, "connectionStatus")) return cmd_connection_status(reply); - if (std.mem.eql(u8, name, "listDatabases")) return cmd_list_databases(ctx, reply); - if (std.mem.eql(u8, name, "listCollections")) return cmd_list_collections(ctx, msg, reply); - if (std.mem.eql(u8, name, "create")) return cmd_create(ctx, msg, reply); - if (std.mem.eql(u8, name, "drop")) return cmd_drop(ctx, msg, reply); - if (std.mem.eql(u8, name, "dropDatabase")) return cmd_drop_database(ctx, msg, reply); - if (std.mem.eql(u8, name, "insert")) return cmd_insert(ctx, msg, reply); - if (std.mem.eql(u8, name, "find")) return cmd_find(ctx, msg, reply); - if (std.mem.eql(u8, name, "update")) return cmd_update(ctx, msg, reply); - if (std.mem.eql(u8, name, "delete")) return cmd_delete(ctx, msg, reply); - if (std.mem.eql(u8, name, "findAndModify")) return cmd_find_and_modify(ctx, msg, reply); - if (std.mem.eql(u8, name, "count")) return cmd_count(ctx, msg, reply); - if (std.mem.eql(u8, name, "aggregate")) return cmd_aggregate(ctx, msg, reply); - if (std.mem.eql(u8, name, "getMore")) return cmd_get_more(reply); - if (std.mem.eql(u8, name, "killCursors")) return cmd_kill_cursors(reply); - - var buf: [256]u8 = undefined; - const errmsg = try std.fmt.bufPrint(&buf, "no such command: '{s}'", .{name}); - return reply.put_error(@intFromEnum(ErrorCode.command_not_found), "CommandNotFound", errmsg); -} - -pub fn reply_error(reply: *wire.Reply, code: i32, code_name: []const u8, message: []const u8) !void { - try reply.put_error(code, code_name, message); -} - // --------------------------------------------------------------------------- // Handshake / administration // --------------------------------------------------------------------------- @@ -133,28 +125,28 @@ fn add_server_info(ctx: *Context, reply: *wire.Reply) !void { const oid = ctx.oid_gen.new(ctx.io); const tv = try reply.arena_alloc().alloc(bson.Pair, 2); - tv[0] = .{ .key = try reply.arena_alloc().dupe(u8, "processId"), .value = .{ .object_id = oid } }; - tv[1] = .{ .key = try reply.arena_alloc().dupe(u8, "counter"), .value = .{ .int64 = 0 } }; + tv[0] = .{ .key = "processId", .value = .{ .object_id = oid } }; + tv[1] = .{ .key = "counter", .value = .{ .int64 = 0 } }; try reply.put("topologyVersion", .{ .doc = tv }); } -fn cmd_hello(ctx: *Context, reply: *wire.Reply) !void { +fn cmd_hello(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void { try add_server_info(ctx, reply); try reply.put_ok(); } -fn cmd_is_master(ctx: *Context, reply: *wire.Reply) !void { +fn cmd_is_master(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void { try add_server_info(ctx, reply); try reply.put("ismaster", .{ .bool = true }); try reply.put("helloOk", .{ .bool = true }); try reply.put_ok(); } -fn cmd_ping(reply: *wire.Reply) !void { +fn cmd_ping(_: *Context, _: *wire.Message, reply: *wire.Reply) !void { try reply.put_ok(); } -fn cmd_build_info(reply: *wire.Reply) !void { +fn cmd_build_info(_: *Context, _: *wire.Message, reply: *wire.Reply) !void { try reply.put("version", .{ .string = "4.4.0" }); try reply.put("gitVersion", .{ .string = "mongo-light" }); try reply.put("versionArray", .{ .array = try int_array(reply, &.{ 4, 4, 0, 0 }) }); @@ -170,79 +162,57 @@ fn cmd_build_info(reply: *wire.Reply) !void { try reply.put_ok(); } -fn cmd_get_parameter(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { - _ = ctx; +fn cmd_get_parameter(_: *Context, msg: *wire.Message, reply: *wire.Reply) !void { // mongosh probes featureCompatibilityVersion; respond per requested key. - const params_value = msg.body.get("getParameter") orelse return reply.put_error( - @intFromEnum(ErrorCode.invalid_argument), - "InvalidArgument", - "getParameter requires a document", - ); - const params = switch (params_value) { - .doc => |d| d, - else => return reply.put_error( - @intFromEnum(ErrorCode.invalid_argument), - "InvalidArgument", - "getParameter requires a document", - ), - }; - var found: bool = false; - for (params) |p| { - if (std.mem.eql(u8, p.key, "featureCompatibilityVersion")) { - const fcv = try reply.arena_alloc().alloc(bson.Pair, 1); - fcv[0] = .{ .key = try reply.arena_alloc().dupe(u8, "version"), .value = .{ .string = "4.4" } }; - try reply.put("featureCompatibilityVersion", .{ .doc = fcv }); - found = true; - } - } - if (!found) { - return reply.put_error( - @intFromEnum(ErrorCode.invalid_argument), - "InvalidArgument", - "no option found to get", - ); + const params = doc_arg(msg.body.get("getParameter")) orelse + return invalid_arg(reply, "getParameter requires a document"); + if (bson.get_pair(params, "featureCompatibilityVersion") == null) { + return invalid_arg(reply, "no option found to get"); } + const fcv = try reply.arena_alloc().alloc(bson.Pair, 1); + fcv[0] = .{ .key = "version", .value = .{ .string = "4.4" } }; + try reply.put("featureCompatibilityVersion", .{ .doc = fcv }); try reply.put_ok(); } -fn cmd_whatsmyuri(ctx: *Context, reply: *wire.Reply) !void { +fn cmd_whatsmyuri(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void { try reply.put("you", .{ .string = ctx.client_desc }); try reply.put_ok(); } -fn cmd_host_info(ctx: *Context, reply: *wire.Reply) !void { +fn cmd_host_info(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void { const now = std.Io.Timestamp.now(ctx.io, .real); const cpu_count: usize = std.Thread.getCpuCount() catch 1; const system = try reply.arena_alloc().alloc(bson.Pair, 6); - system[0] = .{ .key = try reply.arena_alloc().dupe(u8, "currentTime"), .value = .{ .datetime = now.toMilliseconds() } }; - system[1] = .{ .key = try reply.arena_alloc().dupe(u8, "hostname"), .value = .{ .string = "localhost" } }; - system[2] = .{ .key = try reply.arena_alloc().dupe(u8, "cpuAddrSize"), .value = .{ .int32 = 64 } }; - system[3] = .{ .key = try reply.arena_alloc().dupe(u8, "memSizeMB"), .value = .{ .int32 = 0 } }; - system[4] = .{ .key = try reply.arena_alloc().dupe(u8, "numCores"), .value = .{ .int32 = @intCast(cpu_count) } }; - system[5] = .{ .key = try reply.arena_alloc().dupe(u8, "cpuArch"), .value = .{ .string = @tagName(builtin.cpu.arch) } }; + system[0] = .{ .key = "currentTime", .value = .{ .datetime = now.toMilliseconds() } }; + system[1] = .{ .key = "hostname", .value = .{ .string = "localhost" } }; + system[2] = .{ .key = "cpuAddrSize", .value = .{ .int32 = 64 } }; + system[3] = .{ .key = "memSizeMB", .value = .{ .int32 = 0 } }; + system[4] = .{ .key = "numCores", .value = .{ .int32 = @intCast(cpu_count) } }; + system[5] = .{ .key = "cpuArch", .value = .{ .string = @tagName(builtin.cpu.arch) } }; try reply.put("system", .{ .doc = system }); const os = try reply.arena_alloc().alloc(bson.Pair, 3); - os[0] = .{ .key = try reply.arena_alloc().dupe(u8, "type"), .value = .{ .string = @tagName(builtin.os.tag) } }; - os[1] = .{ .key = try reply.arena_alloc().dupe(u8, "name"), .value = .{ .string = @tagName(builtin.os.tag) } }; - os[2] = .{ .key = try reply.arena_alloc().dupe(u8, "version"), .value = .{ .string = "unknown" } }; + os[0] = .{ .key = "type", .value = .{ .string = @tagName(builtin.os.tag) } }; + os[1] = .{ .key = "name", .value = .{ .string = @tagName(builtin.os.tag) } }; + os[2] = .{ .key = "version", .value = .{ .string = "unknown" } }; try reply.put("os", .{ .doc = os }); try reply.put("extra", .{ .doc = &.{} }); try reply.put_ok(); } -fn cmd_get_cmd_line_opts(reply: *wire.Reply) !void { +fn cmd_get_cmd_line_opts(_: *Context, _: *wire.Message, reply: *wire.Reply) !void { const argv = try reply.arena_alloc().alloc(bson.Pair, 2); - argv[0] = .{ .key = try reply.arena_alloc().dupe(u8, "dbpath"), .value = .{ .string = "mongo-light.log" } }; - argv[1] = .{ .key = try reply.arena_alloc().dupe(u8, "port"), .value = .{ .int32 = 27017 } }; + argv[0] = .{ .key = "dbpath", .value = .{ .string = "mongo-light.log" } }; + argv[1] = .{ .key = "port", .value = .{ .int32 = 27017 } }; try reply.put("argv", .{ .array = &.{} }); try reply.put("parsed", .{ .doc = argv }); try reply.put_ok(); } -fn cmd_server_status(ctx: *Context, reply: *wire.Reply) !void { +fn cmd_server_status(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void { const now = std.Io.Timestamp.now(ctx.io, .real); const uptime: i64 = std.Io.Timestamp.durationTo(ctx.server_start, now).toSeconds(); @@ -252,24 +222,24 @@ fn cmd_server_status(ctx: *Context, reply: *wire.Reply) !void { try reply.put("uptime", .{ .double = @floatFromInt(uptime) }); try reply.put("localTime", .{ .datetime = now.toMilliseconds() }); const connections = try reply.arena_alloc().alloc(bson.Pair, 1); - connections[0] = .{ .key = try reply.arena_alloc().dupe(u8, "current"), .value = .{ .int32 = @intCast(ctx.connection_id) } }; + connections[0] = .{ .key = "current", .value = .{ .int32 = @intCast(ctx.connection_id) } }; try reply.put("connections", .{ .doc = connections }); try reply.put_ok(); } -fn cmd_end_sessions(reply: *wire.Reply) !void { +fn cmd_end_sessions(_: *Context, _: *wire.Message, reply: *wire.Reply) !void { try reply.put_ok(); } -fn cmd_connection_status(reply: *wire.Reply) !void { +fn cmd_connection_status(_: *Context, _: *wire.Message, reply: *wire.Reply) !void { const auth_info = try reply.arena_alloc().alloc(bson.Pair, 2); - auth_info[0] = .{ .key = try reply.arena_alloc().dupe(u8, "authenticatedUsers"), .value = .{ .array = &.{} } }; - auth_info[1] = .{ .key = try reply.arena_alloc().dupe(u8, "authenticatedUserRoles"), .value = .{ .array = &.{} } }; + auth_info[0] = .{ .key = "authenticatedUsers", .value = .{ .array = &.{} } }; + auth_info[1] = .{ .key = "authenticatedUserRoles", .value = .{ .array = &.{} } }; try reply.put("authInfo", .{ .doc = auth_info }); try reply.put_ok(); } -fn cmd_list_databases(ctx: *Context, reply: *wire.Reply) !void { +fn cmd_list_databases(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void { var names: std.ArrayListUnmanaged([]const u8) = .empty; defer names.deinit(ctx.gpa); try ctx.engine.database_names(&names); @@ -277,9 +247,9 @@ fn cmd_list_databases(ctx: *Context, reply: *wire.Reply) !void { const values = try reply.arena_alloc().alloc(bson.Value, names.items.len); for (names.items, 0..) |n, i| { const entry = try reply.arena_alloc().alloc(bson.Pair, 3); - entry[0] = .{ .key = try reply.arena_alloc().dupe(u8, "name"), .value = .{ .string = try reply.arena_alloc().dupe(u8, n) } }; - entry[1] = .{ .key = try reply.arena_alloc().dupe(u8, "sizeOnDisk"), .value = .{ .double = 0 } }; - entry[2] = .{ .key = try reply.arena_alloc().dupe(u8, "empty"), .value = .{ .bool = true } }; + entry[0] = .{ .key = "name", .value = .{ .string = try reply.arena_alloc().dupe(u8, n) } }; + entry[1] = .{ .key = "sizeOnDisk", .value = .{ .double = 0 } }; + entry[2] = .{ .key = "empty", .value = .{ .bool = true } }; values[i] = .{ .doc = entry }; } try reply.put("databases", .{ .array = values }); @@ -289,11 +259,7 @@ fn cmd_list_databases(ctx: *Context, reply: *wire.Reply) !void { } fn cmd_list_collections(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { - const db_name = msg.db_name() orelse return reply.put_error( - @intFromEnum(ErrorCode.invalid_argument), - "InvalidArgument", - "listCollections requires $db", - ); + const db_name = msg.db_name() orelse return invalid_arg(reply, "listCollections requires $db"); var names: std.ArrayListUnmanaged([]const u8) = .empty; defer names.deinit(ctx.gpa); try ctx.engine.collection_names(db_name, &names); @@ -301,64 +267,33 @@ fn cmd_list_collections(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) ! const values = try reply.arena_alloc().alloc(bson.Value, names.items.len); for (names.items, 0..) |n, i| { const entry = try reply.arena_alloc().alloc(bson.Pair, 3); - entry[0] = .{ .key = try reply.arena_alloc().dupe(u8, "name"), .value = .{ .string = try reply.arena_alloc().dupe(u8, n) } }; - entry[1] = .{ .key = try reply.arena_alloc().dupe(u8, "type"), .value = .{ .string = "collection" } }; - entry[2] = .{ .key = try reply.arena_alloc().dupe(u8, "options"), .value = .{ .doc = &.{} } }; + entry[0] = .{ .key = "name", .value = .{ .string = try reply.arena_alloc().dupe(u8, n) } }; + entry[1] = .{ .key = "type", .value = .{ .string = "collection" } }; + entry[2] = .{ .key = "options", .value = .{ .doc = &.{} } }; values[i] = .{ .doc = entry }; } - try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, ""), values) }); + try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, ""), "firstBatch", values) }); try reply.put_ok(); } fn cmd_create(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { - const db_name = msg.db_name() orelse return reply.put_error( - @intFromEnum(ErrorCode.invalid_argument), - "InvalidArgument", - "create requires $db", - ); - const coll_name = msg.body.get("create") orelse return reply.put_error( - @intFromEnum(ErrorCode.bad_value), - "BadValue", - "create requires a collection name", - ); - switch (coll_name) { - .string => |s| { - _ = try ctx.engine.get_or_create_collection(db_name, s); - try reply.put_ok(); - }, - else => return reply.put_error(@intFromEnum(ErrorCode.bad_value), "BadValue", "collection name must be a string"), - } + const db_name = msg.db_name() orelse return invalid_arg(reply, "create requires $db"); + const coll_name = str_arg(msg.body.get("create")) orelse return bad_value(reply, "create requires a collection name"); + _ = try ctx.engine.get_or_create_collection(db_name, coll_name); + try reply.put_ok(); } fn cmd_drop(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { - const db_name = msg.db_name() orelse return reply.put_error( - @intFromEnum(ErrorCode.invalid_argument), - "InvalidArgument", - "drop requires $db", - ); - const coll_name = msg.body.get("drop") orelse return reply.put_error( - @intFromEnum(ErrorCode.bad_value), - "BadValue", - "drop requires a collection name", - ); - switch (coll_name) { - .string => |s| { - const dropped = try ctx.engine.drop_collection(db_name, s); - if (!dropped) { - return reply.put_error(@intFromEnum(ErrorCode.namespace_not_found), "NamespaceNotFound", "ns not found"); - } - try reply.put_ok(); - }, - else => return reply.put_error(@intFromEnum(ErrorCode.bad_value), "BadValue", "collection name must be a string"), + const db_name = msg.db_name() orelse return invalid_arg(reply, "drop requires $db"); + const coll_name = str_arg(msg.body.get("drop")) orelse return bad_value(reply, "drop requires a collection name"); + if (!try ctx.engine.drop_collection(db_name, coll_name)) { + return reply.put_error(@intFromEnum(ErrorCode.namespace_not_found), "NamespaceNotFound", "ns not found"); } + try reply.put_ok(); } fn cmd_drop_database(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { - const db_name = msg.db_name() orelse return reply.put_error( - @intFromEnum(ErrorCode.invalid_argument), - "InvalidArgument", - "dropDatabase requires $db", - ); + const db_name = msg.db_name() orelse return invalid_arg(reply, "dropDatabase requires $db"); _ = try ctx.engine.drop_database(db_name); try reply.put("dropped", .{ .string = try reply.arena_alloc().dupe(u8, db_name) }); try reply.put_ok(); @@ -372,26 +307,7 @@ fn cmd_insert(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const db_name = msg.db_name() orelse return invalid_arg(reply, "insert requires $db"); const coll_name = str_arg(msg.body.get("insert")) orelse return bad_value(reply, "insert requires a collection name"); - var docs: []const bson.Document = &.{}; - for (msg.seqs) |seq| { - if (std.mem.eql(u8, seq.name, "documents")) docs = seq.docs; - } - if (docs.len == 0) { - const d = msg.body.get("documents") orelse return bad_value(reply, "insert requires documents"); - switch (d) { - .array => |arr| { - const parsed = try reply.arena_alloc().alloc(bson.Document, arr.len); - for (arr, 0..) |item, i| { - parsed[i] = switch (item) { - .doc => |pairs| bson.Document{ .arena = undefined, .pairs = pairs }, - else => return bad_value(reply, "documents must be documents"), - }; - } - docs = parsed; - }, - else => return bad_value(reply, "documents must be an array"), - } - } + const docs = try batch_arg(msg, reply, "insert", "documents") orelse return; var inserted: i64 = 0; var write_errors: std.ArrayListUnmanaged(bson.Value) = .empty; @@ -404,12 +320,9 @@ fn cmd_insert(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { switch (err) { error.DuplicateKey => { const e = try reply.arena_alloc().alloc(bson.Pair, 3); - e[0] = .{ .key = try reply.arena_alloc().dupe(u8, "index"), .value = .{ .int32 = @intCast(i) } }; - e[1] = .{ .key = try reply.arena_alloc().dupe(u8, "code"), .value = .{ .int32 = @intFromEnum(ErrorCode.duplicate_key) } }; - const id = doc.get("_id") orelse bson.Value.null; - const key = try serialize_value_compact(reply, id); - const msg_text = try std.fmt.allocPrint(reply.arena_alloc(), "E11000 duplicate key error collection: {s}.{s} index: _id_ dup key: {s}", .{ db_name, coll_name, key }); - e[2] = .{ .key = try reply.arena_alloc().dupe(u8, "errmsg"), .value = .{ .string = msg_text } }; + e[0] = .{ .key = "index", .value = .{ .int32 = @intCast(i) } }; + e[1] = .{ .key = "code", .value = .{ .int32 = @intFromEnum(ErrorCode.duplicate_key) } }; + e[2] = .{ .key = "errmsg", .value = .{ .string = try duplicate_key_message(reply, db_name, coll_name, doc) } }; try write_errors.append(reply.arena_alloc(), .{ .doc = e }); }, else => return err, @@ -433,49 +346,58 @@ fn cmd_find(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const coll_name = str_arg(msg.body.get("find")) orelse return bad_value(reply, "find requires a collection name"); const filter = doc_arg(msg.body.get("filter")) orelse return bad_value(reply, "filter must be a document"); - const filter_doc = bson.Document{ .arena = undefined, .pairs = filter }; const sort_keys = try parse_sort_keys(reply, msg.body.get("sort")); const proj_pairs = doc_arg(msg.body.get("projection")); const skip: u64 = int_arg(msg.body.get("skip")) orelse 0; - const limit_raw: i64 = switch (msg.body.get("limit") orelse bson.Value{ .int32 = 0 }) { - .int32 => |i| i, - .int64 => |i| i, - else => 0, - }; - const limit: usize = if (limit_raw < 0) @intCast(-limit_raw) else @intCast(limit_raw); + // A negative limit means "return this many in a single batch"; we always + // reply with one batch, so only the magnitude matters. + const limit: usize = @abs(int_value(msg.body.get("limit")) orelse 0); var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty; defer matched.deinit(ctx.gpa); - - if (ctx.engine.get_collection(db_name, coll_name)) |coll| { - var it = coll.docs.iterator(); - while (it.next()) |entry| { - if (try query.matches(ctx.gpa, &filter_doc, entry.value_ptr.*)) { - try matched.append(ctx.gpa, entry.value_ptr.*); - } - } - } + _ = try scan_matching(ctx, db_name, coll_name, filter, 0, &matched); if (sort_keys.len > 0) { - try query.sort_docs(ctx.gpa, reply.arena_alloc(), matched.items, sort_keys); - } - if (skip < matched.items.len) { - const start = matched.items[skip..]; - const end = if (limit > 0 and limit < start.len) start[0..limit] else start; - try emit_docs(reply, db_name, coll_name, proj_pairs, end); - } else { - try emit_docs(reply, db_name, coll_name, proj_pairs, &.{}); + try query.sort_docs(reply.arena_alloc(), matched.items, sort_keys); } + const rest = if (skip < matched.items.len) matched.items[skip..] else &.{}; + const page = if (limit > 0 and limit < rest.len) rest[0..limit] else rest; + try emit_docs(reply, db_name, coll_name, proj_pairs, page); try reply.put_ok(); } +/// Collect the documents in `db_name.coll_name` matching `filter`, stopping +/// after `limit` matches (0 = unlimited). Returns the number matched; `out` +/// may be null when only the count is wanted. Every command that scans a +/// collection goes through here, so an index would only need one call site. +fn scan_matching( + ctx: *Context, + db_name: []const u8, + coll_name: []const u8, + filter: []const bson.Pair, + limit: usize, + out: ?*std.ArrayListUnmanaged(*const bson.Document), +) !usize { + const coll = ctx.engine.get_collection(db_name, coll_name) orelse return 0; + const filter_doc = bson.Document{ .arena = undefined, .pairs = filter }; + var n: usize = 0; + var it = coll.docs.iterator(); + while (it.next()) |entry| { + if (!try query.matches(ctx.gpa, &filter_doc, entry.value_ptr.*)) continue; + if (out) |list| try list.append(ctx.gpa, entry.value_ptr.*); + n += 1; + if (limit != 0 and n >= limit) break; + } + return n; +} + fn emit_docs(reply: *wire.Reply, db_name: []const u8, coll_name: []const u8, proj_pairs: ?[]const bson.Pair, docs: []const *const bson.Document) !void { const values = try reply.arena_alloc().alloc(bson.Value, docs.len); for (docs, 0..) |d, i| { values[i] = try project_doc(reply, d, proj_pairs); } - try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, coll_name), values) }); + try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, coll_name), "firstBatch", values) }); } /// Project a stored doc (or deep-copy it) into the reply arena. @@ -493,25 +415,7 @@ fn cmd_update(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const db_name = msg.db_name() orelse return invalid_arg(reply, "update requires $db"); const coll_name = str_arg(msg.body.get("update")) orelse return bad_value(reply, "update requires a collection name"); - var specs: []const bson.Document = &.{}; - for (msg.seqs) |seq| { - if (std.mem.eql(u8, seq.name, "updates")) specs = seq.docs; - } - if (specs.len == 0) { - const d = msg.body.get("updates") orelse return bad_value(reply, "update requires updates"); - const arr = switch (d) { - .array => |a| a, - else => return bad_value(reply, "updates must be an array"), - }; - const parsed = try reply.arena_alloc().alloc(bson.Document, arr.len); - for (arr, 0..) |item, i| { - parsed[i] = switch (item) { - .doc => |pairs| bson.Document{ .arena = undefined, .pairs = pairs }, - else => return bad_value(reply, "updates must be documents"), - }; - } - specs = parsed; - } + const specs = try batch_arg(msg, reply, "update", "updates") orelse return; var n_matched: i64 = 0; var n_modified: i64 = 0; @@ -524,29 +428,21 @@ fn cmd_update(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const multi = bool_arg(spec.get("multi")) orelse false; const upsert = bool_arg(spec.get("upsert")) orelse false; - var matched: std.ArrayListUnmanaged(*bson.Document) = .empty; + var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty; defer matched.deinit(ctx.gpa); - if (ctx.engine.get_collection(db_name, coll_name)) |coll| { - var it = coll.docs.iterator(); - while (it.next()) |entry| { - if (try query.matches(ctx.gpa, &.{ .arena = undefined, .pairs = q }, entry.value_ptr.*)) { - try matched.append(ctx.gpa, entry.value_ptr.*); - if (!multi) break; - } - } - } + _ = try scan_matching(ctx, db_name, coll_name, q, if (multi) 0 else 1, &matched); if (matched.items.len == 0) { if (upsert) { - const new_doc = try build_upsert_doc(ctx, reply, q, u_doc); + const new_doc = try build_upsert_doc(reply, q, u_doc); ctx.engine.insert(db_name, coll_name, new_doc, ctx.oid_gen) catch |err| switch (err) { error.DuplicateKey => return duplicate_key_error(reply, db_name, coll_name, new_doc), else => return err, }; const id = new_doc.get("_id") orelse bson.Value.null; const u = try reply.arena_alloc().alloc(bson.Pair, 2); - u[0] = .{ .key = try reply.arena_alloc().dupe(u8, "index"), .value = .{ .int32 = @intCast(si) } }; - u[1] = .{ .key = try reply.arena_alloc().dupe(u8, "_id"), .value = try bson.copy_value(reply.arena_alloc(), id) }; + u[0] = .{ .key = "index", .value = .{ .int32 = @intCast(si) } }; + u[1] = .{ .key = "_id", .value = try bson.copy_value(reply.arena_alloc(), id) }; try upserted.append(reply.arena_alloc(), .{ .key = "u", .value = .{ .doc = u } }); n_matched += 1; } @@ -557,7 +453,7 @@ fn cmd_update(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { for (matched.items) |doc| { // Work on a copy: the log write must precede any visible change, // and a rejected update must not corrupt the stored document. - const copy = try clone_doc(ctx, reply, doc); + const copy = try clone_doc(reply, doc); update.apply(copy, &.{ .arena = undefined, .pairs = u_doc }) catch |err| switch (err) { error.ImmutableId, error.InvalidUpdate => return bad_value(reply, "bad update"), else => return err, @@ -581,50 +477,19 @@ fn cmd_delete(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const db_name = msg.db_name() orelse return invalid_arg(reply, "delete requires $db"); const coll_name = str_arg(msg.body.get("delete")) orelse return bad_value(reply, "delete requires a collection name"); - var specs: []const bson.Document = &.{}; - for (msg.seqs) |seq| { - if (std.mem.eql(u8, seq.name, "deletes")) specs = seq.docs; - } - if (specs.len == 0) { - const d = msg.body.get("deletes") orelse return bad_value(reply, "delete requires deletes"); - const arr = switch (d) { - .array => |a| a, - else => return bad_value(reply, "deletes must be an array"), - }; - const parsed = try reply.arena_alloc().alloc(bson.Document, arr.len); - for (arr, 0..) |item, i| { - parsed[i] = switch (item) { - .doc => |pairs| bson.Document{ .arena = undefined, .pairs = pairs }, - else => return bad_value(reply, "deletes must be documents"), - }; - } - specs = parsed; - } + const specs = try batch_arg(msg, reply, "delete", "deletes") orelse return; var n_deleted: i64 = 0; for (specs) |*spec| { const q = doc_arg(spec.get("q")) orelse return bad_value(reply, "delete spec requires q"); - const limit: i64 = switch (spec.get("limit") orelse bson.Value{ .int32 = 1 }) { - .int32 => |i| i, - .int64 => |i| i, - else => 1, - }; - if (ctx.engine.get_collection(db_name, coll_name)) |coll| { - var to_remove: std.ArrayListUnmanaged([]const u8) = .empty; - defer { - for (to_remove.items) |k| ctx.gpa.free(k); - to_remove.deinit(ctx.gpa); - } - var it = coll.docs.iterator(); - while (it.next()) |entry| { - if (try query.matches(ctx.gpa, &.{ .arena = undefined, .pairs = q }, entry.value_ptr.*)) { - try to_remove.append(ctx.gpa, try ctx.gpa.dupe(u8, entry.key_ptr.*)); - if (limit == 1) break; - } - } - for (to_remove.items) |key| { - if (try ctx.engine.remove(db_name, coll_name, key)) n_deleted += 1; - } + const limit = int_value(spec.get("limit")) orelse 1; + + var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty; + defer matched.deinit(ctx.gpa); + _ = try scan_matching(ctx, db_name, coll_name, q, if (limit == 1) 1 else 0, &matched); + for (matched.items) |doc| { + const id = doc.get("_id") orelse continue; + if (try ctx.engine.remove_by_id(db_name, coll_name, id)) n_deleted += 1; } } @@ -649,80 +514,57 @@ fn cmd_find_and_modify(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !v var matched: std.ArrayListUnmanaged(*const bson.Document) = .empty; defer matched.deinit(ctx.gpa); - if (ctx.engine.get_collection(db_name, coll_name)) |coll| { - var it = coll.docs.iterator(); - while (it.next()) |entry| { - if (try query.matches(ctx.gpa, &.{ .arena = undefined, .pairs = q }, entry.value_ptr.*)) { - try matched.append(ctx.gpa, entry.value_ptr.*); - } - } - } + // Without a sort, only the first match is ever used. + _ = try scan_matching(ctx, db_name, coll_name, q, if (sort_keys.len > 0) 0 else 1, &matched); if (sort_keys.len > 0) { - try query.sort_docs(ctx.gpa, reply.arena_alloc(), matched.items, sort_keys); + try query.sort_docs(reply.arena_alloc(), matched.items, sort_keys); } - var leo: std.ArrayListUnmanaged(bson.Pair) = .empty; - defer leo.deinit(reply.arena_alloc()); + const arena = reply.arena_alloc(); + const target = if (matched.items.len > 0) matched.items[0] else null; - if (matched.items.len == 0 and do_update and upsert) { + // Each branch decides what the reply says; the tail below emits it once. + var n: i32 = 0; + var updated_existing = false; + var upserted_id: ?bson.Value = null; + var value: bson.Value = .null; + + if (target == null and do_update and upsert) { const u_doc = doc_arg(msg.body.get("update")) orelse return bad_value(reply, "update must be a document"); - const new_doc = try build_upsert_doc(ctx, reply, q, u_doc); + const new_doc = try build_upsert_doc(reply, q, u_doc); ctx.engine.insert(db_name, coll_name, new_doc, ctx.oid_gen) catch |err| switch (err) { error.DuplicateKey => return duplicate_key_error(reply, db_name, coll_name, new_doc), else => return err, }; - try leo.append(reply.arena_alloc(), .{ .key = "n", .value = .{ .int32 = 1 } }); - try leo.append(reply.arena_alloc(), .{ .key = "updatedExisting", .value = .{ .bool = false } }); - try leo.append(reply.arena_alloc(), .{ .key = "upserted", .value = try bson.copy_value(reply.arena_alloc(), new_doc.get("_id") orelse bson.Value.null) }); - if (ret_new) { - try reply.put("value", try project_doc(reply, new_doc, proj_pairs)); - } else { - try reply.put("value", .null); - } - try reply.put("lastErrorObject", .{ .doc = try leo.toOwnedSlice(reply.arena_alloc()) }); - try reply.put_ok(); - return; - } - - const target = if (matched.items.len > 0) matched.items[0] else null; - if (target != null and remove) { - const id_key = try bson.serialize_value(ctx.gpa, target.?.get("_id") orelse unreachable); - defer ctx.gpa.free(id_key); - try leo.append(reply.arena_alloc(), .{ .key = "n", .value = .{ .int32 = 1 } }); - try leo.append(reply.arena_alloc(), .{ .key = "updatedExisting", .value = .{ .bool = false } }); - try reply.put("value", try project_doc(reply, target.?, proj_pairs)); - _ = try ctx.engine.remove(db_name, coll_name, id_key); - try reply.put("lastErrorObject", .{ .doc = try leo.toOwnedSlice(reply.arena_alloc()) }); - try reply.put_ok(); - return; - } - - if (target != null and do_update) { + n = 1; + upserted_id = try bson.copy_value(arena, new_doc.get("_id") orelse bson.Value.null); + if (ret_new) value = try project_doc(reply, new_doc, proj_pairs); + } else if (target != null and remove) { + n = 1; + // Project before removing: this reads the stored document. + value = try project_doc(reply, target.?, proj_pairs); + _ = try ctx.engine.remove_by_id(db_name, coll_name, target.?.get("_id") orelse unreachable); + } else if (target != null and do_update) { const u_doc = doc_arg(msg.body.get("update")) orelse return bad_value(reply, "update must be a document"); - const before = try bson.copy_pairs(reply.arena_alloc(), target.?.pairs); - const copy = try clone_doc(ctx, reply, target.?); + const before = try bson.copy_pairs(arena, target.?.pairs); + const copy = try clone_doc(reply, target.?); update.apply(copy, &.{ .arena = undefined, .pairs = u_doc }) catch |err| switch (err) { error.ImmutableId, error.InvalidUpdate => return bad_value(reply, "bad update"), else => return err, }; try ctx.engine.replace(db_name, coll_name, copy, ctx.oid_gen); - try leo.append(reply.arena_alloc(), .{ .key = "n", .value = .{ .int32 = 1 } }); - try leo.append(reply.arena_alloc(), .{ .key = "updatedExisting", .value = .{ .bool = true } }); - if (ret_new) { - try reply.put("value", try project_doc(reply, copy, proj_pairs)); - } else { - try reply.put("value", .{ .doc = before }); - } - try reply.put("lastErrorObject", .{ .doc = try leo.toOwnedSlice(reply.arena_alloc()) }); - try reply.put_ok(); - return; - } + n = 1; + updated_existing = true; + value = if (ret_new) try project_doc(reply, copy, proj_pairs) else .{ .doc = before }; + } // else: no match and no upsert — an empty result - // No match and no upsert - try leo.append(reply.arena_alloc(), .{ .key = "n", .value = .{ .int32 = 0 } }); - try leo.append(reply.arena_alloc(), .{ .key = "updatedExisting", .value = .{ .bool = false } }); - try reply.put("value", .null); - try reply.put("lastErrorObject", .{ .doc = try leo.toOwnedSlice(reply.arena_alloc()) }); + var leo: std.ArrayListUnmanaged(bson.Pair) = .empty; + defer leo.deinit(arena); + try leo.append(arena, .{ .key = "n", .value = .{ .int32 = n } }); + try leo.append(arena, .{ .key = "updatedExisting", .value = .{ .bool = updated_existing } }); + if (upserted_id) |id| try leo.append(arena, .{ .key = "upserted", .value = id }); + try reply.put("value", value); + try reply.put("lastErrorObject", .{ .doc = try leo.toOwnedSlice(arena) }); try reply.put_ok(); } @@ -731,13 +573,7 @@ fn cmd_count(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const coll_name = str_arg(msg.body.get("count")) orelse return bad_value(reply, "count requires a collection name"); const q = doc_arg(msg.body.get("query")) orelse &.{}; - var n: i64 = 0; - if (ctx.engine.get_collection(db_name, coll_name)) |coll| { - var it = coll.docs.iterator(); - while (it.next()) |entry| { - if (try query.matches(ctx.gpa, &.{ .arena = undefined, .pairs = q }, entry.value_ptr.*)) n += 1; - } - } + const n = try scan_matching(ctx, db_name, coll_name, q, 0, null); try reply.put("n", .{ .int32 = @intCast(n) }); try reply.put_ok(); } @@ -791,21 +627,13 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { } else if (std.mem.eql(u8, stage_name, "$sort")) { const keys = try parse_sort_keys(reply, stage[0].value); if (keys.len > 0) { - try query.sort_docs(ctx.gpa, reply.arena_alloc(), stream.items[start..end], keys); + try query.sort_docs(reply.arena_alloc(), stream.items[start..end], keys); } } else if (std.mem.eql(u8, stage_name, "$skip")) { - const n: usize = switch (stage[0].value) { - .int32 => |i| @intCast(@max(0, i)), - .int64 => |i| @intCast(@max(0, i)), - else => return bad_value(reply, "$skip requires a number"), - }; + const n = try stage_count(reply, stage[0].value, "$skip") orelse return; start = @min(start + n, end); } else if (std.mem.eql(u8, stage_name, "$limit")) { - const n: usize = switch (stage[0].value) { - .int32 => |i| @intCast(@max(0, i)), - .int64 => |i| @intCast(@max(0, i)), - else => return bad_value(reply, "$limit requires a number"), - }; + const n = try stage_count(reply, stage[0].value, "$limit") orelse return; end = @min(end, start + n); } else if (std.mem.eql(u8, stage_name, "$project")) { proj_pairs = doc_arg(stage[0].value); @@ -836,7 +664,7 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { c[0] = .{ .key = try reply.arena_alloc().dupe(u8, name), .value = .{ .int32 = @intCast(slice.len) } }; const values = try reply.arena_alloc().alloc(bson.Value, 1); values[0] = .{ .doc = c }; - try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, coll_name), values) }); + try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, coll_name), "firstBatch", values) }); } else { try emit_docs(reply, db_name, coll_name, proj_pairs, slice); } @@ -852,11 +680,11 @@ fn run_group(ctx: *Context, reply: *wire.Reply, group_pairs: []const bson.Pair, return null; }; - var accs: std.ArrayListUnmanaged([]const bson.Pair) = .empty; + var accs: std.ArrayListUnmanaged(bson.Pair) = .empty; defer accs.deinit(arena); for (group_pairs) |p| { if (std.mem.eql(u8, p.key, "_id")) continue; - try accs.append(arena, &.{.{ .key = p.key, .value = p.value }}); + try accs.append(arena, p); } const Group = struct { @@ -872,28 +700,29 @@ fn run_group(ctx: *Context, reply: *wire.Reply, group_pairs: []const bson.Pair, keys_owned.deinit(ctx.gpa); } + var id_key_buf: std.ArrayListUnmanaged(u8) = .empty; + defer id_key_buf.deinit(ctx.gpa); for (docs) |doc| { const id_value: bson.Value = switch (id_expr) { - .null => .null, - else => switch (id_expr) { - .string => |s| if (s.len > 0 and s[0] == '$') query_path_value(doc, s[1..]) orelse .null else id_expr, - else => id_expr, - }, + .string => |s| if (s.len > 0 and s[0] == '$') query_path_value(doc, s[1..]) orelse .null else id_expr, + else => id_expr, }; - var id_key_buf: std.ArrayListUnmanaged(u8) = .empty; - defer id_key_buf.deinit(ctx.gpa); + id_key_buf.clearRetainingCapacity(); try bson.write_value(id_value, ctx.gpa, &id_key_buf); - const key = try ctx.gpa.dupe(u8, id_key_buf.items); - try keys_owned.append(ctx.gpa, key); - const gop = try groups.getOrPut(ctx.gpa, key); + const gop = try groups.getOrPut(ctx.gpa, id_key_buf.items); if (!gop.found_existing) { + // Only a new group needs an owned copy of the key; the map + // borrows it, so keys_owned keeps it alive until we are done. + const key = try ctx.gpa.dupe(u8, id_key_buf.items); + try keys_owned.append(ctx.gpa, key); + gop.key_ptr.* = key; const sums = try ctx.gpa.alloc(f64, accs.items.len); @memset(sums, 0); gop.value_ptr.* = .{ .id_value = id_value, .sums = sums }; } for (accs.items, 0..) |acc, i| { - var expr = acc[0].value; + var expr = acc.value; // Unwrap {$sum: } accumulator documents. if (expr == .doc) { if (bson.get_pair(expr.doc, "$sum")) |inner| { @@ -930,14 +759,14 @@ fn run_group(ctx: *Context, reply: *wire.Reply, group_pairs: []const bson.Pair, while (it.next()) |entry| { const npairs = 1 + accs.items.len; const pairs = try arena.alloc(bson.Pair, npairs); - pairs[0] = .{ .key = try arena.dupe(u8, "_id"), .value = try bson.copy_value(arena, entry.value_ptr.id_value) }; + pairs[0] = .{ .key = "_id", .value = try bson.copy_value(arena, entry.value_ptr.id_value) }; for (accs.items, 0..) |acc, i| { const sum: f64 = entry.value_ptr.sums[i]; const sum_value: bson.Value = if (sum == @floor(sum) and sum <= 2_147_483_647 and sum >= -2_147_483_648) .{ .int32 = @intFromFloat(sum) } else .{ .double = sum }; - pairs[1 + i] = .{ .key = try arena.dupe(u8, acc[0].key), .value = sum_value }; + pairs[1 + i] = .{ .key = acc.key, .value = sum_value }; } const doc = try arena.create(bson.Document); doc.* = bson.Document{ .arena = undefined, .pairs = pairs }; @@ -961,16 +790,13 @@ fn query_path_value(doc: *const bson.Document, path: []const u8) ?bson.Value { return cur; } -fn cmd_get_more(reply: *wire.Reply) !void { - try reply.put("cursor", .{ .doc = &.{ - .{ .key = "id", .value = .{ .int64 = 0 } }, - .{ .key = "ns", .value = .{ .string = "test.$cmd" } }, - .{ .key = "nextBatch", .value = .{ .array = &.{} } }, - } }); +fn cmd_get_more(_: *Context, _: *wire.Message, reply: *wire.Reply) !void { + // Cursors are never left open, so getMore always yields an empty batch. + try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, "test.$cmd", "nextBatch", &.{}) }); try reply.put_ok(); } -fn cmd_kill_cursors(reply: *wire.Reply) !void { +fn cmd_kill_cursors(_: *Context, _: *wire.Message, reply: *wire.Reply) !void { try reply.put("cursorsKilled", .{ .array = &.{} }); try reply.put("cursorsNotFound", .{ .array = &.{} }); try reply.put("cursorsAlive", .{ .array = &.{} }); @@ -983,25 +809,25 @@ fn cmd_kill_cursors(reply: *wire.Reply) !void { /// Deep-copy a stored document into the reply arena so updates can be /// applied off the live doc. -fn clone_doc(ctx: *Context, reply: *wire.Reply, doc: *const bson.Document) !*bson.Document { - var buf: std.ArrayListUnmanaged(u8) = .empty; - defer buf.deinit(ctx.gpa); - try doc.to_bytes(ctx.gpa, &buf); - const owned = try reply.arena_alloc().create(bson.Document); - owned.* = try bson.Document.parse(reply.arena_alloc(), buf.items); +fn clone_doc(reply: *wire.Reply, doc: *const bson.Document) !*bson.Document { + const arena = reply.arena_alloc(); + const owned = try arena.create(bson.Document); + owned.* = .{ + .arena = std.heap.ArenaAllocator.init(arena), + .pairs = try bson.copy_pairs(arena, doc.pairs), + }; return owned; } /// Build the document for an upsert: equality fields from the filter, then /// the update operators applied. Owned by the reply arena. -fn build_upsert_doc(ctx: *Context, reply: *wire.Reply, q: []const bson.Pair, u_doc: []const bson.Pair) !*bson.Document { - _ = ctx; +fn build_upsert_doc(reply: *wire.Reply, q: []const bson.Pair, u_doc: []const bson.Pair) !*bson.Document { const arena = reply.arena_alloc(); var pairs: std.ArrayListUnmanaged(bson.Pair) = .empty; defer pairs.deinit(arena); for (q) |p| { const is_operator = p.key.len > 0 and p.key[0] == '$'; - const is_embedded_operator = p.value == .doc and doc_all_operators(p.value.doc); + const is_embedded_operator = p.value == .doc and query.all_operator_keys(p.value.doc); if (!is_operator and !is_embedded_operator) { try pairs.append(arena, .{ .key = try arena.dupe(u8, p.key), .value = try bson.copy_value(arena, p.value) }); } @@ -1013,18 +839,9 @@ fn build_upsert_doc(ctx: *Context, reply: *wire.Reply, q: []const bson.Pair, u_d error.ImmutableId, error.InvalidUpdate => return error.InvalidUpdate, else => return err, }; - _ = ctx; return owned; } -fn doc_all_operators(pairs: []const bson.Pair) bool { - if (pairs.len == 0) return false; - for (pairs) |p| { - if (p.key.len == 0 or p.key[0] != '$') return false; - } - return true; -} - fn parse_sort_keys(reply: *wire.Reply, value: ?bson.Value) ![]const query.SortKey { const pairs = doc_arg(value) orelse return &.{}; const out = try reply.arena_alloc().alloc(query.SortKey, pairs.len); @@ -1067,14 +884,49 @@ fn bool_arg(v: ?bson.Value) ?bool { } fn int_arg(v: ?bson.Value) ?u64 { + const i = int_value(v) orelse return null; + return if (i < 0) null else @intCast(i); +} + +/// Numeric argument as a signed integer, with no clamping — callers that care +/// about the sign (find's limit, delete's limit, $skip) apply their own rule. +fn int_value(v: ?bson.Value) ?i64 { return switch (v orelse return null) { - .int32 => |i| if (i < 0) null else @intCast(i), - .int64 => |i| if (i < 0) null else @intCast(i), - .double => |d| if (d < 0) null else @intFromFloat(d), + .int32 => |i| i, + .int64 => |i| i, + // lossyCast saturates instead of trapping on out-of-range doubles. + .double => |d| std.math.lossyCast(i64, d), else => null, }; } +/// A `$skip`/`$limit` stage operand: a non-negative document count. Writes +/// the error reply and returns null when the operand is not a number. +fn stage_count(reply: *wire.Reply, v: bson.Value, stage: []const u8) !?usize { + const n = int_value(v) orelse { + const text = try std.fmt.allocPrint(reply.arena_alloc(), "{s} requires a number", .{stage}); + try bad_value(reply, text); + return null; + }; + return @intCast(@max(0, n)); +} + +/// Fetch a batch argument, writing the standard error reply and returning +/// null when it is missing or malformed. +fn batch_arg(msg: *wire.Message, reply: *wire.Reply, cmd: []const u8, name: []const u8) !?[]const bson.Document { + return msg.batch(name) catch |err| { + const arena = reply.arena_alloc(); + const text = switch (err) { + error.MissingBatch => try std.fmt.allocPrint(arena, "{s} requires {s}", .{ cmd, name }), + error.BatchNotArray => try std.fmt.allocPrint(arena, "{s} must be an array", .{name}), + error.BatchElementNotDoc => try std.fmt.allocPrint(arena, "{s} must be documents", .{name}), + error.OutOfMemory => return error.OutOfMemory, + }; + try bad_value(reply, text); + return null; + }; +} + fn invalid_arg(reply: *wire.Reply, msg: []const u8) !void { return reply.put_error(@intFromEnum(ErrorCode.invalid_argument), "InvalidArgument", msg); } @@ -1083,10 +935,19 @@ fn bad_value(reply: *wire.Reply, msg: []const u8) !void { return reply.put_error(@intFromEnum(ErrorCode.bad_value), "BadValue", msg); } +/// The E11000 text, shared by the top-level error reply and the per-document +/// `writeErrors` entries of a batch insert. +fn duplicate_key_message(reply: *wire.Reply, db_name: []const u8, coll_name: []const u8, doc: *const bson.Document) ![]const u8 { + const key = try serialize_value_compact(reply, doc.get("_id") orelse bson.Value.null); + return std.fmt.allocPrint( + reply.arena_alloc(), + "E11000 duplicate key error collection: {s}.{s} index: _id_ dup key: {s}", + .{ db_name, coll_name, key }, + ); +} + fn duplicate_key_error(reply: *wire.Reply, db_name: []const u8, coll_name: []const u8, doc: *const bson.Document) !void { - const id = doc.get("_id") orelse bson.Value.null; - const key = try serialize_value_compact(reply, id); - const msg_text = try std.fmt.allocPrint(reply.arena_alloc(), "E11000 duplicate key error collection: {s}.{s} index: _id_ dup key: {s}", .{ db_name, coll_name, key }); + const msg_text = try duplicate_key_message(reply, db_name, coll_name, doc); return reply.put_error(@intFromEnum(ErrorCode.duplicate_key), "DuplicateKey", msg_text); } @@ -1107,12 +968,12 @@ fn serialize_value_compact(reply: *wire.Reply, v: bson.Value) ![]const u8 { }; } -/// Build a { id: 0, ns: "...", firstBatch: [...] } cursor document. -pub fn cursor_doc(reply: *wire.Reply, cursor_id: i64, ns: []const u8, docs: []const bson.Value) ![]const bson.Pair { +/// Build a { id: , ns: "...", : [...] } cursor document. +pub fn cursor_doc(reply: *wire.Reply, cursor_id: i64, ns: []const u8, batch_key: []const u8, docs: []const bson.Value) ![]const bson.Pair { const c = try reply.arena_alloc().alloc(bson.Pair, 3); - c[0] = .{ .key = try reply.arena_alloc().dupe(u8, "id"), .value = .{ .int64 = cursor_id } }; - c[1] = .{ .key = try reply.arena_alloc().dupe(u8, "ns"), .value = .{ .string = try reply.arena_alloc().dupe(u8, ns) } }; - c[2] = .{ .key = try reply.arena_alloc().dupe(u8, "firstBatch"), .value = .{ .array = docs } }; + c[0] = .{ .key = "id", .value = .{ .int64 = cursor_id } }; + c[1] = .{ .key = "ns", .value = .{ .string = ns } }; + c[2] = .{ .key = batch_key, .value = .{ .array = docs } }; return c; } @@ -1138,26 +999,54 @@ fn str_array(reply: *wire.Reply, values: []const []const u8) ![]const bson.Value const testing = std.testing; +/// Temp-log-backed engine plus the Context a command test dispatches with. +const TestDb = struct { + tmp: std.testing.TmpDir, + path: []u8, + engine: db.Engine, + gen: bson.ObjectIdGen, + + fn init(io: std.Io) !TestDb { + const tmp = std.testing.tmpDir(.{}); + const path = try std.fmt.allocPrint(testing.allocator, ".zig-cache/tmp/{s}/cmd.log", .{tmp.sub_path}); + return .{ + .tmp = tmp, + .path = path, + .engine = try db.Engine.open(testing.allocator, io, path), + .gen = bson.ObjectIdGen.init(io), + }; + } + + fn deinit(self: *TestDb) void { + self.engine.deinit(); + self.tmp.cleanup(); + testing.allocator.free(self.path); + } + + fn ctx(self: *TestDb, io: std.Io) Context { + return test_ctx(io, &self.engine, &self.gen, 1); + } +}; + +fn test_ctx(io: std.Io, engine: *db.Engine, gen: *bson.ObjectIdGen, connection_id: u32) Context { + return .{ + .gpa = testing.allocator, + .io = io, + .oid_gen = gen, + .connection_id = connection_id, + .client_desc = "127.0.0.1:0", + .engine = engine, + .server_start = std.Io.Timestamp.now(io, .real), + }; +} + test "ping and hello replies parse" { var threaded: std.Io.Threaded = .init_single_threaded; 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}/cmd.log", .{tmp.sub_path}); - defer testing.allocator.free(path); - var engine = try db.Engine.open(testing.allocator, io, path); - defer engine.deinit(); - var ctx = Context{ - .gpa = testing.allocator, - .io = io, - .oid_gen = &gen, - .connection_id = 1, - .client_desc = "127.0.0.1:0", - .engine = &engine, - .server_start = std.Io.Timestamp.now(io, .real), - }; + var tdb = try TestDb.init(io); + defer tdb.deinit(); + var ctx = tdb.ctx(io); var reply = wire.Reply.init(testing.allocator); defer reply.deinit(); @@ -1182,22 +1071,9 @@ test "unknown command gives CommandNotFound" { var threaded: std.Io.Threaded = .init_single_threaded; 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}/cmd.log", .{tmp.sub_path}); - defer testing.allocator.free(path); - var engine = try db.Engine.open(testing.allocator, io, path); - defer engine.deinit(); - var ctx = Context{ - .gpa = testing.allocator, - .io = io, - .oid_gen = &gen, - .connection_id = 1, - .client_desc = "127.0.0.1:0", - .engine = &engine, - .server_start = std.Io.Timestamp.now(io, .real), - }; + var tdb = try TestDb.init(io); + defer tdb.deinit(); + var ctx = tdb.ctx(io); var reply = wire.Reply.init(testing.allocator); defer reply.deinit(); var msg = try parse_fake_msg("nonsenseCmd", .null, &.{}); @@ -1212,22 +1088,9 @@ test "getParameter responds for featureCompatibilityVersion" { var threaded: std.Io.Threaded = .init_single_threaded; 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}/cmd.log", .{tmp.sub_path}); - defer testing.allocator.free(path); - var engine = try db.Engine.open(testing.allocator, io, path); - defer engine.deinit(); - var ctx = Context{ - .gpa = testing.allocator, - .io = io, - .oid_gen = &gen, - .connection_id = 1, - .client_desc = "127.0.0.1:0", - .engine = &engine, - .server_start = std.Io.Timestamp.now(io, .real), - }; + var tdb = try TestDb.init(io); + defer tdb.deinit(); + var ctx = tdb.ctx(io); var reply = wire.Reply.init(testing.allocator); defer reply.deinit(); var fcv_doc = try testing.allocator.alloc(bson.Pair, 1); @@ -1244,22 +1107,9 @@ test "insert counts only successful inserts" { var threaded: std.Io.Threaded = .init_single_threaded; 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}/cmd.log", .{tmp.sub_path}); - defer testing.allocator.free(path); - var engine = try db.Engine.open(testing.allocator, io, path); - defer engine.deinit(); - var ctx = Context{ - .gpa = testing.allocator, - .io = io, - .oid_gen = &gen, - .connection_id = 1, - .client_desc = "127.0.0.1:0", - .engine = &engine, - .server_start = std.Io.Timestamp.now(io, .real), - }; + var tdb = try TestDb.init(io); + defer tdb.deinit(); + var ctx = tdb.ctx(io); // documents: [{_id:1}, {_id:1} (dup), {_id:2}] → n: 2 + one writeError. const docs = [_]bson.Value{ @@ -1317,7 +1167,6 @@ test "concurrent insert/find commands on a threaded Io" { 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; @@ -1327,17 +1176,9 @@ test "concurrent insert/find commands on a threaded Io" { 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 { + 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) 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, - }; + var ctx = test_ctx(iow, eng, &wgen, fiber_id); while (true) { const id = id_counter.fetchAdd(1, .monotonic); if (id > total_writes) return; @@ -1353,17 +1194,9 @@ test "concurrent insert/find commands on a threaded Io" { } } - 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 { + fn reader(iow: std.Io, eng: *db.Engine, pending: *std.atomic.Value(usize), fiber_id: u32, total_writes: i32) 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, - }; + var ctx = test_ctx(iow, eng, &rgen, fiber_id); while (pending.load(.acquire) > 0) { var msg = parse_fake_msg("count", .{ .string = "users" }, &.{}) catch return error.Canceled; defer msg.deinit(); @@ -1378,19 +1211,11 @@ test "concurrent insert/find commands on a threaded Io" { 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 }); + for (0..readers) |i| group.async(io, Worker.reader, .{ io, &engine, &remaining, @intCast(i + 1), total }); + for (0..writers) |i| group.async(io, Worker.writer, .{ io, &engine, &next_id, &remaining, @intCast(i + 1), total }); 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 ctx = test_ctx(io, &engine, &gen, 0); var msg = try parse_fake_msg("count", .{ .string = "users" }, &.{}); defer msg.deinit(); var reply = wire.Reply.init(testing.allocator); diff --git a/src/db.zig b/src/db.zig index ffffd14..826b21a 100644 --- a/src/db.zig +++ b/src/db.zig @@ -55,24 +55,43 @@ pub const Engine = struct { pub fn deinit(self: *Engine) void { var db_it = self.dbs.iterator(); while (db_it.next()) |db_entry| { - var coll_it = db_entry.value_ptr.collections.iterator(); - while (coll_it.next()) |coll_entry| { - var doc_it = coll_entry.value_ptr.docs.iterator(); - while (doc_it.next()) |doc_entry| { - doc_entry.value_ptr.*.deinit(); - self.gpa.destroy(doc_entry.value_ptr.*); - self.gpa.free(doc_entry.key_ptr.*); - } - coll_entry.value_ptr.docs.deinit(self.gpa); - self.gpa.free(coll_entry.key_ptr.*); - } - db_entry.value_ptr.collections.deinit(self.gpa); + self.free_db(db_entry.value_ptr); self.gpa.free(db_entry.key_ptr.*); } self.dbs.deinit(self.gpa); self.log.close(); } + /// Free every document in a collection along with its owned _id keys. + fn free_collection(self: *Engine, coll: *Collection) void { + var doc_it = coll.docs.iterator(); + while (doc_it.next()) |doc_entry| { + doc_entry.value_ptr.*.deinit(); + self.gpa.destroy(doc_entry.value_ptr.*); + self.gpa.free(doc_entry.key_ptr.*); + } + coll.docs.deinit(self.gpa); + } + + /// Free every collection in a database along with its owned name keys. + fn free_db(self: *Engine, db: *Db) void { + var coll_it = db.collections.iterator(); + while (coll_it.next()) |coll_entry| { + self.free_collection(coll_entry.value_ptr); + self.gpa.free(coll_entry.key_ptr.*); + } + db.collections.deinit(self.gpa); + } + + /// Drop the document stored under `id_key`, freeing it and its key. + /// No-op when the id is absent. + fn evict_doc(self: *Engine, coll: *Collection, id_key: []const u8) void { + const old = coll.docs.fetchRemove(id_key) orelse return; + old.value.*.deinit(); + self.gpa.destroy(old.value); + self.gpa.free(old.key); + } + // -- commands (callers must hold the matching lock) --------------------- /// Exclusive lock: for commands that mutate the engine. @@ -97,74 +116,76 @@ pub const Engine = struct { /// Insert a document. Fails with error.DuplicateKey if the _id exists. /// Generates an ObjectId _id when absent. pub fn insert(self: *Engine, db_name: []const u8, coll_name: []const u8, doc: *const bson.Document, oid_gen: *bson.ObjectIdGen) !void { - const coll = try self.get_or_create_collection(db_name, coll_name); - const owned = try self.own_with_id(doc, oid_gen); - errdefer { - owned.deinit(); - self.gpa.destroy(owned); - } - - const id_value = owned.get("_id") orelse unreachable; - const id_key = try bson.serialize_value(self.gpa, id_value); - defer self.gpa.free(id_key); - - if (coll.docs.contains(id_key)) return error.DuplicateKey; - - const doc_bytes = try serialize_doc(self.gpa, owned); - defer self.gpa.free(doc_bytes); - self.seq += 1; - try self.log.append_upsert(db_name, coll_name, doc_bytes, self.seq); - - const key_owned = try self.gpa.dupe(u8, id_key); - try coll.docs.put(self.gpa, key_owned, owned); - try self.maybe_compact(); + return self.upsert(db_name, coll_name, doc, oid_gen, .insert); } /// Insert or replace a document by _id (upsert without existence check). pub fn replace(self: *Engine, db_name: []const u8, coll_name: []const u8, doc: *const bson.Document, oid_gen: *bson.ObjectIdGen) !void { + return self.upsert(db_name, coll_name, doc, oid_gen, .replace); + } + + /// Shared body of `insert` and `replace`: they differ only in how an + /// existing _id is treated. Logs (and syncs) the new document before it + /// becomes visible in memory. + fn upsert( + self: *Engine, + db_name: []const u8, + coll_name: []const u8, + doc: *const bson.Document, + oid_gen: *bson.ObjectIdGen, + mode: enum { insert, replace }, + ) !void { const coll = try self.get_or_create_collection(db_name, coll_name); const owned = try self.own_with_id(doc, oid_gen); - errdefer { + const id_value = owned.get("_id") orelse unreachable; + // Ownership of the key moves to the map once `stored` is set; until + // then this frame still owns both it and `owned`. + const id_key = try bson.serialize_value(self.gpa, id_value); + var stored = false; + errdefer if (!stored) { owned.deinit(); self.gpa.destroy(owned); - } + self.gpa.free(id_key); + }; - const id_value = owned.get("_id") orelse unreachable; - const id_key = try bson.serialize_value(self.gpa, id_value); - defer self.gpa.free(id_key); + if (mode == .insert and coll.docs.contains(id_key)) return error.DuplicateKey; const doc_bytes = try serialize_doc(self.gpa, owned); defer self.gpa.free(doc_bytes); self.seq += 1; try self.log.append_upsert(db_name, coll_name, doc_bytes, self.seq); - if (coll.docs.fetchRemove(id_key)) |old| { - old.value.*.deinit(); - self.gpa.destroy(old.value); - self.gpa.free(old.key); - } - const key_owned = try self.gpa.dupe(u8, id_key); - try coll.docs.put(self.gpa, key_owned, owned); + if (mode == .replace) self.evict_doc(coll, id_key); + try coll.docs.put(self.gpa, id_key, owned); + stored = true; try self.maybe_compact(); } - /// Remove a document by _id. Returns true if it existed. - pub fn remove(self: *Engine, db_name: []const u8, coll_name: []const u8, id_key: []const u8) !bool { + /// Remove a document by its `_id` value. Returns true if it existed. + /// The serialized-key encoding stays private to the engine. + pub fn remove_by_id(self: *Engine, db_name: []const u8, coll_name: []const u8, id: bson.Value) !bool { + const id_key = try bson.serialize_value(self.gpa, id); + defer self.gpa.free(id_key); + return self.remove(db_name, coll_name, id_key); + } + + fn remove(self: *Engine, db_name: []const u8, coll_name: []const u8, id_key: []const u8) !bool { const db = self.dbs.get(db_name) orelse return false; const coll = db.collections.getPtr(coll_name) orelse return false; const doc = coll.docs.get(id_key) orelse return false; // Log (and sync) the delete before removing it from memory, so the // log always describes at least as much as the in-memory state. - const doc_bytes = try serialize_doc(self.gpa, doc); - defer self.gpa.free(doc_bytes); + // Replay only reads _id out of a delete record, so log just that + // rather than a copy of the whole document. + const id_pairs = [_]bson.Pair{.{ .key = "_id", .value = doc.get("_id") orelse unreachable }}; + var id_doc: std.ArrayListUnmanaged(u8) = .empty; + defer id_doc.deinit(self.gpa); + try bson.write_doc(&id_pairs, self.gpa, &id_doc); self.seq += 1; - try self.log.append_delete(db_name, coll_name, doc_bytes, self.seq); + try self.log.append_delete(db_name, coll_name, id_doc.items, self.seq); - const removed = coll.docs.fetchRemove(id_key) orelse unreachable; - removed.value.*.deinit(); - self.gpa.destroy(removed.value); - self.gpa.free(removed.key); + self.evict_doc(coll, id_key); return true; } @@ -181,31 +202,14 @@ pub const Engine = struct { pub fn drop_collection(self: *Engine, db_name: []const u8, coll_name: []const u8) !bool { const db = self.dbs.getPtr(db_name) orelse return false; var removed = db.collections.fetchRemove(coll_name) orelse return false; - var doc_it = removed.value.docs.iterator(); - while (doc_it.next()) |doc_entry| { - doc_entry.value_ptr.*.deinit(); - self.gpa.destroy(doc_entry.value_ptr.*); - self.gpa.free(doc_entry.key_ptr.*); - } - removed.value.docs.deinit(self.gpa); + self.free_collection(&removed.value); self.gpa.free(removed.key); return true; } pub fn drop_database(self: *Engine, db_name: []const u8) !bool { var removed = self.dbs.fetchRemove(db_name) orelse return false; - var coll_it = removed.value.collections.iterator(); - while (coll_it.next()) |coll_entry| { - var docs_it = coll_entry.value_ptr.docs.iterator(); - while (docs_it.next()) |doc_entry| { - doc_entry.value_ptr.*.deinit(); - self.gpa.destroy(doc_entry.value_ptr.*); - self.gpa.free(doc_entry.key_ptr.*); - } - coll_entry.value_ptr.docs.deinit(self.gpa); - self.gpa.free(coll_entry.key_ptr.*); - } - removed.value.collections.deinit(self.gpa); + self.free_db(&removed.value); self.gpa.free(removed.key); return true; } @@ -328,28 +332,19 @@ fn apply_record(ctx: *anyopaque, record: storage.Record, doc: *bson.Document) an return; }; const id_key = try bson.serialize_value(self.gpa, id_value); - defer self.gpa.free(id_key); + var key_owned = false; + defer if (!key_owned) self.gpa.free(id_key); const coll = self.get_or_create_collection(record.db, record.coll) catch return; switch (record.type) { storage.record_type_upsert => { - if (coll.docs.fetchRemove(id_key)) |old| { - old.value.*.deinit(); - self.gpa.destroy(old.value); - self.gpa.free(old.key); - } - const key_owned = try self.gpa.dupe(u8, id_key); - try coll.docs.put(self.gpa, key_owned, doc); + self.evict_doc(coll, id_key); + try coll.docs.put(self.gpa, id_key, doc); + key_owned = true; stored = true; }, - storage.record_type_delete => { - if (coll.docs.fetchRemove(id_key)) |old| { - old.value.*.deinit(); - self.gpa.destroy(old.value); - self.gpa.free(old.key); - } - }, + storage.record_type_delete => self.evict_doc(coll, id_key), else => {}, } } @@ -360,21 +355,7 @@ fn apply_record(ctx: *anyopaque, record: storage.Record, doc: *bson.Document) an const testing = std.testing; -const TmpLog = struct { - tmp: std.testing.TmpDir, - path: []u8, - - fn init(gpa: std.mem.Allocator) !TmpLog { - const tmp = std.testing.tmpDir(.{}); - const path = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}/test.log", .{tmp.sub_path}); - return .{ .tmp = tmp, .path = path }; - } - - fn deinit(self: *TmpLog, gpa: std.mem.Allocator) void { - self.tmp.cleanup(); - gpa.free(self.path); - } -}; +const TmpLog = storage.TmpLog; fn test_env(threaded: *std.Io.Threaded) struct { io: std.Io, gen: bson.ObjectIdGen } { const io = threaded.io(); @@ -426,7 +407,7 @@ test "insert, query, remove" { try engine.lock(); const found = engine.get_doc("app", "users", id_key).?; try testing.expectEqualStrings("bob", found.get("name").?.string); - const removed = try engine.remove("app", "users", id_key); + const removed = try engine.remove_by_id("app", "users", .{ .int32 = 2 }); try testing.expect(removed); engine.unlock(); } @@ -450,9 +431,7 @@ test "reopen replays log" { try engine.lock(); try engine.insert("app", "users", &d1, &env.gen); try engine.insert("app", "users", &d2, &env.gen); - const id_key2 = try bson.serialize_value(gpa, bson.Value{ .int32 = 2 }); - defer gpa.free(id_key2); - _ = try engine.remove("app", "users", id_key2); + _ = try engine.remove_by_id("app", "users", .{ .int32 = 2 }); engine.unlock(); } diff --git a/src/index.zig b/src/index.zig new file mode 100644 index 0000000..17bc9ce --- /dev/null +++ b/src/index.zig @@ -0,0 +1,1190 @@ +//! Secondary indexes: per-collection, single-field and compound, with +//! `unique` and `sparse` options, persisted through the log and used by the +//! query planner as candidate-id generators. +//! +//! The governing invariant: an index is only ever used to produce a +//! candidate id set; the full filter is re-applied to every candidate. An +//! index that over-approximates is merely slow, never wrong. The entire +//! correctness risk collapses onto one question — can the index ever +//! under-approximate? Every design decision here answers that with *no*: +//! +//! - Entry generation mirrors field_matches (src/query.zig) exactly, +//! indexing the value at the path plus the elements of any array there, +//! so whole-array equality can never be missed. +//! - A sparse index is never used for a null-valued component ({a: null} +//! would otherwise miss docs whose field the sparse index skipped). +//! - The _id_ fast path (the docs map) is only used when the queried value's +//! compare-equivalence class is serialization-canonical: bson.compare +//! treats int32 1, int64 1 and double 1.0 as equal, but serialize_value +//! produces different map keys, so a hash lookup would miss. +//! - Entry insertion is infallible after the log append (capacity is +//! reserved first), so a document can never be live but unindexed. +//! +//! Entries alias: `Entry.key` values point into the stored document's arena +//! and `Entry.id` aliases the docs map key; only the `[]bson.Value` slice +//! itself is owned (freed on entry removal). Index removal happens at the +//! top of evict_doc (src/db.zig) — the single chokepoint where a document +//! dies — so the aliasing is safe by construction. + +const std = @import("std"); +const bson = @import("bson.zig"); +const query = @import("query.zig"); + +/// MongoDB's compound index field limit. +pub const max_index_keys: usize = 32; + +/// Cap on the $in cartesian product a plan will generate; beyond it the +/// planner falls back to a scan. +const max_combos: u64 = 100; + +pub const ParallelArraysError = error{ ParallelArrays }; + +/// One key in an index spec. `path` is owned by the Index. +pub const IndexKey = struct { + path: []const u8, + descending: bool, +}; + +/// One index entry. `key` is a gpa-owned slice of Values aliasing the +/// stored document's arena; `id` aliases the docs map key. +pub const Entry = struct { + key: []const bson.Value, + id: []const u8, +}; + +/// Entries built for one document, before they are committed to the index. +pub const BuiltEntries = struct { + entries: std.ArrayListUnmanaged(Entry), + multikey: bool, + + pub fn deinit(self: *BuiltEntries, gpa: std.mem.Allocator) void { + for (self.entries.items) |e| gpa.free(e.key); + self.entries.deinit(gpa); + } +}; + +pub const Index = struct { + name: []const u8, + keys: []const IndexKey, + unique: bool, + sparse: bool, + multikey: bool, + entries: std.ArrayListUnmanaged(Entry), + + pub fn init(gpa: std.mem.Allocator, name: []const u8, keys: []const IndexKey, unique: bool, sparse: bool) !Index { + var self: Index = .{ + .name = undefined, + .keys = undefined, + .unique = unique, + .sparse = sparse, + .multikey = false, + .entries = .empty, + }; + self.name = try gpa.dupe(u8, name); + const owned_keys = try gpa.alloc(IndexKey, keys.len); + var n: usize = 0; + errdefer { + for (owned_keys[0..n]) |k| gpa.free(k.path); + gpa.free(owned_keys); + gpa.free(self.name); + } + while (n < keys.len) : (n += 1) { + owned_keys[n] = .{ .path = try gpa.dupe(u8, keys[n].path), .descending = keys[n].descending }; + } + self.keys = owned_keys; + return self; + } + + pub fn deinit(self: *Index, gpa: std.mem.Allocator) void { + for (self.entries.items) |e| gpa.free(e.key); + self.entries.deinit(gpa); + for (self.keys) |k| gpa.free(k.path); + gpa.free(self.keys); + gpa.free(self.name); + } + + // -- entry generation --------------------------------------------------- + + /// Build the entries one document contributes. Mirrors field_matches + /// (src/query.zig): the value at each path plus the elements of any + /// array there, so both `{tags: "a"}` element queries and whole-array + /// equality on `{tags: ["a","b"]}` are covered. Returns an empty list + /// for a sparse index when a path yields no values (the document is + /// skipped); a non-sparse index indexes missing fields as null. + pub fn build_entries(self: *const Index, gpa: std.mem.Allocator, doc: *const bson.Document, id: []const u8) !BuiltEntries { + var per_path: std.ArrayListUnmanaged(std.ArrayListUnmanaged(bson.Value)) = .empty; + defer { + for (per_path.items) |*list| list.deinit(gpa); + per_path.deinit(gpa); + } + var multikey = false; + var multi_paths: usize = 0; + + for (self.keys) |k| { + var values: std.ArrayListUnmanaged(bson.Value) = .empty; + errdefer values.deinit(gpa); + try query.collect_values(gpa, doc.pairs, k.path, &values, 0); + // Index the array itself and each element, like field_matches. + const direct = values.items.len; + var i: usize = 0; + while (i < direct) : (i += 1) { + if (values.items[i] == .array) { + for (values.items[i].array) |elem| try values.append(gpa, elem); + } + } + if (direct > 1) multikey = true; + for (values.items[0..direct]) |v| { + if (v == .array) { + multikey = true; + break; + } + } + if (values.items.len > 1) multi_paths += 1; + if (values.items.len == 0) { + if (self.sparse) return .{ .entries = .empty, .multikey = false }; + try values.append(gpa, .null); + } + try per_path.append(gpa, values); + } + if (multi_paths > 1) return error.ParallelArrays; + + // Cartesian product across paths, then dedupe identical entries. + var out: std.ArrayListUnmanaged(Entry) = .empty; + errdefer { + for (out.items) |e| gpa.free(e.key); + out.deinit(gpa); + } + const nkeys = self.keys.len; + var choice: [max_index_keys]usize = undefined; + @memset(choice[0..nkeys], 0); + while (true) { + const key = try gpa.alloc(bson.Value, nkeys); + errdefer gpa.free(key); + for (0..nkeys) |ci| key[ci] = per_path.items[ci].items[choice[ci]]; + try out.append(gpa, .{ .key = key, .id = id }); + var ci: usize = nkeys; + var carry = true; + while (carry and ci > 0) { + ci -= 1; + choice[ci] += 1; + if (choice[ci] < per_path.items[ci].items.len) carry = false else choice[ci] = 0; + } + if (carry) break; + } + if (out.items.len > 1) { + std.mem.sort(Entry, out.items, {}, entry_less); + var w: usize = 1; + for (out.items[1..]) |e| { + if (compare_entries(out.items[w - 1], e) != .eq) { + out.items[w] = e; + w += 1; + } else { + gpa.free(e.key); + } + } + out.items.len = w; + } + return .{ .entries = out, .multikey = multikey }; + } + + // -- mutation ----------------------------------------------------------- + + /// Ensure the entry array has room for `n` more entries. Called before + /// the log append, so the subsequent insert_entries is infallible. + pub fn reserve_for(self: *Index, gpa: std.mem.Allocator, n: usize) !void { + try self.entries.ensureUnusedCapacity(gpa, n); + } + + /// Insert pre-built entries (maintaining sort order) and drain the + /// batch: ownership of each entry's key slice moves into the index, so + /// the batch's deinit must not free them. Infallible: capacity must + /// already be reserved. + pub fn insert_entries(self: *Index, gpa: std.mem.Allocator, built: *BuiltEntries) void { + _ = gpa; + for (built.entries.items) |e| { + const pos = self.insert_pos(e); + self.entries.insertAssumeCapacity(pos, e); + } + // Key slices now belong to the index; forget them in the batch so + // its deinit only frees the (now empty) ArrayList buffer. + built.entries.items.len = 0; + } + + /// Build, check, and insert entries for one document; the one-shot form + /// used when rebuilding an index on open. + pub fn add_doc(self: *Index, gpa: std.mem.Allocator, doc: *const bson.Document, id: []const u8) !void { + var built = try self.build_entries(gpa, doc, id); + defer built.deinit(gpa); + if (self.unique) try self.check_unique(built.entries.items, id); + if (built.multikey) self.multikey = true; + try self.entries.ensureUnusedCapacity(gpa, built.entries.items.len); + self.insert_entries(gpa, &built); + } + + /// Remove every entry for `id` and free its key slices. Infallible. + pub fn remove_id(self: *Index, gpa: std.mem.Allocator, id: []const u8) void { + var i: usize = 0; + while (i < self.entries.items.len) { + if (std.mem.eql(u8, self.entries.items[i].id, id)) { + gpa.free(self.entries.items[i].key); + _ = self.entries.orderedRemove(i); + } else i += 1; + } + } + + /// Reject when any of `new_entries` has a key already present under a + /// different id. Entries with `exclude_id` (the replacing document's + /// own old entries) are allowed. + pub fn check_unique(self: *const Index, new_entries: []const Entry, exclude_id: []const u8) error{DuplicateKeyIndex}!void { + for (new_entries) |e| { + const start = self.lower_bound_prefix(e.key); + const end = self.upper_bound_prefix(e.key); + for (self.entries.items[start..end]) |existing| { + if (!std.mem.eql(u8, existing.id, exclude_id)) return error.DuplicateKeyIndex; + } + } + } + + fn insert_pos(self: *const Index, e: Entry) usize { + var lo: usize = 0; + var hi: usize = self.entries.items.len; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (compare_entries(self.entries.items[mid], e) == .lt) lo = mid + 1 else hi = mid; + } + return lo; + } + + // -- search ------------------------------------------------------------- + + /// All ids whose key equals `key` (component-wise). For a partial key + /// (fewer components than the index has) this is a prefix search. + pub fn lookup_eq(self: *const Index, gpa: std.mem.Allocator, key: []const bson.Value, out: *std.ArrayListUnmanaged([]const u8)) !void { + const start = self.lower_bound_prefix(key); + const end = self.upper_bound_prefix(key); + var i = start; + while (i < end) : (i += 1) try out.append(gpa, self.entries.items[i].id); + } + + /// All ids whose key starts with `prefix` and whose component at + /// `prefix.len` falls within [lo, hi]. Range bounds apply to the next + /// component after the equality prefix. + pub fn lookup_range( + self: *const Index, + gpa: std.mem.Allocator, + prefix: []const bson.Value, + lo: ?bson.Value, + lo_incl: bool, + hi: ?bson.Value, + hi_incl: bool, + out: *std.ArrayListUnmanaged([]const u8), + ) !void { + const start = self.lower_bound_prefix(prefix); + const end = self.upper_bound_prefix(prefix); + var i = start; + while (i < end) : (i += 1) { + const v = self.entries.items[i].key[prefix.len]; + if (lo) |l| { + const o = bson.compare(v, l); + if (o == .lt or (o == .eq and !lo_incl)) continue; + } + if (hi) |h| { + const o = bson.compare(v, h); + if (o == .gt or (o == .eq and !hi_incl)) continue; + } + try out.append(gpa, self.entries.items[i].id); + } + } + + /// First entry whose first `prefix.len` components are not less than + /// `prefix`. + fn lower_bound_prefix(self: *const Index, prefix: []const bson.Value) usize { + var lo: usize = 0; + var hi: usize = self.entries.items.len; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (prefix_lt(self.entries.items[mid].key, prefix)) lo = mid + 1 else hi = mid; + } + return lo; + } + + /// First entry whose first `prefix.len` components are greater than + /// `prefix`. + fn upper_bound_prefix(self: *const Index, prefix: []const bson.Value) usize { + var lo: usize = 0; + var hi: usize = self.entries.items.len; + while (lo < hi) { + const mid = lo + (hi - lo) / 2; + if (prefix_le(self.entries.items[mid].key, prefix)) lo = mid + 1 else hi = mid; + } + return lo; + } + + // -- serialization ------------------------------------------------------ + + /// The canonical spec document bytes ({v, key, name, unique?, sparse?}) + /// stored in the log and used to rebuild the index on replay. + pub fn write_spec(self: *const Index, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged(u8)) !void { + var pairs: std.ArrayListUnmanaged(bson.Pair) = .empty; + defer pairs.deinit(gpa); + var key_pairs: std.ArrayListUnmanaged(bson.Pair) = .empty; + defer key_pairs.deinit(gpa); + try pairs.append(gpa, .{ .key = "v", .value = .{ .int32 = 2 } }); + for (self.keys) |k| { + try key_pairs.append(gpa, .{ .key = k.path, .value = if (k.descending) .{ .int32 = -1 } else .{ .int32 = 1 } }); + } + try pairs.append(gpa, .{ .key = "key", .value = .{ .doc = key_pairs.items } }); + try pairs.append(gpa, .{ .key = "name", .value = .{ .string = self.name } }); + if (self.unique) try pairs.append(gpa, .{ .key = "unique", .value = .{ .bool = true } }); + if (self.sparse) try pairs.append(gpa, .{ .key = "sparse", .value = .{ .bool = true } }); + try bson.write_doc(pairs.items, gpa, out); + } + + /// The spec as pairs in `arena` (values alias this index's own storage, + /// which outlives any reply). Used by listIndexes. + pub fn spec_pairs(self: *const Index, arena: std.mem.Allocator, out: *std.ArrayListUnmanaged(bson.Pair)) !void { + try out.append(arena, .{ .key = "v", .value = .{ .int32 = 2 } }); + const key_pairs = try arena.alloc(bson.Pair, self.keys.len); + for (self.keys, 0..) |k, i| { + key_pairs[i] = .{ .key = k.path, .value = if (k.descending) .{ .int32 = -1 } else .{ .int32 = 1 } }; + } + try out.append(arena, .{ .key = "key", .value = .{ .doc = key_pairs } }); + try out.append(arena, .{ .key = "name", .value = .{ .string = self.name } }); + if (self.unique) try out.append(arena, .{ .key = "unique", .value = .{ .bool = true } }); + if (self.sparse) try out.append(arena, .{ .key = "sparse", .value = .{ .bool = true } }); + } + + pub fn spec_equal(a: *const Index, b: *const Index) bool { + if (!std.mem.eql(u8, a.name, b.name)) return false; + if (a.unique != b.unique or a.sparse != b.sparse) return false; + if (a.keys.len != b.keys.len) return false; + for (a.keys, b.keys) |ka, kb| { + if (!std.mem.eql(u8, ka.path, kb.path)) return false; + if (ka.descending != kb.descending) return false; + } + return true; + } +}; + +pub const SpecError = error{ InvalidIndexSpec, OutOfMemory }; + +/// Parse {key: {...}, name?, unique?, sparse?} from a spec document — the +/// form drivers send and the form the log stores. +pub fn parse_spec(gpa: std.mem.Allocator, spec: *const bson.Document) SpecError!Index { + const key_value = bson.get_pair(spec.pairs, "key") orelse return error.InvalidIndexSpec; + const key_pairs = switch (key_value) { + .doc => |p| p, + else => return error.InvalidIndexSpec, + }; + if (key_pairs.len == 0 or key_pairs.len > max_index_keys) return error.InvalidIndexSpec; + var keys: [max_index_keys]IndexKey = undefined; + for (key_pairs, 0..) |p, i| { + const ok_mag = switch (p.value) { + .int32 => |n| n == 1 or n == -1, + .int64 => |n| n == 1 or n == -1, + .double => |n| n == 1.0 or n == -1.0, + else => false, + }; + if (!ok_mag) return error.InvalidIndexSpec; + keys[i] = .{ .path = p.key, .descending = descending(p.value) }; + } + var unique = false; + var sparse = false; + if (bson.get_pair(spec.pairs, "unique")) |v| unique = truthy(v); + if (bson.get_pair(spec.pairs, "sparse")) |v| sparse = truthy(v); + const name_value = bson.get_pair(spec.pairs, "name") orelse { + const nm = try default_name(gpa, key_pairs); + defer gpa.free(nm); + return Index.init(gpa, nm, keys[0..key_pairs.len], unique, sparse); + }; + const name = switch (name_value) { + .string => |s| s, + else => return error.InvalidIndexSpec, + }; + return Index.init(gpa, name, keys[0..key_pairs.len], unique, sparse); +} + +fn descending(v: bson.Value) bool { + return switch (v) { + .int32 => |n| n < 0, + .int64 => |n| n < 0, + .double => |n| n < 0, + else => false, + }; +} + +fn truthy(v: bson.Value) bool { + return switch (v) { + .bool => |b| b, + .int32 => |i| i != 0, + .int64 => |i| i != 0, + .double => |d| d != 0, + else => false, + }; +} + +/// MongoDB's default index name: a_1_b_-1. +fn default_name(gpa: std.mem.Allocator, key_pairs: []const bson.Pair) ![]u8 { + var out: std.ArrayListUnmanaged(u8) = .empty; + errdefer out.deinit(gpa); + for (key_pairs, 0..) |p, i| { + if (i > 0) try out.append(gpa, '_'); + try out.appendSlice(gpa, p.key); + try out.append(gpa, '_'); + if (descending(p.value)) { + try out.append(gpa, '-'); + } + try out.append(gpa, '1'); + } + return out.toOwnedSlice(gpa); +} + +// --------------------------------------------------------------------------- +// Comparison +// --------------------------------------------------------------------------- + +/// Component-wise bson.compare, tie-broken by the id bytes. This is the +/// total order the entry array is kept in. +pub fn compare_entries(a: Entry, b: Entry) std.math.Order { + const n = @min(a.key.len, b.key.len); + for (a.key[0..n], b.key[0..n]) |av, bv| { + const o = bson.compare(av, bv); + if (o != .eq) return o; + } + const l = std.math.order(a.key.len, b.key.len); + if (l != .eq) return l; + return std.mem.order(u8, a.id, b.id); +} + +fn entry_less(_: void, a: Entry, b: Entry) bool { + return compare_entries(a, b) == .lt; +} + +/// Whether entry key's first `prefix.len` components are less than `prefix`. +fn prefix_lt(key: []const bson.Value, prefix: []const bson.Value) bool { + for (key[0..prefix.len], prefix) |a, b| { + const o = bson.compare(a, b); + if (o != .eq) return o == .lt; + } + return false; +} + +/// Whether entry key's first `prefix.len` components are <= `prefix`. +fn prefix_le(key: []const bson.Value, prefix: []const bson.Value) bool { + for (key[0..prefix.len], prefix) |a, b| { + const o = bson.compare(a, b); + if (o != .eq) return o == .lt; + } + return true; +} + +fn less_ids(_: void, a: []const u8, b: []const u8) bool { + return std.mem.order(u8, a, b) == .lt; +} + +// --------------------------------------------------------------------------- +// Query planning +// --------------------------------------------------------------------------- + +const Clause = struct { + path: []const u8, + value: bson.Value, +}; + +/// Flatten top-level pairs and $and members into AND-ed predicates. Every +/// other top-level operator ($or, $nor, ...) is skipped: the full filter is +/// re-applied later, so a usable sibling still yields a valid superset. +fn flatten_clauses(gpa: std.mem.Allocator, pairs: []const bson.Pair, out: *std.ArrayListUnmanaged(Clause)) !void { + for (pairs) |p| { + if (p.key.len > 0 and p.key[0] == '$') { + if (std.mem.eql(u8, p.key, "$and")) { + const members = switch (p.value) { + .array => |a| a, + else => continue, + }; + for (members) |m| { + const mp = switch (m) { + .doc => |d| d, + else => continue, + }; + try flatten_clauses(gpa, mp, out); + } + } + continue; + } + try out.append(gpa, .{ .path = p.key, .value = p.value }); + } +} + +/// Per-index-component info extracted from the filter clauses. +const CompInfo = struct { + eq: ?bson.Value = null, + in_values: ?[]const bson.Value = null, + lo: ?bson.Value = null, + lo_incl: bool = false, + hi: ?bson.Value = null, + hi_incl: bool = false, +}; + +/// Extract the usable constraint from one clause value: a bare non-regex +/// equality, or an operator doc whose operators are all in +/// {$eq, $in, $gt, $gte, $lt, $lte}. Anything else leaves the info +/// untouched (unusable — the full filter is re-applied anyway). +fn analyze_clause(value: bson.Value, info: *CompInfo) void { + if (value == .doc) { + const pairs = value.doc; + if (pairs.len > 0 and !query.all_operator_keys(pairs)) { + // Bare document equality: {a: {n: 1}} compares the whole doc. + info.eq = value; + return; + } + for (pairs) |p| { + if (std.mem.eql(u8, p.key, "$eq")) { + info.eq = p.value; + } else if (std.mem.eql(u8, p.key, "$in")) { + info.in_values = switch (p.value) { + .array => |a| a, + else => return, + }; + } else if (std.mem.eql(u8, p.key, "$gt")) { + info.lo = p.value; + info.lo_incl = false; + } else if (std.mem.eql(u8, p.key, "$gte")) { + info.lo = p.value; + info.lo_incl = true; + } else if (std.mem.eql(u8, p.key, "$lt")) { + info.hi = p.value; + info.hi_incl = false; + } else if (std.mem.eql(u8, p.key, "$lte")) { + info.hi = p.value; + info.hi_incl = true; + } else { + info.eq = null; + info.in_values = null; + info.lo = null; + info.hi = null; + return; + } + } + return; + } + if (value == .regex) return; + info.eq = value; +} + +/// A candidate-generation plan for one index: `lookup_keys` is the +/// cartesian product of the leading equality/$in components (each key has +/// `key_len` Values aliasing the filter), plus an optional range on the +/// next component. +pub const Plan = struct { + index: *const Index, + lookup_keys: std.ArrayListUnmanaged([]const bson.Value), + key_len: usize, + lo: ?bson.Value, + lo_incl: bool, + hi: ?bson.Value, + hi_incl: bool, + + pub fn deinit(self: *Plan, gpa: std.mem.Allocator) void { + for (self.lookup_keys.items) |k| gpa.free(k); + self.lookup_keys.deinit(gpa); + } + + /// Collect the candidate ids, sorted and deduplicated. Range scans can + /// return the same id non-adjacently (a doc with {tags: ["a","b"]} + /// contributes two entries inside one range), so adjacent-dup skipping + /// would be wrong. + pub fn search(self: *const Plan, gpa: std.mem.Allocator, out: *std.ArrayListUnmanaged([]const u8)) !void { + for (self.lookup_keys.items) |key| { + if (self.lo == null and self.hi == null) { + try self.index.lookup_eq(gpa, key, out); + } else { + try self.index.lookup_range(gpa, key, self.lo, self.lo_incl, self.hi, self.hi_incl, out); + } + } + if (out.items.len > 1) { + std.mem.sort([]const u8, out.items, {}, less_ids); + var w: usize = 1; + for (out.items[1..]) |id| { + if (!std.mem.eql(u8, id, out.items[w - 1])) { + out.items[w] = id; + w += 1; + } + } + out.items.len = w; + } + } +}; + +/// Pick the index (if any) that can generate a superset of the matching +/// documents: the one covering the longest leading run of equality/$in +/// predicates, optionally with a range on the next key. Returns null when +/// nothing usable remains — the caller scans. +pub fn plan(gpa: std.mem.Allocator, indexes: []const Index, filter: []const bson.Pair) !?Plan { + if (indexes.len == 0) return null; + var clauses: std.ArrayListUnmanaged(Clause) = .empty; + defer clauses.deinit(gpa); + try flatten_clauses(gpa, filter, &clauses); + + var best: ?Plan = null; + for (indexes) |*ix| { + var cand = (try evaluate_index(gpa, ix, clauses.items)) orelse continue; + if (best) |b| { + if (plan_better(&cand, &b)) { + best.?.deinit(gpa); + best = cand; + } else { + cand.deinit(gpa); + } + } else { + best = cand; + } + } + return best; +} + +fn plan_better(a: *const Plan, b: *const Plan) bool { + if (a.key_len != b.key_len) return a.key_len > b.key_len; + const a_range = a.lo != null or a.hi != null; + const b_range = b.lo != null or b.hi != null; + if (a_range != b_range) return a_range; + return a.lookup_keys.items.len < b.lookup_keys.items.len; +} + +fn evaluate_index(gpa: std.mem.Allocator, ix: *const Index, clauses: []const Clause) !?Plan { + const n = ix.keys.len; + var infos: [max_index_keys]CompInfo = undefined; + for (0..n) |i| { + var info = CompInfo{}; + for (clauses) |cl| { + if (std.mem.eql(u8, cl.path, ix.keys[i].path)) analyze_clause(cl.value, &info); + } + infos[i] = info; + } + + // Longest leading run of equality/$in components. + var run: usize = 0; + while (run < n and (infos[run].eq != null or infos[run].in_values != null)) run += 1; + var lo: ?bson.Value = null; + var lo_incl = false; + var hi: ?bson.Value = null; + var hi_incl = false; + if (run < n) { + lo = infos[run].lo; + lo_incl = infos[run].lo_incl; + hi = infos[run].hi; + hi_incl = infos[run].hi_incl; + } + if (run == 0 and lo == null and hi == null) return null; + + // A sparse index skips documents missing a field; {a: null} would + // otherwise miss them. Never use a sparse index for a null component. + if (ix.sparse) { + for (0..run) |i| { + if (infos[i].eq) |v| { + if (v == .null) return null; + } + if (infos[i].in_values) |list| { + for (list) |m| if (m == .null) return null; + } + } + if (lo) |v| if (v == .null) return null; + if (hi) |v| if (v == .null) return null; + } + + var opts: [max_index_keys][]const bson.Value = undefined; + for (0..run) |i| { + if (infos[i].eq) |v| { + opts[i] = &.{v}; + } else opts[i] = infos[i].in_values.?; + } + var combos: u64 = 1; + for (0..run) |i| { + combos *= opts[i].len; + if (combos > max_combos) return null; + } + + var pl = Plan{ + .index = ix, + .lookup_keys = .empty, + .key_len = run, + .lo = lo, + .lo_incl = lo_incl, + .hi = hi, + .hi_incl = hi_incl, + }; + errdefer pl.deinit(gpa); + if (run == 0) { + const empty = try gpa.alloc(bson.Value, 0); + errdefer gpa.free(empty); + try pl.lookup_keys.append(gpa, empty); + } else { + var choice: [max_index_keys]usize = undefined; + @memset(choice[0..run], 0); + while (true) { + const key = try gpa.alloc(bson.Value, run); + errdefer gpa.free(key); + for (0..run) |i| key[i] = opts[i][choice[i]]; + try pl.lookup_keys.append(gpa, key); + var i: usize = run; + var carry = true; + while (carry and i > 0) { + i -= 1; + choice[i] += 1; + if (choice[i] < opts[i].len) carry = false else choice[i] = 0; + } + if (carry) break; + } + } + return pl; +} + +// --------------------------------------------------------------------------- +// _id fast path +// --------------------------------------------------------------------------- + +pub const IdPlan = struct { + values: []const bson.Value, // aliases the filter +}; + +/// Plan for the implicit _id_ index (the docs map). Only applies to a +/// top-level _id equality/$eq/$in clause, and only when the queried value's +/// compare-equivalence class is serialization-canonical. bson.compare calls +/// int32 1, int64 1 and double 1.0 equal (and string/symbol/code "x" equal, +/// and nested variants), but serialize_value produces different map keys — +/// a hash lookup would then miss documents a scan would match. +pub fn plan_id(filter: []const bson.Pair) ?IdPlan { + var clauses: [16]Clause = undefined; + var n: usize = 0; + if (!flatten_id(filter, &clauses, &n)) return null; + for (clauses[0..n]) |cl| { + if (!std.mem.eql(u8, cl.path, "_id")) continue; + if (id_lookup_values(cl.value)) |values| return .{ .values = values }; + } + return null; +} + +fn flatten_id(pairs: []const bson.Pair, out: *[16]Clause, n: *usize) bool { + for (pairs) |p| { + if (p.key.len > 0 and p.key[0] == '$') { + if (std.mem.eql(u8, p.key, "$and")) { + const members = switch (p.value) { + .array => |a| a, + else => continue, + }; + for (members) |m| { + const mp = switch (m) { + .doc => |d| d, + else => continue, + }; + if (!flatten_id(mp, out, n)) return false; + } + } + continue; + } + if (n.* >= out.len) return false; + out[n.*] = .{ .path = p.key, .value = p.value }; + n.* += 1; + } + return true; +} + +fn id_lookup_values(v: bson.Value) ?[]const bson.Value { + if (v == .regex) return null; + if (v != .doc) { + return if (value_fast_path_safe(v)) &.{v} else null; + } + const pairs = v.doc; + if (pairs.len > 0 and !query.all_operator_keys(pairs)) { + // Bare document equality (compare the whole doc). + return if (value_fast_path_safe(v)) &.{v} else null; + } + var eq: ?bson.Value = null; + var in_list: ?[]const bson.Value = null; + for (pairs) |p| { + if (std.mem.eql(u8, p.key, "$eq")) { + eq = p.value; + } else if (std.mem.eql(u8, p.key, "$in")) { + in_list = switch (p.value) { + .array => |a| a, + else => return null, + }; + } else return null; + } + if (eq) |e| { + return if (value_fast_path_safe(e)) &.{e} else null; + } + if (in_list) |list| { + for (list) |m| { + if (!value_fast_path_safe(m)) return null; + } + return list; + } + return null; +} + +/// Whether serialize_value is injective on the value's compare-equivalence +/// class. False for numbers (int32 1 / int64 1 / double 1.0), string/symbol/ +/// code (compared equal, serialized with different type tags), opaque +/// values (compared by payload only), and anything containing those. +fn value_fast_path_safe(v: bson.Value) bool { + return switch (v) { + .double, .int32, .int64, .string, .symbol, .code, .opaque_val => false, + .doc => |pairs| blk: { + for (pairs) |p| { + if (!value_fast_path_safe(p.value)) break :blk false; + } + break :blk true; + }, + .array => |items| blk: { + for (items) |it| { + if (!value_fast_path_safe(it)) break :blk false; + } + break :blk true; + }, + else => true, + }; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +const testing = std.testing; + +fn doc_of(pairs: []const bson.Pair) bson.Document { + return .{ .arena = undefined, .pairs = pairs }; +} + +fn simple_index(gpa: std.mem.Allocator, paths: []const []const u8, unique: bool, sparse: bool) !Index { + var keys: [max_index_keys]IndexKey = undefined; + for (paths, 0..) |p, i| keys[i] = .{ .path = p, .descending = false }; + return Index.init(gpa, "test", keys[0..paths.len], unique, sparse); +} + +/// Look up ids and compare with the expected set. Entry ids alias the +/// caller's storage, so tests pass stable static byte strings as ids. +fn expect_ids(gpa: std.mem.Allocator, ix: *const Index, key: []const bson.Value, expected: []const []const u8) !void { + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, key, &out); + std.mem.sort([]const u8, out.items, {}, less_ids); + try testing.expectEqual(expected.len, out.items.len); + for (out.items, 0..) |id, i| try testing.expectEqualStrings(expected[i], id); +} + +test "entries sort across numeric types and string/null/objectid" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{"a"}, false, false); + defer ix.deinit(gpa); + + const d_int = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "a", .value = .{ .int32 = 5 } } }); + const d_dbl = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 2 } }, .{ .key = "a", .value = .{ .double = 5.0 } } }); + const d_str = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 3 } }, .{ .key = "a", .value = .{ .string = "b" } } }); + const d_nul = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 4 } }, .{ .key = "a", .value = .null } }); + const d_oid = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 5 } }, .{ .key = "a", .value = .{ .object_id = [_]u8{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 } } } }); + try ix.add_doc(gpa, &d_int, "i1"); + try ix.add_doc(gpa, &d_dbl, "i2"); + try ix.add_doc(gpa, &d_str, "i3"); + try ix.add_doc(gpa, &d_nul, "i4"); + try ix.add_doc(gpa, &d_oid, "i5"); + + // An int64 query finds both the int32 and double entries: compare-equal. + try expect_ids(gpa, &ix, &.{.{ .int64 = 5 }}, &.{ "i1", "i2" }); + try expect_ids(gpa, &ix, &.{.null}, &.{"i4"}); + try expect_ids(gpa, &ix, &.{.{ .string = "b" }}, &.{"i3"}); + try expect_ids(gpa, &ix, &.{.{ .object_id = [_]u8{ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11 } }}, &.{"i5"}); + + // Full order: null, 5, "b", oid — sorted entries respect it. + try testing.expect(compare_entries(ix.entries.items[0], ix.entries.items[1]) == .lt); + try testing.expect(compare_entries(ix.entries.items[1], ix.entries.items[2]) == .lt); + try testing.expect(compare_entries(ix.entries.items[2], ix.entries.items[3]) == .lt); +} + +test "missing field is indexed as null; sparse skips the document" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{"a"}, false, false); + defer ix.deinit(gpa); + const d = doc_of(&.{.{ .key = "_id", .value = .{ .int32 = 1 } }}); + try ix.add_doc(gpa, &d, "m1"); + try testing.expectEqual(@as(usize, 1), ix.entries.items.len); + try expect_ids(gpa, &ix, &.{.null}, &.{"m1"}); + + var sp = try simple_index(gpa, &.{"a"}, false, true); + defer sp.deinit(gpa); + try sp.add_doc(gpa, &d, "m2"); + try testing.expectEqual(@as(usize, 0), sp.entries.items.len); +} + +test "multikey expansion indexes the array and its elements" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{"tags"}, false, false); + defer ix.deinit(gpa); + + const d = doc_of(&.{.{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "tags", .value = .{ .array = &.{ .{ .string = "a" }, .{ .string = "b" } } } }}); + try ix.add_doc(gpa, &d, "mk1"); + + // 3 entries: the array itself, "a", "b". + try testing.expectEqual(@as(usize, 3), ix.entries.items.len); + // Element query. + try expect_ids(gpa, &ix, &.{.{ .string = "a" }}, &.{"mk1"}); + try expect_ids(gpa, &ix, &.{.{ .string = "b" }}, &.{"mk1"}); + // Whole-array equality — the reason arrays are indexed twice. + try expect_ids(gpa, &ix, &.{.{ .array = &.{ .{ .string = "a" }, .{ .string = "b" } } }}, &.{"mk1"}); + try expect_ids(gpa, &ix, &.{.{ .array = &.{ .{ .string = "b" }, .{ .string = "a" } } }}, &.{}); +} + +test "per-document dedup keeps {a: [1,1]} under a unique index" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{"a"}, true, false); + defer ix.deinit(gpa); + const d = doc_of(&.{.{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "a", .value = .{ .array = &.{ .{ .int32 = 1 }, .{ .int32 = 1 } } } }}); + try ix.add_doc(gpa, &d, "d1"); + // Entries after dedup: the array itself and one element. + try testing.expectEqual(@as(usize, 2), ix.entries.items.len); +} + +test "parallel arrays are rejected" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{ "a", "b" }, false, false); + defer ix.deinit(gpa); + const d = doc_of(&.{ + .{ .key = "_id", .value = .{ .int32 = 1 } }, + .{ .key = "a", .value = .{ .array = &.{ .{ .int32 = 1 }, .{ .int32 = 2 } } } }, + .{ .key = "b", .value = .{ .array = &.{ .{ .int32 = 3 }, .{ .int32 = 4 } } } }, + }); + try testing.expectError(error.ParallelArrays, ix.add_doc(gpa, &d, "p1")); + + // One array path is fine. + const ok = doc_of(&.{ + .{ .key = "_id", .value = .{ .int32 = 2 } }, + .{ .key = "a", .value = .{ .array = &.{ .{ .int32 = 1 }, .{ .int32 = 2 } } } }, + .{ .key = "b", .value = .{ .int32 = 3 } }, + }); + try ix.add_doc(gpa, &ok, "p2"); + try testing.expectEqual(@as(usize, 3), ix.entries.items.len); +} + +test "unique conflict across documents, replace of own entries allowed" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{"a"}, true, false); + defer ix.deinit(gpa); + + const d1 = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "a", .value = .{ .int32 = 10 } } }); + try ix.add_doc(gpa, &d1, "u1"); + + const d2 = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 2 } }, .{ .key = "a", .value = .{ .int32 = 10 } } }); + try testing.expectError(error.DuplicateKeyIndex, ix.add_doc(gpa, &d2, "u2")); + + // A replace keeps its own key: remove old entries first (the engine's + // evict_doc does this), then add the new ones. + const d1b = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "a", .value = .{ .int32 = 20 } } }); + ix.remove_id(gpa, "u1"); + try ix.add_doc(gpa, &d1b, "u1"); + try expect_ids(gpa, &ix, &.{.{ .int32 = 20 }}, &.{"u1"}); + try expect_ids(gpa, &ix, &.{.{ .int32 = 10 }}, &.{}); +} + +test "range bounds inclusive and exclusive" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{"a"}, false, false); + defer ix.deinit(gpa); + const docs = [_]struct { id: []const u8, a: i32 }{ + .{ .id = "r1", .a = 1 }, + .{ .id = "r2", .a = 2 }, + .{ .id = "r3", .a = 3 }, + .{ .id = "r4", .a = 4 }, + .{ .id = "r5", .a = 5 }, + }; + for (docs) |s| { + const d = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = s.a } }, .{ .key = "a", .value = .{ .int32 = s.a } } }); + try ix.add_doc(gpa, &d, s.id); + } + + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + // 2 <= a < 4 → {2, 3} + try ix.lookup_range(gpa, &.{}, .{ .int32 = 2 }, true, .{ .int32 = 4 }, false, &out); + try testing.expectEqual(@as(usize, 2), out.items.len); + try testing.expect(std.mem.eql(u8, out.items[0], "r2") or std.mem.eql(u8, out.items[0], "r3")); + out.clearRetainingCapacity(); + // 2 < a <= 3 → {3} + try ix.lookup_range(gpa, &.{}, .{ .int32 = 2 }, false, .{ .int32 = 3 }, true, &out); + try testing.expectEqual(@as(usize, 1), out.items.len); + try testing.expectEqualStrings("r3", out.items[0]); + out.clearRetainingCapacity(); + // a > 4 → {5} + try ix.lookup_range(gpa, &.{}, .{ .int32 = 4 }, false, null, false, &out); + try testing.expectEqual(@as(usize, 1), out.items.len); + try testing.expectEqualStrings("r5", out.items[0]); +} + +test "empty index and remove_id" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{"a"}, false, false); + defer ix.deinit(gpa); + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_eq(gpa, &.{.{ .int32 = 1 }}, &out); + try testing.expectEqual(@as(usize, 0), out.items.len); + + const d = doc_of(&.{ .{ .key = "_id", .value = .{ .int32 = 1 } }, .{ .key = "a", .value = .{ .int32 = 1 } } }); + try ix.add_doc(gpa, &d, "e1"); + ix.remove_id(gpa, "e1"); + try testing.expectEqual(@as(usize, 0), ix.entries.items.len); + try ix.lookup_eq(gpa, &.{.{ .int32 = 1 }}, &out); + try testing.expectEqual(@as(usize, 0), out.items.len); +} + +test "compound index prefix search and range on the next key" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{ "a", "b" }, false, false); + defer ix.deinit(gpa); + + const id1 = "c1"; + const id2 = "c2"; + const id3 = "c3"; + const specs = [_]struct { id: []const u8, a: i32, b: i32 }{ + .{ .id = id1, .a = 1, .b = 2 }, + .{ .id = id2, .a = 1, .b = 3 }, + .{ .id = id3, .a = 2, .b = 1 }, + }; + for (specs) |s| { + const d = doc_of(&.{ + .{ .key = "_id", .value = .{ .int32 = s.a } }, + .{ .key = "a", .value = .{ .int32 = s.a } }, + .{ .key = "b", .value = .{ .int32 = s.b } }, + }); + try ix.add_doc(gpa, &d, s.id); + } + + // Prefix on a only. + try expect_ids(gpa, &ix, &.{.{ .int32 = 1 }}, &.{ "c1", "c2" }); + // Full key. + try expect_ids(gpa, &ix, &.{ .{ .int32 = 1 }, .{ .int32 = 3 } }, &.{"c2"}); + try expect_ids(gpa, &ix, &.{ .{ .int32 = 2 }, .{ .int32 = 1 } }, &.{"c3"}); + // Range on the next key: a == 1 and 2 <= b <= 3. + var out: std.ArrayListUnmanaged([]const u8) = .empty; + defer out.deinit(gpa); + try ix.lookup_range(gpa, &.{.{ .int32 = 1 }}, .{ .int32 = 2 }, true, .{ .int32 = 3 }, true, &out); + try testing.expectEqual(@as(usize, 2), out.items.len); + try testing.expect(std.mem.eql(u8, out.items[0], "c1") or std.mem.eql(u8, out.items[0], "c2")); +} + +test "id fast path guards and $in" { + // Numbers never use the fast path (compare-equal but serialize-different). + try testing.expect(plan_id(&.{.{ .key = "_id", .value = .{ .int32 = 1 } }}) == null); + // Strings are skipped too: a stored symbol/code _id compare-equals a + // string query but serializes differently. + try testing.expect(plan_id(&.{.{ .key = "_id", .value = .{ .string = "x" } }}) == null); + // Truly canonical values (bool, ObjectId) do use it. + try testing.expect(plan_id(&.{.{ .key = "_id", .value = .{ .bool = true } }}) != null); + const oid = [_]u8{ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12 }; + try testing.expect(plan_id(&.{.{ .key = "_id", .value = .{ .object_id = oid } }}) != null); + // $in with a number member is skipped. + const mixed = [_]bson.Pair{.{ .key = "_id", .value = .{ .doc = &.{ + .{ .key = "$in", .value = .{ .array = &.{ .{ .bool = true }, .{ .int32 = 1 } } } }, + } } }}; + try testing.expect(plan_id(&mixed) == null); + const safe_in = [_]bson.Pair{.{ .key = "_id", .value = .{ .doc = &.{ + .{ .key = "$in", .value = .{ .array = &.{ .{ .bool = true }, .{ .bool = false } } } }, + } } }}; + try testing.expect(plan_id(&safe_in) != null); + // $and members count as top-level. + const and_f = [_]bson.Pair{.{ .key = "$and", .value = .{ .array = &.{ + .{ .doc = &.{.{ .key = "_id", .value = .{ .bool = true } }} }, + } } }}; + try testing.expect(plan_id(&and_f) != null); + // $or is not usable for the fast path. + const or_f = [_]bson.Pair{.{ .key = "$or", .value = .{ .array = &.{ + .{ .doc = &.{.{ .key = "_id", .value = .{ .bool = true } }} }, + } } }}; + try testing.expect(plan_id(&or_f) == null); + // A doc containing a number is not fast-path safe. + const doc_id = [_]bson.Pair{.{ .key = "_id", .value = .{ .doc = &.{.{ .key = "n", .value = .{ .int32 = 1 } }} } }}; + try testing.expect(plan_id(&doc_id) == null); + // A doc containing a string is unsafe too (string/symbol/code class). + const doc_str = [_]bson.Pair{.{ .key = "_id", .value = .{ .doc = &.{.{ .key = "n", .value = .{ .string = "x" } }} } }}; + try testing.expect(plan_id(&doc_str) == null); + // A doc of canonical values is safe. + const doc_safe = [_]bson.Pair{.{ .key = "_id", .value = .{ .doc = &.{.{ .key = "n", .value = .{ .bool = true } }} } }}; + try testing.expect(plan_id(&doc_safe) != null); +} + +test "planner picks eq run, ranges, and bails on sparse null" { + const gpa = testing.allocator; + var ix = try simple_index(gpa, &.{ "a", "b" }, false, false); + defer ix.deinit(gpa); + var sp = try simple_index(gpa, &.{ "a", "b" }, false, true); + defer sp.deinit(gpa); + + // {a: 1, b: 2} → full-key equality. + { + const f = [_]bson.Pair{ + .{ .key = "a", .value = .{ .int32 = 1 } }, + .{ .key = "b", .value = .{ .int32 = 2 } }, + }; + var p = (try plan(gpa, &.{ix}, &f)).?; + defer p.deinit(gpa); + try testing.expectEqual(@as(usize, 2), p.key_len); + try testing.expect(p.lo == null and p.hi == null); + } + // {a: 1, b: {$gt: 2}} → equality run of 1 + range on the next key. + { + const f = [_]bson.Pair{ + .{ .key = "a", .value = .{ .int32 = 1 } }, + .{ .key = "b", .value = .{ .doc = &.{.{ .key = "$gt", .value = .{ .int32 = 2 } }} } }, + }; + var p = (try plan(gpa, &.{ix}, &f)).?; + defer p.deinit(gpa); + try testing.expectEqual(@as(usize, 1), p.key_len); + try testing.expect(p.hi == null and p.lo != null and !p.lo_incl); + } + // {a: 1} only → prefix run of 1. + { + const f = [_]bson.Pair{.{ .key = "a", .value = .{ .int32 = 1 } }}; + var p = (try plan(gpa, &.{ix}, &f)).?; + defer p.deinit(gpa); + try testing.expectEqual(@as(usize, 1), p.key_len); + } + // Pure range on the first key → key_len 0 with a bound. + { + const f = [_]bson.Pair{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$gte", .value = .{ .int32 = 1 } }} } }}; + var p = (try plan(gpa, &.{ix}, &f)).?; + defer p.deinit(gpa); + try testing.expectEqual(@as(usize, 0), p.key_len); + try testing.expect(p.lo != null and p.lo_incl); + } + // Unusable filter → no plan. + { + const f = [_]bson.Pair{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$regex", .value = .{ .string = "^x" } }} } }}; + try testing.expect((try plan(gpa, &.{ix}, &f)) == null); + const or_f = [_]bson.Pair{.{ .key = "$or", .value = .{ .array = &.{ + .{ .doc = &.{.{ .key = "a", .value = .{ .int32 = 1 } }} }, + } } }}; + try testing.expect((try plan(gpa, &.{ix}, &or_f)) == null); + } + // Sparse index bails on a null component. + { + const f = [_]bson.Pair{.{ .key = "a", .value = .null }}; + try testing.expect((try plan(gpa, &.{sp}, &f)) == null); + // Non-sparse is fine with null. + var p = (try plan(gpa, &.{ix}, &f)).?; + defer p.deinit(gpa); + try testing.expect(p.key_len == 1); + // A null inside $in bails too. + const fin = [_]bson.Pair{.{ .key = "a", .value = .{ .doc = &.{.{ .key = "$in", .value = .{ .array = &.{ .{ .int32 = 1 }, .null } } }} } }}; + try testing.expect((try plan(gpa, &.{sp}, &fin)) == null); + } + // $in cartesian product is capped. + { + var members: [20]bson.Value = undefined; + for (0..20) |j| members[j] = .{ .int32 = @intCast(j) }; + const f = [_]bson.Pair{ + .{ .key = "a", .value = .{ .doc = &.{.{ .key = "$in", .value = .{ .array = &members } }} } }, + .{ .key = "b", .value = .{ .doc = &.{.{ .key = "$in", .value = .{ .array = &members } }} } }, + }; + // 20 * 20 = 400 > 100 → fall back to a scan. + try testing.expect((try plan(gpa, &.{ix}, &f)) == null); + } +} diff --git a/src/lib.zig b/src/lib.zig index b8ee928..b483f22 100644 --- a/src/lib.zig +++ b/src/lib.zig @@ -8,6 +8,7 @@ pub const storage = @import("storage.zig"); pub const db = @import("db.zig"); pub const query = @import("query.zig"); pub const update = @import("update.zig"); +pub const index = @import("index.zig"); test { _ = @import("bson.zig"); @@ -18,4 +19,5 @@ test { _ = @import("db.zig"); _ = @import("query.zig"); _ = @import("update.zig"); + _ = @import("index.zig"); } diff --git a/src/query.zig b/src/query.zig index 93ea36d..c5aff04 100644 --- a/src/query.zig +++ b/src/query.zig @@ -57,16 +57,26 @@ fn match_top_level(gpa: std.mem.Allocator, op: []const u8, value: bson.Value, do return false; } +/// True when every key is a `$`-prefixed operator. The one definition of +/// what an operator document looks like — filters, $elemMatch, $pull and +/// upsert-document construction all ask this question. +pub fn all_operator_keys(pairs: []const bson.Pair) bool { + if (pairs.len == 0) return false; + for (pairs) |p| { + if (p.key.len == 0 or p.key[0] != '$') return false; + } + return true; +} + +/// The operator pairs of a `{$op: ...}` value, or null if it is not one. An +/// empty document counts as an (vacuously satisfied) operator document. fn is_operator_doc(value: bson.Value) ?[]const bson.Pair { - return switch (value) { - .doc => |pairs| blk: { - for (pairs) |p| { - if (p.key.len == 0 or p.key[0] != '$') break :blk null; - } - break :blk pairs; - }, - else => null, + const pairs = switch (value) { + .doc => |p| p, + else => return null, }; + if (pairs.len > 0 and !all_operator_keys(pairs)) return null; + return pairs; } fn field_matches(gpa: std.mem.Allocator, path: []const u8, expected: bson.Value, doc: *const bson.Document) QueryError!bool { @@ -220,13 +230,7 @@ fn match_operator(gpa: std.mem.Allocator, op: []const u8, value: bson.Value, act .doc => |pairs| pairs, else => return false, }; - var all_operators = operand.len > 0; - for (operand) |p| { - if (p.key.len == 0 or p.key[0] != '$') { - all_operators = false; - break; - } - } + const all_operators = all_operator_keys(operand); for (actuals) |a| { if (a != .array) continue; for (a.array) |elem| { @@ -258,8 +262,9 @@ fn match_operator(gpa: std.mem.Allocator, op: []const u8, value: bson.Value, act /// Collect values reachable at `path` (dot-separated), descending into /// documents and, per MongoDB multikey semantics, into arrays of documents. /// Appends into `out`; on OOM, collection stops early (the engine is -/// already failing at that point). -fn collect_values(gpa: std.mem.Allocator, pairs: []const bson.Pair, path: []const u8, out: *std.ArrayListUnmanaged(bson.Value), depth: usize) QueryError!void { +/// already failing at that point). Public because index entry generation +/// must mirror field_matches exactly (src/index.zig). +pub fn collect_values(gpa: std.mem.Allocator, pairs: []const bson.Pair, path: []const u8, out: *std.ArrayListUnmanaged(bson.Value), depth: usize) QueryError!void { var it = std.mem.splitScalar(u8, path, '.'); const first = it.next() orelse return; @@ -550,7 +555,7 @@ pub const SortKey = struct { /// Sort `docs` in place by `keys`. Candidate values are collected up front /// (allocations happen before the sort), so the comparator itself is pure /// and cannot fail — OOM during collection propagates as QueryError. -pub fn sort_docs(gpa: std.mem.Allocator, arena: std.mem.Allocator, docs: []*const bson.Document, keys: []const SortKey) QueryError!void { +pub fn sort_docs(arena: std.mem.Allocator, docs: []*const bson.Document, keys: []const SortKey) QueryError!void { if (keys.len == 0 or docs.len < 2) return; const SortedDoc = struct { @@ -583,14 +588,13 @@ pub fn sort_docs(gpa: std.mem.Allocator, arena: std.mem.Allocator, docs: []*cons std.mem.sort(SortedDoc, entries, Ctx{ .keys = keys }, Ctx.lessThan); for (entries, 0..) |e, i| docs[i] = e.doc; - _ = gpa; } // --------------------------------------------------------------------------- // Projection // --------------------------------------------------------------------------- -pub const ProjectionError = error{ OutOfMemory, InvalidProjection }; +pub const ProjectionError = std.mem.Allocator.Error; /// Apply a projection document, writing resulting pairs into `out` (which /// should use the caller's arena so strings are owned). @@ -614,7 +618,7 @@ pub fn project(arena: std.mem.Allocator, doc: *const bson.Document, proj: *const } if (include_id) { if (bson.get_pair(doc.pairs, "_id")) |idv| { - try out.append(arena, .{ .key = try arena.dupe(u8, "_id"), .value = try bson.copy_value(arena, idv) }); + try out.append(arena, .{ .key = "_id", .value = try bson.copy_value(arena, idv) }); } } for (proj.pairs) |p| { @@ -659,10 +663,7 @@ fn exclude_doc(arena: std.mem.Allocator, pairs: []const bson.Pair, proj: *const } fn is_excluded(proj: *const bson.Document, key: []const u8) bool { - for (proj.pairs) |pp| { - if (std.mem.eql(u8, pp.key, key)) return true; - } - return false; + return bson.get_pair(proj.pairs, key) != null; } fn has_deeper_exclusion(proj: *const bson.Document, key: []const u8) bool { @@ -875,17 +876,17 @@ test "sort compares by BSON order" { var docs = [_]*const bson.Document{ &b, &a }; const asc = [_]SortKey{.{ .path = "n", .descending = false }}; - try sort_docs(testing.allocator, arena.allocator(), &docs, &asc); + try sort_docs(arena.allocator(), &docs, &asc); try testing.expect(docs[0] == &a); try testing.expect(docs[1] == &b); var docs2 = [_]*const bson.Document{ &a, &b }; const desc = [_]SortKey{.{ .path = "n", .descending = true }}; - try sort_docs(testing.allocator, arena.allocator(), &docs2, &desc); + try sort_docs(arena.allocator(), &docs2, &desc); try testing.expect(docs2[0] == &b); const missing = [_]SortKey{.{ .path = "zz", .descending = false }}; - try sort_docs(testing.allocator, arena.allocator(), &docs2, &missing); + try sort_docs(arena.allocator(), &docs2, &missing); try testing.expect(docs2[0] == &b); // stable-ish: order untouched by missing key } diff --git a/src/server.zig b/src/server.zig index 94c6b42..1f6f35b 100644 --- a/src/server.zig +++ b/src/server.zig @@ -120,17 +120,14 @@ fn handle_connection_inner(io: std.Io, stream: std.Io.net.Stream, server: *Serve }; out_buf.clearRetainingCapacity(); - if (msg.op_code == wire.op_code_query) { - wire.write_reply_query(server.gpa, reply_request_id, msg.request_id, reply.pairs.items, &out_buf) catch |err| { - std.debug.print("mongo-light: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) }); - return; - }; - } else { - reply.build(server.gpa, reply_request_id, msg.request_id, &out_buf) catch |err| { - std.debug.print("mongo-light: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) }); - return; - }; - } + const built = if (msg.op_code == wire.op_code_query) + wire.write_reply_query(server.gpa, reply_request_id, msg.request_id, reply.pairs.items, &out_buf) + else + reply.build(server.gpa, reply_request_id, msg.request_id, &out_buf); + built catch |err| { + std.debug.print("mongo-light: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) }); + return; + }; reply_request_id +%= 1; writer.interface.writeAll(out_buf.items) catch |err| { std.debug.print("mongo-light: write error on conn {d}: {s}\n", .{ connection_id, @errorName(err) }); diff --git a/src/storage.zig b/src/storage.zig index f086ad3..fc4a1c7 100644 --- a/src/storage.zig +++ b/src/storage.zig @@ -14,6 +14,8 @@ const bson = @import("bson.zig"); pub const record_type_upsert: u8 = 1; pub const record_type_delete: u8 = 2; +pub const record_type_index_create: u8 = 3; +pub const record_type_index_drop: u8 = 4; pub const header_len: usize = 20; // len + crc + seq + type + reserved @@ -26,7 +28,6 @@ pub const Record = struct { type: u8, db: []const u8, // transient: valid only during replay callback coll: []const u8, - doc: []const u8, // raw bson bytes }; pub const Error = error{ @@ -44,6 +45,9 @@ pub const Log = struct { path: []const u8, end_pos: u64, log_bytes: u64, // bytes written since the log was last rewritten + // Reused record-framing buffer. Appends are single-writer (the engine's + // exclusive lock), so one buffer avoids a realloc cycle per record. + scratch: std.ArrayListUnmanaged(u8), pub fn open(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Log { // Resolve to an absolute path so compaction can rename the file @@ -68,10 +72,12 @@ pub const Log = struct { .path = abs_path, .end_pos = 0, .log_bytes = 0, + .scratch = .empty, }; } pub fn close(self: *Log) void { + self.scratch.deinit(self.gpa); self.file.close(self.io); self.gpa.free(self.path); } @@ -134,7 +140,6 @@ pub const Log = struct { .type = rtype, .db = db, .coll = coll, - .doc = doc_bytes, }, doc); pos += total; @@ -150,12 +155,23 @@ pub const Log = struct { try self.append(record_type_delete, db, coll, doc, seq); } + /// The payload is the canonical index spec document ({v, key, name, + /// unique?, sparse?}); only apply_record interprets it. + pub fn append_index_create(self: *Log, db: []const u8, coll: []const u8, doc: []const u8, seq: u64) !void { + try self.append(record_type_index_create, db, coll, doc, seq); + } + + /// The payload is {name: "..."}; only apply_record interprets it. + pub fn append_index_drop(self: *Log, db: []const u8, coll: []const u8, doc: []const u8, seq: u64) !void { + try self.append(record_type_index_drop, db, coll, doc, seq); + } + fn append(self: *Log, rtype: u8, db: []const u8, coll: []const u8, doc: []const u8, seq: u64) !void { if (std.mem.indexOfScalar(u8, db, 0) != null or std.mem.indexOfScalar(u8, coll, 0) != null) { return error.NulInName; } - var buf: std.ArrayListUnmanaged(u8) = .empty; - defer buf.deinit(self.gpa); + const buf = &self.scratch; + buf.clearRetainingCapacity(); try buf.appendNTimes(self.gpa, 0, header_len); std.mem.writeInt(u64, buf.items[8..16], seq, .little); buf.items[16] = rtype; @@ -189,17 +205,19 @@ pub const Log = struct { const testing = std.testing; -const TmpLog = struct { +/// Throwaway log file in the test temp dir, shared by the storage and +/// engine test suites. +pub const TmpLog = struct { tmp: std.testing.TmpDir, path: []u8, - fn init(gpa: std.mem.Allocator) !TmpLog { + pub fn init(gpa: std.mem.Allocator) !TmpLog { const tmp = std.testing.tmpDir(.{}); const path = try std.fmt.allocPrint(gpa, ".zig-cache/tmp/{s}/test.log", .{tmp.sub_path}); return .{ .tmp = tmp, .path = path }; } - fn deinit(self: *TmpLog, gpa: std.mem.Allocator) void { + pub fn deinit(self: *TmpLog, gpa: std.mem.Allocator) void { self.tmp.cleanup(); gpa.free(self.path); } diff --git a/src/update.zig b/src/update.zig index fc4f4d0..447fc62 100644 --- a/src/update.zig +++ b/src/update.zig @@ -13,7 +13,7 @@ const max_path_segments = 16; /// Apply an update document (whose fields are operator documents) to `doc`. pub fn apply(doc: *bson.Document, update: *const bson.Document) UpdateError!void { const arena = doc.arena.allocator(); - var pairs = try copy_pairs_to_list(arena, doc.pairs); + var pairs = try copy_to_list(bson.Pair, arena, doc.pairs); for (update.pairs) |op| { if (op.key.len == 0 or op.key[0] != '$') return error.InvalidUpdate; try apply_operator(arena, &pairs, op.key, op.value); @@ -48,7 +48,7 @@ fn apply_operator(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson. const n = split_path(p.key, &segs) orelse return error.InvalidUpdate; const current = get_value(pairs.items, segs[0..n]) orelse bson.Value{ .int32 = 0 }; if (!current.is_number() or !p.value.is_number()) return error.InvalidUpdate; - const sum = try numeric_add(arena, current, p.value); + const sum = try numeric_add(current, p.value); try set_path(arena, pairs, segs[0..n], sum); } return; @@ -145,19 +145,15 @@ fn parse_index(seg: []const u8) ?usize { return std.fmt.parseInt(usize, seg, 10) catch null; } -fn copy_pairs_to_list(arena: std.mem.Allocator, pairs: []const bson.Pair) UpdateError!std.ArrayListUnmanaged(bson.Pair) { - var out: std.ArrayListUnmanaged(bson.Pair) = .empty; +/// Shallow-copy a slice into a growable list backed by `arena`. +fn copy_to_list(comptime T: type, arena: std.mem.Allocator, items: []const T) UpdateError!std.ArrayListUnmanaged(T) { + var out: std.ArrayListUnmanaged(T) = .empty; errdefer out.deinit(arena); - try out.appendSlice(arena, pairs); + try out.appendSlice(arena, items); return out; } -fn find_pair(pairs: []const bson.Pair, key: []const u8) ?usize { - for (pairs, 0..) |p, i| { - if (std.mem.eql(u8, p.key, key)) return i; - } - return null; -} +const find_pair = bson.get_pair_index; fn get_value(pairs: []const bson.Pair, segs: []const []const u8) ?bson.Value { const idx = find_pair(pairs, segs[0]) orelse return null; @@ -193,7 +189,7 @@ fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), }; switch (pairs.items[idx].value) { .doc => |sub| { - var sub_pairs = try copy_pairs_to_list(arena, sub); + var sub_pairs = try copy_to_list(bson.Pair, arena, sub); defer sub_pairs.deinit(arena); try set_path(arena, &sub_pairs, segs[1..], value); pairs.items[idx].value = .{ .doc = try sub_pairs.toOwnedSlice(arena) }; @@ -207,7 +203,7 @@ fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), pairs.items[idx].value = .{ .doc = try sub_pairs.toOwnedSlice(arena) }; return; }; - var items = try copy_array_to_list(arena, arr); + var items = try copy_to_list(bson.Value, arena, arr); defer items.deinit(arena); if (index >= items.items.len) { try items.appendNTimes(arena, .null, index + 1 - items.items.len); @@ -217,7 +213,7 @@ fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), } else { switch (items.items[index]) { .doc => |sub| { - var sub_pairs = try copy_pairs_to_list(arena, sub); + var sub_pairs = try copy_to_list(bson.Pair, arena, sub); defer sub_pairs.deinit(arena); try set_path(arena, &sub_pairs, segs[2..], value); items.items[index] = .{ .doc = try sub_pairs.toOwnedSlice(arena) }; @@ -241,13 +237,6 @@ fn set_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), } } -fn copy_array_to_list(arena: std.mem.Allocator, arr: []const bson.Value) UpdateError!std.ArrayListUnmanaged(bson.Value) { - var out: std.ArrayListUnmanaged(bson.Value) = .empty; - errdefer out.deinit(arena); - try out.appendSlice(arena, arr); - return out; -} - fn unset_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair), segs: []const []const u8) void { if (segs.len == 1) { if (find_pair(pairs.items, segs[0])) |idx| { @@ -258,7 +247,7 @@ fn unset_path(arena: std.mem.Allocator, pairs: *std.ArrayListUnmanaged(bson.Pair const idx = find_pair(pairs.items, segs[0]) orelse return; switch (pairs.items[idx].value) { .doc => |sub| { - var sub_pairs = copy_pairs_to_list(arena, sub) catch return; + var sub_pairs = copy_to_list(bson.Pair, arena, sub) catch return; unset_path(arena, &sub_pairs, segs[1..]); pairs.items[idx].value = .{ .doc = sub_pairs.items }; }, @@ -273,14 +262,7 @@ fn pull_matches(arena: std.mem.Allocator, condition: bson.Value, elem: bson.Valu .doc => |pairs| pairs, else => return false, }; - var all_operators = cond_pairs.len > 0; - for (cond_pairs) |p| { - if (p.key.len == 0 or p.key[0] != '$') { - all_operators = false; - break; - } - } - if (all_operators) { + if (query.all_operator_keys(cond_pairs)) { // Operator condition against the element's value at each // operator's field — treat element doc as the doc. var ok = true; @@ -297,8 +279,7 @@ fn pull_matches(arena: std.mem.Allocator, condition: bson.Value, elem: bson.Valu } } -fn numeric_add(arena: std.mem.Allocator, a: bson.Value, b: bson.Value) UpdateError!bson.Value { - _ = arena; +fn numeric_add(a: bson.Value, b: bson.Value) UpdateError!bson.Value { if (a == .double or b == .double) { const sum: f64 = @floatCast(a.as_f128() + b.as_f128()); return .{ .double = sum }; diff --git a/src/wire.zig b/src/wire.zig index 63bbec6..f4ed4bc 100644 --- a/src/wire.zig +++ b/src/wire.zig @@ -13,9 +13,6 @@ pub const op_code_reply: i32 = 2001; pub const max_message_size: usize = 48 * 1024 * 1024; pub const max_bson_object_size: i32 = 16 * 1024 * 1024; -pub const flag_checksum_present: u32 = 1 << 1; -pub const flag_more_to_come: u32 = 1 << 0; - pub const Seq = struct { name: []const u8, docs: []bson.Document, @@ -175,6 +172,37 @@ pub const Message = struct { else => null, }; } + + /// Documents of a batch argument (`documents`, `updates`, `deletes`). + /// Drivers send them either as an OP_MSG document sequence or as an array + /// inside the command body; callers should not have to care which. The + /// result borrows this message's storage — never deinit the documents. + pub fn batch(self: *Message, name: []const u8) BatchError![]const bson.Document { + for (self.seqs) |seq| { + if (std.mem.eql(u8, seq.name, name) and seq.docs.len > 0) return seq.docs; + } + const arr = switch (bson.get_pair(self.body.pairs, name) orelse return error.MissingBatch) { + .array => |a| a, + else => return error.BatchNotArray, + }; + const docs = try self.arena.allocator().alloc(bson.Document, arr.len); + for (arr, 0..) |item, i| { + docs[i] = switch (item) { + // Pairs are borrowed from the body document, which owns the + // arena; these views must never be deinited. + .doc => |pairs| .{ .arena = undefined, .pairs = pairs }, + else => return error.BatchElementNotDoc, + }; + } + return docs; + } +}; + +pub const BatchError = error{ + MissingBatch, // no sequence and no body field of that name + BatchNotArray, // body field is present but not an array + BatchElementNotDoc, // an array element is not a document + OutOfMemory, }; /// Builder for a command reply document. Strings for keys and values must @@ -221,6 +249,32 @@ pub const Reply = struct { } }; +/// Write the 16-byte message header with a zero length placeholder. Returns +/// the offset `end_message` needs to patch the length in. +fn begin_message( + gpa: std.mem.Allocator, + out: *std.ArrayListUnmanaged(u8), + request_id: u32, + response_to: u32, + op_code: i32, +) !usize { + const len_pos = out.items.len; + var header: [16]u8 = undefined; + std.mem.writeInt(u32, header[0..4], 0, .little); // patched by end_message + std.mem.writeInt(u32, header[4..8], request_id, .little); + std.mem.writeInt(u32, header[8..12], response_to, .little); + std.mem.writeInt(i32, header[12..16], op_code, .little); + try out.appendSlice(gpa, &header); + return len_pos; +} + +/// Patch in the total length of the message started at `len_pos`. +fn end_message(out: *std.ArrayListUnmanaged(u8), len_pos: usize) !void { + const total = out.items.len - len_pos; + if (total > std.math.maxInt(u32)) return error.MessageTooLarge; + std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little); +} + /// Serialize an OP_MSG reply: header + flags + single body section. pub fn write_message( gpa: std.mem.Allocator, @@ -230,18 +284,13 @@ pub fn write_message( body: []const bson.Pair, out: *std.ArrayListUnmanaged(u8), ) !void { - const len_pos = out.items.len; - var header: [20]u8 = undefined; - std.mem.writeInt(u32, header[4..8], request_id, .little); - std.mem.writeInt(u32, header[8..12], response_to, .little); - std.mem.writeInt(i32, header[12..16], op_code_msg, .little); - std.mem.writeInt(u32, header[16..20], flags, .little); - try out.appendSlice(gpa, &header); + const len_pos = try begin_message(gpa, out, request_id, response_to, op_code_msg); + var flag_bytes: [4]u8 = undefined; + std.mem.writeInt(u32, &flag_bytes, flags, .little); + try out.appendSlice(gpa, &flag_bytes); try out.append(gpa, 0x00); // single body section try bson.write_doc(body, gpa, out); - const total = out.items.len - len_pos; - if (total > std.math.maxInt(u32)) return error.MessageTooLarge; - std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little); + try end_message(out, len_pos); } /// Serialize an OP_REPLY (legacy): header + responseFlags | cursorID | @@ -254,19 +303,12 @@ pub fn write_reply_query( body: []const bson.Pair, out: *std.ArrayListUnmanaged(u8), ) !void { - const len_pos = out.items.len; - var header: [16]u8 = undefined; - std.mem.writeInt(u32, header[4..8], request_id, .little); - std.mem.writeInt(u32, header[8..12], response_to, .little); - std.mem.writeInt(i32, header[12..16], op_code_reply, .little); - try out.appendSlice(gpa, &header); + const len_pos = try begin_message(gpa, out, request_id, response_to, op_code_reply); var reply_fields: [20]u8 = [_]u8{0} ** 20; // responseFlags | cursorID | startingFrom | numberReturned std.mem.writeInt(i32, reply_fields[16..20], 1, .little); // numberReturned = 1 try out.appendSlice(gpa, &reply_fields); try bson.write_doc(body, gpa, out); - const total = out.items.len - len_pos; - if (total > std.math.maxInt(u32)) return error.MessageTooLarge; - std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little); + try end_message(out, len_pos); } // ---------------------------------------------------------------------------