zig: basic tcp echo server w/ stdlib event loop

main.zig

const std = @import("std");

const os = std.os;
const net = std.net;
const mem = std.mem;
const log = std.log.scoped(.server);

const assert = std.debug.assert;

pub const io_mode = .evented;
pub const log_level = .debug;

pub fn main() !void {
    var gpa: std.heap.GeneralPurposeAllocator(.{}) = .{};
    defer assert(!gpa.deinit());

    var listener = net.StreamServer.init(.{});
    defer listener.deinit();

    try listener.listen(net.Address.initIp4(.{ 0, 0, 0, 0 }, 9000));

    log.info("listening for clients at: {}", .{listener.listen_address});

    var server: Server = .{ .listener = listener };
    defer {
        server.shutdown();
        server.deinit(&gpa.allocator);
    }

    var server_frame = async server.serve(&gpa.allocator);
    defer await server_frame catch |err| log.warn("server error: {}", .{err});

    // usually you'd wait for a Ctrl+C or interrupt signal here
}

pub const Server = struct {
    lock: std.Thread.Mutex = .{},
    listener: net.StreamServer,
    clients: std.AutoHashMapUnmanaged(*@Frame(Server.serveClient), net.StreamServer.Connection) = .{},

    pub fn deinit(self: *Server, gpa: *mem.Allocator) void {
        self.clients.deinit(gpa);
    }

    pub fn shutdown(self: *Server) void {
        std.os.shutdown(self.listener.sockfd.?, .recv) catch {};

        const held = self.lock.acquire();
        defer held.release();

        var it = self.clients.valueIterator();
        while (it.next()) |conn| {
            std.os.shutdown(conn.stream.handle, .recv) catch {};
        }
    }

    pub fn serve(self: *Server, gpa: *mem.Allocator) !void {
        while (true) {
            std.event.Loop.instance.?.yield();

            const conn = self.listener.accept() catch |err| switch (err) {
                error.ProcessFdQuotaExceeded, error.SystemFdQuotaExceeded, error.SystemResources => continue,
                error.SocketNotListening => return,
                else => return err,
            };

            const frame = gpa.create(@Frame(Server.serveClient)) catch {
                conn.stream.close();
                continue;
            };

            {
                const held = self.lock.acquire();
                defer held.release();

                self.clients.putNoClobber(gpa, frame, conn) catch {
                    gpa.destroy(frame);
                    conn.stream.close();
                    continue;
                };
            }

            frame.* = async self.serveClient(gpa, conn);
        }
    }

    pub fn serveClient(self: *Server, gpa: *mem.Allocator, conn: net.StreamServer.Connection) !void {
        defer {
            conn.stream.close();
            suspend self.deallocateClient(gpa, @frame());
        }

        log.info("peer connected: {}", .{conn.address});
        defer log.info("peer disconnected: {}", .{conn.address});

        var buffer: [1024]u8 = undefined;
        while (true) {
            std.event.Loop.instance.?.yield();

            const num_bytes = try conn.stream.read(&buffer);
            if (num_bytes == 0) return;
            try conn.stream.writer().writeAll(buffer[0..num_bytes]);
        }
    }

    pub fn deallocateClient(self: *Server, gpa: *mem.Allocator, frame: *@Frame(Server.serveClient)) void {
        const held = self.lock.acquire();
        defer held.release();

        assert(self.clients.remove(frame));
        gpa.destroy(frame);
    }
};