Chapter 15: Complete Example Program

  const process     = import("std/process");
const io_mod      = import("std/async");
const net         = import("std/net");
const sync        = import("std/sync");
const collections = import("std/collections");
const fmt         = import("std/fmt");

-- the store: a thread-safe HashMap
type Store = struct {
    data: sync.Mutex[collections.HashMap[String, String]],
};

impl Store {
    static fn new(alloc: Allocator): !Self {
        return .{
            .data = sync.Mutex[collections.HashMap[String, String]]::new(
                collections.HashMap[String, String]::new()
            ),
        };
    }

    fn get(*const self, key: []const u8): ?String {
        var guard = self.data.lock();
        return case guard.*.get(key) {
            .Some { val } => .Some { .value = val.clone() },
            .None         => .None,
        };
    }

    fn set(*self, alloc: Allocator, key: String, val: String): !void {
        var guard = self.data.lock();
        try guard.*.set(alloc, key, val);
    }

    fn delete(*self, key: []const u8): bool {
        var guard = self.data.lock();
        return guard.*.remove(key) != .None;
    }
}

-- parse and handle one command
fn handleCommand(
    io:    *Io,
    store: *Store,
    alloc: Allocator,
    line:  []const u8,
    conn:  *net.TcpStream,
): !void {
    var parts = splitLine(line);

    case parts[0] {
        "GET" if parts.len = 2 => {
            if store.get(parts[1]) -> val {
                defer val.deinit(alloc);
                try conn.write(io, val.asSlice());
                try conn.write(io, "\n");
            } else {
                try conn.write(io, "NOT_FOUND\n");
            }
        },
        "SET" if parts.len = 3 => {
            var key = try String::from(alloc, parts[1]);
            errdefer key.deinit(alloc);
            var val = try String::from(alloc, parts[2]);
            errdefer val.deinit(alloc);

            try store.set(alloc, key, val);
            try conn.write(io, "OK\n");
        },
        "DEL" if parts.len = 2 => {
            if store.delete(parts[1]) {
                try conn.write(io, "DELETED\n");
            } else {
                try conn.write(io, "NOT_FOUND\n");
            }
        },
        _ => {
            try conn.write(io, "ERR unknown command\n");
        },
    }
}

-- handle one connection
fn handleConn(
    io:    *Io,
    store: *Store,
    alloc: Allocator,
    conn:  net.TcpStream,
): !void {
    defer conn.deinit();

    var buf: [4096]u8 = #undefined;

    loop {
        var n = try conn.read(io, &buf);
        if n = 0 { break; }    -- connection closed

        var line = trim(buf[0..n]);
        handleCommand(io, store, alloc, line, &conn) catch |e| {
            var msg = try fmt.format(alloc, "ERR {}\n", .{e});
            defer alloc.free(msg);
            conn.write(io, msg) catch {};
        };
    }
}

fn main(ctx: process.Context): !void {
    using ctx.alloc;

    var io = try io_mod.IoUringScheduler::new(ctx.alloc, 1024);
    defer io.deinit(ctx.alloc);

    -- shared store, Arc for thread safety
    var store = Arc[Store]::new(ctx.alloc, try Store::new(ctx.alloc));
    defer store.deinit(ctx.alloc);

    var listener = try net.TcpListener::bind(&io, "0.0.0.0:6379");
    defer listener.deinit();

    debug.print("Listening on :6379\n");

    -- structured concurrency: all connections joined on exit
    try io.scope([](s: *IoScope) {
        while true {
            var conn = try listener.accept(&io);
            var s_clone = store.clone(ctx.alloc);

            try s.spawn([move conn, move s_clone](_: *IoScope) {
                handleConn(&io, s_clone.*, ctx.alloc, conn) catch |e| {
                    debug.print("connection error: {}\n", .{e});
                };
                s_clone.deinit(ctx.alloc);
            });
        }
    });
}