diff --git a/AGENTS.md b/AGENTS.md index 79b386e..fac2755 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -13,9 +13,11 @@ indexes, and TTL/unique/sparse/compound index support. **Forward plan**: the project's direction — a full-fledged embedded, tens-of-GB, maximally MongoDB-compatible database — is decided and written -in **[PLAN.md](PLAN.md)**. The current milestone is **M0 (mmap + WAL -storage foundation)**. Before starting any work, read PLAN.md; its decision -record (D1-D9) and ground rules are binding. +in **[PLAN.md](PLAN.md)**. M0 (mmap + WAL storage foundation) has landed; +the current milestone is **M1 (cursors + wire polish)**, whose cursor work is +done — see `src/cursor.zig` and `tests/e2e/e2e7.js`. Before starting any +work, read PLAN.md; its decision record (D1-D9) and ground rules are +binding. ## Read first, in order @@ -71,6 +73,7 @@ node tests/e2e/e2e2.js crash-b # restart, verify all 50 survived node tests/e2e/e2e3.js # secondary indexes node tests/e2e/e2e4.js # TTL indexes (server must run --ttl-sweep-secs 1) node tests/e2e/e2e6.js # self-contained full lifecycle (spawns its own server, incl. kill -9) +node tests/e2e/e2e7.js # self-contained cursors (spawns its own servers; needs no server running) ``` Which suites to run for a given change: @@ -78,6 +81,7 @@ Which suites to run for a given change: - anything touching the write path or log format → the crash pair (e2e2 crash-a/b) and e2e6 - anything touching indexes → e2e3.js and e2e4.js +- anything touching cursors, batching or the reply size → e2e7.js - everything → all of the above `tests/e2e/README.md` has the full matrix, ports, and harness docs @@ -180,7 +184,8 @@ src/server.zig TCP accept loop, per-connection handlers, TTL sweep monitor src/db.zig engine: db → collection → _id → document maps, slab storage src/storage.zig append-only log: blocks, LZ4, XxHash3, replay, compaction src/query.zig filter matcher, regex engine, sort, projection -src/index.zig B+tree indexes: entries, search, query planner +src/index.zig B+tree indexes: entries, search, query planner, scan resume +src/cursor.zig server-side cursor state: registry, batch policy, expiry src/update.zig update operators with dot-path navigation src/main.zig CLI: --port, --bind, --db, --ttl-sweep-secs, --compact-threshold ``` @@ -193,6 +198,8 @@ src/main.zig CLI: --port, --bind, --db, --ttl-sweep-secs, --compact-threshol baseline. 3. Commit scorecard and benchmark results with each milestone (PLAN D9) so progress stays verifiable across sessions. -4. Deferred designs (cursors, aggregation, transactions, change streams, - C API) are deliberately *not* specified yet — grill the design with the - human before implementing (PLAN section 6). +4. Deferred designs (aggregation, transactions, change streams, C API) are + deliberately *not* specified yet — grill the design with the human before + implementing (PLAN section 6). Cursors are no longer among them: the + design was settled and implemented in M1, and `src/cursor.zig`'s module + comment is where it is written down. diff --git a/PLAN.md b/PLAN.md index 561b99e..e0f2270 100644 --- a/PLAN.md +++ b/PLAN.md @@ -635,9 +635,47 @@ it. ## 6. Deferred designs (grill each at its milestone) -- **M1 cursors**: cursor id allocation, idle expiration, batchSize - semantics, getMore against a lagging/compactable engine, cursor state - lifecycle across compaction. +- **M1 cursors** — *settled and implemented.* The design lives in + `src/cursor.zig`'s module comment; the decisions it records, and how each + was reached: + - **Cursor ids** are `(nonce << 20) | slot`, always positive, never 0. The + nonce is not decoration: without it a recycled slot serves one client + another's documents, which is the worst failure this feature could have. + - **batchSize semantics** were *measured against mongod 8.3.7*, not + recalled, and three assumptions were wrong: a bare `getMore` does **not** + inherit the find's batchSize (4998 of 5000 documents come back), a + namespace mismatch is `Unauthorized` (13) rather than `CursorNotFound`, + and `CursorInUse` is 143 rather than the 12051 an earlier note claimed. + `internalQueryFindCommandBatchSize` is 101, `cursorTimeoutMillis` + 600000, `clientCursorMonitorFrequencySecs` 4. + - **Never look ahead**: a batch that met its target leaves the cursor open + even when the source is in fact exhausted, so four documents at + `batchSize: 2` take three commands. The pinned suites assert that count. + - **Idle expiration** is a second monitor fiber, separate from the TTL one: + the cadences differ by an order of magnitude, and a TTL sweep failure + must not stop cursors being reclaimed. The registry is fixed-capacity and + evicts the least recently used cursor, whose client sees the same + `CursorNotFound` an idle timeout gives. + - **Against a lagging/compactable engine**, what survives depends on what + the cursor remembers, so the check is per-source: a repack changes no + key, so a streaming cursor resumes; slab offsets all move, so an offsets + cursor is killed with `QueryPlanKilled`; a snapshot needs no collection at + all. `Collection.layout_epoch` and `Index.epoch` are the tokens. + - **Resume** anchors on `(key, off)` plus a position hint gated on + `Index.epoch`, with an exact-order band walk bounded by + `resume_walk_max`. Two hazards found while implementing: a deleted anchor + must resume at its band position or the rest of an equal-key band is + silently dropped, and on a *unique* index a same-key entry can only be + the anchor rewritten — resuming at it returned updated documents twice, + caught by draining a collection being updated underneath. + + Still open in M1: the doc-level free list, sessions plumbing (`lsid` + accepted), and command-monitoring (`expectEvents`) in the spec runner. + **A prerequisite the free list must honour**, recorded here while it is + still being designed: *an offset that was ever a record start must remain a + record start.* `doc_bytes` reads a `u32` length prefix in place, so an + offset landing mid-record after a re-split is a garbage-length read rather + than a wrong answer — and an offsets cursor holds exactly such offsets. - **M2 aggregation**: stage/expression tiers, which spec-test files are the gate, whether $lookup/$unwind/facet make the first cut. - **M4 transactions**: snapshot isolation over mmap (COW vs undo), read diff --git a/README.md b/README.md index e8eaf07..980fe9e 100644 --- a/README.md +++ b/README.md @@ -12,8 +12,9 @@ maximally MongoDB-compatible database — its decision record, milestones and gates live in [PLAN.md](PLAN.md). Milestone 0 (mmap + WAL storage foundation) has landed; its measured gate results are in [`tests/e2e/results/m0-gates.txt`](tests/e2e/results/m0-gates.txt). -Milestone 1 (cursors, and the doc-level free list the churn gate showed is -needed) is next. +Milestone 1 is in progress: server-side cursors have landed (see +**Cursors** below); the doc-level free list the churn gate showed is needed +is still open. ## Quick start @@ -145,10 +146,43 @@ whole pass, so the interval is the tuning knob: expiry is never more precise than `--ttl-sweep-secs`, and a very large TTL index wants a longer one. +## Cursors + +`find`, `aggregate`, `listCollections` and `listIndexes` return real cursor +ids, and `getMore`/`killCursors` work. Batching follows MongoDB: a first +batch of 101 documents unless `batchSize` says otherwise, a `getMore` with +no `batchSize` bounded only by the 16 MiB batch cap, `batchSize: 0` as an +empty batch with a live cursor, and `limit` honoured across batches. Every +default here was measured against a real `mongod` rather than assumed. + +A cursor holds no lock between requests, so what it remembers has to survive +arbitrary concurrent writes. Three shapes, picked by the query: + +| query | what the cursor keeps | +| --- | --- | +| a whole-index walk (`find({})`, or a sort an index provides) | the last key and offset it yielded — O(key), whatever the collection size | +| a narrowed index plan | the matching offsets, 8 bytes each | +| a sort no index provides, or aggregate/listing output | a snapshot of the remaining documents | + +The first is what lets a cursor walk a collection larger than memory. It +also survives a compaction, because a repack changes no key; the offsets +form cannot, and says so with `QueryPlanKilled` rather than returning +documents from the wrong place. + +Cursors are not pinned to the connection that created them, so a `getMore` +may arrive on any connection — which is what the driver specification +allows. They are reclaimed when exhausted, when killed, or after +`--cursor-timeout-ms` idle (default 600000, MongoDB's own +`cursorTimeoutMillis`); `--max-open-cursors` bounds the registry and evicts +the least recently used cursor at capacity, whose client then sees the same +`CursorNotFound` an idle timeout gives. + ## Not (yet) implemented - Authentication (SCRAM) — run without credentials -- Real cursors (all results are returned in one batch, cursor id 0) +- Tailable/awaitData cursors, which need capped collections; a tailable + `find` is rejected, exactly as MongoDB rejects one on a non-capped + collection - Transactions, change streams, replicasets - Compression (OP_COMPRESSED) - `collMod`, so an index's `expireAfterSeconds` cannot be changed in diff --git a/src/commands.zig b/src/commands.zig index f2d1b0a..4f92d5b 100644 --- a/src/commands.zig +++ b/src/commands.zig @@ -10,6 +10,7 @@ const Collection = db.Collection; const query = @import("query.zig"); const update = @import("update.zig"); const index = @import("index.zig"); +const cursor = @import("cursor.zig"); // Always active, including in the default ReleaseFast build -- see assert.zig. const assert = @import("assert.zig").assert; const assert_msg = @import("assert.zig").assert_msg; @@ -27,17 +28,30 @@ pub const Context = struct { pub const ErrorCode = enum(i32) { command_not_found = 59, bad_value = 2, - invalid_argument = 72, + /// 72 is MongoDB's `InvalidOptions`. There is no `InvalidArgument` in its + /// table at all, so that is the name this used to send. + invalid_options = 72, namespace_not_found = 26, index_not_found = 27, duplicate_key = 11000, namespace_exists = 48, failed_to_parse = 9, internal_error = 1, - invalid_pipeline_operator = 40324, + /// "Unrecognized pipeline stage name". A `Location` code, so mongod names it + /// `Location40324` rather than after any symbol. + location_unrecognized_stage = 40324, index_options_conflict = 85, cannot_create_index = 67, invalid_index_specification_option = 197, + // Cursor codes. Taken from MongoDB's own error_codes.js rather than + // recalled -- CursorInUse in particular is 143, not the 12051 that turns up + // in older notes. + cursor_not_found = 43, + cursor_in_use = 143, + query_plan_killed = 175, + unauthorized = 13, + type_mismatch = 14, + operation_failed = 96, }; /// Which lock (if any) a command needs on the engine. Contract: only @@ -45,13 +59,18 @@ pub const ErrorCode = enum(i32) { /// remove, drop*, get_or_create_collection, compact); `.read` commands may /// only read (`get_collection`, `database_names`, `collection_names`); /// `.none` commands must not touch the engine at all. +/// +/// This describes engine *data* only. The cursor store is a separate resource +/// with its own leaf mutex, and any kind may mutate it: `getMore` is `.read` +/// because it only reads documents, even though it advances cursor state, and +/// `killCursors` is `.none` because the store is all it touches. const CommandKind = enum { none, read, write }; /// Lock shape for one command, acquired by dispatch: the catalog lock mode /// and whether the command's target collection is locked (shared for reads, /// exclusive for writes/DDL). The target collection is the message field /// named after the command (find/count/insert/...), which every -/// collection-targeting command uses. +/// collection-targeting command uses except `getMore` -- see `Command.coll_field`. const LockShape = struct { catalog: enum { none, shared, exclusive } = .none, coll: enum { none, shared, exclusive } = .none, @@ -61,6 +80,9 @@ const Command = struct { name: []const u8, kind: CommandKind, locks: LockShape = .{}, + /// Body field naming the target collection, when it is not the command's + /// own field. Only `getMore` needs it: its own value is an int64 cursor id. + coll_field: ?[]const u8 = null, handler: *const fn (*Context, *wire.Message, *wire.Reply) anyerror!void, }; @@ -80,10 +102,20 @@ const command_table = [_]Command{ .{ .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 }, + // killCursors touches only the cursor store, so it needs no lock -- and its + // own field really is the collection name, unlike getMore's. .{ .name = "killCursors", .kind = .none, .handler = cmd_kill_cursors }, // Read-only: scan the engine without mutating it. .{ .name = "find", .kind = .read, .locks = .{ .catalog = .shared, .coll = .shared }, .handler = cmd_find }, + // getMore is the one command whose collection is not its own field: its + // value is an int64 cursor id, so dispatch reads `collection` instead. + .{ + .name = "getMore", + .kind = .read, + .locks = .{ .catalog = .shared, .coll = .shared }, + .coll_field = "collection", + .handler = cmd_get_more, + }, .{ .name = "count", .kind = .read, .locks = .{ .catalog = .shared, .coll = .shared }, .handler = cmd_count }, .{ .name = "aggregate", .kind = .read, .locks = .{ .catalog = .shared, .coll = .shared }, .handler = cmd_aggregate }, .{ .name = "listDatabases", .kind = .read, .locks = .{ .catalog = .shared }, .handler = cmd_list_databases }, @@ -140,7 +172,8 @@ pub fn dispatch(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { if (cmd.locks.coll != .none) { const db_name = msg.db_name() orelse return bad_value(reply, "command requires a $db"); - const coll_name = str_arg(msg.body.get(name)) orelse + const field = cmd.coll_field orelse name; + const coll_name = str_arg(msg.body.get(field)) orelse return bad_value(reply, "command requires a collection name"); ns = .{ .db = db_name, .coll = coll_name }; } @@ -368,21 +401,30 @@ fn cmd_list_databases(ctx: *Context, _: *wire.Message, reply: *wire.Reply) !void try reply.put_ok(); } +/// The collection part of the namespace a `listCollections` cursor reports. +/// mongod uses this pseudo-collection, and the exact string matters: the previous +/// `"."` had an *empty* collection part, and the driver throws client-side +/// when it tries to build a getMore or killCursors from a namespace like that -- +/// so the moment such a cursor stopped being id 0 it would have broken. +const list_collections_ns: []const u8 = "$cmd.listCollections"; + fn cmd_list_collections(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const db_name = msg.db_name() orelse return invalid_arg(reply, "listCollections requires $db"); + const batch_size = (try aggregate_batch_size(reply, msg) orelse return).value; var names: std.ArrayListUnmanaged([]const u8) = .empty; defer names.deinit(ctx.gpa); try ctx.engine.collection_names(db_name, &names); - const values = try reply.arena_alloc().alloc(bson.Value, names.items.len); + const arena = reply.arena_alloc(); + const docs = try arena.alloc(*const bson.Document, names.items.len); for (names.items, 0..) |n, i| { - const entry = try reply.arena_alloc().alloc(bson.Pair, 3); - entry[0] = .{ .key = "name", .value = .{ .string = try reply.arena_alloc().dupe(u8, n) } }; + const entry = try arena.alloc(bson.Pair, 3); + entry[0] = .{ .key = "name", .value = .{ .string = try arena.dupe(u8, n) } }; entry[1] = .{ .key = "type", .value = .{ .string = "collection" } }; entry[2] = .{ .key = "options", .value = .{ .doc = &.{} } }; - values[i] = .{ .doc = entry }; + docs[i] = try doc_from_pairs(arena, entry); } - try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, ""), "firstBatch", values) }); + try emit_first_batch(ctx, reply, db_name, list_collections_ns, null, docs, batch_size); try reply.put_ok(); } @@ -506,21 +548,22 @@ fn cmd_list_indexes(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void const coll = ctx.engine.get_collection(db_name, coll_name) orelse return reply.put_error(@intFromEnum(ErrorCode.namespace_not_found), "NamespaceNotFound", "ns not found"); + const batch_size = (try aggregate_batch_size(reply, msg) orelse return).value; // The _id_ index first, then the secondaries. - const n = coll.indexes.items.len + 1; - const values = try reply.arena_alloc().alloc(bson.Value, n); - const id_pairs = try reply.arena_alloc().alloc(bson.Pair, 2); + const arena = reply.arena_alloc(); + const values = try arena.alloc(*const bson.Document, coll.indexes.items.len + 1); + const id_pairs = try arena.alloc(bson.Pair, 2); id_pairs[0] = .{ .key = "v", .value = .{ .int32 = 2 } }; id_pairs[1] = .{ .key = "key", .value = .{ .doc = &.{.{ .key = "_id", .value = .{ .int32 = 1 } }} } }; - values[0] = .{ .doc = try index_pairs_append(reply, id_pairs, "_id_") }; + values[0] = try doc_from_pairs(arena, try index_pairs_append(reply, id_pairs, "_id_")); for (coll.indexes.items, 0..) |ix, i| { - // The pairs live in the reply arena (freed with it); the values - // array below references them. + // The pairs live in the reply arena (freed with it); the docs array + // below references them. var pairs: std.ArrayListUnmanaged(bson.Pair) = .empty; - try ix.spec_pairs(reply.arena_alloc(), &pairs); - values[1 + i] = .{ .doc = pairs.items }; + try ix.spec_pairs(arena, &pairs); + values[1 + i] = try doc_from_pairs(arena, pairs.items); } - try reply.put("cursor", .{ .doc = try cursor_doc(reply, 0, try format_namespace(reply, db_name, coll_name), "firstBatch", values) }); + try emit_first_batch(ctx, reply, db_name, coll_name, null, values, batch_size); try reply.put_ok(); } @@ -656,68 +699,307 @@ fn cmd_insert(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { try reply.put_ok(); } +/// Parse a `batchSize`-shaped option. Null means the error reply is already +/// written. `zero_is_default` distinguishes `find`, where 0 is a real request for +/// an empty batch, from `getMore`, where mongod reads it as "no document target" +/// -- which is also what an absent field means. +fn batch_size_arg( + reply: *wire.Reply, + v: bson.Value, + label: []const u8, + zero_is_default: bool, +) !??u32 { + const n = int_value(v) orelse { + const text = try std.fmt.allocPrint(reply.arena_alloc(), "{s} must be a number", .{label}); + try bad_value(reply, text); + return null; + }; + if (n < 0) { + const text = try std.fmt.allocPrint( + reply.arena_alloc(), + "{s} value must be non-negative", + .{label}, + ); + try bad_value(reply, text); + return null; + } + if (n == 0 and zero_is_default) return @as(?u32, null); + return @as(?u32, std.math.cast(u32, n) orelse std.math.maxInt(u32)); +} + +/// The `find` options that shape the cursor rather than the query. +const CursorOpts = struct { + batch_size: ?u32 = null, + /// One batch and no cursor. Set explicitly, and also by a negative `limit`. + single_batch: bool = false, + no_timeout: bool = false, +}; + +/// Parse the cursor-shaping options, or write an error reply and return null. +fn parse_cursor_opts(reply: *wire.Reply, msg: *wire.Message) !?CursorOpts { + // Every tailable form is refused, and that is parity rather than a gap: + // mongod rejects a tailable cursor on a non-capped collection, and this + // engine has no capped collections at all. Silently ignoring the flag would + // be worse than erroring -- the cursor would report EOF and a driver's tail + // loop would exit, which reads to the application as data loss. + if (bool_arg(msg.body.get("tailable")) orelse false) { + try bad_value(reply, "tailable cursor requested on non capped collection"); + return null; + } + if (bool_arg(msg.body.get("awaitData")) orelse false) { + try bad_value(reply, "Cannot set 'awaitData' without also setting 'tailable'"); + return null; + } + + var opts = CursorOpts{}; + if (msg.body.get("batchSize")) |v| { + opts.batch_size = try batch_size_arg(reply, v, "batchSize", false) orelse return null; + } + opts.single_batch = bool_arg(msg.body.get("singleBatch")) orelse false; + opts.no_timeout = bool_arg(msg.body.get("noCursorTimeout")) orelse false; + return opts; +} + +/// Whether a source may outlive the request that built it. +/// +/// Only a snapshot can fail this: it is the one arm that pins bytes proportional +/// to the result set, and holding a large one for the idle timeout times the +/// number of live cursors is exactly what `cursor_buffer_max` exists to prevent. +/// Over the bound the caller emits everything in one batch instead -- which is +/// what this server did before cursors existed. +fn source_keepable(source: cursor.Source) bool { + const buffered = switch (source) { + .buffered => |b| b, + else => return true, + }; + var total: u64 = 0; + for (buffered.docs) |d| total += d.len; + return total <= cursor_buffer_max; +} + fn cmd_find(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const db_name = msg.db_name() orelse return invalid_arg(reply, "find requires $db"); 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 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; - // 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); + const limit_raw = int_value(msg.body.get("limit")) orelse 0; + const opts = try parse_cursor_opts(reply, msg) orelse return; + // A negative limit is the historical `ntoreturn` shape: this many documents + // in exactly one batch. Drivers translate it before sending, but the wire + // form is still legal. + const single_batch = opts.single_batch or limit_raw < 0; + const limit: u64 = @abs(limit_raw); + const target = cursor.batch_target(opts.batch_size, true); + + // A find on a namespace that does not exist is an empty cursor, not an + // error, and the ordinary path below says exactly that. + const coll = ctx.engine.get_collection(db_name, coll_name); + + var arena = std.heap.ArenaAllocator.init(ctx.gpa); + var arena_owned = false; + defer if (!arena_owned) arena.deinit(); + + // Whether this reply is the whole answer, which is what makes the top-k + // sort shortcut legal -- it leaves everything past k unspecified, and an + // open cursor would later need those documents. + const closes_here = single_batch or (limit > 0 and limit <= (target orelse 0)); + + var feed = try find_feed(ctx, reply, &arena, coll, .{ + .db_name = db_name, + .coll_name = coll_name, + .filter = filter, + .sort_keys = sort_keys, + .skip = skip, + .limit = limit, + .closes_here = closes_here, + }); + + var values: std.ArrayListUnmanaged(bson.Value) = .empty; + const exhausted = try fill_batch(ctx, reply, coll, &feed, filter, proj_pairs, target, &values); + + var cursor_id: i64 = 0; + if (!exhausted and !single_batch and source_keepable(feed.source)) { + cursor_id = keep_find_cursor(ctx, &arena, coll, &feed, .{ + .db = db_name, + .coll = coll_name, + }, filter, proj_pairs, opts); + arena_owned = cursor_id != 0; + } + // No cursor, but documents still to come: one batch holding the rest. + if (cursor_id == 0 and !exhausted and !single_batch) { + _ = try fill_batch(ctx, reply, coll, &feed, filter, proj_pairs, null, &values); + } + + const ns = try format_namespace(reply, db_name, coll_name); + const batch = try cursor_doc(reply, cursor_id, ns, "firstBatch", values.items); + try reply.put("cursor", .{ .doc = batch }); + try reply.put_ok(); +} + +/// Register the remainder of a `find` as a cursor, returning its id or 0. +fn keep_find_cursor( + ctx: *Context, + arena: *std.heap.ArenaAllocator, + coll: ?*Collection, + feed: *const Feed, + ns: cursor.Ns, + filter: []const bson.Pair, + proj_pairs: ?[]const bson.Pair, + opts: CursorOpts, +) i64 { + // Serialized before the arena is handed over: `open_cursor` takes it by + // value, so anything allocated after that call would be invisible to the copy + // the cursor keeps. + const a = arena.allocator(); + const filter_bytes = serialize_pairs(a, filter) catch return 0; + const proj_bytes = if (proj_pairs) |pp| (serialize_pairs(a, pp) catch return 0) else ""; + return open_cursor(ctx, arena.*, .{ + .ns = ns, + .layout_epoch = if (coll) |c| c.layout_epoch else 0, + .filter_bytes = filter_bytes, + .proj_bytes = proj_bytes, + .remaining_limit = feed.remaining_limit, + .batch_size = opts.batch_size, + .no_timeout = opts.no_timeout, + // feed.source, not the source it started as: the first batch advanced it, + // and for a stream that advance *is* the resume point. + .source = feed.source, + }); +} + +/// Everything `find_feed` needs that is not a lock or an arena. +const FindRequest = struct { + db_name: []const u8, + coll_name: []const u8, + filter: []const bson.Pair, + sort_keys: []const query.SortKey, + skip: u64, + limit: u64, + closes_here: bool, +}; + +/// Choose the source for a `find` and open the feed that fills its first batch. +/// +/// A whole-index walk is served without materializing anything: the cursor +/// remembers a key and an offset, so this is what makes `find({})` over a +/// collection larger than memory possible at all. Every other shape collects its +/// matches first, exactly as before cursors existed. +fn find_feed( + ctx: *Context, + reply: *wire.Reply, + arena: *std.heap.ArenaAllocator, + coll: ?*Collection, + req: FindRequest, +) !Feed { + const remaining: ?u64 = if (req.limit == 0) null else req.limit; + + // One plan for the whole request: `stream_shape` reads it to decide whether a + // resumable walk is possible, and the fallback scan reuses it instead of + // planning the same query a second time. + var plan_opt = if (coll) |c| + try index.plan(ctx.gpa, &c.id_index, c.indexes.items, req.filter, req.sort_keys) + else + null; + defer if (plan_opt) |*p| p.deinit(ctx.gpa); + const plan: ?*const index.Plan = if (plan_opt) |*p| p else null; + + if (coll) |c| { + if (stream_shape(c, plan, req.sort_keys)) |st| { + var feed = Feed{ .source = .{ .stream = st }, .remaining_limit = remaining }; + feed.scan = open_scan(c, &feed.source.stream); + // `skip` is consumed once, here, through the same anchor the batch + // uses -- which is what lets a cursor whose entire first batch falls + // inside the skipped prefix still resume from the right place. + try stream_skip(ctx, c, &feed, req.filter, req.skip); + return feed; + } + } + + // Documents needed to fill the page, counting the skipped prefix; 0 means + // unbounded. Only an index-ordered scan may stop there. + const page_end: usize = if (req.limit == 0) 0 else blk: { + const skip_usize = std.math.cast(usize, req.skip) orelse break :blk 0; + break :blk skip_usize +| req.limit; + }; + // Lives only until `find_source` copies what it needs into `arena`. var matched: std.ArrayListUnmanaged(u64) = .empty; defer matched.deinit(ctx.gpa); - // Documents needed to fill the page, counting the skipped prefix; 0 - // means unbounded. - const page_end: usize = if (limit == 0) 0 else blk: { - const skip_usize = std.math.cast(usize, skip) orelse break :blk 0; - break :blk skip_usize +| limit; - }; - // An index whose order already is the requested one lets the scan stop - // at the page boundary and skip sorting entirely. Otherwise a sort has - // to see every match before it can tell which ones the page contains. var index_sorted = false; - _ = try scan_sorted(ctx, db_name, coll_name, filter, page_end, &matched, sort_keys, &index_sorted); + _ = try scan_planned( + ctx, + req.db_name, + req.coll_name, + req.filter, + page_end, + &matched, + req.sort_keys, + &index_sorted, + plan, + ); + if (coll == null) { + assert_msg(matched.items.len == 0, "find matched in a collection that does not exist"); + } + const source = try find_source(reply, arena, coll, matched.items, .{ + .sort_keys = req.sort_keys, + .index_sorted = index_sorted, + .skip = req.skip, + .page_end = page_end, + .closes_here = req.closes_here, + }); + return .{ .source = source, .remaining_limit = remaining }; +} - // Sorting and emitting need the documents as trees; materialize the - // matched page into the reply arena (the slab itself is never copied). - // A find on a namespace that does not exist is an empty cursor, not an - // error and not a reply missing `ok`: the scan above matched nothing, so - // falling through to the emit at the end of this function says exactly - // that without a second exit path to keep in step with it. - const coll = ctx.engine.get_collection(db_name, coll_name); - // The scan above ran against this same collection with the catalog lock - // held, so a missing collection means nothing matched. Asserted rather than - // left implicit: if that ever stops holding, the loop below silently emits - // an empty page for a query that did match, which is the hardest kind of - // wrong answer to notice. - if (coll == null) assert_msg(matched.items.len == 0, "find matched documents in a collection that does not exist"); - // Lives in the reply arena; freed with it. +/// How `find_source` should turn a match list into a source. +const SourceShape = struct { + sort_keys: []const query.SortKey, + index_sorted: bool, + skip: u64, + page_end: usize, + closes_here: bool, +}; + +/// Build the source `find` pulls from, with `skip` already consumed. +/// +/// The query shape decides which source is possible. An index-ordered scan +/// yields offsets -- 8 bytes apiece, and valid until the collection is rebuilt. +/// A sort no index provides had to materialize and order every match, so there +/// is no ordered offset list to point at and the remainder is snapshotted. +fn find_source( + reply: *wire.Reply, + arena: *std.heap.ArenaAllocator, + coll: ?*Collection, + matched: []const u64, + shape: SourceShape, +) !cursor.Source { + const c = coll orelse return .{ .offsets = .{ .items = &.{} } }; + if (shape.sort_keys.len == 0 or shape.index_sorted) { + // Already in the order the client asked for, so skip is a slice. + const rest = if (shape.skip < matched.len) matched[shape.skip..] else &.{}; + return .{ .offsets = .{ .items = try arena.allocator().dupe(u64, rest) } }; + } + + // Ordering needs the values: materialize into the reply arena and sort. + const reply_arena = reply.arena_alloc(); var tree_docs: std.ArrayListUnmanaged(*const bson.Document) = .empty; - const arena = reply.arena_alloc(); - if (coll) |c| { - for (matched.items) |off| { - try tree_docs.append(arena, try doc_tree(arena, c, off)); - } + for (matched) |off| try tree_docs.append(reply_arena, try doc_tree(reply_arena, c, off)); + if (shape.closes_here and shape.page_end > 0 and shape.page_end *| 4 <= tree_docs.items.len) { + try query.sort_docs_top_k(reply_arena, tree_docs.items, shape.sort_keys, shape.page_end); + } else { + try query.sort_docs(reply_arena, tree_docs.items, shape.sort_keys); } - if (sort_keys.len > 0 and !index_sorted) { - // Selecting the page is much cheaper than ordering everything when - // the page is a small fraction of the matches. Above that fraction - // the heap's bookkeeping stops paying for itself. - if (page_end > 0 and page_end *| 4 <= tree_docs.items.len) { - try query.sort_docs_top_k(arena, tree_docs.items, sort_keys, page_end); - } else { - try query.sort_docs(arena, tree_docs.items, sort_keys); - } - } - const rest = if (skip < tree_docs.items.len) tree_docs.items[skip..] else &.{}; - const page = if (limit > 0 and limit < rest.len) rest[0..limit] else rest; - try emit_docs_tree(reply, db_name, coll_name, proj_pairs, page); - try reply.put_ok(); + const rest = if (shape.skip < tree_docs.items.len) tree_docs.items[shape.skip..] else &.{}; + return buffered_source(arena.allocator(), rest); +} + +/// Serialize `pairs` into `arena`. A cursor cannot keep the parsed form: those +/// pairs point into the per-request message arena. +fn serialize_pairs(arena: std.mem.Allocator, pairs: []const bson.Pair) ![]const u8 { + var buf: std.ArrayListUnmanaged(u8) = .empty; + try bson.write_doc(pairs, arena, &buf); + return buf.items; } /// Collect the documents in `db_name.coll_name` matching `filter`, stopping @@ -754,6 +1036,27 @@ fn scan_sorted( out: ?*std.ArrayListUnmanaged(u64), sort: []const query.SortKey, sorted: ?*bool, +) !usize { + return scan_planned(ctx, db_name, coll_name, filter, limit, out, sort, sorted, null); +} + +/// `scan_sorted` with the plan supplied. `find` decides between a resumable walk +/// and a materializing scan by looking at the plan, so without this it would plan +/// once to choose and `scan_sorted` would plan the identical query again -- +/// `index.plan` flattens the filter's clauses and evaluates every index, both +/// allocating, on the hot read path. +/// +/// `prebuilt` is borrowed: the caller keeps ownership, including its `deinit`. +fn scan_planned( + ctx: *Context, + db_name: []const u8, + coll_name: []const u8, + filter: []const bson.Pair, + limit: usize, + out: ?*std.ArrayListUnmanaged(u64), + sort: []const query.SortKey, + sorted: ?*bool, + prebuilt: ?*const index.Plan, ) !usize { if (sorted) |flag| flag.* = false; // Stopping early is only meaningful when the candidates come out in the @@ -781,10 +1084,16 @@ fn scan_sorted( var offs: std.ArrayListUnmanaged(u64) = .empty; defer offs.deinit(ctx.gpa); var cands: index.Candidates = undefined; - var plan_opt = try index.plan(ctx.gpa, &coll.id_index, coll.indexes.items, filter, sort); - defer if (plan_opt) |*p| p.deinit(ctx.gpa); + // Only plan here when the caller did not; theirs is borrowed, so only ours + // is freed. + var owned_plan = if (prebuilt != null) + null + else + try index.plan(ctx.gpa, &coll.id_index, coll.indexes.items, filter, sort); + defer if (owned_plan) |*p| p.deinit(ctx.gpa); + const plan_opt: ?*const index.Plan = prebuilt orelse if (owned_plan) |*p| p else null; - if (plan_opt) |*plan| { + if (plan_opt) |plan| { if (sorted) |flag| flag.* = plan.provides_sort; if (plan.provides_sort) lim = limit; if (plan.full_scan()) { @@ -818,26 +1127,384 @@ fn scan_sorted( /// the pair/value skeleton is allocated. The arena owns the skeleton, so /// the result is never deinit'd — the reply arena frees it with the reply. fn doc_tree(arena: std.mem.Allocator, coll: *const Collection, off: u64) !*const bson.Document { + return doc_tree_bytes(arena, coll.doc_bytes(off)); +} + +/// The same borrowed spine over bytes that are already in hand -- the slab for a +/// live scan, or a cursor's own snapshot for a buffered one. +fn doc_tree_bytes(arena: std.mem.Allocator, bytes: []const u8) !*const bson.Document { const doc = try arena.create(bson.Document); - doc.* = bson.Document{ .arena = undefined, .pairs = try bson.spine(arena, coll.doc_bytes(off)) }; + doc.* = bson.Document{ .arena = undefined, .pairs = try bson.spine(arena, bytes) }; return doc; } -fn emit_docs_tree( - 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); +// --------------------------------------------------------------------------- +// Cursors: filling a batch, and the sources a batch pulls from +// --------------------------------------------------------------------------- + +/// Largest snapshot a cursor will copy into its own arena. Mirrors MongoDB's +/// 32 MiB in-memory sort limit, and only the shapes that had to materialize +/// anyway can reach it. +/// +/// Over the bound the cursor is *declined* and the whole result goes out in one +/// batch -- exactly what this server did before cursors existed. Declining is +/// the right direction to fail: the alternative, holding the snapshot anyway, +/// pins it for the idle timeout times the number of live cursors. +const cursor_buffer_max: u64 = 32 * 1024 * 1024; + +/// The part of a cursor that a batch consumes. Kept separate from +/// `cursor.Cursor` so `find` can fill its first batch and only then decide +/// whether a cursor needs to exist at all. +const Feed = struct { + source: cursor.Source, + /// Documents still owed across every remaining batch; null is unbounded. + remaining_limit: ?u64, + /// A `.stream` source's live position in the tree, valid for **this request + /// only**. It is deliberately not part of `cursor.Source`: a tree position + /// must never outlive the collection lock that made it safe to hold, which is + /// the whole reason the stored form is a value-typed anchor instead. + scan: ?Scan = null, + + fn limit_exhausted(self: *const Feed) bool { + const rem = self.remaining_limit orelse return false; + return rem == 0; + } +}; + +/// A walk over one index, in one direction, for the duration of one request. +const Scan = struct { + walk: union(enum) { + fwd: index.Index.Iter, + rev: index.Index.RevIter, + }, + /// The entry `peek_bytes` has produced but the batch has not yet accepted. + /// Held because a batch that turns out to be full must not consume it. + pending: ?index.Index.Positioned = null, + + fn next(self: *Scan) ?index.Index.Positioned { + return switch (self.walk) { + .fwd => |*it| it.positioned(), + .rev => |*it| it.positioned(), + }; + } +}; + +/// The next document the source will yield, *without* consuming it -- a batch +/// that turns out to be full must not swallow a document it cannot carry. +/// +/// Candidates whose document no longer matches are consumed and skipped here. +/// That re-check is the index invariant (`src/index.zig`) applied per batch, and +/// it is also what makes a saved offset safe once documents start being +/// recycled: a reused offset either fails the filter or resolves to a document +/// that genuinely matches it. +fn peek_bytes( + ctx: *Context, + coll: ?*Collection, + filter: []const bson.Pair, + feed: *Feed, +) !?[]const u8 { + switch (feed.source) { + // A snapshot of documents that already matched, and that nothing can + // mutate underneath us -- so it is not re-filtered. + .buffered => |*b| { + if (b.next >= b.docs.len) return null; + return b.docs[b.next]; + }, + .offsets => |*o| { + const c = coll orelse return null; + while (o.next < o.items.len) { + const bytes = c.doc_bytes(o.items[o.next]); + if (try query.matches_bytes(ctx.gpa, filter, bytes)) return bytes; + o.next += 1; + } + return null; + }, + .stream => { + const c = coll orelse return null; + const sc = &(feed.scan orelse return null); + if (sc.pending) |p| return c.doc_bytes(p.off); + while (sc.next()) |p| { + const bytes = c.doc_bytes(p.off); + if (!try query.matches_bytes(ctx.gpa, filter, bytes)) { + // A candidate the filter rejects is still progress: the + // anchor must move past it, or a resume would walk it again. + feed.source.stream.advance(p.key, p.off, p.leaf, p.slot); + continue; + } + sc.pending = p; + return bytes; + } + return null; + }, } - 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. +fn consume_one(feed: *Feed) void { + switch (feed.source) { + .buffered => |*b| b.next += 1, + .offsets => |*o| o.next += 1, + .stream => |*st| { + const sc = &(feed.scan orelse return); + const p = sc.pending orelse return; + st.advance(p.key, p.off, p.leaf, p.slot); + sc.pending = null; + }, + } +} + +/// Resolve the index a `.stream` cursor names. Empty is the implicit `_id_`, +/// which is deliberately kept out of `Collection.indexes` and so is not findable +/// by name. +fn stream_index(coll: *Collection, name: []const u8) ?*const index.Index { + if (name.len == 0) return &coll.id_index; + return coll.find_index(name); +} + +/// Open this request's walk over the index a `.stream` cursor is following. +/// +/// Null means the cursor can never produce another document: the index it was +/// following is gone. The caller turns that into `QueryPlanKilled` rather than +/// an empty batch, because an empty batch would claim the result set ended. +fn open_scan(coll: *Collection, st: *cursor.Stream) ?Scan { + const ix = stream_index(coll, st.index_name()) orelse return null; + // A hint is a node id plus a slot, and `reset_tree` recycles ids 0 and 1 as + // different nodes -- so the hint is only meaningful at the epoch it was + // taken. The *anchor* is unaffected: it is key bytes and an offset, both + // values, so an epoch change costs a band walk rather than correctness. + const trusted = st.index_epoch == ix.epoch; + st.index_epoch = ix.epoch; + + if (!st.started()) { + // `batchSize: 0` leaves a cursor with no anchor yet, so a first getMore + // starts the walk from the beginning rather than resuming. + return .{ .walk = if (st.backward) + .{ .rev = ix.iter_reverse() } + else + .{ .fwd = ix.iter() } }; + } + if (st.backward) { + const r = ix.resume_reverse( + st.anchor_key(), + st.anchor_off, + st.hint_leaf, + st.hint_slot, + trusted, + ); + if (r.capped) return null; + return .{ .walk = .{ .rev = r.it } }; + } + const r = ix.resume_forward( + st.anchor_key(), + st.anchor_off, + st.band_index, + st.hint_leaf, + st.hint_slot, + trusted, + ); + if (r.capped) return null; + return .{ .walk = .{ .fwd = r.it } }; +} + +/// Whether this query can be served by walking one index end to end, and if so +/// which index and in which direction. +/// +/// This is the shape that lets a cursor outlive its result set: it holds a key +/// and an offset instead of a list, so `find({})` over a collection larger than +/// memory costs O(key) of cursor state rather than 8 bytes per matching +/// document. Everything else keeps the materialized sources. +/// +/// Two shapes qualify. A query the planner declines outright (`{}` with no sort) +/// walks the `_id_` index forward -- every document has an `_id` and the index is +/// not sparse, so a full walk cannot miss one. And a plan whose `full_scan()` +/// holds is by construction a whole-index walk in the requested order; note that +/// `full_scan()` implies `provides_sort`, since a plan with no run, no range and +/// no sort direction is declined before it is built. +/// +/// A narrowed plan is excluded on purpose: it dedupes `$in` and multikey +/// candidates across the whole set, which a stream cannot do without remembering +/// what it has already emitted. +fn stream_shape( + coll: *Collection, + plan_opt: ?*const index.Plan, + sort: []const query.SortKey, +) ?cursor.Stream { + var st = cursor.Stream{}; + if (plan_opt) |plan| { + if (!plan.full_scan()) return null; + st.backward = plan.backward; + if (plan.index != &coll.id_index) { + if (plan.index.name.len > cursor.index_name_max) return null; + @memcpy(st.index_name_buf[0..plan.index.name.len], plan.index.name); + st.index_name_len = @intCast(plan.index.name.len); + } + st.index_epoch = plan.index.epoch; + return st; + } + // No usable predicate and no ordering to honour: every document in _id + // order, which is what the old materializing fallback did too. + if (sort.len != 0) return null; + st.index_epoch = coll.id_index.epoch; + return st; +} + +/// Consume `skip` matching documents without emitting them, advancing the anchor +/// as it goes so a resume does not walk them again. +fn stream_skip( + ctx: *Context, + coll: *Collection, + feed: *Feed, + filter: []const bson.Pair, + skip: u64, +) !void { + var left = skip; + while (left > 0) { + _ = try peek_bytes(ctx, coll, filter, feed) orelse return; + consume_one(feed); + left -= 1; + } +} + +/// Fill one batch from `feed`, projecting into the reply arena. Returns whether +/// the source is **exhausted**. +/// +/// Exhaustion is observed, never predicted: a batch that reached its document +/// target returns false even when the source happens to have nothing left, so +/// the cursor stays open and the client gets one more (possibly empty) batch. +/// Predicting it here would close the cursor a round trip early and break the +/// command-count assertions in the pinned spec suites. +fn fill_batch( + ctx: *Context, + reply: *wire.Reply, + coll: ?*Collection, + feed: *Feed, + filter: []const bson.Pair, + proj_pairs: ?[]const bson.Pair, + target: ?u32, + out: *std.ArrayListUnmanaged(bson.Value), +) !bool { + var builder = cursor.BatchBuilder.init(target); + const arena = reply.arena_alloc(); + while (true) { + // Checked before the target so the batch that takes the last document + // the limit allows is itself the one that closes the cursor. This is + // what lets `batchSize == limit` finish in a single round trip. + if (feed.limit_exhausted()) return true; + if (builder.full()) return false; + const bytes = try peek_bytes(ctx, coll, filter, feed) orelse return true; + // The stored length is exact with no projection and an upper bound with + // one, since `query.project` only ever drops fields. + if (builder.offer(bytes.len) == .batch_full) return false; + const doc = try doc_tree_bytes(arena, bytes); + try out.append(arena, try project_doc(reply, doc, proj_pairs)); + consume_one(feed); + if (feed.remaining_limit) |rem| feed.remaining_limit = rem - 1; + } +} + +/// Register `feed`'s remainder as a cursor and return its id, or 0 when no +/// cursor is needed or one could not be had. +/// +/// Every "could not" path degrades to `id: 0` rather than to an error: the +/// client then has the batch it was given and no cursor, which is precisely how +/// this server behaved before cursors existed. +fn open_cursor(ctx: *Context, arena: std.heap.ArenaAllocator, spec: cursor.OpenSpec) i64 { + var owned = arena; + return ctx.engine.cursors.open(ctx.io, now_ms(ctx), owned, spec) catch |err| switch (err) { + // A namespace too long for the fixed buffers, a store whose every slot + // is pinned, or an allocation failure. None is worth failing a query + // whose documents are already in the reply. + error.NameTooLong, error.TooManyCursors, error.OutOfMemory, error.Canceled => { + owned.deinit(); + return 0; + }, + }; +} + +fn now_ms(ctx: *Context) i64 { + return std.Io.Timestamp.now(ctx.io, .real).toMilliseconds(); +} + +/// Copy `docs` into `arena` as canonical BSON, for a result with no stable +/// backing store to point at. +fn buffered_source( + arena: std.mem.Allocator, + docs: []const *const bson.Document, +) !cursor.Source { + const out = try arena.alloc([]const u8, docs.len); + for (docs, 0..) |d, i| out[i] = try serialize_pairs(arena, d.pairs); + return .{ .buffered = .{ .docs = out } }; +} + +/// Wrap pairs already living in an arena as a borrowed document, so generated +/// results can go through the same batch path as stored ones. +fn doc_from_pairs(arena: std.mem.Allocator, pairs: []const bson.Pair) !*const bson.Document { + const doc = try arena.create(bson.Document); + doc.* = bson.Document{ .arena = undefined, .pairs = pairs }; + return doc; +} + +/// Emit a first batch from documents already materialized in the reply arena, +/// registering a cursor for whatever does not fit in it. +/// +/// The snapshot source is the only one available here: these documents are either +/// generated (`$group`, a listing) or the output of a pipeline that has already +/// materialized its window, so there is no stable structure to point back into. +/// That also makes the resulting cursor independent of its collection, which is +/// what lets a listing hold a cursor over a `$cmd.*` namespace. +/// +/// `ns_coll` is the collection part of the reported namespace and the one a +/// `getMore` must name -- for a listing that is `$cmd.listCollections`, not the +/// empty string that used to be reported. An empty collection part makes the +/// driver throw client-side before it even sends the getMore. +fn emit_first_batch( + ctx: *Context, + reply: *wire.Reply, + ns_db: []const u8, + ns_coll: []const u8, + proj_pairs: ?[]const bson.Pair, + docs: []const *const bson.Document, + batch_size: ?u32, +) !void { + var arena = std.heap.ArenaAllocator.init(ctx.gpa); + var arena_owned = false; + defer if (!arena_owned) arena.deinit(); + + const source = try buffered_source(arena.allocator(), docs); + var feed = Feed{ .source = source, .remaining_limit = null }; + var values: std.ArrayListUnmanaged(bson.Value) = .empty; + const target = cursor.batch_target(batch_size, true); + const exhausted = try fill_batch(ctx, reply, null, &feed, &.{}, proj_pairs, target, &values); + + var cursor_id: i64 = 0; + if (!exhausted and source_keepable(feed.source)) { + cursor_id = open_cursor(ctx, arena, .{ + .ns = .{ .db = ns_db, .coll = ns_coll }, + // Nothing about this cursor depends on the collection's layout. + .layout_epoch = 0, + .batch_size = batch_size, + .source = feed.source, + }); + arena_owned = cursor_id != 0; + } + if (cursor_id == 0 and !exhausted) { + _ = try fill_batch(ctx, reply, null, &feed, &.{}, proj_pairs, null, &values); + } + + const ns = try format_namespace(reply, ns_db, ns_coll); + const batch = try cursor_doc(reply, cursor_id, ns, "firstBatch", values.items); + try reply.put("cursor", .{ .doc = batch }); +} + +/// `batchSize` out of an `aggregate`'s `cursor` option. A bare `cursor: {}` means +/// the default; a missing `cursor` is accepted as the same thing, which is looser +/// than mongod (it requires the field) but cannot surprise any driver. +fn aggregate_batch_size(reply: *wire.Reply, msg: *wire.Message) !?struct { value: ?u32 } { + const spec = doc_arg(msg.body.get("cursor")) orelse return .{ .value = null }; + const v = bson.get_pair(spec, "batchSize") orelse return .{ .value = null }; + const parsed = try batch_size_arg(reply, v, "cursor.batchSize", false) orelse return null; + return .{ .value = parsed }; +} + fn project_doc( reply: *wire.Reply, doc: *const bson.Document, @@ -1065,8 +1732,22 @@ fn cmd_count(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const db_name = msg.db_name() orelse return invalid_arg(reply, "count requires $db"); 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 &.{}; + // `count` takes skip and limit like `find` does, and ignoring them was a + // silent wrong answer for `countDocuments(f, {limit})`. + const skip: u64 = int_arg(msg.body.get("skip")) orelse 0; + const limit: u64 = @abs(int_value(msg.body.get("limit")) orelse 0); - const n = try scan_matching(ctx, db_name, coll_name, q, 0, null); + // Counting only needs to know whether the matches reach skip + limit, so the + // scan may stop there. Unlike `find` this needs no index to be an early stop: + // the *count* of a window does not depend on which documents fall in it. + const ceiling: usize = if (limit == 0) 0 else blk: { + const s = std.math.cast(usize, skip) orelse break :blk 0; + const l = std.math.cast(usize, limit) orelse break :blk 0; + break :blk s +| l; + }; + const matched = try scan_matching(ctx, db_name, coll_name, q, ceiling, null); + const after_skip = matched -| (std.math.cast(usize, skip) orelse matched); + const n = if (limit == 0) after_skip else @min(after_skip, limit); try reply.put("n", .{ .int32 = @intCast(n) }); try reply.put_ok(); } @@ -1075,6 +1756,7 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const db_name = msg.db_name() orelse return invalid_arg(reply, "aggregate requires $db"); const coll_name = str_arg(msg.body.get("aggregate")) orelse return bad_value(reply, "aggregate requires a collection name"); + const batch_size = (try aggregate_batch_size(reply, msg) orelse return).value; const pipeline_value = msg.body.get("pipeline") orelse return bad_value(reply, "aggregate requires pipeline"); var stages = switch (pipeline_value) { .array => |a| a, @@ -1106,13 +1788,12 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { .{ .double = sum }, }; } - const doc = try arena.create(bson.Document); - doc.* = bson.Document{ .arena = undefined, .pairs = pairs }; + const doc = try doc_from_pairs(arena, pairs); const one = try arena.alloc(*const bson.Document, 1); one[0] = doc; docs = one; } - try emit_docs_tree(reply, db_name, coll_name, null, docs); + try emit_first_batch(ctx, reply, db_name, coll_name, null, docs, batch_size); return reply.put_ok(); } @@ -1132,7 +1813,7 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { // collection. MongoDB answers an aggregate over a missing collection with // an empty cursor. const coll = ctx.engine.get_collection(db_name, coll_name) orelse { - try emit_docs_tree(reply, db_name, coll_name, null, &.{}); + try emit_first_batch(ctx, reply, db_name, coll_name, null, &.{}, batch_size); return reply.put_ok(); }; // A leading $match is pushed down into an indexed candidate scan; the @@ -1249,7 +1930,12 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { }; } else { const msg_text = try std.fmt.allocPrint(reply.arena_alloc(), "Unrecognized pipeline stage name: '{s}'", .{stage_name}); - return reply.put_error(@intFromEnum(ErrorCode.invalid_pipeline_operator), "InvalidPipelineOperator", msg_text); + // 40324 is right for "unrecognized stage" but its name is not + // `InvalidPipelineOperator` (that is 168). mongod reports numeric + // Location codes under a `Location` name -- verified by asking a + // real mongod for an unknown stage. + const code = @intFromEnum(ErrorCode.location_unrecognized_stage); + return reply.put_error(code, "Location40324", msg_text); } } @@ -1257,17 +1943,19 @@ fn cmd_aggregate(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { const len = if (in_trees) trees.items[start..end].len else offs.items[start..end].len; const c = try reply.arena_alloc().alloc(bson.Pair, 1); c[0] = .{ .key = try reply.arena_alloc().dupe(u8, name), .value = .{ .int32 = @intCast(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), "firstBatch", values) }); + const one = try reply.arena_alloc().alloc(*const bson.Document, 1); + one[0] = try doc_from_pairs(reply.arena_alloc(), c); + try emit_first_batch(ctx, reply, db_name, coll_name, null, one, batch_size); } else { const arena = reply.arena_alloc(); if (in_trees) { - try emit_docs_tree(reply, db_name, coll_name, proj_pairs, trees.items[start..end]); + const window = trees.items[start..end]; + try emit_first_batch(ctx, reply, db_name, coll_name, proj_pairs, window, batch_size); } else { var page: std.ArrayListUnmanaged(*const bson.Document) = .empty; for (offs.items[start..end]) |off| try page.append(arena, try doc_tree(arena, coll, off)); - try emit_docs_tree(reply, db_name, coll_name, proj_pairs, page.items); + const window = page.items; + try emit_first_batch(ctx, reply, db_name, coll_name, proj_pairs, window, batch_size); } } try reply.put_ok(); @@ -1471,16 +2159,227 @@ fn query_path_value_bytes( return cur; } -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", &.{}) }); +/// A cursor id from the wire. Accepts int32 as well as int64, so a hand-written +/// `runCommand` is not rejected on a technicality; the driver always sends a +/// long. +fn cursor_id_arg(v: ?bson.Value) ?i64 { + return switch (v orelse return null) { + .int64 => |i| i, + .int32 => |i| i, + else => null, + }; +} + +fn cursor_not_found(reply: *wire.Reply, id: i64) !void { + const text = try std.fmt.allocPrint(reply.arena_alloc(), "cursor id {d} not found", .{id}); + return reply.put_error(@intFromEnum(ErrorCode.cursor_not_found), "CursorNotFound", text); +} + +fn query_plan_killed(reply: *wire.Reply, why: []const u8) !void { + const text = try std.fmt.allocPrint( + reply.arena_alloc(), + "query plan killed :: caused by :: {s}", + .{why}, + ); + const code = @intFromEnum(ErrorCode.query_plan_killed); + return reply.put_error(code, "QueryPlanKilled", text); +} + +/// Whether this cursor can still be answered from `coll`, or why not. +/// +/// What a rebuild since the cursor was created costs depends entirely on what the +/// cursor remembers. `.offsets` holds slab offsets and a rebuild moved every +/// document, so those offsets now name unrelated bytes. `.stream` holds key +/// bytes, and a repack changes no key -- only the anchor's offset and position +/// hint go stale, and both are checked before they are believed, so the walk +/// resumes at the right key and the offsets it yields come fresh out of the tree. +/// `.buffered` holds copies of the documents and needs no collection at all, +/// which is what lets a listing hold a cursor over a `$cmd.*` namespace. +fn cursor_still_valid(c: *cursor.Cursor, coll: ?*Collection) ?[]const u8 { + const live = coll orelse { + // A snapshot needs no collection at all, which is what lets a listing + // hold a cursor over a `$cmd.*` namespace nothing backs. + return if (c.source == .buffered) null else "collection dropped"; + }; + if (live.layout_epoch == c.layout_epoch) return null; + if (c.source == .offsets) return "collection rebuilt"; + // Survived the rebuild: adopt the new layout so the next getMore does not + // re-examine it. + c.layout_epoch = live.layout_epoch; + return null; +} + +/// The stored filter and projection, reparsed for this request. They were +/// serialized into the cursor's own arena because the parsed forms pointed into +/// the request that created them; `spine` parses the structure and borrows the +/// leaf bytes, so this costs no copy of the filter's strings. +fn cursor_query(reply: *wire.Reply, c: *const cursor.Cursor) !struct { + filter: []const bson.Pair, + proj: ?[]const bson.Pair, +} { + const arena = reply.arena_alloc(); + return .{ + .filter = if (c.filter_bytes.len == 0) &.{} else try bson.spine(arena, c.filter_bytes), + .proj = if (c.proj_bytes.len == 0) null else try bson.spine(arena, c.proj_bytes), + }; +} + +/// The shape of a `getMore` request, or null once the error reply is written. +fn parse_get_more(reply: *wire.Reply, msg: *wire.Message) !?struct { id: i64, batch_size: ?u32 } { + const id = cursor_id_arg(msg.body.get("getMore")) orelse { + const text = "BSON field 'getMore.getMore' is the wrong type"; + try reply.put_error(@intFromEnum(ErrorCode.type_mismatch), "TypeMismatch", text); + return null; + }; + var batch_size: ?u32 = null; + if (msg.body.get("batchSize")) |v| { + batch_size = try batch_size_arg(reply, v, "batchSize", true) orelse return null; + } + return .{ .id = id, .batch_size = batch_size }; +} + +/// Claim the cursor for this request, or write the error reply and return null. +/// +/// The namespace check inside `pin` is load-bearing rather than cosmetic: +/// dispatch locks the collection the *message* names, so a getMore quoting one +/// cursor's id and another collection's name would otherwise iterate the first +/// collection's index while holding the second collection's lock. +fn pin_cursor(ctx: *Context, reply: *wire.Reply, id: i64, ns: cursor.Ns) !?*cursor.Cursor { + return ctx.engine.cursors.pin(ctx.io, id, ns, now_ms(ctx)) catch |err| switch (err) { + error.CursorNotFound => { + try cursor_not_found(reply, id); + return null; + }, + error.CursorNamespaceMismatch => { + var found: cursor.NsBuf = .{}; + const arena = reply.arena_alloc(); + const text = if (ctx.engine.cursors.ns_of(ctx.io, id, &found)) + try std.fmt.allocPrint( + arena, + "Requested getMore on namespace '{s}.{s}', but cursor belongs to" ++ + " a different namespace {s}.{s}", + .{ ns.db, ns.coll, found.ns().db, found.ns().coll }, + ) + else + try std.fmt.allocPrint( + arena, + "Requested getMore on namespace '{s}.{s}', but cursor belongs to" ++ + " a different namespace", + .{ ns.db, ns.coll }, + ); + try reply.put_error(@intFromEnum(ErrorCode.unauthorized), "Unauthorized", text); + return null; + }, + error.CursorInUse => { + const text = try std.fmt.allocPrint( + reply.arena_alloc(), + "cursor id {d} is already in use", + .{id}, + ); + try reply.put_error(@intFromEnum(ErrorCode.cursor_in_use), "CursorInUse", text); + return null; + }, + error.Canceled => return err, + }; +} + +fn cmd_get_more(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { + const db_name = msg.db_name() orelse return invalid_arg(reply, "getMore requires $db"); + const coll_name = str_arg(msg.body.get("collection")) orelse + return bad_value(reply, "Field 'collection' must be of type string"); + if (coll_name.len == 0) return bad_value(reply, "Collection names cannot be empty"); + const req = try parse_get_more(reply, msg) orelse return; + const id = req.id; + + const ns = cursor.Ns{ .db = db_name, .coll = coll_name }; + const c = try pin_cursor(ctx, reply, id, ns) orelse return; + // Set to true by every path that must not leave the cursor behind, including + // the error paths below: a cursor whose collection is gone can never answer + // again, so keeping it would only occupy a slot until the idle sweep. + var done = false; + defer ctx.engine.cursors.release(ctx.io, c, now_ms(ctx), done); + + const coll = ctx.engine.get_collection(db_name, coll_name); + if (cursor_still_valid(c, coll)) |why| { + done = true; + return query_plan_killed(reply, why); + } + const q = try cursor_query(reply, c); + + var feed = Feed{ .source = c.source, .remaining_limit = c.remaining_limit }; + // Persist on every exit, not just the success path. For a stream the advance + // *is* the resume anchor, so an early return that skipped this would report a + // batch and then hand the same documents out again on the next getMore. + // Registered after `release`'s defer, so it runs before it. + defer { + c.source = feed.source; + c.remaining_limit = feed.remaining_limit; + } + if (feed.source == .stream) { + // Only a stream needs the collection here; `open_scan` walks its index. + assert_msg(coll != null, "a streaming cursor reached getMore with no collection"); + // Reopening the walk is where a resume actually happens. Null means the + // index the cursor was following is gone, or its anchor could not be + // located within the walk bound -- either way there is no honest way to + // continue, and an empty batch would falsely claim the result ended. + feed.scan = open_scan(coll.?, &feed.source.stream); + if (feed.scan == null) { + done = true; + return query_plan_killed(reply, "the index this cursor was following is gone"); + } + } + var values: std.ArrayListUnmanaged(bson.Value) = .empty; + // Deliberately not `batch_size orelse c.batch_size`: mongod does not carry + // the find's batchSize into a bare getMore. Measured -- find with + // batchSize 2 then a bare getMore returns 4998 of 5000 documents. + const target = cursor.batch_target(req.batch_size, false); + done = try fill_batch(ctx, reply, coll, &feed, q.filter, q.proj, target, &values); + + const ns_str = try format_namespace(reply, db_name, coll_name); + const live_id: i64 = if (done) 0 else id; + const batch = try cursor_doc(reply, live_id, ns_str, "nextBatch", values.items); + try reply.put("cursor", .{ .doc = batch }); try reply.put_ok(); } -fn cmd_kill_cursors(_: *Context, _: *wire.Message, reply: *wire.Reply) !void { - try reply.put("cursorsKilled", .{ .array = &.{} }); - try reply.put("cursorsNotFound", .{ .array = &.{} }); +fn cmd_kill_cursors(ctx: *Context, msg: *wire.Message, reply: *wire.Reply) !void { + const db_name = msg.db_name() orelse return invalid_arg(reply, "killCursors requires $db"); + // Unlike getMore, killCursors' own field really is the collection name. + const coll_name = str_arg(msg.body.get("killCursors")) orelse + return bad_value(reply, "killCursors requires a collection name"); + const cursors_arg = msg.body.get("cursors") orelse + return bad_value(reply, "killCursors requires a cursors array"); + const ids = switch (cursors_arg) { + .array => |a| a, + else => return bad_value(reply, "cursors must be an array"), + }; + + const arena = reply.arena_alloc(); + var killed: std.ArrayListUnmanaged(bson.Value) = .empty; + var not_found: std.ArrayListUnmanaged(bson.Value) = .empty; + const ns = cursor.Ns{ .db = db_name, .coll = coll_name }; + for (ids) |v| { + const id = cursor_id_arg(v) orelse { + try not_found.append(arena, v); + continue; + }; + // A namespace mismatch reports not-found rather than erroring: + // killCursors is best-effort by design, and the driver ignores its reply + // entirely. + switch (ctx.engine.cursors.kill(ctx.io, id, ns)) { + .killed => try killed.append(arena, .{ .int64 = id }), + .not_found => try not_found.append(arena, .{ .int64 = id }), + } + } + + try reply.put("cursorsKilled", .{ .array = killed.items }); + try reply.put("cursorsNotFound", .{ .array = not_found.items }); + // Empty by construction: a cursor pinned by an in-flight getMore is marked + // and reported killed, since the client's intent is satisfied and the + // request frees it on release. `cursorsUnknown` exists for shape -- every + // outcome here is classified. try reply.put("cursorsAlive", .{ .array = &.{} }); + try reply.put("cursorsUnknown", .{ .array = &.{} }); try reply.put_ok(); } @@ -1618,7 +2517,7 @@ fn batch_arg( } fn invalid_arg(reply: *wire.Reply, msg: []const u8) !void { - return reply.put_error(@intFromEnum(ErrorCode.invalid_argument), "InvalidArgument", msg); + return reply.put_error(@intFromEnum(ErrorCode.invalid_options), "InvalidOptions", msg); } fn bad_value(reply: *wire.Reply, msg: []const u8) !void { @@ -2075,8 +2974,8 @@ fn dispatch_find_ids( defer reply.deinit(); try dispatch(&ctx, &msg, &reply); try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double); - const cursor = bson.get_pair(reply.pairs.items, "cursor") orelse return error.TestUnexpectedResult; - const batch = switch (bson.get_pair(cursor.doc, "firstBatch") orelse return error.TestUnexpectedResult) { + const cur = bson.get_pair(reply.pairs.items, "cursor") orelse return error.TestUnexpectedResult; + const batch = switch (bson.get_pair(cur.doc, "firstBatch") orelse return error.TestUnexpectedResult) { .array => |a| a, else => return error.TestUnexpectedResult, }; @@ -2132,8 +3031,8 @@ test "createIndexes, listIndexes, dropIndexes, and idempotent re-create" { defer reply.deinit(); try dispatch(&ctx, &msg, &reply); try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double); - const cursor = bson.get_pair(reply.pairs.items, "cursor").?; - const batch = bson.get_pair(cursor.doc, "firstBatch").?.array; + const cur = bson.get_pair(reply.pairs.items, "cursor").?; + const batch = bson.get_pair(cur.doc, "firstBatch").?.array; try testing.expectEqual(@as(usize, 2), batch.len); try testing.expectEqualStrings("_id_", bson.get_pair(batch[0].doc, "name").?.string); try testing.expectEqualStrings("email_1", bson.get_pair(batch[1].doc, "name").?.string); @@ -2210,8 +3109,8 @@ test "TTL index round-trips through createIndexes/listIndexes; bad specs give 67 var reply = wire.Reply.init(testing.allocator); defer reply.deinit(); try dispatch(&ctx, &msg, &reply); - const cursor = bson.get_pair(reply.pairs.items, "cursor").?; - const batch = bson.get_pair(cursor.doc, "firstBatch").?.array; + const cur = bson.get_pair(reply.pairs.items, "cursor").?; + const batch = bson.get_pair(cur.doc, "firstBatch").?.array; try testing.expectEqual(@as(usize, 2), batch.len); try testing.expectEqual(@as(i32, 60), bson.get_pair(batch[1].doc, "expireAfterSeconds").?.int32); // The _id_ entry never carries one. @@ -2380,8 +3279,8 @@ test "aggregate $sort without a preceding $group sorts and frees correctly" { try dispatch(&ctx, &msg, &reply); try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double); - const cursor = bson.get_pair(reply.pairs.items, "cursor") orelse return error.TestUnexpectedResult; - const batch = switch (bson.get_pair(cursor.doc, "firstBatch") orelse return error.TestUnexpectedResult) { + const cur = bson.get_pair(reply.pairs.items, "cursor") orelse return error.TestUnexpectedResult; + const batch = switch (bson.get_pair(cur.doc, "firstBatch") orelse return error.TestUnexpectedResult) { .array => |a| a, else => return error.TestUnexpectedResult, }; @@ -2489,8 +3388,8 @@ test "a sorted full scan over a multikey index returns each document once" { defer reply.deinit(); try dispatch(&ctx, &msg, &reply); try testing.expectEqual(@as(f64, 1.0), bson.get_pair(reply.pairs.items, "ok").?.double); - const cursor = bson.get_pair(reply.pairs.items, "cursor") orelse return error.TestUnexpectedResult; - const batch = switch (bson.get_pair(cursor.doc, "firstBatch") orelse return error.TestUnexpectedResult) { + const cur = bson.get_pair(reply.pairs.items, "cursor") orelse return error.TestUnexpectedResult; + const batch = switch (bson.get_pair(cur.doc, "firstBatch") orelse return error.TestUnexpectedResult) { .array => |arr| arr, else => return error.TestUnexpectedResult, }; @@ -2801,3 +3700,136 @@ test "indexed queries are equivalent to scans over a mixed corpus" { for (scanned.items, indexed.items) |a, b| try testing.expectEqualSlices(u8, a, b); } } + +/// Insert `n` documents `{_id: 1..n, a: i % 5, pad}` in one batch. +/// +/// Goes through `dispatch_insert` rather than dispatching itself, because that +/// helper checks `ok`, `writeErrors` *and* `n` -- and its comment records what +/// omitting those cost last time: a corpus silently lost a document and every +/// test over it still passed. A cursor test whose premise is "60 documents" must +/// not be able to become "0 documents" quietly. +const seed_pad = "0123456789012345678901234567890123456789"; + +fn seed_docs(tdb: *TestDb, io: std.Io, coll: []const u8, n: i32) !void { + const arena = testing.allocator; + const docs = try arena.alloc(bson.Value, @intCast(n)); + defer { + for (docs) |d| arena.free(d.doc); + arena.free(docs); + } + for (docs, 0..) |*d, i| { + const pairs = try arena.alloc(bson.Pair, 3); + pairs[0] = .{ .key = "_id", .value = .{ .int32 = @intCast(i + 1) } }; + pairs[1] = .{ .key = "a", .value = .{ .int32 = @intCast(@mod(i + 1, 5)) } }; + pairs[2] = .{ .key = "pad", .value = .{ .string = seed_pad } }; + d.* = .{ .doc = pairs }; + } + try dispatch_insert(tdb, io, coll, docs); +} + +/// Run `find` and return (cursor id, first-batch length). +fn dispatch_find( + ctx: *Context, + coll: []const u8, + filter: []const bson.Pair, + batch_size: i32, +) !struct { id: i64, n: usize } { + var reply = wire.Reply.init(testing.allocator); + defer reply.deinit(); + var msg = try parse_fake_msg("find", .{ .string = coll }, &.{ + .{ .key = "filter", .value = .{ .doc = filter } }, + .{ .key = "batchSize", .value = .{ .int32 = batch_size } }, + }); + defer msg.deinit(); + try dispatch(ctx, &msg, &reply); + const cur = bson.get_pair(reply.pairs.items, "cursor") orelse return error.TestUnexpectedResult; + const id = bson.get_pair(cur.doc, "id").?.int64; + const batch = bson.get_pair(cur.doc, "firstBatch").?.array; + return .{ .id = id, .n = batch.len }; +} + +/// Run `getMore` and return the error code, or 0 on success. +fn dispatch_get_more(ctx: *Context, coll: []const u8, id: i64) !i32 { + var reply = wire.Reply.init(testing.allocator); + defer reply.deinit(); + var msg = try parse_fake_msg("getMore", .{ .int64 = id }, &.{ + .{ .key = "collection", .value = .{ .string = coll } }, + .{ .key = "batchSize", .value = .{ .int32 = 5 } }, + }); + defer msg.deinit(); + try dispatch(ctx, &msg, &reply); + if (bson.get_pair(reply.pairs.items, "code")) |c| return c.int32; + return 0; +} + +test "a rebuild kills an offsets cursor and spares a streaming one" { + // This is also the test that proves the streaming source is *selected*: a + // whole-index walk and an indexed-predicate scan are given different sources, + // and a rebuild is exactly what tells them apart. If `find({})` quietly fell + // back to materializing offsets, both would die here. + // + // A rebuild is triggered directly rather than through churn, because whether + // churn crosses the compaction threshold is not something a test should have + // to guess at. + var threaded: std.Io.Threaded = .init_single_threaded; + defer threaded.deinit(); + const io = threaded.io(); + var tdb = try TestDb.init(io); + defer tdb.deinit(); + var ctx = tdb.ctx(io); + + try seed_docs(&tdb, io, "c", 60); + + // A whole-index walk: no predicate, no sort. + const walk = try dispatch_find(&ctx, "c", &.{}, 5); + try testing.expectEqual(@as(usize, 5), walk.n); + try testing.expect(walk.id != 0); + // A narrowed plan over the _id_ index, which materializes its candidates. + const narrowed = try dispatch_find(&ctx, "c", &.{ + .{ .key = "_id", .value = .{ .doc = &.{.{ .key = "$gte", .value = .{ .int32 = 1 } }} } }, + }, 5); + try testing.expect(narrowed.id != 0); + + // Both still work before the rebuild, so the difference below is the rebuild. + try testing.expectEqual(@as(i32, 0), try dispatch_get_more(&ctx, "c", walk.id)); + try testing.expectEqual(@as(i32, 0), try dispatch_get_more(&ctx, "c", narrowed.id)); + + try ctx.engine.compact(); + + // The stream remembers key bytes, which a repack does not change. + try testing.expectEqual(@as(i32, 0), try dispatch_get_more(&ctx, "c", walk.id)); + // The offsets name bytes that have moved, so continuing would be a wrong + // answer; QueryPlanKilled says so instead. + try testing.expectEqual( + @intFromEnum(ErrorCode.query_plan_killed), + try dispatch_get_more(&ctx, "c", narrowed.id), + ); + // And a second getMore on the killed cursor no longer knows it at all. + try testing.expectEqual( + @intFromEnum(ErrorCode.cursor_not_found), + try dispatch_get_more(&ctx, "c", narrowed.id), + ); + + // The surviving stream drains to exactly the 60 documents, once each. + // Three batches of 5 are already out: the first batch, the getMore before the + // rebuild, and the one just after it. + var seen: u32 = 5 + 5 + 5; + var id = walk.id; + var rounds: u32 = 0; + while (rounds < 100) : (rounds += 1) { + var reply = wire.Reply.init(testing.allocator); + defer reply.deinit(); + var msg = try parse_fake_msg("getMore", .{ .int64 = id }, &.{ + .{ .key = "collection", .value = .{ .string = "c" } }, + .{ .key = "batchSize", .value = .{ .int32 = 7 } }, + }); + defer msg.deinit(); + try dispatch(&ctx, &msg, &reply); + const cur = bson.get_pair(reply.pairs.items, "cursor").?; + seen += @intCast(bson.get_pair(cur.doc, "nextBatch").?.array.len); + id = bson.get_pair(cur.doc, "id").?.int64; + if (id == 0) break; + } + try testing.expectEqual(@as(i64, 0), id); + try testing.expectEqual(@as(u32, 60), seen); +} diff --git a/src/cursor.zig b/src/cursor.zig new file mode 100644 index 0000000..039319f --- /dev/null +++ b/src/cursor.zig @@ -0,0 +1,979 @@ +//! Server-side cursor state: what a `find`/`aggregate` leaves behind so a later +//! `getMore` can carry on, and the fixed-capacity registry that holds it. +//! +//! This module is deliberately *pure*: it owns state and policy, never +//! execution. It does not import `db.zig` or `commands.zig`, so `db.Engine` can +//! own a `Store` with no import cycle, and the batch policy below is testable +//! with no engine, no socket and no allocator. Filling a batch stays in +//! `commands.zig`, which already owns orchestration the way `index.zig` owns +//! planning. +//! +//! **The one rule the whole batching protocol follows: never look ahead.** A +//! batch ends either because it reached its target -- and the cursor stays open +//! -- or because the source reported EOF, and then the cursor closes with +//! `id: 0` in that same reply. A batch that reached its target leaves the cursor +//! open *even when the source happens to be exhausted*. So four documents at +//! `batchSize: 2` need a third command answering `nextBatch: []` with `id: 0`; +//! that empty terminal batch is correct, not a bug, and the pinned spec suites +//! assert exactly that command count. +//! +//! ## What a cursor is allowed to remember +//! +//! A cursor holds no lock between requests, so everything it saves must survive +//! arbitrary concurrent mutation. Nothing here is a pointer, and the two things +//! that look like stable addresses are not: +//! +//! - A tree position `(leaf, slot)` is invalidated by `Index.reset_tree`, +//! which clears the node table so ids 0 and 1 become a live but *unrelated* +//! root and leaf. Guarded by `Stream.index_epoch`. +//! - A slab offset is invalidated by `rebuild_collection`, which moves every +//! document. Guarded by `layout_epoch`. +//! +//! Both are checked as error returns rather than assertions, because a client +//! can reach either one by keeping a cursor open across maintenance. + +const std = @import("std"); +const bson = @import("bson.zig"); +const index = @import("index.zig"); +// Always active, including in the default ReleaseFast build -- see assert.zig. +const assert = @import("assert.zig").assert; +const assert_msg = @import("assert.zig").assert_msg; + +// --------------------------------------------------------------------------- +// Bounds +// --------------------------------------------------------------------------- + +/// Longest index key a `.stream` cursor will anchor on. Tied to the B+tree's +/// own "this record is normal" threshold rather than picked: a key past it has +/// already spilled to the overflow slab, so the tree itself considers it +/// exceptional. Also what makes the anchor a fixed inline array instead of an +/// allocation. +/// +/// Unbounded, this is a memory denial of service and not a subtle one: +/// `bson.encode_key` escapes NULs, so a 16 MB string doubles, and a compound +/// index may carry 32 of them. +pub const anchor_key_max: usize = 1024; + +comptime { + // The bound is only defensible if it really is the tree's spill threshold. + // If the page size or the spill fraction ever changes, this fails to + // compile rather than silently becoming an arbitrary number. + std.debug.assert(anchor_key_max == index.inline_limit); +} + +pub const ns_db_max: usize = 64; +pub const ns_coll_max: usize = 192; +pub const index_name_max: usize = 128; + +/// Documents in a first batch when the client named no `batchSize`. MongoDB's +/// own default (`internalQueryFindCommandBatchSize`). +pub const default_first_batch: u32 = 101; + +/// Cap on a batch's document payload: `maxBsonObjectSize`, which is also what +/// leaves room for the reply envelope inside the 48 MiB message limit. +pub const batch_bytes_max: u64 = 16 * 1024 * 1024; + +/// Idle milliseconds before the sweep reaps a cursor. MongoDB's +/// `cursorTimeoutMillis`. +pub const default_idle_timeout_ms: i64 = 10 * 60 * 1000; + +/// Slots in the registry unless configured otherwise. +pub const default_capacity: u32 = 4096; + +/// Low bits of a cursor id that address its slot; the rest is the nonce. +const slot_bits: u6 = 20; +const slot_mask: u64 = (@as(u64, 1) << slot_bits) - 1; +/// Nonce width, leaving the sign bit clear so every id is a positive i64. +const nonce_mask: u64 = (@as(u64, 1) << (63 - slot_bits)) - 1; + +pub const max_capacity: u32 = @intCast(slot_mask); + +// --------------------------------------------------------------------------- +// Cursor state +// --------------------------------------------------------------------------- + +/// A namespace, by value. A cursor cannot hold a `*Collection`: `drop` frees +/// it, and the pointer would dangle exactly the way the M0 notes on +/// heap-allocating collections describe. +pub const Ns = struct { + db: []const u8, + coll: []const u8, +}; + +/// Where the remaining documents come from. +pub const Source = union(enum) { + /// An index-ordered scan, resumed from a value-typed anchor. O(key) + /// memory, so this is the shape that lets a cursor walk a collection far + /// larger than memory -- the reason M0 made whole-index scans stream. + stream: Stream, + /// Matched slab offsets, 8 bytes each, which `scan_sorted` has already + /// materialized for a narrowed plan. + /// + /// Safe against an offset the coming document free list has recycled, + /// because every batch re-applies the full filter -- the index invariant. + /// A recycled offset is therefore either rejected or resolves to a + /// document that genuinely matches. It needs one guarantee from the free + /// list, recorded in PLAN: an offset that was ever a record start must + /// stay a record start, since `doc_bytes` reads a length prefix in place. + offsets: struct { items: []u64, next: u32 = 0 }, + /// Canonical BSON bytes owned by the cursor's arena, for results with no + /// stable backing store to point at: a sort no index provides, and + /// aggregate/listCollections/listIndexes output. + buffered: struct { docs: []const []const u8, next: u32 = 0 }, +}; + +/// A resumable index scan. Every field is a value; nothing here is a pointer +/// into the tree, the slab or the request that created it. +pub const Stream = struct { + /// Empty means the implicit `_id_` index. Re-resolved by name on every + /// `getMore`, so a `dropIndexes` cannot leave a dangling `*Index`. + index_name_buf: [index_name_max]u8 = undefined, + index_name_len: u8 = 0, + /// Bumped by `reset_tree`/`replace_root_with_leaf`; if it moved, the hint + /// below addresses a different tree and must not be trusted. + index_epoch: u64 = 0, + backward: bool = false, + anchor_buf: [anchor_key_max]u8 = undefined, + anchor_len: u16 = 0, + anchor_off: u64 = 0, + /// Entries sharing the anchor's key that this cursor has already yielded. + /// Without it, an anchor whose document was deleted between batches would + /// resume past the entire equal-key band -- on a three-value index that is + /// millions of documents silently missing. + band_index: u64 = 0, + /// Last known position of the anchor. A hint, never trusted without + /// re-reading the entry there: it turns resume from a walk down the + /// equal-key band into O(1), which is what keeps a low-cardinality + /// `sort({status: 1})` from costing O(band) per batch. + hint_leaf: u32 = 0, + hint_slot: u32 = 0, + + pub fn index_name(self: *const Stream) []const u8 { + return self.index_name_buf[0..self.index_name_len]; + } + + pub fn anchor_key(self: *const Stream) []const u8 { + return self.anchor_buf[0..self.anchor_len]; + } + + /// Whether anything has been yielded yet. Derived rather than stored: an + /// encoded index key always begins with `bson.encode_key`'s rank byte, so it + /// is never empty, and a separate `started` flag would be a second field that + /// has to agree with this one. + /// + /// It can legitimately be false on a live cursor: `batchSize: 0` returns an + /// empty first batch without consuming anything, and such a cursor starts at + /// `iter()`/`iter_reverse()` rather than resuming. + pub fn started(self: *const Stream) bool { + return self.anchor_len > 0; + } + + /// Record the entry just yielded as the point to resume after. + /// + /// Asserts the anchor advances in scan order. This is the single check most + /// likely to catch a resume bug: going backwards duplicates documents, + /// standing still makes `getMore` loop forever, and both are far easier to + /// see here than in a client's result set. Equal keys are legal (a + /// duplicate band), which is exactly why `band_index` also has to move. + pub fn advance(self: *Stream, key: []const u8, off: u64, leaf: u32, slot: u32) void { + assert(key.len <= anchor_key_max); + // What makes `started()` derivable, so pin it here rather than trust it. + assert_msg(key.len > 0, "an encoded index key is never empty"); + if (self.started()) { + const order = std.mem.order(u8, key, self.anchor_key()); + if (self.backward) { + assert_msg(order != .gt, "a reverse cursor's anchor moved forward"); + } else { + assert_msg(order != .lt, "a forward cursor's anchor moved backward"); + } + if (order == .eq) { + assert_msg( + off != self.anchor_off or self.band_index > 0, + "a cursor re-anchored on the entry it just yielded", + ); + // Still inside the anchor's band, so the position within it has + // to move or a resume could not tell the two entries apart. + self.band_index += 1; + } else { + self.band_index = 0; + } + } + @memcpy(self.anchor_buf[0..key.len], key); + self.anchor_len = @intCast(key.len); + self.anchor_off = off; + self.hint_leaf = leaf; + self.hint_slot = slot; + } +}; + +pub const Cursor = struct { + /// Positive and never 0: `id: 0` is "no cursor" on the wire. + id: i64, + ns_db_buf: [ns_db_max]u8 = undefined, + ns_db_len: u8 = 0, + ns_coll_buf: [ns_coll_max]u8 = undefined, + ns_coll_len: u8 = 0, + /// Bumped when a rebuild moves documents, so a saved offset or anchor + /// offset is stale. Also the drop detector. + layout_epoch: u64 = 0, + /// Serialized so they outlive the request that parsed them: a parsed + /// `[]bson.Pair` points into the per-request message arena, and the reply + /// arena is reset on every request. + filter_bytes: []const u8 = &.{}, + proj_bytes: []const u8 = &.{}, + /// Documents still owed across all remaining batches; null is unbounded. + /// Reaching 0 is an EOF *source*, which is what closes the cursor in the + /// very batch that exhausts the limit rather than one round trip later. + /// Optional rather than "0 means unbounded" precisely because 0 has to keep + /// its literal meaning here. + remaining_limit: ?u64 = null, + /// The client's `batchSize`, reused when a `getMore` names none. + batch_size: ?u32 = null, + /// Exempt from the idle sweep. Still killable by `killCursors` and by + /// eviction -- a fixed-capacity registry cannot promise "never expires". + no_timeout: bool = false, + /// A request is using this cursor right now. Concurrent use is rejected + /// rather than queued: queueing lets one client turn a single cursor into a + /// connection-count denial of service. + pinned: bool = false, + /// `killCursors` arrived while pinned; the in-flight request frees it. + kill_requested: bool = false, + last_use_ms: i64 = 0, + arena: std.heap.ArenaAllocator, + source: Source, + + pub fn ns(self: *const Cursor) Ns { + return .{ + .db = self.ns_db_buf[0..self.ns_db_len], + .coll = self.ns_coll_buf[0..self.ns_coll_len], + }; + } + + pub fn ns_matches(self: *const Cursor, other: Ns) bool { + const own = self.ns(); + return std.mem.eql(u8, own.db, other.db) and std.mem.eql(u8, own.coll, other.coll); + } +}; + +/// Everything a caller must decide before a cursor can exist. Grouped so +/// `open` cannot be called with an argument silently in the wrong position. +pub const OpenSpec = struct { + ns: Ns, + layout_epoch: u64, + filter_bytes: []const u8 = &.{}, + proj_bytes: []const u8 = &.{}, + remaining_limit: ?u64 = null, + batch_size: ?u32 = null, + no_timeout: bool = false, + source: Source, +}; + +pub const OpenError = error{ + /// The namespace does not fit the fixed buffers. Callers degrade to a + /// single batch rather than failing the query. + NameTooLong, + /// Every slot is pinned by an in-flight request. + TooManyCursors, + OutOfMemory, + /// Taking the store mutex was cancelled (shutdown). + Canceled, +}; + +pub const PinError = error{ + CursorNotFound, + /// The id exists but belongs to another namespace. Distinct from + /// `CursorNotFound` because mongod answers this with `Unauthorized` (13), + /// not 43, and leaves the cursor alive -- the request is wrong, not the + /// cursor. + CursorNamespaceMismatch, + CursorInUse, + Canceled, +}; + +/// Owned storage for a namespace copied out of the store, so an error message +/// can name a cursor's namespace without holding the store's mutex or a pointer +/// into its slots. +pub const NsBuf = struct { + db_buf: [ns_db_max]u8 = undefined, + db_len: u8 = 0, + coll_buf: [ns_coll_max]u8 = undefined, + coll_len: u8 = 0, + + pub fn ns(self: *const NsBuf) Ns { + return .{ .db = self.db_buf[0..self.db_len], .coll = self.coll_buf[0..self.coll_len] }; + } + + fn set(self: *NsBuf, from: Ns) void { + @memcpy(self.db_buf[0..from.db.len], from.db); + self.db_len = @intCast(from.db.len); + @memcpy(self.coll_buf[0..from.coll.len], from.coll); + self.coll_len = @intCast(from.coll.len); + } +}; + +pub const KillOutcome = enum { killed, not_found }; + +// --------------------------------------------------------------------------- +// The registry +// --------------------------------------------------------------------------- + +pub const Store = struct { + /// Guards every field below. A **leaf** lock: no other lock -- catalog, + /// collection, log -- is ever acquired while it is held, so it cannot + /// participate in a cycle. In particular a `getMore` copies what it needs + /// out, releases this, and only then iterates under the collection lock; + /// otherwise the reaper would block behind a full scan. + mutex: std.Io.Mutex = .init, + /// Boxed, not inline. A `Cursor` inlines its anchor and namespace buffers and + /// so is ~1.5 KiB; at the default capacity an inline table would be 6.2 MiB + /// allocated and zeroed at *every* `Engine.open` -- paid by every embedded + /// user and by all ~47 engine opens in the unit suite, to hold zero cursors. + /// A pointer table is 32 KiB and the cursor itself is allocated when one + /// actually exists, which is also when its arena is created anyway. + slots: []?*Cursor, + /// Mixed into every id so ids are not guessable across processes, and so a + /// reused slot rejects the previous id exactly. Without this a stale + /// `getMore` can address a recycled slot and read another client's cursor. + nonce: u64, + live: u32 = 0, + idle_timeout_ms: i64 = default_idle_timeout_ms, + gpa: std.mem.Allocator, + + pub fn init( + gpa: std.mem.Allocator, + io: std.Io, + capacity: u32, + idle_timeout_ms: i64, + ) !Store { + assert(capacity > 0 and capacity <= max_capacity); + var seed: [8]u8 = undefined; + io.random(&seed); + const slots = try gpa.alloc(?*Cursor, capacity); + @memset(slots, null); + return .{ + .slots = slots, + // A zero nonce would make the first slot's id equal to its index, + // and slot 0's id would be 0 -- which means "no cursor". + .nonce = std.mem.readInt(u64, &seed, .little) | 1, + .idle_timeout_ms = idle_timeout_ms, + .gpa = gpa, + }; + } + + pub fn deinit(self: *Store) void { + for (self.slots) |maybe| { + if (maybe) |c| destroy_cursor(self.gpa, c); + } + self.gpa.free(self.slots); + self.slots = &.{}; + } + + /// Free a cursor: its arena first, then the box the slot pointed at. + fn destroy_cursor(gpa: std.mem.Allocator, c: *Cursor) void { + c.arena.deinit(); + gpa.destroy(c); + } + + fn slot_of(id: i64) usize { + return @intCast(@as(u64, @bitCast(id)) & slot_mask); + } + + /// Build the id for `slot` at the store's current nonce, then advance the + /// nonce so the next cursor in this slot gets a different id. + fn mint(self: *Store, slot: usize) i64 { + const n = self.nonce & nonce_mask; + self.nonce +%= 1; + const raw = (n << slot_bits) | @as(u64, @intCast(slot)); + const id: i64 = @intCast(raw & ~(@as(u64, 1) << 63)); + // Both properties are load-bearing on the wire and in lookup. + assert_msg(id > 0, "a cursor id must be a positive int64"); + assert_msg(slot_of(id) == slot, "a cursor id must address its own slot"); + return id; + } + + /// Register a cursor and return its id, or null when the caller should + /// answer in a single batch instead (`NameTooLong` is not worth failing a + /// query over -- the degradation is exactly today's behaviour). + /// + /// Takes ownership of `spec.source` and of the arena backing it. + pub fn open( + self: *Store, + io: std.Io, + now_ms: i64, + arena: std.heap.ArenaAllocator, + spec: OpenSpec, + ) OpenError!i64 { + if (spec.ns.db.len > ns_db_max or spec.ns.coll.len > ns_coll_max) { + return error.NameTooLong; + } + try self.mutex.lock(io); + defer self.mutex.unlock(io); + + const slot = self.free_slot(now_ms) orelse return error.TooManyCursors; + assert(self.slots[slot] == null); + + const c = try self.gpa.create(Cursor); + errdefer self.gpa.destroy(c); + c.* = .{ + .id = self.mint(slot), + .layout_epoch = spec.layout_epoch, + .filter_bytes = spec.filter_bytes, + .proj_bytes = spec.proj_bytes, + .remaining_limit = spec.remaining_limit, + .batch_size = spec.batch_size, + .no_timeout = spec.no_timeout, + .last_use_ms = now_ms, + .arena = arena, + .source = spec.source, + }; + @memcpy(c.ns_db_buf[0..spec.ns.db.len], spec.ns.db); + c.ns_db_len = @intCast(spec.ns.db.len); + @memcpy(c.ns_coll_buf[0..spec.ns.coll.len], spec.ns.coll); + c.ns_coll_len = @intCast(spec.ns.coll.len); + + self.slots[slot] = c; + self.live += 1; + return c.id; + } + + /// An empty slot: a genuinely free one, else the least-recently-used + /// unpinned cursor. Evicting is legal and cheap to reason about, because + /// the victim's client gets `CursorNotFound` on its next `getMore` -- the + /// same answer an idle timeout gives, which every driver already handles. + /// Caller holds the mutex. + fn free_slot(self: *Store, now_ms: i64) ?usize { + var lru: ?usize = null; + var lru_ms: i64 = std.math.maxInt(i64); + for (self.slots, 0..) |maybe, i| { + const c = maybe orelse return i; + // Reap on the way past, so a store that has gone quiet does not + // wait for the sweep tick to reclaim what already expired. + if (self.expired(c, now_ms)) { + self.destroy(i); + return i; + } + if (c.pinned) continue; + if (c.last_use_ms < lru_ms) { + lru_ms = c.last_use_ms; + lru = i; + } + } + if (lru) |i| { + self.destroy(i); + return i; + } + return null; + } + + /// Caller holds the mutex. + fn expired(self: *const Store, c: *const Cursor, now_ms: i64) bool { + if (c.pinned or c.no_timeout or self.idle_timeout_ms <= 0) return false; + return now_ms -| c.last_use_ms >= self.idle_timeout_ms; + } + + /// Caller holds the mutex. + fn destroy(self: *Store, slot: usize) void { + const c = self.slots[slot] orelse return; + assert_msg(!c.pinned, "a pinned cursor must not be destroyed under its user"); + destroy_cursor(self.gpa, c); + self.slots[slot] = null; + self.live -= 1; + } + + /// Claim a cursor for one request. The returned pointer is stable only + /// until `release`, and only the pinning request may touch it. + /// + /// The namespace check is not cosmetic: `dispatch` locks the collection + /// named in the *message*, so a `getMore` quoting one cursor's id and + /// another collection's name would otherwise iterate the first collection's + /// index while holding the second collection's lock. A mismatch leaves the + /// cursor alive -- it is the request that is wrong, not the cursor. + pub fn pin(self: *Store, io: std.Io, id: i64, ns: Ns, now_ms: i64) PinError!*Cursor { + try self.mutex.lock(io); + defer self.mutex.unlock(io); + if (id <= 0) return error.CursorNotFound; + const slot = slot_of(id); + if (slot >= self.slots.len) return error.CursorNotFound; + const c = self.slots[slot] orelse return error.CursorNotFound; + // Compare the whole id, not just the slot: this is what makes a + // recycled slot reject its predecessor's id. + if (c.id != id) return error.CursorNotFound; + if (self.expired(c, now_ms)) { + self.destroy(slot); + return error.CursorNotFound; + } + if (!c.ns_matches(ns)) return error.CursorNamespaceMismatch; + if (c.pinned) return error.CursorInUse; + c.pinned = true; + c.last_use_ms = now_ms; + return c; + } + + /// The namespace a live cursor belongs to, copied out. Only used to build + /// the namespace-mismatch error message, so a second lock acquisition on an + /// error path is the right trade for not threading an out-parameter through + /// the success path. + pub fn ns_of(self: *Store, io: std.Io, id: i64, out: *NsBuf) bool { + self.mutex.lock(io) catch return false; + defer self.mutex.unlock(io); + if (id <= 0) return false; + const slot = slot_of(id); + if (slot >= self.slots.len) return false; + const c = self.slots[slot] orelse return false; + if (c.id != id) return false; + out.set(c.ns()); + return true; + } + + /// Hand a pinned cursor back. `exhausted` destroys it, and so does a + /// `killCursors` that arrived while it was pinned. + pub fn release(self: *Store, io: std.Io, c: *Cursor, now_ms: i64, exhausted: bool) void { + self.mutex.lock(io) catch { + // Cancellation while returning a cursor would otherwise leave it + // pinned forever, unreachable and un-reapable. Unpinning without + // the lock is the lesser evil: the field is only ever written by + // the one request that owns the pin. + c.pinned = false; + return; + }; + defer self.mutex.unlock(io); + assert_msg(c.pinned, "released a cursor that was not pinned"); + c.pinned = false; + c.last_use_ms = now_ms; + if (exhausted or c.kill_requested) { + const slot = slot_of(c.id); + assert(self.slots[slot].? == c); + self.destroy(slot); + } + } + + /// `killCursors` for one id. A pinned cursor is marked and reported killed: + /// the client's intent is satisfied, and the in-flight request frees it on + /// release. Storage is never freed under a running request. + pub fn kill(self: *Store, io: std.Io, id: i64, ns: Ns) KillOutcome { + self.mutex.lock(io) catch return .not_found; + defer self.mutex.unlock(io); + if (id <= 0) return .not_found; + const slot = slot_of(id); + if (slot >= self.slots.len) return .not_found; + const c = self.slots[slot] orelse return .not_found; + if (c.id != id or !c.ns_matches(ns)) return .not_found; + if (c.pinned) { + c.kill_requested = true; + return .killed; + } + self.destroy(slot); + return .killed; + } + + /// Kill every cursor on a namespace. Called when the collection or its + /// database is dropped: a later `getMore` would fail anyway, since the + /// cursor holds names rather than a pointer, but reaping here frees the + /// slots at once and keeps the open-cursor metric honest. + pub fn kill_namespace( + self: *Store, + io: std.Io, + db_name: []const u8, + coll_name: ?[]const u8, + ) u32 { + self.mutex.lock(io) catch return 0; + defer self.mutex.unlock(io); + var n: u32 = 0; + for (self.slots, 0..) |maybe, i| { + const c = maybe orelse continue; + const own = c.ns(); + if (!std.mem.eql(u8, own.db, db_name)) continue; + if (coll_name) |name| { + if (!std.mem.eql(u8, own.coll, name)) continue; + } + if (c.pinned) { + c.kill_requested = true; + } else { + self.destroy(i); + } + n += 1; + } + return n; + } + + /// Reap idle cursors. Returns how many went. + pub fn sweep(self: *Store, io: std.Io, now_ms: i64) u32 { + self.mutex.lock(io) catch return 0; + defer self.mutex.unlock(io); + if (self.live == 0) return 0; + var n: u32 = 0; + for (self.slots, 0..) |maybe, i| { + const c = maybe orelse continue; + if (!self.expired(c, now_ms)) continue; + self.destroy(i); + n += 1; + } + return n; + } +}; + +// --------------------------------------------------------------------------- +// Batch policy +// --------------------------------------------------------------------------- + +/// What `offer` decided about one document. +pub const Offered = enum { + appended, + /// The batch is full. **Not** EOF: the cursor stays open, and this document + /// has not been consumed -- the caller must hand it to the next batch. + batch_full, +}; + +/// Accumulates one batch and owns the two limits that end it. +/// +/// Split out from the emit path so the whole policy is a pure function of +/// `(target, emitted, bytes, size)` and can be unit-tested without an engine, +/// a socket or an allocator. The subtle parts are all here: a target of 0 is +/// unbounded (a `getMore` naming no `batchSize`), a `batchSize: 0` first batch +/// is a target that is *reached immediately*, and the byte cap must still let +/// the first document through or an oversized document would wedge the cursor +/// forever, returning empty batches with no progress. +pub const BatchBuilder = struct { + /// Documents wanted; null means no document target, fill to the byte cap. + /// Optional rather than "0 means unbounded" because `batchSize: 0` is a real + /// request for an empty batch, and conflating the two returned the whole + /// collection where mongod returns nothing. + target: ?u32, + bytes_max: u64 = batch_bytes_max, + emitted: u32 = 0, + bytes: u64 = 0, + + pub fn init(target: ?u32) BatchBuilder { + return .{ .target = target }; + } + + /// Whether the batch has already met its document target, checked before + /// pulling from the source so a full batch never consumes a document it + /// cannot carry. + pub fn full(self: *const BatchBuilder) bool { + const t = self.target orelse return false; + return self.emitted >= t; + } + + /// Account for a document of `size` serialized bytes. + pub fn offer(self: *BatchBuilder, size: u64) Offered { + assert_msg(!self.full(), "offered a document to a batch that was already full"); + // The at-least-one rule: an empty batch takes the document whatever it + // measures. Stored documents cannot exceed the cap (inserts enforce + // 16 MiB), so this only arises for a generated one. + if (self.emitted > 0 and self.bytes + size > self.bytes_max) return .batch_full; + self.emitted += 1; + self.bytes += size; + return .appended; + } +}; + +/// The document target for a batch: the client's `batchSize` if it named one, +/// otherwise 101 for a first batch and *no* document target for a `getMore`. +/// +/// Both defaults are measured against mongod 8.3.7 rather than assumed. +/// `internalQueryFindCommandBatchSize` reports 101, and a `getMore` carrying no +/// `batchSize` after a `find` with `batchSize: 2` returns 4998 of 5000 +/// documents -- so a bare `getMore` is bounded by bytes alone and does *not* +/// inherit the `batchSize` the cursor was created with. +pub fn batch_target(batch_size: ?u32, first: bool) ?u32 { + if (batch_size) |n| return n; + return if (first) default_first_batch else null; +} + +// --------------------------------------------------------------------------- +// Tests +// --------------------------------------------------------------------------- + +const testing = std.testing; + +/// A Store for tests, with a threaded Io so the mutex is real. +const TestStore = struct { + threaded: std.Io.Threaded, + store: Store, + + fn init(capacity: u32, idle_timeout_ms: i64) !TestStore { + var self: TestStore = undefined; + self.threaded = std.Io.Threaded.init(testing.allocator, .{}); + self.store = try Store.init( + testing.allocator, + self.threaded.io(), + capacity, + idle_timeout_ms, + ); + return self; + } + + fn io(self: *TestStore) std.Io { + return self.threaded.io(); + } + + fn deinit(self: *TestStore) void { + self.store.deinit(); + self.threaded.deinit(); + } + + fn open_one(self: *TestStore, coll: []const u8, now_ms: i64) !i64 { + const arena = std.heap.ArenaAllocator.init(testing.allocator); + return self.store.open(self.io(), now_ms, arena, .{ + .ns = .{ .db = "t", .coll = coll }, + .layout_epoch = 0, + .source = .{ .buffered = .{ .docs = &.{} } }, + }); + } +}; + +test "cursor ids are positive, address their slot, and never repeat" { + var ts = try TestStore.init(4, default_idle_timeout_ms); + defer ts.deinit(); + + var seen: [16]i64 = undefined; + for (0..16) |i| { + const id = try ts.open_one("c", 0); + try testing.expect(id > 0); + // Freeing the slot immediately means the next open reuses it, which is + // exactly the case the nonce has to survive. + const killed = ts.store.kill(ts.io(), id, .{ .db = "t", .coll = "c" }); + try testing.expectEqual(KillOutcome.killed, killed); + seen[i] = id; + } + for (seen, 0..) |a, i| { + for (seen[i + 1 ..]) |b| try testing.expect(a != b); + } +} + +test "a recycled slot rejects the id it used to hold" { + // The guard that keeps one client from reading another's cursor. Without + // the nonce in the id, `slot_of(stale) == slot_of(fresh)` and the stale + // getMore would be served the new cursor's documents. + var ts = try TestStore.init(1, default_idle_timeout_ms); + defer ts.deinit(); + const ns = Ns{ .db = "t", .coll = "c" }; + + const stale = try ts.open_one("c", 0); + try testing.expectEqual(KillOutcome.killed, ts.store.kill(ts.io(), stale, ns)); + const fresh = try ts.open_one("c", 0); + + try testing.expectEqual(Store.slot_of(stale), Store.slot_of(fresh)); + try testing.expect(stale != fresh); + try testing.expectError(error.CursorNotFound, ts.store.pin(ts.io(), stale, ns, 0)); + _ = try ts.store.pin(ts.io(), fresh, ns, 0); +} + +test "pin rejects a wrong namespace and a second holder, and leaves the cursor alive" { + var ts = try TestStore.init(4, default_idle_timeout_ms); + defer ts.deinit(); + const ns = Ns{ .db = "t", .coll = "c" }; + const id = try ts.open_one("c", 0); + + // A wrong namespace must not kill the cursor: the request is wrong, not + // the cursor, and the client is allowed to retry correctly. Reported apart + // from CursorNotFound because mongod answers it with Unauthorized (13). + const wrong_coll = Ns{ .db = "t", .coll = "other" }; + const wrong_db = Ns{ .db = "other", .coll = "c" }; + const mismatch = error.CursorNamespaceMismatch; + try testing.expectError(mismatch, ts.store.pin(ts.io(), id, wrong_coll, 0)); + try testing.expectError(mismatch, ts.store.pin(ts.io(), id, wrong_db, 0)); + + var found: NsBuf = .{}; + try testing.expect(ts.store.ns_of(ts.io(), id, &found)); + try testing.expectEqualStrings("t", found.ns().db); + try testing.expectEqualStrings("c", found.ns().coll); + + const c = try ts.store.pin(ts.io(), id, ns, 0); + try testing.expectError(error.CursorInUse, ts.store.pin(ts.io(), id, ns, 0)); + ts.store.release(ts.io(), c, 1, false); + // Released, so it can be pinned again. + const again = try ts.store.pin(ts.io(), id, ns, 2); + ts.store.release(ts.io(), again, 3, true); + try testing.expectError(error.CursorNotFound, ts.store.pin(ts.io(), id, ns, 4)); +} + +test "a full store evicts the least recently used unpinned cursor" { + var ts = try TestStore.init(3, default_idle_timeout_ms); + defer ts.deinit(); + const ns = Ns{ .db = "t", .coll = "c" }; + + const a = try ts.open_one("c", 100); + const b = try ts.open_one("c", 200); + const c = try ts.open_one("c", 300); + // Touch `a` so `b` becomes the least recently used. + const pinned_a = try ts.store.pin(ts.io(), a, ns, 400); + ts.store.release(ts.io(), pinned_a, 400, false); + + const d = try ts.open_one("c", 500); + try testing.expectError(error.CursorNotFound, ts.store.pin(ts.io(), b, ns, 500)); + for ([_]i64{ a, c, d }) |id| { + const live = try ts.store.pin(ts.io(), id, ns, 500); + ts.store.release(ts.io(), live, 500, false); + } +} + +test "a store whose every slot is pinned refuses rather than evicting" { + var ts = try TestStore.init(2, default_idle_timeout_ms); + defer ts.deinit(); + const ns = Ns{ .db = "t", .coll = "c" }; + const a = try ts.open_one("c", 0); + const b = try ts.open_one("c", 0); + _ = try ts.store.pin(ts.io(), a, ns, 0); + _ = try ts.store.pin(ts.io(), b, ns, 0); + try testing.expectError(error.TooManyCursors, ts.open_one("c", 0)); +} + +test "the sweep reaps idle cursors and spares noCursorTimeout" { + var ts = try TestStore.init(4, 1000); + defer ts.deinit(); + const ns = Ns{ .db = "t", .coll = "c" }; + + const perishable = try ts.open_one("c", 0); + const arena = std.heap.ArenaAllocator.init(testing.allocator); + const immortal = try ts.store.open(ts.io(), 0, arena, .{ + .ns = ns, + .layout_epoch = 0, + .no_timeout = true, + .source = .{ .buffered = .{ .docs = &.{} } }, + }); + + // Just short of the timeout: nothing goes. + try testing.expectEqual(@as(u32, 0), ts.store.sweep(ts.io(), 999)); + try testing.expectEqual(@as(u32, 1), ts.store.sweep(ts.io(), 1000)); + try testing.expectError(error.CursorNotFound, ts.store.pin(ts.io(), perishable, ns, 1000)); + + // The exempt one survives an interval it would otherwise have died in... + try testing.expectEqual(@as(u32, 0), ts.store.sweep(ts.io(), 100_000)); + const live = try ts.store.pin(ts.io(), immortal, ns, 100_000); + ts.store.release(ts.io(), live, 100_000, false); + // ...but is still killable explicitly. + try testing.expectEqual(KillOutcome.killed, ts.store.kill(ts.io(), immortal, ns)); +} + +test "killCursors reports a pinned cursor killed and frees it on release" { + var ts = try TestStore.init(4, default_idle_timeout_ms); + defer ts.deinit(); + const ns = Ns{ .db = "t", .coll = "c" }; + const id = try ts.open_one("c", 0); + const c = try ts.store.pin(ts.io(), id, ns, 0); + + try testing.expectEqual(KillOutcome.killed, ts.store.kill(ts.io(), id, ns)); + // Still pinned, so its storage must not have been freed under the request. + try testing.expect(c.kill_requested); + // Not exhausted, but the pending kill wins. + ts.store.release(ts.io(), c, 1, false); + try testing.expectError(error.CursorNotFound, ts.store.pin(ts.io(), id, ns, 2)); + try testing.expectEqual(KillOutcome.not_found, ts.store.kill(ts.io(), id, ns)); +} + +test "kill_namespace reaps a collection's cursors and leaves the rest" { + var ts = try TestStore.init(8, default_idle_timeout_ms); + defer ts.deinit(); + const doomed = try ts.open_one("doomed", 0); + const spared = try ts.open_one("spared", 0); + + try testing.expectEqual(@as(u32, 1), ts.store.kill_namespace(ts.io(), "t", "doomed")); + const doomed_ns = Ns{ .db = "t", .coll = "doomed" }; + try testing.expectError(error.CursorNotFound, ts.store.pin(ts.io(), doomed, doomed_ns, 0)); + const live = try ts.store.pin(ts.io(), spared, .{ .db = "t", .coll = "spared" }, 0); + ts.store.release(ts.io(), live, 0, false); + + // Whole-database form. + try testing.expectEqual(@as(u32, 1), ts.store.kill_namespace(ts.io(), "t", null)); + try testing.expectEqual(@as(u32, 0), ts.store.live); +} + +test "a namespace too long for the fixed buffers declines rather than failing" { + var ts = try TestStore.init(2, default_idle_timeout_ms); + defer ts.deinit(); + const long = "c" ** (ns_coll_max + 1); + try testing.expectError(error.NameTooLong, ts.open_one(long, 0)); +} + +test "batch_target: 101 for a first batch, unbounded for a getMore, honoured when given" { + try testing.expectEqual(@as(?u32, default_first_batch), batch_target(null, true)); + // A bare getMore has no document target at all. Measured against mongod: + // it does not inherit the batchSize the cursor was created with. + try testing.expectEqual(@as(?u32, null), batch_target(null, false)); + try testing.expectEqual(@as(?u32, 7), batch_target(7, true)); + + // batchSize: 0 is a real target of zero -- an empty first batch with a live + // cursor, which drivers use to obtain a cursor cheaply. It must NOT read as + // "unbounded": conflating the two returns the whole collection where mongod + // returns nothing, which is exactly the bug this optional prevents. + try testing.expectEqual(@as(?u32, 0), batch_target(0, true)); + var zero = BatchBuilder.init(batch_target(0, true)); + try testing.expect(zero.full()); + var bare = BatchBuilder.init(batch_target(null, false)); + try testing.expect(!bare.full()); +} + +test "BatchBuilder stops at its document target" { + var b = BatchBuilder.init(2); + try testing.expect(!b.full()); + try testing.expectEqual(Offered.appended, b.offer(10)); + try testing.expect(!b.full()); + try testing.expectEqual(Offered.appended, b.offer(10)); + try testing.expect(b.full()); + try testing.expectEqual(@as(u32, 2), b.emitted); +} + +test "BatchBuilder: a null target is unbounded by documents" { + var b = BatchBuilder.init(null); + for (0..5000) |_| { + try testing.expect(!b.full()); + try testing.expectEqual(Offered.appended, b.offer(1)); + } + try testing.expect(!b.full()); +} + +test "BatchBuilder stops on bytes, but always takes at least one document" { + // Hitting the byte cap must not read as EOF, or the cursor would close and + // silently drop the rest of the result. + var b = BatchBuilder.init(null); + b.bytes_max = 100; + try testing.expectEqual(Offered.appended, b.offer(60)); + try testing.expectEqual(Offered.batch_full, b.offer(60)); + // The refused document was not accounted for, so the caller can hand it to + // the next batch. + try testing.expectEqual(@as(u32, 1), b.emitted); + try testing.expectEqual(@as(u64, 60), b.bytes); + + // An oversized document on an empty batch goes through anyway: refusing it + // would wedge the cursor, returning empty batches and never progressing. + var solo = BatchBuilder.init(null); + solo.bytes_max = 100; + try testing.expectEqual(Offered.appended, solo.offer(1_000_000)); + try testing.expectEqual(@as(u32, 1), solo.emitted); + try testing.expectEqual(Offered.batch_full, solo.offer(1)); +} + +test "Stream.advance records the anchor and counts an equal-key band" { + var s = Stream{}; + try testing.expect(!s.started()); + + s.advance("aaa", 10, 3, 4); + try testing.expect(s.started()); + try testing.expectEqualStrings("aaa", s.anchor_key()); + try testing.expectEqual(@as(u64, 10), s.anchor_off); + try testing.expectEqual(@as(u32, 3), s.hint_leaf); + try testing.expectEqual(@as(u32, 4), s.hint_slot); + try testing.expectEqual(@as(u64, 0), s.band_index); + + // Same key, different document: still inside the band, so the position + // within it has to advance or a resume could not tell them apart. + s.advance("aaa", 11, 3, 5); + try testing.expectEqual(@as(u64, 1), s.band_index); + s.advance("aaa", 12, 3, 6); + try testing.expectEqual(@as(u64, 2), s.band_index); + + // A new key ends the band. + s.advance("bbb", 13, 3, 7); + try testing.expectEqual(@as(u64, 0), s.band_index); + try testing.expectEqualStrings("bbb", s.anchor_key()); +} + +test "Stream.advance accepts a reverse cursor moving down" { + var s = Stream{ .backward = true }; + s.advance("ccc", 1, 1, 5); + s.advance("bbb", 2, 1, 4); + s.advance("aaa", 3, 1, 3); + try testing.expectEqualStrings("aaa", s.anchor_key()); +} diff --git a/src/db.zig b/src/db.zig index 771a0e9..57ba93e 100644 --- a/src/db.zig +++ b/src/db.zig @@ -22,6 +22,7 @@ const bson = @import("bson.zig"); const storage = @import("storage.zig"); const index = @import("index.zig"); const pgr = @import("pager.zig"); +const cursor = @import("cursor.zig"); // Always active, including in the default ReleaseFast build -- see assert.zig // for why std.debug.assert is the wrong tool for these invariants. const assert = @import("assert.zig").assert; @@ -97,8 +98,21 @@ pub const Collection = struct { /// replaces the old serialization-guarded docs-map fast path for /// integer/string/etc. _id lookups. id_index: index.Index, + /// Identity-and-layout token for open cursors. Drawn from + /// `Engine.layout_epoch_seq`, so it is unique across the engine's life and + /// bumped again by every rebuild. + /// + /// It answers two questions a cursor cannot answer any other way. A rebuild + /// moves every document, so a saved slab offset (or a saved index anchor's + /// offset) is stale -- and the keys surviving unchanged makes that *worse*, + /// because a lookup then succeeds and quietly resolves to the wrong bytes. + /// And a cursor holds namespace *strings*, not a `*Collection`, so a + /// drop-and-recreate under the same name would otherwise be invisible to it; + /// drawing from an engine-wide sequence rather than starting each collection + /// at zero is what makes the recreated one compare unequal. + layout_epoch: u64, - fn init(gpa: std.mem.Allocator, pager: *pgr.Pager) !Collection { + fn init(gpa: std.mem.Allocator, pager: *pgr.Pager, layout_epoch: u64) !Collection { var self: Collection = .{ .doc_count = 0, .pager = pager, @@ -110,6 +124,7 @@ pub const Collection = struct { .hold = .{}, .indexes = .empty, .id_index = undefined, + .layout_epoch = layout_epoch, }; const keys = [_]index.IndexKey{.{ .path = "_id", .descending = false }}; // unique: the tree, not the docs map, is what enforces _id uniqueness @@ -290,6 +305,15 @@ pub const Engine = struct { /// rewrite is worth doing — see `note_compact`. live_docs: u64 = 0, dead_docs: u64 = 0, + /// Hands out `Collection.layout_epoch` values. Monotonic and never reset, so + /// no two collection instances -- including a drop followed by a recreate + /// under the same name -- ever share one. + layout_epoch_seq: u64 = 0, + /// Open cursors. Lives on the engine rather than the server because the C + /// API seam (PLAN D1) lists cursor iteration, and because the unit tests + /// build an Engine with no server at all. Its mutex is a leaf: see + /// `cursor.Store`. + cursors: cursor.Store, /// The same question in bytes, about the *data file* rather than the log. /// Once a checkpoint truncates the log, the log no longer holds the garbage /// -- the doc slab does, and only a rebuild reclaims it. These are what @@ -324,6 +348,12 @@ pub const Engine = struct { /// command reads it while still holding the write lock. dup_index: ?[]const u8 = null, + /// The registry an embedded caller gets without configuring anything; the + /// CLI replaces it through `reconfigure_cursors`. + fn default_cursor_store(gpa: std.mem.Allocator, io: std.Io) !cursor.Store { + return cursor.Store.init(gpa, io, cursor.default_capacity, cursor.default_idle_timeout_ms); + } + pub fn open(gpa: std.mem.Allocator, io: std.Io, path: []const u8) !Engine { var log = try storage.Log.open(gpa, io, path); errdefer log.close(); @@ -346,8 +376,10 @@ pub const Engine = struct { .dbs = .empty, .seq = 0, .compact_threshold = 16 * 1024 * 1024, + .cursors = try default_cursor_store(gpa, io), }; errdefer { + engine.cursors.deinit(); engine.pager.deinit(); engine.dbs.deinit(gpa); } @@ -393,6 +425,17 @@ pub const Engine = struct { return engine; } + /// Replace the cursor registry with one of a different shape. Only legal + /// before the server starts accepting connections, because it drops every + /// cursor -- asserted rather than left to the comment, since the method is + /// public and a later caller would otherwise get silent data loss. + pub fn reconfigure_cursors(self: *Engine, capacity: u32, idle_timeout_ms: i64) !void { + assert_msg(self.cursors.live == 0, "reconfigured the cursor registry with cursors open"); + const fresh = try cursor.Store.init(self.gpa, self.io, capacity, idle_timeout_ms); + self.cursors.deinit(); + self.cursors = fresh; + } + pub fn deinit(self: *Engine) void { var db_it = self.dbs.iterator(); while (db_it.next()) |db_entry| { @@ -400,6 +443,11 @@ pub const Engine = struct { self.gpa.free(db_entry.key_ptr.*); } self.dbs.deinit(self.gpa); + // Before the pager: a cursor's arena is its own, but freeing cursors + // first keeps the teardown order the same as the construction order + // reversed, which is the only order that stays obviously correct as + // cursors grow to hold more. + self.cursors.deinit(); self.pager.deinit(); self.gpa.destroy(self.pager); self.log.close(); @@ -936,6 +984,11 @@ pub const Engine = struct { const removed = db.collections.fetchRemove(coll_name) orelse return false; self.free_collection(removed.value); self.gpa.free(removed.key); + // A cursor on this namespace is already safe -- it holds names, so its + // next getMore finds nothing to lock -- but reaping here frees the slots + // now instead of at the idle timeout, and keeps the open-cursor metric + // describing cursors that can still return something. + _ = self.cursors.kill_namespace(self.io, db_name, coll_name); return true; } @@ -943,6 +996,7 @@ pub const Engine = struct { var removed = self.dbs.fetchRemove(db_name) orelse return false; self.free_db(&removed.value); self.gpa.free(removed.key); + _ = self.cursors.kill_namespace(self.io, db_name, null); return true; } @@ -1156,7 +1210,8 @@ pub const Engine = struct { errdefer self.gpa.free(coll_key); const new_coll = try self.gpa.create(Collection); errdefer self.gpa.destroy(new_coll); - new_coll.* = try Collection.init(self.gpa, self.pager); + self.layout_epoch_seq += 1; + new_coll.* = try Collection.init(self.gpa, self.pager, self.layout_epoch_seq); errdefer new_coll.id_index.deinit(self.gpa); try db.collections.put(self.gpa, coll_key, new_coll); return new_coll; @@ -1389,6 +1444,13 @@ pub const Engine = struct { for (coll.indexes.items) |ix| try self.repack_index(coll, ix, moved.items); for (old_extents) |e| try self.pager.free_pages(e.first, e.pages); + + // Every document has moved, so every offset an open cursor is holding + // now names different bytes. Bumped last, after the rebuild can no + // longer fail: a cursor invalidated by a rebuild that then errored out + // would have been invalidated for nothing. + self.layout_epoch_seq += 1; + coll.layout_epoch = self.layout_epoch_seq; } fn repack_index( @@ -3886,3 +3948,64 @@ const Reader = struct { return self.take(n); } }; + +test "the epochs that invalidate a cursor move exactly when they must" { + // Three separate promises, each one load-bearing for an open cursor: + // + // - a rebuild moves every document, so a saved slab offset is stale; + // - a drop-and-recreate under the same name is a different collection, + // which a cursor holding only namespace strings cannot otherwise see; + // - `Index.reset_tree` re-creates node ids 0 and 1 as different nodes, so a + // saved (leaf, slot) position becomes valid-and-wrong rather than absent. + // + // A cursor's whole safety story is these three bumps, so assert them here + // rather than inferring them from cursor behaviour later. + 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(); + + try engine.lock(); + var i: i32 = 0; + while (i < 40) : (i += 1) { + var d = try make_doc(gpa, i, "payload"); + defer d.deinit(); + try engine.insert("app", "c", &d, &env.gen); + } + const before = engine.get_collection("app", "c").?.layout_epoch; + engine.unlock(); + try testing.expect(before != 0); + + // A rebuild moves documents, so the epoch must move with them. + try engine.compact(); + try engine.lock(); + const after_rebuild = engine.get_collection("app", "c").?.layout_epoch; + engine.unlock(); + try testing.expect(after_rebuild != before); + + // A recreated collection must not be mistaken for the one that was + // dropped. Starting each collection's epoch at zero would fail here. + try engine.lock(); + try testing.expect(try engine.drop_collection("app", "c")); + var fresh_doc = try make_doc(gpa, 1, "fresh"); + defer fresh_doc.deinit(); + try engine.insert("app", "c", &fresh_doc, &env.gen); + const after_recreate = engine.get_collection("app", "c").?.layout_epoch; + engine.unlock(); + try testing.expect(after_recreate != after_rebuild); + try testing.expect(after_recreate != before); + + // And the index-level token, which guards the position hint. + try engine.lock(); + const coll = engine.get_collection("app", "c").?; + const index_before = coll.id_index.epoch; + try coll.id_index.reset_tree(gpa); + try testing.expect(coll.id_index.epoch != index_before); + engine.unlock(); +} diff --git a/src/index.zig b/src/index.zig index afda333..a06897e 100644 --- a/src/index.zig +++ b/src/index.zig @@ -106,7 +106,9 @@ const page_size = 4096; /// Bytes of node payload: a 32-byte header plus the slotted region. const page_data = page_size - 32; /// Records longer than a quarter of a node spill to the overflow slab. -const inline_limit = page_size / 4; +/// Public because it is also the bound a resumable cursor anchors within: a key +/// past it has already spilled, so the tree itself treats it as exceptional. +pub const inline_limit = page_size / 4; /// Upper bound on the slots one node can hold, since every slot costs at /// least its own size. Bounds the split scratch. const max_slots = page_data / slot_size; @@ -234,6 +236,19 @@ pub const Index = struct { depth: u32, /// Total entries, maintained incrementally. entry_count: usize, + /// Bumped whenever a node id stops meaning what it meant, which is the one + /// thing that makes a saved `(leaf, slot)` position dangerous rather than + /// merely stale. Node ids are otherwise append-only (`alloc_node`, and + /// `drop_child` abandons a page without recycling its id), and `page()` + /// resolves ids through `node_pages`, so copy-on-write and checkpoints move + /// pages without disturbing ids. Only `reset_tree` and + /// `replace_root_with_leaf` reuse an id for different contents. + /// + /// A resumable cursor keeps a position hint to avoid walking an equal-key + /// band on every `getMore`; it must compare this first. Without it the hint + /// would address a live but unrelated leaf after a compaction and the cursor + /// would iterate a tree that no longer exists. + epoch: u64, /// Repack scratch: any single node's record bytes fit here. scratch: [page_data]u8, /// Promoted-key scratch: inline keys being propagated up a split are @@ -268,6 +283,7 @@ pub const Index = struct { .leaf_count = 0, .depth = 0, .entry_count = 0, + .epoch = 0, .scratch = undefined, .promo = undefined, }; @@ -614,6 +630,10 @@ pub const Index = struct { self.depth = 0; self.entry_count = 0; self.multikey = false; + // Node ids 0 and 1 were just re-created as different nodes, so every + // position anyone saved into the old tree now points somewhere valid + // and wrong. This is the bump that tells them apart. + self.epoch += 1; } /// Remove every entry for `id`, in one pass over the leaves. Infallible. @@ -777,6 +797,14 @@ pub const Index = struct { leaf: u32, slot: u32, + /// `next`, plus the position of the entry it yielded. Exact because + /// `next` leaves `leaf` alone on the call that yields and has already + /// incremented `slot` past the entry. + pub fn positioned(self: *Iter) ?Positioned { + const e = self.next() orelse return null; + return .{ .key = e.key, .off = e.off, .leaf = self.leaf, .slot = self.slot - 1 }; + } + pub fn next(self: *Iter) ?EntryRef { const ix = self.ix; while (self.leaf != 0) { @@ -872,6 +900,13 @@ pub const Index = struct { /// One past the slot to yield next, so 0 means this leaf is done. slot: u32, + /// As `Iter.positioned`, but `RevIter.next` decrements *onto* the entry + /// it yields, so the slot needs no adjustment. + pub fn positioned(self: *RevIter) ?Positioned { + const e = self.next() orelse return null; + return .{ .key = e.key, .off = e.off, .leaf = self.leaf, .slot = self.slot }; + } + pub fn next(self: *RevIter) ?EntryRef { const ix = self.ix; while (self.leaf != 0) { @@ -919,6 +954,174 @@ pub const Index = struct { return .{ .ix = self, .leaf = b.leaf, .slot = b.slot }; } + // -- resuming an interrupted scan --------------------------------------- + + /// Entries a resume will walk past before giving up and reporting `capped`. + /// A bound rather than a hope: `seek` lands at the *start* of an equal-key + /// band, so without one a key with millions of duplicates would make every + /// batch cost O(band) and a full drain quadratic. + pub const resume_walk_max: u32 = 1 << 16; + + /// One entry, with enough of its position to resume after it next time. + pub const Positioned = struct { + key: []const u8, + off: u64, + leaf: u32, + slot: u32, + }; + + /// A resumed forward walk. `capped` means the anchor could not be located + /// within `resume_walk_max` steps, so the position is not trustworthy and + /// the caller must fail rather than return documents from the wrong place. + pub const Resumed = struct { it: Iter, capped: bool = false }; + pub const ResumedRev = struct { it: RevIter, capped: bool = false }; + + /// Does `(leaf, slot)` still hold exactly `(key, off)`? + /// + /// A hint is never believed, only checked, and the checks are ordered so the + /// cheap structural ones run first: `off_of` asserts `is_leaf` with + /// `std.debug.assert`, which in ReleaseFast is a promise to the optimizer + /// rather than a check, so `is_leaf` must be tested for real beforehand. + /// + /// Node ids are append-only, so a stale id is always in bounds; what makes a + /// hint dangerous rather than merely wrong is `reset_tree` re-creating ids 0 + /// and 1 as different nodes, and `Index.epoch` is what the caller compares + /// for that. + fn hint_holds(self: *const Index, leaf: u32, slot: u32, key: []const u8, off: u64) bool { + if (leaf == 0 or leaf >= self.node_pages.items.len) return false; + const node = self.page(leaf); + if (node.is_leaf != 1) return false; + if (slot >= node.count) return false; + if (self.off_of(leaf, slot) != off) return false; + return std.mem.eql(u8, self.key_of(leaf, slot), key); + } + + /// Locate the anchor `(key, off)` by walking its equal-key band. + /// + /// Comparison is `std.mem.order`, not `cmp_prefix`, because the band is + /// defined as the entries whose key is byte-equal to the anchor's and prefix + /// semantics would call `"ab"` and `"abc"` equal. In fairness the two happen + /// to agree on where this function resumes -- the fallback is positional, and + /// the first out-of-band entry is the same entry either way -- so this is a + /// clarity choice, not a bug fix; an attempted mutation to `cmp_prefix` does + /// not change any observable result. What does matter is that `lower_bound` + /// uses prefix semantics and therefore errs *before* the band, never past it, + /// so the walk cannot start beyond the anchor and skip it. + /// + /// When the anchor is gone, "gone" turns out to mean two different things and + /// they want opposite answers: + /// + /// - **Deleted.** A sibling has moved up into the anchor's band position, and + /// that sibling has not been returned yet. Resume *at* band position + /// `band_index`. Resuming after the whole band instead would silently drop + /// every remaining member, which on a three-value index is most of the + /// collection. + /// - **Updated.** The document was rewritten, so its key is unchanged but its + /// offset moved. The entry at the anchor's band position *is* the anchor, + /// already returned. Resume *after* it. + /// + /// The index cannot tell these apart in general -- both look like "same key, + /// different offset". On a **unique** index it can: two entries cannot share a + /// key, so a same-key entry is necessarily the same document, hence the update + /// case, hence resume past the band. That covers `_id_` and so every unsorted + /// scan and every `_id` sort, which is where an update-during-drain otherwise + /// returns a document twice -- observed as duplicate `_id`s draining a + /// collection that was being updated underneath. + /// + /// On a non-unique index the positional fallback stands, so an updated document + /// may come back a second time. That is legal: MongoDB documents that a + /// non-snapshot cursor may return a document more than once if an intervening + /// write moves it. + fn band_resume(self: *const Index, key: []const u8, off: u64, band_index: u64) Resumed { + var it = self.seek(key); + var fallback: ?Iter = null; + var pos: u64 = 0; + var steps: u32 = 0; + while (steps < resume_walk_max) : (steps += 1) { + // The iterator state that would yield the entry we are about to + // look at, i.e. "resume *at* this entry". + const before = it; + const e = it.next() orelse break; + if (std.mem.order(u8, e.key, key) != .eq) { + // Past the band. Prefer the fallback if the band held one. + return .{ .it = fallback orelse before }; + } + if (e.off == off) return .{ .it = it }; // resume just after the anchor + // On a unique index a same-key entry can only be the anchor itself, + // rewritten, so there is no sibling to fall back to. + if (!self.unique and pos == band_index and fallback == null) fallback = before; + pos += 1; + } + if (steps == resume_walk_max) return .{ .it = it, .capped = true }; + return .{ .it = fallback orelse it }; + } + + /// A forward walk positioned just after `(key, off)`. + /// + /// O(1) whenever the hint still holds, which is the case unless something + /// wrote to that exact leaf between batches. The band walk is the fallback, + /// and it is what the walk bound exists to contain. + pub fn resume_forward( + self: *const Index, + key: []const u8, + off: u64, + band_index: u64, + hint_leaf: u32, + hint_slot: u32, + hint_trusted: bool, + ) Resumed { + if (hint_trusted and self.hint_holds(hint_leaf, hint_slot, key, off)) { + return .{ .it = .{ .ix = self, .leaf = hint_leaf, .slot = hint_slot + 1 } }; + } + return self.band_resume(key, off, band_index); + } + + /// A reverse walk positioned just before `(key, off)` in key order, i.e. at + /// the next entry a descending scan owes. + /// + /// `RevIter` decrements before yielding, so slot `s` yields `s - 1` -- the + /// entry immediately below the anchor -- and crosses into `prev` when the + /// anchor sat at slot 0. + /// + /// Known limitation, and it is a deliberate trade. When the anchor is gone + /// *and* it had duplicates, this resumes below the whole band rather than at + /// the anchor's position within it, so the band's remaining members are not + /// returned. Placing a reverse fallback exactly would need the band's length, + /// which is only known after walking it, hence a second walk on a path that + /// requires a descending scan over a duplicate-heavy index whose anchor was + /// deleted mid-cursor. Forward resumes -- every unsorted scan and every + /// ascending sort -- use `band_index` and have no such gap. + pub fn resume_reverse( + self: *const Index, + key: []const u8, + off: u64, + hint_leaf: u32, + hint_slot: u32, + hint_trusted: bool, + ) ResumedRev { + if (hint_trusted and self.hint_holds(hint_leaf, hint_slot, key, off)) { + return .{ .it = .{ .ix = self, .leaf = hint_leaf, .slot = hint_slot } }; + } + // Find the anchor by walking forward, then turn around on it. + var it = self.seek(key); + var steps: u32 = 0; + while (steps < resume_walk_max) : (steps += 1) { + const e = it.next() orelse break; + if (std.mem.order(u8, e.key, key) != .eq) break; // past the band + if (e.off == off) { + // `it` has already stepped past the anchor, so the anchor sat at + // `it.slot - 1` and a RevIter there yields the entry below it. + return .{ .it = .{ .ix = self, .leaf = it.leaf, .slot = it.slot - 1 } }; + } + } + if (steps == resume_walk_max) { + return .{ .it = .{ .ix = self, .leaf = 0, .slot = 0 }, .capped = true }; + } + // Anchor gone: resume below the band. + const b = self.lower_bound(key); + return .{ .it = .{ .ix = self, .leaf = b.leaf, .slot = b.slot } }; + } + // -- serialization ------------------------------------------------------ /// The canonical spec document bytes @@ -1710,6 +1913,9 @@ pub const Index = struct { self.first_leaf = self.root; self.leaf_count = 1; self.depth = 0; + // The root's id is unchanged but it is a leaf now, so a saved position + // that named it as an internal node describes a different tree shape. + self.epoch += 1; } /// Pack the sorted staging array into a fresh tree: leaves filled in @@ -3672,3 +3878,258 @@ test "planner picks eq run, ranges, and bails on sparse null" { try testing.expect((try plan(gpa, null, &.{&ix}, &f, &.{})) == null); } } + +/// Drain `ix` by resuming every `stride` entries, the way a cursor with that +/// batch size would, and return the offsets in the order they came out. +fn drain_resuming( + gpa: std.mem.Allocator, + ix: *const Index, + stride: u32, + backward: bool, + out: *std.ArrayListUnmanaged(u64), +) !void { + var started = false; + var anchor: std.ArrayListUnmanaged(u8) = .empty; + defer anchor.deinit(gpa); + var anchor_off: u64 = 0; + var band_index: u64 = 0; + var hint_leaf: u32 = 0; + var hint_slot: u32 = 0; + + while (true) { + // One "batch": open a walk where the last one stopped. + var fwd: Index.Iter = undefined; + var rev: Index.RevIter = undefined; + if (!started) { + if (backward) rev = ix.iter_reverse() else fwd = ix.iter(); + } else if (backward) { + const r = ix.resume_reverse(anchor.items, anchor_off, hint_leaf, hint_slot, true); + try testing.expect(!r.capped); + rev = r.it; + } else { + const r = ix.resume_forward( + anchor.items, + anchor_off, + band_index, + hint_leaf, + hint_slot, + true, + ); + try testing.expect(!r.capped); + fwd = r.it; + } + + var n: u32 = 0; + while (n < stride) : (n += 1) { + const e = if (backward) + rev.positioned() + else + fwd.positioned(); + const got = e orelse return; + try out.append(gpa, got.off); + if (started and std.mem.eql(u8, anchor.items, got.key)) { + band_index += 1; + } else { + band_index = 0; + } + anchor.clearRetainingCapacity(); + try anchor.appendSlice(gpa, got.key); + anchor_off = got.off; + hint_leaf = got.leaf; + hint_slot = got.slot; + started = true; + } + } +} + +test "a resumed walk yields exactly what an uninterrupted one does" { + // The property the whole streaming cursor rests on: stopping and restarting + // a scan changes nothing about what it returns, in either direction, at any + // batch size, including one entry at a time. + // + // Mutation-checked: changing `resume_forward`'s `hint_slot + 1` to + // `hint_slot` makes every batch boundary repeat an entry, and this test goes + // red. (A third mutation was tried and rejected as meaningless: swapping + // `std.mem.order` for `cmp_prefix` inside `band_resume` changes nothing + // observable, so no test can catch it -- see the note there.) + // + // Note this test always resumes from a *valid* hint, since nothing mutates + // the tree between its batches. The band walk is covered by the two tests + // below, which invalidate the hint on purpose. + const gpa = testing.allocator; + + // Three corpora, each hard for a different reason: distinct keys spanning + // several leaves and an interior level; a low-cardinality index whose bands + // span leaves; and *variable-length string keys in prefix relationships* + // ("a" < "ab" < "abc"), which is the only shape that can tell `std.mem.order` + // apart from `cmp_prefix` -- fixed-width integer keys never differ, so an + // integer-only corpus cannot catch that mistake at all. + const Shape = enum { distinct, duplicates, prefixes }; + for ([_]Shape{ .distinct, .duplicates, .prefixes }) |shape| { + var ix = try simple_index(gpa, test_pager(), &.{"a"}, false, false); + defer ix.deinit(gpa); + + const n = 400; + var key_buf: [40]u8 = undefined; + for (0..n) |i| { + const value: bson.Value = switch (shape) { + .distinct => .{ .int32 = @intCast(i + 1) }, + .duplicates => .{ .int32 = @intCast(i % 3) }, + // Every key is a prefix of the next in its group of eight, so + // each band start is also a proper prefix of later keys. + .prefixes => blk: { + const written = try std.fmt.bufPrint(&key_buf, "k{d}", .{i / 8}); + const depth = (i % 8) + 1; + @memset(key_buf[written.len .. written.len + depth], 'x'); + break :blk .{ .string = key_buf[0 .. written.len + depth] }; + }, + }; + const d = try bytes_of(gpa, &.{ + .{ .key = "_id", .value = .{ .int32 = @intCast(i + 1) } }, + .{ .key = "a", .value = value }, + }); + defer gpa.free(d); + _ = try ix.add_doc(gpa, d, @intCast(i + 1), false); + } + try testing.expect(ix.depth >= 1); + + for ([_]bool{ false, true }) |backward| { + var whole: std.ArrayListUnmanaged(u64) = .empty; + defer whole.deinit(gpa); + if (backward) { + var it = ix.iter_reverse(); + while (it.next()) |e| try whole.append(gpa, e.off); + } else { + var it = ix.iter(); + while (it.next()) |e| try whole.append(gpa, e.off); + } + try testing.expectEqual(@as(usize, n), whole.items.len); + + for ([_]u32{ 1, 2, 7, 101, 399, 400, 1000 }) |stride| { + var resumed: std.ArrayListUnmanaged(u64) = .empty; + defer resumed.deinit(gpa); + try drain_resuming(gpa, &ix, stride, backward, &resumed); + try testing.expectEqualSlices(u64, whole.items, resumed.items); + } + } + } +} + +test "a resume survives a split between batches" { + // A cursor holds no lock, so the tree it comes back to is not the tree it + // left. Inserting mid-drain moves entries between leaves and invalidates the + // position hint, which is exactly what the anchor is for. + const gpa = testing.allocator; + var ix = try simple_index(gpa, test_pager(), &.{"a"}, false, false); + defer ix.deinit(gpa); + + const n = 200; + for (0..n) |i| { + // Even keys only, so the inserts below land between existing entries. + const d = try bytes_of(gpa, &.{ + .{ .key = "_id", .value = .{ .int32 = @intCast(i + 1) } }, + .{ .key = "a", .value = .{ .int32 = @intCast((i + 1) * 2) } }, + }); + defer gpa.free(d); + _ = try ix.add_doc(gpa, d, @intCast(i + 1), false); + } + + var seen: std.ArrayListUnmanaged(u64) = .empty; + defer seen.deinit(gpa); + var anchor: std.ArrayListUnmanaged(u8) = .empty; + defer anchor.deinit(gpa); + var anchor_off: u64 = 0; + var hint_leaf: u32 = 0; + var hint_slot: u32 = 0; + var started = false; + var next_id: i32 = 10_000; + + while (true) { + var it = if (!started) ix.iter() else blk: { + const r = ix.resume_forward(anchor.items, anchor_off, 0, hint_leaf, hint_slot, true); + try testing.expect(!r.capped); + break :blk r.it; + }; + var n_in_batch: u32 = 0; + while (n_in_batch < 5) : (n_in_batch += 1) { + const got = it.positioned() orelse break; + try seen.append(gpa, got.off); + anchor.clearRetainingCapacity(); + try anchor.appendSlice(gpa, got.key); + anchor_off = got.off; + hint_leaf = got.leaf; + hint_slot = got.slot; + started = true; + } + if (n_in_batch < 5) break; + + // Between batches, insert odd keys across the whole range: guaranteed to + // split leaves and to appear both before and after the anchor. + for (0..20) |k| { + next_id += 1; + const d = try bytes_of(gpa, &.{ + .{ .key = "_id", .value = .{ .int32 = next_id } }, + .{ .key = "a", .value = .{ .int32 = @intCast(k * 19 + 1) } }, + }); + defer gpa.free(d); + _ = try ix.add_doc(gpa, d, @intCast(next_id), false); + } + } + + // The 200 originals must each appear exactly once. Entries inserted behind + // the cursor may or may not show up -- that is ordinary non-snapshot cursor + // behaviour -- but nothing may be duplicated or lost. + var originals: u32 = 0; + var counts = std.AutoHashMap(u64, u32).init(gpa); + defer counts.deinit(); + for (seen.items) |off| { + const e = try counts.getOrPutValue(off, 0); + e.value_ptr.* += 1; + try testing.expectEqual(@as(u32, 1), e.value_ptr.*); // no duplicates + if (off <= n) originals += 1; + } + try testing.expectEqual(@as(u32, n), originals); +} + +test "a resume whose anchor was deleted keeps the rest of its band" { + // The failure this guards against is silent and large: with the anchor gone, + // resuming after the whole equal-key band drops every remaining member, and + // on a low-cardinality index that is most of the collection. + // + // Mutation-checked: `pos == band_index + 1` in `band_resume` shifts the + // resume by one entry and this test goes red. + const gpa = testing.allocator; + var ix = try simple_index(gpa, test_pager(), &.{"a"}, false, false); + defer ix.deinit(gpa); + + // One key, 50 documents: a single band. + for (0..50) |i| { + const d = try bytes_of(gpa, &.{ + .{ .key = "_id", .value = .{ .int32 = @intCast(i + 1) } }, + .{ .key = "a", .value = .{ .int32 = 7 } }, + }); + defer gpa.free(d); + _ = try ix.add_doc(gpa, d, @intCast(i + 1), false); + } + + // Yield three, then delete the third -- the anchor itself. + var it = ix.iter(); + var third: Index.Positioned = undefined; + var band_index: u64 = 0; + for (0..3) |i| { + third = it.positioned().?; + if (i > 0) band_index += 1; + } + const anchor_key = try gpa.dupe(u8, third.key); + defer gpa.free(anchor_key); + ix.remove_off(third.off); + + const r = ix.resume_forward(anchor_key, third.off, band_index, third.leaf, third.slot, true); + try testing.expect(!r.capped); + var rest: u32 = 0; + var walk = r.it; + while (walk.next()) |_| rest += 1; + + // 50 inserted, 1 deleted, 2 already returned before the anchor: 47 left. + try testing.expectEqual(@as(u32, 47), rest); +} diff --git a/src/lib.zig b/src/lib.zig index 406c10e..ff69beb 100644 --- a/src/lib.zig +++ b/src/lib.zig @@ -10,6 +10,7 @@ pub const db = @import("db.zig"); pub const query = @import("query.zig"); pub const update = @import("update.zig"); pub const index = @import("index.zig"); +pub const cursor = @import("cursor.zig"); pub const pager = @import("pager.zig"); test { @@ -23,5 +24,6 @@ test { _ = @import("query.zig"); _ = @import("update.zig"); _ = @import("index.zig"); + _ = @import("cursor.zig"); _ = @import("pager.zig"); } diff --git a/src/main.zig b/src/main.zig index fe57d94..e47c049 100644 --- a/src/main.zig +++ b/src/main.zig @@ -10,6 +10,18 @@ const usage = \\ --db database file (default multiforadb.log) \\ --ttl-sweep-secs \\ seconds between TTL index sweeps (default 60, 0 disables) + \\ --cursor-timeout-ms + \\ idle milliseconds before an open cursor is reaped + \\ (default 600000, matching MongoDB's cursorTimeoutMillis; + \\ 0 disables expiry) + \\ --cursor-sweep-secs + \\ seconds between idle-cursor sweeps (default 4, matching + \\ MongoDB's clientCursorMonitorFrequencySecs; 0 disables) + \\ --max-open-cursors + \\ cursor registry capacity (default 4096). At capacity the + \\ least recently used cursor is evicted, and its client + \\ sees CursorNotFound -- the same answer an idle timeout + \\ gives, which every driver already handles. \\ --compact-threshold \\ minimum log bytes between compactions; suffixes k/m/g \\ (default 16m). The actual trigger also scales with the @@ -44,34 +56,77 @@ fn parse_size_suffix(v: []const u8) ?u64 { return n * mult; } -pub fn main(init: std.process.Init) !void { - var port: u16 = 27017; - var bind_ip: []const u8 = "127.0.0.1"; - var db_path: []const u8 = "multiforadb.log"; - var ttl_sweep_secs: i64 = 60; - var compact_threshold: u64 = 16 * 1024 * 1024; +/// Everything the CLI can set. Parsed apart from `main` so the option table has +/// room to grow without main outgrowing the 70-line limit. +const Options = struct { + port: u16 = 27017, + bind_ip: []const u8 = "127.0.0.1", + db_path: []const u8 = "multiforadb.log", + ttl_sweep_secs: i64 = 60, + compact_threshold: u64 = 16 * 1024 * 1024, + cursor_timeout_ms: i64 = mongo.cursor.default_idle_timeout_ms, + cursor_sweep_secs: i64 = 4, + max_open_cursors: u32 = mongo.cursor.default_capacity, + /// Set when --help was given: print usage and exit without opening anything. + help: bool = false, +}; +pub fn main(init: std.process.Init) !void { + const opts = try parse_args(init) orelse { + try std.Io.File.writeStreamingAll(.stdout(), init.io, usage); + return; + }; + + const oid_gen = mongo.bson.ObjectIdGen.init(init.io); + var engine = try mongo.db.Engine.open(init.gpa, init.io, opts.db_path); + defer engine.deinit(); + engine.compact_threshold = opts.compact_threshold; + // The engine builds its registry with the defaults so an embedded caller + // needs no configuration; the CLI replaces it when asked for something else. + try engine.reconfigure_cursors(opts.max_open_cursors, opts.cursor_timeout_ms); + std.debug.print("multiforadb: opened database '{s}' (compact threshold {d})\n", .{ + opts.db_path, + opts.compact_threshold, + }); + + var server = mongo.server.Server{ + .gpa = init.gpa, + .port = opts.port, + .bind_ip = opts.bind_ip, + .oid_gen = oid_gen, + .connection_counter = .init(1), + .engine = &engine, + .start_time = std.Io.Timestamp.now(init.io, .real), + .ttl_sweep_secs = opts.ttl_sweep_secs, + .cursor_sweep_secs = opts.cursor_sweep_secs, + }; + try server.run(); +} + +/// Null means --help: the caller prints usage and exits. +fn parse_args(init: std.process.Init) !?Options { + var o = Options{}; var it = std.process.Args.Iterator.init(init.minimal.args); defer it.deinit(); _ = it.next(); // program name while (it.next()) |arg| { if (std.mem.eql(u8, arg, "--port")) { const v = it.next() orelse return error.MissingValue; - port = std.fmt.parseInt(u16, v, 10) catch { + o.port = std.fmt.parseInt(u16, v, 10) catch { std.debug.print("multiforadb: invalid port '{s}'\n", .{v}); return error.InvalidPort; }; } else if (std.mem.eql(u8, arg, "--bind")) { - bind_ip = it.next() orelse return error.MissingValue; + o.bind_ip = it.next() orelse return error.MissingValue; } else if (std.mem.eql(u8, arg, "--db")) { - db_path = it.next() orelse return error.MissingValue; + o.db_path = it.next() orelse return error.MissingValue; } else if (std.mem.eql(u8, arg, "--ttl-sweep-secs")) { const v = it.next() orelse return error.MissingValue; - // i64 is the width std.Io.Duration.fromSeconds takes, so the - // value reaches the sweeper without a cast; negatives are the - // only thing parseInt would otherwise let through. - ttl_sweep_secs = std.fmt.parseInt(i64, v, 10) catch -1; - if (ttl_sweep_secs < 0) { + // i64 is the width std.Io.Duration.fromSeconds takes, so the value + // reaches the sweeper without a cast; negatives are the only thing + // parseInt would otherwise let through. + o.ttl_sweep_secs = std.fmt.parseInt(i64, v, 10) catch -1; + if (o.ttl_sweep_secs < 0) { std.debug.print("multiforadb: invalid ttl sweep interval '{s}'\n", .{v}); return error.InvalidTtlSweepSecs; } @@ -85,34 +140,45 @@ pub fn main(init: std.process.Init) !void { std.debug.print("multiforadb: compact threshold must be at least 1m\n", .{}); return error.InvalidCompactThreshold; } - compact_threshold = parsed; + o.compact_threshold = parsed; } else if (std.mem.eql(u8, arg, "--help") or std.mem.eql(u8, arg, "-h")) { - try std.Io.File.writeStreamingAll(.stdout(), init.io, usage); - return; + return null; + } else if (try parse_cursor_flag(arg, &it, &o)) { + // Handled: one of the cursor-registry flags. } else { std.debug.print("multiforadb: unknown option '{s}'\n{s}", .{ arg, usage }); return error.UnknownOption; } } - - const oid_gen = mongo.bson.ObjectIdGen.init(init.io); - var engine = try mongo.db.Engine.open(init.gpa, init.io, db_path); - defer engine.deinit(); - engine.compact_threshold = compact_threshold; - std.debug.print("multiforadb: opened database '{s}' (compact threshold {d})\n", .{ - db_path, - compact_threshold, - }); - - var server = mongo.server.Server{ - .gpa = init.gpa, - .port = port, - .bind_ip = bind_ip, - .oid_gen = oid_gen, - .connection_counter = .init(1), - .engine = &engine, - .start_time = std.Io.Timestamp.now(init.io, .real), - .ttl_sweep_secs = ttl_sweep_secs, - }; - try server.run(); + return o; +} + +/// The cursor-registry flags, grouped so `parse_args` stays one flat table of +/// options. Returns whether `arg` was one of them; consumes its value if so. +fn parse_cursor_flag(arg: []const u8, it: *std.process.Args.Iterator, o: *Options) !bool { + if (std.mem.eql(u8, arg, "--cursor-timeout-ms")) { + const v = it.next() orelse return error.MissingValue; + o.cursor_timeout_ms = std.fmt.parseInt(i64, v, 10) catch -1; + if (o.cursor_timeout_ms < 0) { + std.debug.print("multiforadb: invalid cursor timeout '{s}'\n", .{v}); + return error.InvalidCursorTimeout; + } + } else if (std.mem.eql(u8, arg, "--cursor-sweep-secs")) { + const v = it.next() orelse return error.MissingValue; + o.cursor_sweep_secs = std.fmt.parseInt(i64, v, 10) catch -1; + if (o.cursor_sweep_secs < 0) { + std.debug.print("multiforadb: invalid cursor sweep interval '{s}'\n", .{v}); + return error.InvalidCursorSweepSecs; + } + } else if (std.mem.eql(u8, arg, "--max-open-cursors")) { + const v = it.next() orelse return error.MissingValue; + o.max_open_cursors = std.fmt.parseInt(u32, v, 10) catch 0; + if (o.max_open_cursors == 0 or o.max_open_cursors > mongo.cursor.max_capacity) { + std.debug.print("multiforadb: invalid max open cursors '{s}'\n", .{v}); + return error.InvalidMaxOpenCursors; + } + } else { + return false; + } + return true; } diff --git a/src/server.zig b/src/server.zig index c5b8065..cf21d9e 100644 --- a/src/server.zig +++ b/src/server.zig @@ -20,6 +20,9 @@ pub const Server = struct { /// because that is what std.Io.Duration.fromSeconds takes — the CLI /// rejects negatives. ttl_sweep_secs: i64, + /// Seconds between idle-cursor sweeps; 0 leaves that monitor unspawned. + /// mongod's own `clientCursorMonitorFrequencySecs` default is 4. + cursor_sweep_secs: i64, pub fn run(self: *Server) !void { // Unbounded async limit: connection handlers otherwise fall back to @@ -43,6 +46,12 @@ pub const Server = struct { // The TTL monitor is just another member of the connection group, so // the `group.cancel` above stops it with everything else. if (self.ttl_sweep_secs > 0) group.async(io, ttl_monitor, .{ io, self }); + // A separate fiber rather than a branch inside ttl_monitor, for two + // reasons: the cadences differ by more than an order of magnitude (4 s + // against 60 s), and a TTL sweep that fails must not stop cursors being + // reclaimed. It is also spawned when TTL sweeping is disabled entirely, + // which is the configuration the spec runner uses. + if (self.cursor_sweep_secs > 0) group.async(io, cursor_monitor, .{ io, self }); while (true) { const stream = listener.accept(io) catch |err| switch (err) { @@ -87,6 +96,20 @@ fn ttl_monitor(io: std.Io, server: *Server) error{Canceled}!void { } } +/// Reap cursors nobody has touched for `cursor_timeout_ms`, until the group is +/// canceled. Takes no engine lock: a cursor owns its own arena, and the store's +/// mutex is a leaf. +fn cursor_monitor(io: std.Io, server: *Server) error{Canceled}!void { + const interval: std.Io.Duration = .fromSeconds(server.cursor_sweep_secs); + while (true) { + // Sleep first, for the same reason the TTL monitor does: at startup + // there is nothing to reap and the listener wants the CPU. + try std.Io.sleep(io, interval, .awake); + const now_ms = std.Io.Timestamp.now(io, .real).toMilliseconds(); + _ = server.engine.cursors.sweep(io, now_ms); + } +} + /// Entry point required by `Group.async`: must return only `error.Canceled`. fn handle_connection(io: std.Io, stream: std.Io.net.Stream, server: *Server) error{Canceled}!void { handle_connection_inner(io, stream, server) catch {}; diff --git a/src/wire.zig b/src/wire.zig index 0173340..df6cc9d 100644 --- a/src/wire.zig +++ b/src/wire.zig @@ -299,9 +299,20 @@ fn begin_message( } /// Patch in the total length of the message started at `len_pos`. +/// +/// The bound is `max_message_size`, the same 48 MiB we advertise to drivers as +/// `maxMessageSizeBytes`, not `maxInt(u32)`. A reply past what we told the +/// client to expect is not a large reply, it is a desynchronized connection: +/// the driver reads the length, refuses or mis-frames it, and every later +/// command on that socket reads the wrong bytes. Failing here turns that into +/// one honest error on the request that caused it. +/// +/// Reachable today: nothing caps how many documents a `find` puts in its single +/// batch, so ~3000 documents of 16 KiB clears 48 MB. The cursor batch budget +/// makes it unreachable, which is the point of keeping this as the backstop. 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; + if (total > max_message_size) return error.MessageTooLarge; std.mem.writeInt(u32, out.items[len_pos..][0..4], @intCast(total), .little); } @@ -411,3 +422,31 @@ test "parse OP_QUERY handshake" { try testing.expectEqual(@as(i32, op_code_query), msg.op_code); try testing.expectEqualStrings("isMaster", msg.command_name()); } + +test "a reply past the advertised message size fails to build" { + // The guard exists because exceeding it desynchronizes the connection + // rather than merely making one reply large, so it must be an error return + // and not a truncation. One oversized string is the cheapest way past it + // without allocating 48 MB of documents. + const gpa = testing.allocator; + var reply = Reply.init(gpa); + defer reply.deinit(); + + const big = try reply.arena_alloc().alloc(u8, max_message_size + 1); + @memset(big, 'x'); + try reply.put_ok(); + try reply.put("payload", .{ .string = big }); + + var out: std.ArrayListUnmanaged(u8) = .empty; + defer out.deinit(gpa); + try testing.expectError(error.MessageTooLarge, reply.build(gpa, 1, 1, &out)); + + // And a reply comfortably inside the bound still builds, so the guard is + // not simply rejecting everything. + var small = Reply.init(gpa); + defer small.deinit(); + try small.put_ok(); + out.clearRetainingCapacity(); + try small.build(gpa, 1, 1, &out); + try testing.expect(out.items.len < max_message_size); +} diff --git a/tests/e2e/README.md b/tests/e2e/README.md index 4a1db62..1e04d19 100644 --- a/tests/e2e/README.md +++ b/tests/e2e/README.md @@ -41,6 +41,19 @@ node tests/e2e/e2e6.js # 73 checks, ~15 s, needs no running server E2E6_PORT=27300 node tests/e2e/e2e6.js # different port if 27220 is taken ``` +`e2e7.js` is the cursor suite, self-contained for a different reason: cursor +behaviour is only observable with non-default flags. It spawns three servers in +turn -- default flags for batching/streaming/aggregate, then +`--cursor-timeout-ms 800 --cursor-sweep-secs 1 --max-open-cursors 4` for idle +expiry and registry capacity, then a restart on the same database to confirm a +cursor does not survive one. Most of it uses raw `runCommand`, because the +driver hides `cursor.id` and that is the thing under test: + +```sh +node tests/e2e/e2e7.js # 86 checks, needs no running server +E2E7_PORT=27310 node tests/e2e/e2e7.js # different port if 27230 is taken +``` + Rebuild with `zig build` after any change under `src/` before restarting the server: `zig build test` compiles the test binary only and leaves `zig-out/bin/multiforadb` stale, so the suites keep running against the old diff --git a/tests/e2e/e2e7.js b/tests/e2e/e2e7.js new file mode 100644 index 0000000..69b4293 --- /dev/null +++ b/tests/e2e/e2e7.js @@ -0,0 +1,474 @@ +// E2E part 7: server-side cursors, self-contained. +// +// Spawns its own multiforadb servers, because cursor behaviour is only +// observable with non-default flags (a short idle timeout, a tiny registry) and +// with raw `runCommand` — the driver hides `cursor.id`, which is the thing under +// test. +// +// node tests/e2e/e2e7.js +// +// Env: E2E7_PORT listen port (default 27230) +// MFDB_BIN server binary (default ../../zig-out/bin/multiforadb) +// E2E7_KEEP keep the log files after the run +// +// The one rule most of this file is about: **never look ahead.** A batch ends +// either because it reached its target — cursor stays open — or because the +// source reported EOF, and only then does the cursor close with `id: 0`. So four +// documents at `batchSize: 2` require a third command answering an empty +// `nextBatch` with `id: 0`. That empty terminal batch is correct, and it is what +// real mongod does (measured, not assumed — see checks 5 and 6). +const { MongoClient, Long } = require('mongodb'); +const { spawn } = require('child_process'); +const fs = require('fs'); +const path = require('path'); + +const PORT = Number(process.env.E2E7_PORT || 27230); +const BIN = process.env.MFDB_BIN || path.resolve(__dirname, '../../zig-out/bin/multiforadb'); +const DBFILE = path.resolve(__dirname, '../../.zig-cache/e2e7-cursors.log'); +const URL = `mongodb://127.0.0.1:${PORT}`; + +const results = []; +function check(name, cond, detail = '') { + results.push({ name, ok: !!cond, detail: String(detail) }); + if (!cond) console.error(` x ${name} ${detail}`); +} +function eq(name, got, want) { + const ok = JSON.stringify(got) === JSON.stringify(want); + check(name, ok, ok ? '' : `got ${JSON.stringify(got)} want ${JSON.stringify(want)}`); +} +const sleep = (ms) => new Promise((r) => setTimeout(r, ms)); + +/// The error code a command produces, or 0 when it succeeds. +async function codeOf(fn) { + try { + await fn(); + return 0; + } catch (e) { + return e.code === undefined ? -1 : e.code; + } +} + +let server = null; +let serverLog = ''; +let serverDead = false; + +function cleanup() { + if (server && !serverDead) { + try { server.kill('SIGKILL'); } catch {} + } +} +process.on('exit', cleanup); +process.on('SIGINT', () => { cleanup(); process.exit(130); }); +process.on('SIGTERM', () => { cleanup(); process.exit(143); }); + +function startServer(args, fresh = true) { + return new Promise((resolve, reject) => { + if (fresh) { + fs.rmSync(DBFILE, { force: true }); + fs.rmSync(DBFILE + '.data', { force: true }); + } + serverDead = false; + serverLog = ''; + server = spawn(BIN, ['--port', String(PORT), '--db', DBFILE, ...args], { + stdio: ['ignore', 'pipe', 'pipe'], + }); + server.stdout.on('data', (d) => (serverLog += d)); + server.stderr.on('data', (d) => (serverLog += d)); + server.on('error', (e) => reject(new Error(`cannot start ${BIN}: ${e.message}`))); + server.on('exit', (code, sig) => { + serverDead = true; + if (code !== null && sig === null) serverLog += `\n[child exited rc=${code}]`; + }); + const deadline = Date.now() + 15000; + (async () => { + while (Date.now() < deadline) { + if (serverDead) { + reject(new Error(`server exited during start (port ${PORT} busy?)\n${serverLog}`)); + return; + } + const c = new MongoClient(URL, { serverSelectionTimeoutMS: 1000 }); + try { + await c.connect(); + await c.db('admin').command({ ping: 1 }); + await c.close(); + return resolve(); + } catch { + try { await c.close(); } catch {} + await sleep(100); + } + } + reject(new Error(`server did not come up on :${PORT}\n${serverLog}`)); + })(); + }); +} + +async function stopServer(sig = 'SIGTERM') { + if (!server) return; + const exited = new Promise((r) => server.once('exit', r)); + server.kill(sig); + await Promise.race([exited, sleep(5000)]); + serverDead = true; + server = null; +} + +// --------------------------------------------------------------------------- +// Phase A — batching, lifecycle and errors, on default flags +// --------------------------------------------------------------------------- + +async function phaseA(db) { + const col = db.collection('c'); + await col.deleteMany({}); + await col.insertMany([...Array(250)].map((_, i) => ({ _id: i + 1, x: i }))); + + // 1. A cursor is a real cursor: nonzero id, a namespace with both parts. + let r = await db.command({ find: 'c', filter: {}, batchSize: 2 }); + eq('1 batchSize 2 returns 2', r.cursor.firstBatch.length, 2); + check('1 cursor id is nonzero', r.cursor.id > 0, r.cursor.id); + eq('1 ns is db.coll', r.cursor.ns, 'e2e7.c'); + + // 2. The default first batch is 101, MongoDB's own + // internalQueryFindCommandBatchSize. + eq('2 default first batch is 101', (await db.command({ find: 'c', filter: {} })).cursor.firstBatch.length, 101); + + // 3. A getMore naming no batchSize is bounded by bytes, not by the batchSize + // the cursor was created with. Measured against mongod 8.3.7: find with + // batchSize 2 then a bare getMore returns 4998 of 5000 documents. + const id3 = (await db.command({ find: 'c', filter: {}, batchSize: 2 })).cursor.id; + let g = await db.command({ getMore: id3, collection: 'c' }); + eq('3 a bare getMore drains the rest', g.cursor.nextBatch.length, 248); + eq('3 and closes at EOF', String(g.cursor.id), '0'); + + // 4. A getMore's batchSize applies to that batch only. + const id4 = (await db.command({ find: 'c', filter: {}, batchSize: 2 })).cursor.id; + g = await db.command({ getMore: id4, collection: 'c', batchSize: 3 }); + eq('4 getMore batchSize 3', g.cursor.nextBatch.map((d) => d._id), [3, 4, 5]); + check('4 still open', g.cursor.id > 0); + + // 5. No look-ahead: 4 documents at batchSize 2 needs a third command whose + // nextBatch is empty. Closing on "batch full and source dry" would break + // the command counts the pinned spec suites assert. + const four = db.collection('four'); + await four.deleteMany({}); + await four.insertMany([1, 2, 3, 4].map((i) => ({ _id: i }))); + r = await db.command({ find: 'four', filter: {}, batchSize: 2 }); + g = await db.command({ getMore: r.cursor.id, collection: 'four', batchSize: 2 }); + eq('5 second full batch is 2 documents', g.cursor.nextBatch.length, 2); + check('5 and leaves the cursor open', g.cursor.id > 0); + g = await db.command({ getMore: g.cursor.id, collection: 'four', batchSize: 2 }); + eq('5 terminal batch is empty and closed', [g.cursor.nextBatch.length, String(g.cursor.id)], [0, '0']); + + // 6. limit is an EOF source, so the batch that exhausts it also closes the + // cursor. This is why the driver sends batchSize = limit + 1. + r = await db.command({ find: 'c', filter: {}, sort: { _id: 1 }, limit: 4, batchSize: 5 }); + eq('6 limit 4 batchSize 5 closes in one reply', String(r.cursor.id), '0'); + eq('6 and returns exactly the limit', r.cursor.firstBatch.map((d) => d._id), [1, 2, 3, 4]); + r = await db.command({ find: 'c', filter: {}, sort: { _id: 1 }, limit: 4, batchSize: 2 }); + check('6 limit 4 batchSize 2 stays open', r.cursor.id > 0); + g = await db.command({ getMore: r.cursor.id, collection: 'c', batchSize: 2 }); + eq('6 the batch reaching the limit closes', String(g.cursor.id), '0'); + eq('6 across batches, limit still honoured', g.cursor.nextBatch.map((d) => d._id), [3, 4]); + + // 7. skip is consumed once, at creation, and never re-applied. + r = await db.command({ find: 'c', filter: {}, skip: 20, batchSize: 3 }); + eq('7 skip 20 starts at 21', r.cursor.firstBatch.map((d) => d._id), [21, 22, 23]); + g = await db.command({ getMore: r.cursor.id, collection: 'c', batchSize: 3 }); + eq('7 skip not re-applied on getMore', g.cursor.nextBatch.map((d) => d._id), [24, 25, 26]); + + // 8. batchSize 0 is a real request for an empty batch with a live cursor, not + // "unbounded". Drivers use it to obtain a cursor cheaply. Nothing may be + // consumed. + r = await db.command({ find: 'c', filter: {}, batchSize: 0 }); + eq('8 batchSize 0 returns nothing', r.cursor.firstBatch.length, 0); + check('8 but a live cursor', r.cursor.id > 0); + g = await db.command({ getMore: r.cursor.id, collection: 'c', batchSize: 1 }); + eq('8 nothing was consumed', g.cursor.nextBatch.map((d) => d._id), [1]); + + // 9. singleBatch, and its wire-legacy form, a negative limit. + eq('9 singleBatch closes', String((await db.command({ find: 'c', filter: {}, batchSize: 2, singleBatch: true })).cursor.id), '0'); + r = await db.command({ find: 'c', filter: {}, limit: -3 }); + eq('9 negative limit is one batch', [r.cursor.firstBatch.length, String(r.cursor.id)], [3, '0']); + + // 10. A cursor is not pinned to the connection that created it: the driver + // spec allows a getMore from any connection to the same server. + const other = new MongoClient(URL); + await other.connect(); + const shared = (await db.command({ find: 'c', filter: {}, batchSize: 2 })).cursor.id; + eq('10 getMore from another connection', await codeOf(() => other.db('e2e7').command({ getMore: shared, collection: 'c', batchSize: 2 })), 0); + await other.close(); + + // 11. A getMore naming the wrong collection is Unauthorized (13), not 43, and + // leaves the cursor alive — the request is wrong, not the cursor. + // Measured against mongod, which answers exactly this code. + const live = (await db.command({ find: 'c', filter: {}, batchSize: 2 })).cursor.id; + eq('11 wrong collection is Unauthorized 13', await codeOf(() => db.command({ getMore: live, collection: 'four' })), 13); + eq('11 the cursor survived it', (await db.command({ getMore: live, collection: 'c', batchSize: 1 })).cursor.nextBatch.length, 1); + + // 12. killCursors: all four arrays, and the right partitioning. + let k = await db.command({ killCursors: 'c', cursors: [live] }); + eq('12 a live cursor is killed', k.cursorsKilled.map(String), [String(live)]); + check('12 all four arrays present', ['cursorsKilled', 'cursorsNotFound', 'cursorsAlive', 'cursorsUnknown'].every((f) => Array.isArray(k[f])), Object.keys(k).join(',')); + k = await db.command({ killCursors: 'c', cursors: [live] }); + eq('12 killing it twice reports notFound', k.cursorsNotFound.map(String), [String(live)]); + const other_ns = (await db.command({ find: 'four', filter: {}, batchSize: 1 })).cursor.id; + k = await db.command({ killCursors: 'c', cursors: [other_ns] }); + eq('12 a wrong-namespace id reports notFound', k.cursorsNotFound.map(String), [String(other_ns)]); + eq('12 and that cursor still lives', await codeOf(() => db.command({ getMore: other_ns, collection: 'four', batchSize: 1 })), 0); + + // 13. Malformed and unknown ids. + eq('13 getMore after kill is 43', await codeOf(() => db.command({ getMore: live, collection: 'c' })), 43); + eq('13 id 0 is 43', await codeOf(() => db.command({ getMore: Long.fromNumber(0), collection: 'c' })), 43); + eq('13 an unknown id is 43', await codeOf(() => db.command({ getMore: Long.fromString('987654321'), collection: 'c' })), 43); + eq('13 a non-numeric id is TypeMismatch 14', await codeOf(() => db.command({ getMore: 'nope', collection: 'c' })), 14); + eq('13 a missing collection is BadValue 2', await codeOf(() => db.command({ getMore: Long.fromNumber(1) })), 2); + + // 14. tailable is refused, which is parity: mongod rejects it on a non-capped + // collection and this engine has none. Ignoring it would make a driver's + // tail loop exit, which the application reads as data loss. + eq('14 tailable is BadValue 2', await codeOf(() => db.command({ find: 'c', filter: {}, tailable: true })), 2); + eq('14 awaitData alone is BadValue 2', await codeOf(() => db.command({ find: 'c', filter: {}, awaitData: true })), 2); + + // 15. The driver's own iteration, which is the point of all of the above. + const ids = (await col.find({}).batchSize(7).toArray()).map((d) => d._id); + eq('15 driver drains 250 at batchSize 7', ids.length, 250); + eq('15 no duplicates and no gaps', [new Set(ids).size, Math.min(...ids), Math.max(...ids)], [250, 1, 250]); + eq('15 sort+skip+limit unchanged', (await col.find({ _id: { $gt: 2 } }, { sort: { _id: 1 }, skip: 2, limit: 2 }).toArray()).map((d) => d._id), [5, 6]); + // A sort no index provides must materialize; it still has to drain correctly. + eq('15 an unindexed sort drains in order', (await col.find({}, { sort: { x: -1 } }).batchSize(10).toArray()).map((d) => d.x)[0], 249); +} + +// --------------------------------------------------------------------------- +// Phase B — streaming cursors: resume across writes, and what survives a rebuild +// --------------------------------------------------------------------------- + +async function phaseB(db) { + const col = db.collection('s'); + await col.deleteMany({}); + await col.insertMany([...Array(300)].map((_, i) => ({ _id: i + 1, a: i % 5, pad: 'q'.repeat(200) }))); + + // 16. A whole-index walk holds a key, not a list, so it resumes across writes + // that move documents. The bug this caught: an update rewrites a document + // to a new offset, and resuming by band position returned it twice. + let r = await db.command({ find: 's', filter: {}, batchSize: 10 }); + const seen = new Set(r.cursor.firstBatch.map((d) => d._id)); + let dupes = 0; + let id = r.cursor.id; + let rounds = 0; + let errored = 0; + while (String(id) !== '0' && rounds++ < 200) { + // Churn between every batch: updates rewrite documents, which both moves + // them in the slab and can split leaves. + await col.updateMany({ _id: { $lt: 60 } }, { $inc: { n: 1 } }); + let g; + try { + g = await db.command({ getMore: id, collection: 's', batchSize: 10 }); + } catch (e) { + errored = e.code; + break; + } + for (const d of g.cursor.nextBatch) { + if (seen.has(d._id)) dupes++; + seen.add(d._id); + } + id = g.cursor.id; + } + eq('16 draining across churn did not error', errored, 0); + eq('16 no document came back twice', dupes, 0); + eq('16 every document was returned', seen.size, 300); + check('16 and nothing outside the collection', [...seen].every((v) => v >= 1 && v <= 300)); + + // 17. Both directions stream, over the _id_ index and a secondary one. + eq('17 ascending _id sort drains', (await col.find({}, { sort: { _id: 1 } }).batchSize(9).toArray()).length, 300); + const desc = (await col.find({}, { sort: { _id: -1 } }).batchSize(9).toArray()).map((d) => d._id); + eq('17 descending drains in order', [desc.length, desc[0], desc[299]], [300, 300, 1]); + await col.createIndex({ a: 1 }); + const bya = await col.find({}, { sort: { a: 1 } }).batchSize(11).toArray(); + eq('17 a secondary-index sort drains', bya.length, 300); + check('17 and in the index order', bya.every((d, i) => i === 0 || bya[i - 1].a <= d.a)); + + // 18. Dropping the index a stream is following cannot be resumed — the walk + // has nothing left to walk. That must be a clean error, not garbage. + await col.createIndex({ b: 1 }); + const onb = (await db.command({ find: 's', filter: {}, sort: { b: 1 }, batchSize: 3 })).cursor.id; + await col.dropIndex('b_1'); + eq('18 dropping the streamed index is 175', await codeOf(() => db.command({ getMore: onb, collection: 's', batchSize: 3 })), 175); + + // 19. Dropping the collection kills every kind of cursor. + const doomed = (await db.command({ find: 's', filter: {}, batchSize: 3 })).cursor.id; + await col.drop(); + const dc = await codeOf(() => db.command({ getMore: doomed, collection: 's', batchSize: 3 })); + check('19 dropping the collection kills the cursor', dc === 175 || dc === 43, dc); +} + +// --------------------------------------------------------------------------- +// Phase C — aggregate, the listing commands, and count +// --------------------------------------------------------------------------- + +async function phaseC(db) { + const col = db.collection('g'); + await col.deleteMany({}); + await col.insertMany([...Array(250)].map((_, i) => ({ _id: i + 1, g: i % 40, v: i }))); + + // 20. aggregate batches through cursor.batchSize; a bare cursor is the default. + let r = await db.command({ aggregate: 'g', pipeline: [], cursor: { batchSize: 3 } }); + eq('20 aggregate batchSize 3', r.cursor.firstBatch.length, 3); + check('20 aggregate cursor is real', r.cursor.id > 0); + eq('20 aggregate ns', r.cursor.ns, 'e2e7.g'); + const g20 = await db.command({ getMore: r.cursor.id, collection: 'g', batchSize: 5 }); + eq('20 aggregate getMore continues', g20.cursor.nextBatch.map((d) => d._id), [4, 5, 6, 7, 8]); + eq('20 bare cursor defaults to 101', (await db.command({ aggregate: 'g', pipeline: [], cursor: {} })).cursor.firstBatch.length, 101); + eq('20 driver aggregate drains', (await col.aggregate([], { batchSize: 7 }).toArray()).length, 250); + const groups = await col.aggregate([{ $group: { _id: '$g', n: { $sum: 1 } } }], { batchSize: 6 }).toArray(); + eq('20 $group drains across batches', [groups.length, groups.reduce((a, x) => a + x.n, 0)], [40, 250]); + eq('20 $count stage', await col.aggregate([{ $count: 'total' }]).toArray(), [{ total: 250 }]); + + // 21. listCollections' namespace. It used to be "." with an empty + // collection part, and the driver throws client-side on a namespace like + // that — so the moment the cursor stopped being id 0 it would have broken. + for (let i = 0; i < 12; i++) await db.createCollection('k' + i); + r = await db.command({ listCollections: 1, cursor: { batchSize: 4 } }); + eq('21 listCollections ns has a collection part', r.cursor.ns, 'e2e7.$cmd.listCollections'); + eq('21 listCollections honours batchSize', r.cursor.firstBatch.length, 4); + check('21 listCollections cursor is real', r.cursor.id > 0); + const g21 = await db.command({ getMore: r.cursor.id, collection: '$cmd.listCollections', batchSize: 100 }); + check('21 its getMore works', g21.cursor.nextBatch.length >= 8, g21.cursor.nextBatch.length); + const listed = await db.listCollections({}, { batchSize: 3 }).toArray(); + check('21 driver listCollections drains', listed.length >= 13, listed.length); + + // 22. listIndexes. + await col.createIndex({ v: 1 }); + await col.createIndex({ g: 1 }); + await col.createIndex({ v: -1, g: 1 }); + r = await db.command({ listIndexes: 'g', cursor: { batchSize: 2 } }); + eq('22 listIndexes honours batchSize', r.cursor.firstBatch.length, 2); + eq('22 listIndexes ns', r.cursor.ns, 'e2e7.g'); + eq('22 driver listIndexes drains', (await col.listIndexes({ batchSize: 1 }).toArray()).length, 4); + + // 23. count honoured neither skip nor limit before, which made + // countDocuments(f, {limit}) a silent wrong answer. + eq('23 count plain', (await db.command({ count: 'g' })).n, 250); + eq('23 count limit', (await db.command({ count: 'g', limit: 10 })).n, 10); + eq('23 count skip', (await db.command({ count: 'g', skip: 240 })).n, 10); + eq('23 count skip and limit', (await db.command({ count: 'g', skip: 245, limit: 10 })).n, 5); + eq('23 count skip past the end', (await db.command({ count: 'g', skip: 1000 })).n, 0); + eq('23 count with a query and limit', (await db.command({ count: 'g', query: { g: 0 }, limit: 3 })).n, 3); + eq('23 driver countDocuments limit', await col.countDocuments({}, { limit: 7 }), 7); + + // 24. A batch is capped by bytes as well as by documents, so a large-document + // result splits instead of building a reply past the advertised message + // size. 40 documents of ~1 MiB cannot all fit one 16 MiB batch. + const big = db.collection('big'); + await big.deleteMany({}); + const pad = 'p'.repeat(1024 * 1024 - 64); + for (let i = 0; i < 40; i++) await big.insertOne({ _id: i + 1, pad }); + r = await db.command({ find: 'big', filter: {}, batchSize: 40 }); + check('24 the byte cap split the batch', r.cursor.firstBatch.length >= 1 && r.cursor.firstBatch.length <= 16, r.cursor.firstBatch.length); + check('24 and left the cursor open', r.cursor.id > 0); + eq('24 the whole result still drains', (await big.find({}).batchSize(40).toArray()).length, 40); +} + +// --------------------------------------------------------------------------- +// Phase D — expiry, capacity, and what a restart does +// --------------------------------------------------------------------------- + +async function phaseD(db) { + const col = db.collection('e'); + await col.deleteMany({}); + await col.insertMany([...Array(50)].map((_, i) => ({ _id: i + 1 }))); + + // 25. An idle cursor is reaped; noCursorTimeout exempts one from that. + const perishable = (await db.command({ find: 'e', filter: {}, batchSize: 2 })).cursor.id; + const immortal = (await db.command({ find: 'e', filter: {}, batchSize: 2, noCursorTimeout: true })).cursor.id; + await sleep(300); + eq('25 before the timeout it is alive', await codeOf(() => db.command({ getMore: perishable, collection: 'e', batchSize: 1 })), 0); + await sleep(2500); + eq('25 an idle cursor is reaped', await codeOf(() => db.command({ getMore: perishable, collection: 'e', batchSize: 1 })), 43); + eq('25 noCursorTimeout survives', await codeOf(() => db.command({ getMore: immortal, collection: 'e', batchSize: 1 })), 0); + const k = await db.command({ killCursors: 'e', cursors: [immortal] }); + eq('25 but is still killable', k.cursorsKilled.map(String), [String(immortal)]); + + // 26. A full registry evicts the least recently used cursor rather than + // refusing the new one. The victim sees the same 43 an idle timeout gives, + // which every driver already handles. + const ids = []; + for (let i = 0; i < 5; i++) ids.push((await db.command({ find: 'e', filter: {}, batchSize: 1 })).cursor.id); + eq('26 the oldest was evicted', await codeOf(() => db.command({ getMore: ids[0], collection: 'e', batchSize: 1 })), 43); + const alive = []; + for (const id of ids.slice(1)) alive.push(await codeOf(() => db.command({ getMore: id, collection: 'e', batchSize: 1 }))); + eq('26 the newest four are alive', alive, [0, 0, 0, 0]); +} + +async function phaseE(db, staleId) { + // 27. Cursors do not survive a restart, and a stale id must be a clean 43 — + // not a hang, and not an empty batch claiming the result ended. + eq('27 a cursor from before the restart is 43', await codeOf(() => db.command({ getMore: staleId, collection: 'e', batchSize: 1 })), 43); + const fresh = await db.command({ find: 'e', filter: {}, batchSize: 2 }); + check('27 and new cursors work after a restart', fresh.cursor.id > 0); +} + +async function main() { + // The same guard e2e6.js and big.js carry: without it a missing binary + // surfaces as a generic spawn error instead of saying what to do about it. + if (!fs.existsSync(BIN)) { + console.error(`E2E7_FAIL server binary not found: ${BIN}\n run: zig build`); + process.exit(1); + } + + // Phase A-C on default cursor flags. + await startServer(['--ttl-sweep-secs', '0', '--compact-threshold', '1m'], true); + let client = new MongoClient(URL); + await client.connect(); + let db = client.db('e2e7'); + console.log('phase A: batching, lifecycle, errors'); + await phaseA(db); + console.log('phase B: streaming cursors across writes'); + await phaseB(db); + console.log('phase C: aggregate, listings, count'); + await phaseC(db); + await client.close(); + await stopServer('SIGTERM'); + + // Phase D needs a short timeout and a tiny registry. + console.log('phase D: idle expiry and registry capacity'); + await startServer( + ['--ttl-sweep-secs', '0', '--cursor-timeout-ms', '800', '--cursor-sweep-secs', '1', '--max-open-cursors', '4'], + true, + ); + client = new MongoClient(URL); + await client.connect(); + db = client.db('e2e7'); + await phaseD(db); + const staleId = (await db.command({ find: 'e', filter: {}, batchSize: 1 })).cursor.id; + await client.close(); + await stopServer('SIGTERM'); + + // Phase E: the same database, a new process. + console.log('phase E: a cursor does not survive a restart'); + await startServer(['--ttl-sweep-secs', '0'], false); + client = new MongoClient(URL); + await client.connect(); + await phaseE(client.db('e2e7'), staleId); + await client.close(); + + if (process.env.E2E7_KEEP !== '1') { + fs.rmSync(DBFILE, { force: true }); + fs.rmSync(DBFILE + '.data', { force: true }); + } + await stopServer('SIGTERM'); + + const failed = results.filter((r) => !r.ok); + console.log(`\n${results.length - failed.length}/${results.length} checks passed`); + if (failed.length) { + console.log('FAILED:', failed.map((f) => f.name).join(', ')); + console.log('--- server log tail ---'); + console.log(serverLog.split('\n').slice(-30).join('\n')); + process.exit(1); + } + console.log('E2E7_OK'); +} + +main().catch((e) => { + console.error('E2E7_FAIL', e); + console.log('--- server log tail ---'); + console.log(serverLog.split('\n').slice(-40).join('\n')); + process.exit(1); +}); diff --git a/tests/spec/run.js b/tests/spec/run.js index 831dbe3..fda126b 100644 --- a/tests/spec/run.js +++ b/tests/spec/run.js @@ -314,11 +314,38 @@ class Unsupported extends Error {} // Argument keys the spec passes positionally rather than as driver options. const POSITIONAL = new Set(['filter', 'document', 'documents', 'update', 'replacement', 'pipeline', 'fieldName', 'models', 'requests', 'keys', 'name', 'indexes', 'command', 'session', 'entity', 'to']); +// Driver options the driver only honours as a JavaScript number. The suites are +// parsed with `EJSON.parse(text, {relaxed: false})` so that `$numberLong` and +// friends keep their exact BSON type in *data* -- but that also turns a plain +// JSON `2` in an *option* into a BSON Int32 object, and the driver gates every +// one of these on `typeof options.skip === 'number'` +// (node_modules/mongodb/lib/operations/find.js:68-95). A BSON wrapper therefore +// failed the check and the option was dropped on the floor: `skip`, `limit` and +// `batchSize` never reached the wire at all, and three find.json cases failed +// with the *unclipped* match count while the engine was applying both correctly. +// Read as an engine bug for a whole milestone. Coerce by name, not by shape: +// unwrapping every numeric-looking value would rewrite the wire type of the +// `comment` and `hint` values that other suites assert on. +const NUMERIC_OPTIONS = new Set([ + 'skip', + 'limit', + 'batchSize', + 'maxTimeMS', + 'maxAwaitTimeMS', + 'expireAfterSeconds', +]); + +function numeric_option(v) { + if (v === null || typeof v !== 'object' || typeof v.valueOf !== 'function') return v; + const n = v.valueOf(); + return typeof n === 'number' ? n : v; +} + function options(args, drop = []) { const o = {}; for (const [k, v] of Object.entries(args || {})) { if (POSITIONAL.has(k) || drop.includes(k)) continue; - o[k] = v; + o[k] = NUMERIC_OPTIONS.has(k) ? numeric_option(v) : v; } return Object.keys(o).length ? o : undefined; } diff --git a/tests/spec/scorecard.txt b/tests/spec/scorecard.txt index a182ab2..a622491 100644 --- a/tests/spec/scorecard.txt +++ b/tests/spec/scorecard.txt @@ -13,7 +13,7 @@ # semantics; ignoring them makes some cases pass that a full runner would # fail, so treat `pass` as an upper bound until M1 wires events up. -total 163 pass 129 fail 195 skip 175 files 0 errored +total 168 pass 124 fail 195 skip 175 files 0 errored # per-file: name pass fail skip aggregate-allowdiskuse.json 3 0 0 @@ -78,7 +78,7 @@ client-bulkWrite-updateOne-sort.json 0 0 1 count-collation.json 1 0 1 count-empty.json 2 0 1 count-rawdata.json 0 0 2 -count.json 3 1 3 +count.json 4 0 3 countDocuments-comment.json 2 0 1 countDocuments-rawdata.json 1 0 1 create-null-ids.json 0 6 1 @@ -116,8 +116,8 @@ find-collation.json 0 1 0 find-comment.json 1 2 2 find-let.json 0 1 1 find-rawdata.json 1 0 1 -find.json 2 3 0 -findOne.json 1 1 0 +find.json 5 0 0 +findOne.json 2 0 0 findOneAndDelete-collation.json 0 1 0 findOneAndDelete-comment.json 2 0 1 findOneAndDelete-hint-serverError.json 0 0 2 @@ -292,7 +292,6 @@ count-collation.json SKIP Deprecated count with collation runner: operation coun count-empty.json SKIP Deprecated count with empty collection runner: operation count count-rawdata.json SKIP Deprecated count with rawData option needs server >= 8.2.0 count-rawdata.json SKIP Deprecated count with rawData option on less than 8.2.0 - ignore argument runner: operation count -count.json FAIL Count documents with skip and limit countDocuments: expected 2, got 3 count.json SKIP Deprecated count without a filter runner: operation count count.json SKIP Deprecated count with a filter runner: operation count count.json SKIP Deprecated count with skip and limit runner: operation count @@ -354,10 +353,6 @@ find-comment.json SKIP find with comment does not set comment on getMore - pre 4 find-let.json SKIP Find with let option needs server >= 5.0 find-let.json FAIL Find with let option unsupported (server-side error) find: expected an error, the operation succeeded find-rawdata.json SKIP Find with rawData option needs server >= 8.2.0 -find.json FAIL Find with filter, sort, skip, and limit find: expected 2 elements, got 4 -find.json FAIL Find with limit, sort, and batchsize find: expected 4 elements, got 6 -find.json FAIL Find with batchSize equal to limit find: expected 4 elements, got 5 -findOne.json FAIL FindOne with filter, sort, and skip findOne._id: expected 5, got 3 findOneAndDelete-collation.json FAIL FindOneAndDelete when one document matches with collation findOneAndDelete: expected a document, got null findOneAndDelete-comment.json SKIP findOneAndDelete with comment - pre 4.4 needs server <= 4.2.99 findOneAndDelete-hint-serverError.json SKIP * needs server <= 4.3.3