MultiforaDB/src/db.zig

//! In-memory database engine backed by the append-only log. Maps
//! db -> collection -> _id(serialized) -> owned Document. All mutations are
//! logged and synced before they become visible in memory, so a crash never
//! loses a committed write. Callers must hold the write lock (`lock`) around
//! any command that mutates state, and the read lock (`lock_read`) around
//! read-only commands so reads overlap with each other.

const std = @import("std");
const bson = @import("bson.zig");
const storage = @import("storage.zig");
const index = @import("index.zig");

pub const Collection = struct {
    docs: std.StringHashMapUnmanaged(*bson.Document),
    indexes: std.ArrayListUnmanaged(index.Index),

    fn init() Collection {
        return .{ .docs = .empty, .indexes = .empty };
    }
};

pub const Db = struct {
    collections: std.StringHashMapUnmanaged(Collection),
};

pub const Engine = struct {
    gpa: std.mem.Allocator,
    io: std.Io,
    // One writer at a time (log append + fsync, map mutation); many
    // concurrent readers (find/count/aggregate scans). Writer-preferring:
    // a queued writer blocks new readers rather than starving.
    rwlock: std.Io.RwLock,
    log: storage.Log,
    dbs: std.StringHashMapUnmanaged(Db),
    seq: u64,
    compact_threshold: u64,
    /// Set to the failing index's own stable name when an upsert is
    /// rejected by a unique secondary index (error.DuplicateKeyIndex). The
    /// command reads it while still holding the write lock.
    dup_index: ?[]const u8 = null,

    pub fn open(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Engine {
        var engine = Engine{
            .gpa = gpa,
            .io = io,
            .rwlock = .init,
            .log = try storage.Log.open(gpa, io, path),
            .dbs = .empty,
            .seq = 0,
            .compact_threshold = 16 * 1024 * 1024,
        };
        errdefer {
            engine.log.close();
            engine.dbs.deinit(gpa);
        }

        try engine.log.replay(&engine, apply_record);
        // Replay registers empty indexes; build them from the live docs
        // once replay completes (order-independent).
        try engine.build_all_indexes();
        return engine;
    }

    pub fn deinit(self: *Engine) void {
        var db_it = self.dbs.iterator();
        while (db_it.next()) |db_entry| {
            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
    /// and secondary indexes (whose entries alias the documents — freed
    /// first).
    fn free_collection(self: *Engine, coll: *Collection) void {
        for (coll.indexes.items) |*ix| ix.deinit(self.gpa);
        coll.indexes.deinit(self.gpa);
        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. This is the single chokepoint where a
    /// document dies, so index entries are removed here — before
    /// old.value.deinit() and gpa.free(old.key) — keeping the entry aliasing
    /// (values into the document arena, id into the docs map key) safe.
    fn evict_doc(self: *Engine, coll: *Collection, id_key: []const u8) void {
        for (coll.indexes.items) |*ix| ix.remove_id(self.gpa, id_key);
        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.
    pub fn lock(self: *Engine) !void {
        try self.rwlock.lock(self.io);
    }

    pub fn unlock(self: *Engine) void {
        self.rwlock.unlock(self.io);
    }

    /// Shared lock: for read-only commands (find, count, aggregate, list*).
    /// Multiple readers may hold it simultaneously; writers wait for them.
    pub fn lock_read(self: *Engine) !void {
        try self.rwlock.lockShared(self.io);
    }

    pub fn unlock_read(self: *Engine) void {
        self.rwlock.unlockShared(self.io);
    }

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

    /// One document's built entries for one index, tracked so a failure
    /// anywhere before the log append frees them all.
    const Built = struct {
        built: index.BuiltEntries,
        ix: *index.Index,
    };

    /// 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);
        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);
        };
        self.dup_index = null;

        // 1. Build entries for every index. ParallelArrays escapes here,
        //    before anything is logged or mutated.
        var built_list: std.ArrayListUnmanaged(Built) = .empty;
        defer {
            for (built_list.items) |*b| b.built.deinit(self.gpa);
            built_list.deinit(self.gpa);
        }
        for (coll.indexes.items) |*ix| {
            var built = try ix.build_entries(self.gpa, owned, id_key);
            built_list.append(self.gpa, .{ .built = built, .ix = ix }) catch |err| {
                built.deinit(self.gpa);
                return err;
            };
        }

        // 2. The _id check, mirroring the pre-index behavior.
        if (mode == .insert and coll.docs.contains(id_key)) return error.DuplicateKey;

        // 3. Unique secondary-index checks; a rejected write never reaches
        //    the log.
        for (built_list.items) |*b| {
            if (!b.ix.unique) continue;
            b.ix.check_unique(b.built.entries.items, id_key) catch {
                self.dup_index = b.ix.name;
                return error.DuplicateKeyIndex;
            };
        }

        // 4. Reserve entry capacity — the last fallible step, so the entry
        //    insertion after the log append is infallible.
        for (built_list.items) |*b| {
            if (b.built.entries.items.len == 0) continue;
            try b.ix.entries.ensureUnusedCapacity(self.gpa, b.built.entries.items.len);
        }

        // 5. Log (and sync) before anything becomes visible.
        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);

        // 6. Replace drops the old document (and its index entries).
        if (mode == .replace) self.evict_doc(coll, id_key);

        // 7. Publish the document and its entries.
        try coll.docs.put(self.gpa, id_key, owned);
        for (built_list.items) |*b| {
            if (b.built.multikey) b.ix.multikey = true;
            b.ix.insert_entries(self.gpa, &b.built);
        }
        stored = true;
        try self.maybe_compact();
    }

    /// 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.
        // 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, id_doc.items, self.seq);

        self.evict_doc(coll, id_key);
        return true;
    }

    pub fn get_collection(self: *Engine, db_name: []const u8, coll_name: []const u8) ?*Collection {
        const db = self.dbs.get(db_name) orelse return null;
        return db.collections.getPtr(coll_name);
    }

    pub fn get_doc(self: *Engine, db_name: []const u8, coll_name: []const u8, id_key: []const u8) ?*const bson.Document {
        const coll = self.get_collection(db_name, coll_name) orelse return null;
        return coll.docs.get(id_key);
    }

    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;
        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;
        self.free_db(&removed.value);
        self.gpa.free(removed.key);
        return true;
    }

    /// Build and register a secondary index from a spec document
    /// ({key, name, unique?, sparse?}). The create record is written only
    /// after the index builds over the existing documents and passes
    /// uniqueness, so a rejected create persists nothing. Returns the new
    /// index (or the existing one when the spec matches — idempotent).
    pub fn create_index(self: *Engine, db_name: []const u8, coll_name: []const u8, spec_doc: *const bson.Document) !*index.Index {
        const coll = try self.get_or_create_collection(db_name, coll_name);
        var ix = try index.parse_spec(self.gpa, spec_doc);
        var committed = false;
        // Runs on every return path (including the idempotent no-op): the
        // parsed spec is only owned by the collection once committed.
        defer if (!committed) ix.deinit(self.gpa);

        for (coll.indexes.items) |*existing| {
            if (std.mem.eql(u8, existing.name, ix.name)) {
                if (index.Index.spec_equal(existing, &ix)) return existing;
                return error.IndexOptionsConflict;
            }
        }

        // Build entries over the existing documents, checking uniqueness as
        // we go (the index is not exposed until the end, so mutating it is
        // safe). Each batch is inserted into the index immediately (which
        // drains it), so on any later failure the errdefer ix.deinit frees
        // every inserted entry key; a batch that fails before insertion is
        // freed by its own errdefer. Nothing is persisted on failure.
        var doc_it = coll.docs.iterator();
        while (doc_it.next()) |entry| {
            var built = try ix.build_entries(self.gpa, entry.value_ptr.*, entry.key_ptr.*);
            // Runs on success too: insert_entries drains the keys, leaving
            // only the (now empty) ArrayList buffer to free.
            defer built.deinit(self.gpa);
            if (built.multikey) ix.multikey = true;
            if (ix.unique) {
                try ix.check_unique(built.entries.items, entry.key_ptr.*);
            }
            if (built.entries.items.len == 0) continue;
            try ix.entries.ensureUnusedCapacity(self.gpa, built.entries.items.len);
            ix.insert_entries(self.gpa, &built);
        }

        // Reserve the collection slot, then persist and publish.
        try coll.indexes.ensureUnusedCapacity(self.gpa, 1);
        var spec_bytes: std.ArrayListUnmanaged(u8) = .empty;
        defer spec_bytes.deinit(self.gpa);
        try ix.write_spec(self.gpa, &spec_bytes);
        self.seq += 1;
        try self.log.append_index_create(db_name, coll_name, spec_bytes.items, self.seq);

        coll.indexes.appendAssumeCapacity(ix);
        committed = true;
        return &coll.indexes.items[coll.indexes.items.len - 1];
    }

    /// Remove a secondary index by name, persisting a drop record first.
    /// Returns false when no such index exists.
    pub fn drop_index(self: *Engine, db_name: []const u8, coll_name: []const u8, index_name: []const u8) !bool {
        const db = self.dbs.get(db_name) orelse return false;
        const coll = db.collections.getPtr(coll_name) orelse return false;
        var found = false;
        for (coll.indexes.items) |ix| {
            if (std.mem.eql(u8, ix.name, index_name)) {
                found = true;
                break;
            }
        }
        if (!found) return false;

        const name_pairs = [_]bson.Pair{.{ .key = "name", .value = .{ .string = index_name } }};
        var name_doc: std.ArrayListUnmanaged(u8) = .empty;
        defer name_doc.deinit(self.gpa);
        try bson.write_doc(&name_pairs, self.gpa, &name_doc);
        self.seq += 1;
        try self.log.append_index_drop(db_name, coll_name, name_doc.items, self.seq);

        var i: usize = 0;
        while (i < coll.indexes.items.len) {
            if (std.mem.eql(u8, coll.indexes.items[i].name, index_name)) {
                var removed = coll.indexes.orderedRemove(i);
                removed.deinit(self.gpa);
            } else i += 1;
        }
        return true;
    }

    pub fn database_names(self: *Engine, out: *std.ArrayListUnmanaged([]const u8)) !void {
        var it = self.dbs.iterator();
        while (it.next()) |entry| try out.append(self.gpa, entry.key_ptr.*);
    }

    pub fn collection_names(self: *Engine, db_name: []const u8, out: *std.ArrayListUnmanaged([]const u8)) !void {
        const db = self.dbs.get(db_name) orelse return;
        var it = db.collections.iterator();
        while (it.next()) |entry| try out.append(self.gpa, entry.key_ptr.*);
    }

    // -- internals -----------------------------------------------------------

    pub fn get_or_create_collection(self: *Engine, db_name: []const u8, coll_name: []const u8) !*Collection {
        const db = self.dbs.getPtr(db_name) orelse {
            const db_key = try self.gpa.dupe(u8, db_name);
            errdefer self.gpa.free(db_key);
            try self.dbs.put(self.gpa, db_key, .{ .collections = .empty });
            return self.get_or_create_collection(db_name, coll_name);
        };
        if (db.collections.getPtr(coll_name)) |coll| return coll;
        const coll_key = try self.gpa.dupe(u8, coll_name);
        errdefer self.gpa.free(coll_key);
        try db.collections.put(self.gpa, coll_key, Collection.init());
        return db.collections.getPtr(coll_name) orelse unreachable;
    }

    /// Deep-copy a document into engine-owned storage, prepending a
    /// generated ObjectId `_id` when absent.
    fn own_with_id(self: *Engine, doc: *const bson.Document, oid_gen: *bson.ObjectIdGen) !*bson.Document {
        var pairs: std.ArrayListUnmanaged(bson.Pair) = .empty;
        defer pairs.deinit(self.gpa);
        if (doc.get("_id") == null) {
            const oid = oid_gen.new(self.io);
            try pairs.append(self.gpa, .{ .key = "_id", .value = .{ .object_id = oid } });
        }
        try pairs.appendSlice(self.gpa, doc.pairs);

        var out: std.ArrayListUnmanaged(u8) = .empty;
        defer out.deinit(self.gpa);
        try bson.write_doc(pairs.items, self.gpa, &out);
        const owned = try self.gpa.create(bson.Document);
        errdefer self.gpa.destroy(owned);
        owned.* = try bson.Document.parse(self.gpa, out.items);
        return owned;
    }

    fn maybe_compact(self: *Engine) !void {
        if (self.log.log_bytes < self.compact_threshold) return;
        try self.compact();
    }

    /// Rewrite the log with only live documents, atomically swapping the file.
    /// Callers must hold the write lock.
    pub fn compact(self: *Engine) !void {
        const tmp_path = try std.fmt.allocPrint(self.gpa, "{s}.tmp", .{self.log.path});
        defer self.gpa.free(tmp_path);
        std.Io.Dir.cwd().deleteFile(self.io, tmp_path) catch {};
        var new_log = try storage.Log.open(self.gpa, self.io, tmp_path);
        defer new_log.close();

        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| {
                // Re-emit the index definitions first: a compacted log that
                // dropped them would resurrect the collections without
                // indexes on replay.
                for (coll_entry.value_ptr.indexes.items) |*ix| {
                    var spec_bytes: std.ArrayListUnmanaged(u8) = .empty;
                    defer spec_bytes.deinit(self.gpa);
                    try ix.write_spec(self.gpa, &spec_bytes);
                    try new_log.append_index_create(db_entry.key_ptr.*, coll_entry.key_ptr.*, spec_bytes.items, self.seq);
                }
                var doc_it = coll_entry.value_ptr.docs.iterator();
                while (doc_it.next()) |doc_entry| {
                    const doc_bytes = try serialize_doc(self.gpa, doc_entry.value_ptr.*);
                    defer self.gpa.free(doc_bytes);
                    try new_log.append_upsert(db_entry.key_ptr.*, coll_entry.key_ptr.*, doc_bytes, self.seq);
                }
            }
        }
        const new_end_pos = new_log.end_pos;

        try std.Io.Dir.renameAbsolute(tmp_path, self.log.path, self.io);
        // Persist the rename: fsync the parent directory so the new
        // directory entry survives a power loss right after compaction.
        const parent = parent_dir(self.log.path);
        var dir_file = try std.Io.Dir.cwd().openFile(self.io, parent, .{ .mode = .read_only, .allow_directory = true });
        defer dir_file.close(self.io);
        try dir_file.sync(self.io);

        const old_path = try self.gpa.dupe(u8, self.log.path);
        self.log.close();
        self.log = try storage.Log.open(self.gpa, self.io, old_path);
        // Log.open starts at end_pos 0 and does not replay; continue appending
        // where the compacted file actually ends.
        self.log.end_pos = new_end_pos;
        self.gpa.free(old_path);
    }

    /// Rebuild every empty index from the live documents. Runs after replay
    /// completes, so it is order-independent: a create record, the documents
    /// it indexes, and any drop record all replay first. A duplicate under a
    /// unique index logs a loud warning and keeps the index (still correct
    /// as a candidate generator; future writes are still enforced) — the
    /// database always opens, leaving dropIndexes as an in-band recovery
    /// path.
    fn build_all_indexes(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| {
                for (coll_entry.value_ptr.indexes.items) |*ix| {
                    if (ix.entries.items.len > 0) continue; // defensive
                    var doc_it = coll_entry.value_ptr.docs.iterator();
                    while (doc_it.next()) |doc_entry| {
                        var built = ix.build_entries(self.gpa, doc_entry.value_ptr.*, doc_entry.key_ptr.*) catch |err| switch (err) {
                            error.ParallelArrays => {
                                std.debug.print("mongo-light: WARNING: index '{s}' cannot index an existing document; entry skipped\n", .{ix.name});
                                continue;
                            },
                            else => return err,
                        };
                        defer built.deinit(self.gpa);
                        if (built.multikey) ix.multikey = true;
                        if (ix.unique) {
                            ix.check_unique(built.entries.items, doc_entry.key_ptr.*) catch {
                                std.debug.print("mongo-light: WARNING: unique index '{s}' has duplicate keys in existing data; duplicates not enforced for existing documents\n", .{ix.name});
                            };
                        }
                        if (built.entries.items.len == 0) continue;
                        try ix.entries.ensureUnusedCapacity(self.gpa, built.entries.items.len);
                        ix.insert_entries(self.gpa, &built);
                    }
                }
            }
        }
    }

    /// Register an (empty) index from a persisted spec document. A repeated
    /// create record for the same name is an idempotent no-op.
    fn register_index_from_spec(self: *Engine, coll: *Collection, spec_doc: *const bson.Document) !void {
        var ix = try index.parse_spec(self.gpa, spec_doc);
        var committed = false;
        defer if (!committed) ix.deinit(self.gpa);
        for (coll.indexes.items) |existing| {
            if (std.mem.eql(u8, existing.name, ix.name)) return;
        }
        try coll.indexes.append(self.gpa, ix);
        committed = true;
    }
};

fn parent_dir(path: []const u8) []const u8 {
    const last = std.mem.lastIndexOfScalar(u8, path, '/') orelse return ".";
    if (last == 0) return "/";
    return path[0..last];
}

fn serialize_doc(gpa: std.mem.Allocator, doc: *const bson.Document) ![]u8 {
    var out: std.ArrayListUnmanaged(u8) = .empty;
    errdefer out.deinit(gpa);
    try doc.to_bytes(gpa, &out);
    return out.toOwnedSlice(gpa);
}

fn apply_record(ctx: *anyopaque, record: storage.Record, doc: *bson.Document) anyerror!void {
    const self: *Engine = @ptrCast(@alignCast(ctx));
    var stored = false;
    defer if (!stored) {
        doc.deinit();
        self.gpa.destroy(doc);
    };

    const coll = self.get_or_create_collection(record.db, record.coll) catch return;

    // Index records carry no _id — handle them before the lookup. Replay
    // registers indexes empty; Engine.open builds them from the live docs
    // after replay completes.
    switch (record.type) {
        storage.record_type_index_create => {
            self.register_index_from_spec(coll, doc) catch |err| {
                std.debug.print("mongo-light: index create record failed to apply: {s}\n", .{@errorName(err)});
                return;
            };
            return;
        },
        storage.record_type_index_drop => {
            const name_value = doc.get("name") orelse return;
            const name = switch (name_value) {
                .string => |s| s,
                else => return,
            };
            var i: usize = 0;
            while (i < coll.indexes.items.len) {
                if (std.mem.eql(u8, coll.indexes.items[i].name, name)) {
                    var removed = coll.indexes.orderedRemove(i);
                    removed.deinit(self.gpa);
                } else i += 1;
            }
            return;
        },
        else => {},
    }

    const id_value = doc.get("_id") orelse {
        std.debug.print("mongo-light: log record without _id, skipping\n", .{});
        return;
    };
    const id_key = try bson.serialize_value(self.gpa, id_value);
    var key_owned = false;
    defer if (!key_owned) self.gpa.free(id_key);

    switch (record.type) {
        storage.record_type_upsert => {
            self.evict_doc(coll, id_key);
            try coll.docs.put(self.gpa, id_key, doc);
            key_owned = true;
            stored = true;
        },
        storage.record_type_delete => self.evict_doc(coll, id_key),
        else => {},
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

const testing = std.testing;

const TmpLog = storage.TmpLog;

fn test_env(threaded: *std.Io.Threaded) struct { io: std.Io, gen: bson.ObjectIdGen } {
    const io = threaded.io();
    const gen = bson.ObjectIdGen.init(io);
    return .{ .io = io, .gen = gen };
}

fn make_doc(gpa: std.mem.Allocator, id: i32, name: []const u8) !bson.Document {
    var arena = std.heap.ArenaAllocator.init(gpa);
    errdefer arena.deinit();
    const pairs = try arena.allocator().alloc(bson.Pair, 2);
    pairs[0] = .{ .key = try arena.allocator().dupe(u8, "_id"), .value = .{ .int32 = id } };
    pairs[1] = .{ .key = try arena.allocator().dupe(u8, "name"), .value = .{ .string = try arena.allocator().dupe(u8, name) } };
    return .{ .arena = arena, .pairs = pairs };
}

test "insert, query, remove" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    var engine = try Engine.open(gpa, io, tmp.path);
    defer engine.deinit();

    var d1 = try make_doc(gpa, 1, "alice");
    defer d1.deinit();
    var d2 = try make_doc(gpa, 2, "bob");
    defer d2.deinit();

    try engine.lock();
    try engine.insert("app", "users", &d1, &env.gen);
    try engine.insert("app", "users", &d2, &env.gen);
    engine.unlock();

    // duplicate key
    var d3 = try make_doc(gpa, 1, "alice2");
    defer d3.deinit();
    try engine.lock();
    try testing.expectError(error.DuplicateKey, engine.insert("app", "users", &d3, &env.gen));
    engine.unlock();

    // find by id
    const id_key = try bson.serialize_value(gpa, bson.Value{ .int32 = 2 });
    defer gpa.free(id_key);
    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_by_id("app", "users", .{ .int32 = 2 });
    try testing.expect(removed);
    engine.unlock();
}

test "reopen replays log" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    {
        var engine = try Engine.open(gpa, io, tmp.path);
        defer engine.deinit();
        var d1 = try make_doc(gpa, 1, "alice");
        defer d1.deinit();
        var d2 = try make_doc(gpa, 2, "bob");
        defer d2.deinit();
        try engine.lock();
        try engine.insert("app", "users", &d1, &env.gen);
        try engine.insert("app", "users", &d2, &env.gen);
        _ = try engine.remove_by_id("app", "users", .{ .int32 = 2 });
        engine.unlock();
    }

    var engine2 = try Engine.open(gpa, io, tmp.path);
    defer engine2.deinit();
    try engine2.lock();
    const id_key = try bson.serialize_value(gpa, bson.Value{ .int32 = 2 });
    defer gpa.free(id_key);
    try testing.expect(engine2.get_doc("app", "users", id_key) == null);
    const id_key1 = try bson.serialize_value(gpa, bson.Value{ .int32 = 1 });
    defer gpa.free(id_key1);
    try testing.expectEqualStrings("alice", engine2.get_doc("app", "users", id_key1).?.get("name").?.string);
    engine2.unlock();
}

test "auto _id generation survives reopen" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    {
        var engine = try Engine.open(gpa, io, tmp.path);
        defer engine.deinit();

        var doc = try make_doc(gpa, 0, "no-id-here");
        defer doc.deinit();
        // strip _id
        const stripped = doc.pairs[1..];

        var arena = std.heap.ArenaAllocator.init(gpa);
        defer arena.deinit();
        var d2 = try bson.Document.alloc(gpa, try arena.allocator().dupe(bson.Pair, stripped));
        defer d2.deinit();

        try engine.lock();
        try engine.insert("app", "no_ids", &d2, &env.gen);
        engine.unlock();
    }

    var engine2 = try Engine.open(gpa, io, tmp.path);
    defer engine2.deinit();
    const coll = engine2.get_collection("app", "no_ids").?;
    var it = coll.docs.iterator();
    var count: usize = 0;
    while (it.next()) |entry| {
        count += 1;
        try testing.expect(entry.value_ptr.*.get("_id").?.object_id.len == 12);
    }
    try testing.expectEqual(@as(usize, 1), count);
}

test "compaction rewrites log and keeps data" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    {
        var engine = try Engine.open(gpa, io, tmp.path);
        engine.compact_threshold = 1; // always compact
        defer engine.deinit();

        var docs: [4]bson.Document = undefined;
        defer for (&docs) |*d| d.deinit();
        try engine.lock();
        for (0..4) |i| {
            docs[i] = try make_doc(gpa, @intCast(i + 1), "user-{d}");
            try engine.insert("app", "users", &docs[i], &env.gen);
        }
        engine.unlock();
    }

    // Reopen after compaction and keep writing: with the log reopened at
    // end_pos 0, appends would clobber the compacted records.
    var engine2 = try Engine.open(gpa, io, tmp.path);
    defer engine2.deinit();
    try engine2.lock();
    var extra = try make_doc(gpa, 5, "eve");
    defer extra.deinit();
    try engine2.insert("app", "users", &extra, &env.gen);
    engine2.unlock();

    var engine3 = try Engine.open(gpa, io, tmp.path);
    defer engine3.deinit();
    try engine3.lock();
    for (1..6) |i| {
        const id_key = try bson.serialize_value(gpa, bson.Value{ .int32 = @intCast(i) });
        defer gpa.free(id_key);
        try testing.expect(engine3.get_doc("app", "users", id_key) != null);
    }
    engine3.unlock();
}

test "concurrent readers and writers on a threaded Io" {
    // Real worker threads: writers hold the exclusive lock, readers the
    // shared lock. Proves the RwLock split keeps committed writes visible
    // to concurrent readers and never corrupts the maps.
    const gpa = testing.allocator;
    var threaded: std.Io.Threaded = .init(gpa, .{});
    defer threaded.deinit();
    const io = threaded.io();

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    var engine = try Engine.open(gpa, io, tmp.path);
    defer engine.deinit();

    const writers = 4;
    const readers = 4;
    const per_writer: i32 = 200;
    const total: i32 = writers * per_writer;
    var next_id = std.atomic.Value(i32).init(1);
    var remaining = std.atomic.Value(usize).init(@intCast(total));

    const Worker = struct {
        fn writer(e: *Engine, id_counter: *std.atomic.Value(i32), pending: *std.atomic.Value(usize), alloc: std.mem.Allocator) error{Canceled}!void {
            while (true) {
                const id = id_counter.fetchAdd(1, .monotonic);
                if (id > total) return;
                var doc = make_doc(alloc, id, "user") catch return error.Canceled;
                defer doc.deinit();
                e.lock() catch return error.Canceled;
                defer e.unlock();
                e.insert("app", "users", &doc, undefined) catch return error.Canceled;
                _ = pending.fetchSub(1, .monotonic);
            }
        }

        fn reader(e: *Engine, pending: *std.atomic.Value(usize)) error{Canceled}!void {
            while (pending.load(.acquire) > 0) {
                e.lock_read() catch return error.Canceled;
                defer e.unlock_read();
                if (e.get_collection("app", "users")) |coll| {
                    var n: usize = 0;
                    var it = coll.docs.iterator();
                    while (it.next()) |_| n += 1;
                    // A reader must never observe more docs than can exist.
                    if (n > @as(usize, @intCast(total))) return error.Canceled;
                }
            }
        }
    };

    var group: std.Io.Group = .init;
    defer group.cancel(io);
    for (0..readers) |_| group.async(io, Worker.reader, .{ &engine, &remaining });
    for (0..writers) |_| group.async(io, Worker.writer, .{ &engine, &next_id, &remaining, gpa });
    try group.await(io);

    // Every committed write must be visible once all writers finish.
    try engine.lock_read();
    defer engine.unlock_read();
    const coll = engine.get_collection("app", "users") orelse return error.TestUnexpectedResult;
    try testing.expectEqual(@as(usize, @intCast(total)), coll.docs.count());
    for (1..total + 1) |i| {
        const id_key = try bson.serialize_value(gpa, bson.Value{ .int32 = @intCast(i) });
        defer gpa.free(id_key);
        try testing.expect(engine.get_doc("app", "users", id_key) != null);
    }
}

// -- index tests -----------------------------------------------------------

/// A spec document for a single-path index, built by serializing and
/// re-parsing so the pairs are arena-owned.
fn index_spec(gpa: std.mem.Allocator, path: []const u8, name: []const u8, unique: bool, sparse: bool) !bson.Document {
    var out: std.ArrayListUnmanaged(u8) = .empty;
    defer out.deinit(gpa);
    const pairs = [_]bson.Pair{
        .{ .key = "key", .value = .{ .doc = &.{.{ .key = path, .value = .{ .int32 = 1 } }} } },
        .{ .key = "name", .value = .{ .string = name } },
        .{ .key = "unique", .value = .{ .bool = unique } },
        .{ .key = "sparse", .value = .{ .bool = sparse } },
    };
    try bson.write_doc(&pairs, gpa, &out);
    return bson.Document.parse(gpa, out.items);
}

/// Number of entries the named index has for a single-value equality key.
fn index_count(gpa: std.mem.Allocator, engine: *Engine, db_name: []const u8, coll_name: []const u8, name: []const u8, key_value: bson.Value) !usize {
    const coll = engine.get_collection(db_name, coll_name) orelse return 0;
    for (coll.indexes.items) |*ix| {
        if (std.mem.eql(u8, ix.name, name)) {
            var out: std.ArrayListUnmanaged([]const u8) = .empty;
            defer out.deinit(gpa);
            try ix.lookup_eq(gpa, &.{key_value}, &out);
            return out.items.len;
        }
    }
    return 0;
}

fn make_user(gpa: std.mem.Allocator, id: i32, email: []const u8) !bson.Document {
    var arena = std.heap.ArenaAllocator.init(gpa);
    errdefer arena.deinit();
    const pairs = try arena.allocator().alloc(bson.Pair, 2);
    pairs[0] = .{ .key = try arena.allocator().dupe(u8, "_id"), .value = .{ .int32 = id } };
    pairs[1] = .{ .key = try arena.allocator().dupe(u8, "email"), .value = .{ .string = try arena.allocator().dupe(u8, email) } };
    return .{ .arena = arena, .pairs = pairs };
}

test "unique index enforced on insert, replace, and upsert-conflict" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    var engine = try Engine.open(gpa, io, tmp.path);
    defer engine.deinit();

    var spec = try index_spec(gpa, "email", "email_1", true, false);
    defer spec.deinit();
    try engine.lock();
    _ = try engine.create_index("app", "users", &spec);

    var d1 = try make_user(gpa, 1, "a@x.io");
    defer d1.deinit();
    try engine.insert("app", "users", &d1, &env.gen);

    // A second doc with the same email is rejected and never logged.
    var d2 = try make_user(gpa, 2, "a@x.io");
    defer d2.deinit();
    try testing.expectError(error.DuplicateKeyIndex, engine.insert("app", "users", &d2, &env.gen));
    try testing.expectEqualStrings("email_1", engine.dup_index.?);

    // A replace that keeps its own email is fine (own entries excluded).
    var d1b = try make_user(gpa, 1, "a@x.io");
    defer d1b.deinit();
    try engine.replace("app", "users", &d1b, &env.gen);
    try testing.expectEqual(@as(usize, 1), try index_count(gpa, &engine, "app", "users", "email_1", .{ .string = "a@x.io" }));

    // An update that would collide is rejected.
    var d2b = try make_user(gpa, 2, "a@x.io");
    defer d2b.deinit();
    try testing.expectError(error.DuplicateKeyIndex, engine.replace("app", "users", &d2b, &env.gen));

    // A different email still inserts.
    var d3 = try make_user(gpa, 3, "b@x.io");
    defer d3.deinit();
    try engine.insert("app", "users", &d3, &env.gen);
    engine.unlock();
}

test "index maintained across update and delete" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    var engine = try Engine.open(gpa, io, tmp.path);
    defer engine.deinit();

    var spec = try index_spec(gpa, "a", "a_1", false, false);
    defer spec.deinit();
    try engine.lock();
    _ = try engine.create_index("app", "items", &spec);

    var d1 = try doc_with_a(gpa, 1, 10);
    defer d1.deinit();
    var d2 = try doc_with_a(gpa, 2, 20);
    defer d2.deinit();
    try engine.insert("app", "items", &d1, &env.gen);
    try engine.insert("app", "items", &d2, &env.gen);
    try testing.expectEqual(@as(usize, 1), try index_count(gpa, &engine, "app", "items", "a_1", .{ .int32 = 20 }));

    // Replace doc 1 with a new value: old entry gone, new entry present.
    var d1b = try doc_with_a(gpa, 1, 30);
    defer d1b.deinit();
    try engine.replace("app", "items", &d1b, &env.gen);
    try testing.expectEqual(@as(usize, 0), try index_count(gpa, &engine, "app", "items", "a_1", .{ .int32 = 10 }));
    try testing.expectEqual(@as(usize, 1), try index_count(gpa, &engine, "app", "items", "a_1", .{ .int32 = 30 }));

    // Delete doc 2: its entry is removed.
    _ = try engine.remove_by_id("app", "items", .{ .int32 = 2 });
    try testing.expectEqual(@as(usize, 0), try index_count(gpa, &engine, "app", "items", "a_1", .{ .int32 = 20 }));
    engine.unlock();
}

/// A document with an integer `a` field (on top of _id + name).
fn doc_with_a(gpa: std.mem.Allocator, id: i32, a: i32) !bson.Document {
    var arena = std.heap.ArenaAllocator.init(gpa);
    errdefer arena.deinit();
    const pairs = try arena.allocator().alloc(bson.Pair, 3);
    pairs[0] = .{ .key = try arena.allocator().dupe(u8, "_id"), .value = .{ .int32 = id } };
    pairs[1] = .{ .key = try arena.allocator().dupe(u8, "name"), .value = .{ .string = try arena.allocator().dupe(u8, "x") } };
    pairs[2] = .{ .key = try arena.allocator().dupe(u8, "a"), .value = .{ .int32 = a } };
    return .{ .arena = arena, .pairs = pairs };
}

test "index survives reopen and compaction" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    {
        var engine = try Engine.open(gpa, io, tmp.path);
        defer engine.deinit();
        engine.compact_threshold = 1; // every write compacts
        var spec = try index_spec(gpa, "email", "email_1", false, false);
        defer spec.deinit();
        try engine.lock();
        _ = try engine.create_index("app", "users", &spec);
        var d1 = try make_user(gpa, 1, "a@x.io");
        defer d1.deinit();
        var d2 = try make_user(gpa, 2, "b@x.io");
        defer d2.deinit();
        try engine.insert("app", "users", &d1, &env.gen);
        try engine.insert("app", "users", &d2, &env.gen);
        engine.unlock();
    }

    // Reopen: the index (rebuilt from the compacted log) still finds docs.
    var engine2 = try Engine.open(gpa, io, tmp.path);
    defer engine2.deinit();
    try engine2.lock();
    try testing.expectEqual(@as(usize, 1), try index_count(gpa, &engine2, "app", "users", "email_1", .{ .string = "b@x.io" }));
    engine2.unlock();
}

test "index drop survives reopen" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    {
        var engine = try Engine.open(gpa, io, tmp.path);
        defer engine.deinit();
        var spec = try index_spec(gpa, "email", "email_1", false, false);
        defer spec.deinit();
        try engine.lock();
        _ = try engine.create_index("app", "users", &spec);
        var d1 = try make_user(gpa, 1, "a@x.io");
        defer d1.deinit();
        try engine.insert("app", "users", &d1, &env.gen);
        try testing.expect(try engine.drop_index("app", "users", "email_1"));
        engine.unlock();
    }

    var engine2 = try Engine.open(gpa, io, tmp.path);
    defer engine2.deinit();
    try engine2.lock();
    try testing.expectEqual(@as(usize, 0), engine2.get_collection("app", "users").?.indexes.items.len);
    engine2.unlock();
}

test "drop_collection frees indexes; log without index records replays" {
    var threaded: std.Io.Threaded = .init_single_threaded;
    defer threaded.deinit();
    var env = test_env(&threaded);
    const io = env.io;
    const gpa = testing.allocator;

    var tmp = try TmpLog.init(gpa);
    defer tmp.deinit(gpa);
    {
        var engine = try Engine.open(gpa, io, tmp.path);
        defer engine.deinit();
        var spec = try index_spec(gpa, "email", "email_1", false, false);
        defer spec.deinit();
        try engine.lock();
        _ = try engine.create_index("app", "users", &spec);
        var d1 = try make_user(gpa, 1, "a@x.io");
        defer d1.deinit();
        try engine.insert("app", "users", &d1, &env.gen);
        // Dropped in memory; free_collection releases the index memory
        // (verified by testing.allocator at engine.deinit).
        try testing.expect(try engine.drop_collection("app", "users"));
        try testing.expect(engine.get_collection("app", "users") == null);
        // A log that only ever contained plain upserts replays fine.
        var d2 = try make_doc(gpa, 2, "bob");
        defer d2.deinit();
        try engine.insert("app", "plain", &d2, &env.gen);
        engine.unlock();
    }

    var engine2 = try Engine.open(gpa, io, tmp.path);
    defer engine2.deinit();
    try engine2.lock();
    const id_key = try bson.serialize_value(gpa, bson.Value{ .int32 = 2 });
    defer gpa.free(id_key);
    try testing.expect(engine2.get_doc("app", "plain", id_key) != null);
    // Pre-existing limitation (documented in the README): drop_collection
    // writes no log record, so the collection and its index resurrect.
    const users = engine2.get_collection("app", "users").?;
    try testing.expectEqual(@as(usize, 1), users.indexes.items.len);
    try testing.expectEqual(@as(usize, 1), try index_count(gpa, &engine2, "app", "users", "email_1", .{ .string = "a@x.io" }));
    engine2.unlock();
}