index: secondary index core — entries, search, planner, _id fast path

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.
This commit is contained in:
2026-08-02 12:21:22 +03:00
parent 662df9b121
commit a38ddc2f50
10 changed files with 1824 additions and 782 deletions

View File

@@ -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`

File diff suppressed because it is too large Load Diff

View File

@@ -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();
}

1190
src/index.zig Normal file

File diff suppressed because it is too large Load Diff

View File

@@ -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");
}

View File

@@ -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
}

View File

@@ -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) });

View File

@@ -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);
}

View File

@@ -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 };

View File

@@ -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);
}
// ---------------------------------------------------------------------------