MultiforaDB/src/server.zig

//! 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;
        };
    }
}