//! TCP server speaking the MongoDB wire protocol. Accept loop dispatches each //! connection onto the Io worker pool; a connection is handled until the peer //! closes it or a protocol error occurs. const std = @import("std"); const bson = @import("bson.zig"); const wire = @import("wire.zig"); const commands = @import("commands.zig"); const db = @import("db.zig"); pub const Server = struct { gpa: std.mem.Allocator, port: u16, bind_ip: []const u8, oid_gen: bson.ObjectIdGen, connection_counter: std.atomic.Value(u32), engine: *db.Engine, start_time: std.Io.Timestamp, /// Seconds between TTL sweeps; 0 leaves the monitor unspawned. Signed /// because that is what std.Io.Duration.fromSeconds takes — the CLI /// rejects negatives. ttl_sweep_secs: i64, pub fn run(self: *Server) !void { // Unbounded async limit: connection handlers otherwise fall back to // running inline on the accept-loop fiber once busy_count hits the // default cpu_count-1, which blocks accept() for the handler's // lifetime and stalls new connections (handshake timeouts). With an // unlimited limit the pool spawns a thread per live connection. var threaded: std.Io.Threaded = std.Io.Threaded.init(self.gpa, .{ .async_limit = .unlimited }); defer threaded.deinit(); const io = threaded.io(); var addr = try std.Io.net.IpAddress.parse(self.bind_ip, self.port); var listener = try addr.listen(io, .{ .reuse_address = true }); defer listener.deinit(io); std.debug.print("mongo-lite: listening on {s}:{d}\n", .{ self.bind_ip, self.port }); var group: std.Io.Group = .init; defer group.cancel(io); // 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 }); while (true) { const stream = listener.accept(io) catch |err| switch (err) { error.Canceled => return, else => { std.debug.print("mongo-lite: accept error: {s}\n", .{@errorName(err)}); continue; }, }; group.async(io, handle_connection, .{ io, stream, self }); } } }; /// Expire documents under TTL indexes every `ttl_sweep_secs` seconds, until /// the group is canceled. Sweeping takes the engine's write lock, so it is /// serialized with commands exactly like any other write; a sweep failure is /// logged rather than fatal, since the next one will retry. fn ttl_monitor(io: std.Io, server: *Server) error{Canceled}!void { const interval: std.Io.Duration = .fromSeconds(server.ttl_sweep_secs); while (true) { // Sleep first: at startup the engine has just replayed the log, and // an immediate sweep would race the listener's first connections for // the write lock. try std.Io.sleep(io, interval, .awake); try server.engine.lock(); defer server.engine.unlock(); const now_ms = std.Io.Timestamp.now(io, .real).toMilliseconds(); _ = server.engine.ttl_sweep(now_ms) catch |err| { std.debug.print("mongo-lite: TTL sweep failed: {s}\n", .{@errorName(err)}); continue; }; } } /// 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 {}; } fn handle_connection_inner(io: std.Io, stream: std.Io.net.Stream, server: *Server) !void { defer stream.close(io); const connection_id = server.connection_counter.fetchAdd(1, .monotonic); var read_scratch: [16 * 1024]u8 = undefined; var write_scratch: [16 * 1024]u8 = undefined; var reader = stream.reader(io, &read_scratch); var writer = stream.writer(io, &write_scratch); var msg_buf: std.ArrayListUnmanaged(u8) = .empty; defer msg_buf.deinit(server.gpa); var out_buf: std.ArrayListUnmanaged(u8) = .empty; defer out_buf.deinit(server.gpa); var reply_request_id: u32 = 1; // One reply for the whole connection, reset per request: its arena // keeps its pages instead of being rebuilt for every command. var reply = wire.Reply.init(server.gpa); defer reply.deinit(); var ctx = commands.Context{ .gpa = server.gpa, .io = io, .oid_gen = &server.oid_gen, .connection_id = connection_id, .client_desc = "127.0.0.1:0", .engine = server.engine, .server_start = server.start_time, }; while (true) { var len_bytes: [4]u8 = undefined; reader.interface.readSliceAll(&len_bytes) catch return; // clean client disconnect (EOF or RST) const total: u32 = std.mem.readInt(u32, &len_bytes, .little); if (total < 16 or total > wire.max_message_size) { std.debug.print("mongo-lite: bad message length {d} on conn {d}\n", .{ total, connection_id }); return; } msg_buf.clearRetainingCapacity(); try msg_buf.ensureTotalCapacity(server.gpa, total); msg_buf.items.len = total; std.mem.writeInt(u32, msg_buf.items[0..4], total, .little); reader.interface.readSliceAll(msg_buf.items[4..]) catch |err| { std.debug.print("mongo-lite: read error on conn {d}: {s} (body, len {d})\n", .{ connection_id, @errorName(err), total }); return; }; var msg = wire.Message.parse(server.gpa, msg_buf.items) catch |err| { // Unparseable request: close the connection. const op: i32 = if (msg_buf.items.len >= 16) std.mem.readInt(i32, msg_buf.items[12..16], .little) else 0; std.debug.print("mongo-lite: bad message on conn {d}: {s} (opCode {d})\n", .{ connection_id, @errorName(err), op }); return; }; defer msg.deinit(); reply.reset(); commands.dispatch(&ctx, &msg, &reply) catch { // Discard any partial reply (the client would read the first // ok field, which may already say 1) and send a clean error. reply.pairs.clearRetainingCapacity(); reply.put_error( @intFromEnum(commands.ErrorCode.internal_error), "InternalError", "internal error", ) catch return; }; out_buf.clearRetainingCapacity(); const built = if (msg.op_code == wire.op_code_query) wire.write_reply_query(server.gpa, reply_request_id, msg.request_id, reply.pairs.items, &out_buf) else reply.build(server.gpa, reply_request_id, msg.request_id, &out_buf); built catch |err| { std.debug.print("mongo-lite: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) }); return; }; reply_request_id +%= 1; writer.interface.writeAll(out_buf.items) catch |err| { std.debug.print("mongo-lite: write error on conn {d}: {s}\n", .{ connection_id, @errorName(err) }); return; }; writer.interface.flush() catch |err| { std.debug.print("mongo-lite: flush error on conn {d}: {s}\n", .{ connection_id, @errorName(err) }); return; }; } }