baseline: mongo-light working tree before concurrency refactor

src/server.zig

@@ -0,0 +1,145 @@
+//! 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,
+
+    pub fn run(self: *Server) !void {
+        var threaded: std.Io.Threaded = std.Io.Threaded.init(self.gpa, .{});
+        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-light: listening on {s}:{d}\n", .{ self.bind_ip, self.port });
+
+        var group: std.Io.Group = .init;
+        defer group.cancel(io);
+
+        while (true) {
+            const stream = listener.accept(io) catch |err| switch (err) {
+                error.Canceled => return,
+                else => {
+                    std.debug.print("mongo-light: accept error: {s}\n", .{@errorName(err)});
+                    continue;
+                },
+            };
+            group.async(io, handle_connection, .{ io, stream, self });
+        }
+    }
+};
+
+/// 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;
+    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 |err| switch (err) {
+            error.EndOfStream => return, // clean client disconnect
+            else => {
+                std.debug.print("mongo-light: read error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
+                return;
+            },
+        };
+        const total: u32 = std.mem.readInt(u32, &len_bytes, .little);
+        if (total < 16 or total > wire.max_message_size) {
+            std.debug.print("mongo-light: 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-light: read error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
+            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-light: bad message on conn {d}: {s} (opCode {d})\n", .{ connection_id, @errorName(err), op });
+            return;
+        };
+        defer msg.deinit();
+
+        var reply = wire.Reply.init(server.gpa);
+        defer reply.deinit();
+
+        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();
+        if (msg.op_code == wire.op_code_query) {
+            wire.write_reply_query(server.gpa, reply_request_id, msg.request_id, reply.pairs.items, &out_buf) catch |err| {
+                std.debug.print("mongo-light: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
+                return;
+            };
+        } else {
+            reply.build(server.gpa, reply_request_id, msg.request_id, &out_buf) catch |err| {
+                std.debug.print("mongo-light: reply build error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
+                return;
+            };
+        }
+        reply_request_id +%= 1;
+        writer.interface.writeAll(out_buf.items) catch |err| {
+            std.debug.print("mongo-light: write error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
+            return;
+        };
+        writer.interface.flush() catch |err| {
+            std.debug.print("mongo-light: flush error on conn {d}: {s}\n", .{ connection_id, @errorName(err) });
+            return;
+        };
+    }
+}