Made progress, but not perfect.

the message isn't crossing container boundaries, but it works in the
test!

3 files changed, 151 insertions, 52 deletions

diff --git a/src/server/client.zig b/src/server/client.zig
index b4bc55b..f8d4c93 100644
--- a/src/server/client.zig
+++ b/src/server/client.zig
@@ -3,7 +3,7 @@ const std = @import("std");
 
 pub const ClientState = struct {
     id: usize,
-    connect: Message.Connect,
+    connect: Message.AllocatedConnect,
     /// Messages that this client is trying to send.
     send_queue: std.Io.Queue(Message),
     send_queue_buffer: [1024]Message = undefined,
@@ -11,17 +11,16 @@ pub const ClientState = struct {
     recv_queue: std.Io.Queue(Message),
     recv_queue_buffer: [1024]Message = undefined,
 
+    io_group: std.Io.Group = .init,
+
     from_client: *std.Io.Reader,
     to_client: *std.Io.Writer,
 
-    write_task: std.Io.Future(void),
-    read_task: std.Io.Future(void),
-
     pub fn init(
         io: std.Io,
         allocator: std.mem.Allocator,
         id: usize,
-        connect: Message.Connect,
+        connect: Message.AllocatedConnect,
         in: *std.Io.Reader,
         out: *std.Io.Writer,
     ) !ClientState {
@@ -32,15 +31,11 @@ pub const ClientState = struct {
             .recv_queue = undefined,
             .from_client = in,
             .to_client = out,
-            .write_task = undefined,
-            .read_task = undefined,
         };
         res.send_queue = .init(&res.send_queue_buffer);
         res.recv_queue = .init(&res.recv_queue_buffer);
-        // res.send_queue = .init(&.{});
-        // res.recv_queue = .init(&.{});
-        res.write_task = try io.concurrent(processWrite, .{ &res, io, out });
-        res.read_task = try io.concurrent(processRead, .{ &res, io, allocator, in });
+        try res.io_group.concurrent(io, processWrite, .{ &res, io, out });
+        try res.io_group.concurrent(io, processRead, .{ &res, io, allocator, in });
 
         return res;
     }
@@ -78,24 +73,30 @@ pub const ClientState = struct {
         allocator: std.mem.Allocator,
         in: *std.Io.Reader,
     ) void {
-        io.sleep(.fromMilliseconds(100), .real) catch @panic("couldn't sleep");
+        // io.sleep(.fromMilliseconds(100), .real) catch @panic("couldn't sleep");
         while (true) {
-            std.debug.print("waiting for message\n", .{});
+            // std.debug.print("waiting for message\n", .{});
             const next_message = Message.next(allocator, in) catch |err| switch (err) {
-                error.EndOfStream => break,
+                error.EndOfStream => {
+                    self.send_queue.putOne(io, .eos) catch {};
+                    break;
+                },
                 else => {
                     std.debug.panic("guh: {any}\n", .{err});
                     break;
-                    // return err;
                 },
             };
-            std.debug.print("got message {}\n", .{next_message});
-            // std.debug.print("queue: {any}\n", .{self.send_queue});
-            self.send_queue.putOneUncancelable(io, next_message); //catch {
-            //     std.debug.print("in catch\n\n\n", .{});
+            std.debug.print("next_message: {any}\n", .{next_message});
 
-            // std.debug.print("queue: {any}\n", .{self.send_queue});
-            // };
+            switch (next_message) {
+                .eos => {
+                    self.send_queue.putOne(io, next_message) catch {};
+                    break;
+                },
+                else => {
+                    self.send_queue.putOne(io, next_message) catch break;
+                },
+            }
         }
 
         std.debug.print("no more messages\n", .{});
@@ -171,21 +172,99 @@ test {
     var from_client_buf: [1024]Message = undefined;
     var from_client_queue: std.Io.Queue(Message) = .init(&from_client_buf);
 
-    while (Message.next(gpa, &from_client)) |msg| {
-        try from_client_queue.putOne(io, msg);
-    } else |_| {}
-
-    for (0..2) |_| {
-        var msg = try from_client_queue.getOne(io);
-        std.debug.print("Message: {any}\n", .{msg});
-        switch (msg) {
-            .connect => |*c| {
-                c.deinit();
-            },
-            else => {},
-        }
+    {
+        while (Message.next(gpa, &from_client)) |msg| {
+            switch (msg) {
+                .eos => {
+                    try from_client_queue.putOne(io, msg);
+                    break;
+                },
+                else => {
+                    try from_client_queue.putOne(io, msg);
+                },
+            }
+        } else |_| {}
+
+        while (from_client_queue.getOne(io)) |msg| {
+            switch (msg) {
+                .eos => {
+                    break;
+                },
+                .connect => |*c| {
+                    std.debug.print("Message: {any}\n", .{msg});
+                    c.deinit();
+                },
+                else => {
+                    std.debug.print("Message: {any}\n", .{msg});
+                },
+            }
+        } else |_| {}
     }
 
+    from_client_queue = .init(&from_client_buf);
+    // Reset the reader to process it again.
+    from_client.seek = 0;
+
+    {
+        const SemiClient = struct {
+            q: std.Io.Queue(Message),
+
+            fn parseClientInput(self: *@This(), ioh: std.Io, in: *std.Io.Reader) void {
+                defer std.debug.print("done parse\n", .{});
+                while (Message.next(gpa, in)) |msg| {
+                    switch (msg) {
+                        .eos => {
+                            self.q.putOne(ioh, msg) catch return;
+                            break;
+                        },
+                        else => {
+                            self.q.putOne(ioh, msg) catch return;
+                        },
+                    }
+                } else |_| {}
+            }
+
+            fn next(self: *@This(), ioh: std.Io) !Message {
+                return self.q.getOne(ioh);
+            }
+
+            fn printAll(self: *@This(), ioh: std.Io) void {
+                defer std.debug.print("done print\n", .{});
+                while (self.next(ioh)) |*msg| {
+                    switch (msg.*) {
+                        .eos => |_| {
+                            break;
+                        },
+                        else => {},
+                    }
+                    std.debug.print("Client msg: {any}\n", .{msg});
+                    switch (msg.*) {
+                        .connect => |c| {
+                            // c.allocator.deinit();
+                            c.deinit();
+                            // @constCast(c).deinit();
+                        },
+                        else => {},
+                    }
+                } else |_| {}
+            }
+        };
+
+        var c: SemiClient = .{ .q = from_client_queue };
+        var group: std.Io.Group = .init;
+        defer group.wait(io);
+
+        group.concurrent(io, SemiClient.printAll, .{ &c, io }) catch {
+            @panic("could not start printAll\n");
+        };
+
+        group.concurrent(io, SemiClient.parseClientInput, .{ &c, io, &from_client }) catch {
+            @panic("could not start printAll\n");
+        };
+    }
+
+    ////////
+
     // const connect = (Message.next(gpa, &from_client) catch unreachable).connect;
 
     // var to_client_alloc: std.Io.Writer.Allocating = .init(gpa);
diff --git a/src/server/main.zig b/src/server/main.zig
index 7f9b9a3..155a455 100644
--- a/src/server/main.zig
+++ b/src/server/main.zig
@@ -6,7 +6,7 @@ const ClientState = @import("./client.zig").ClientState;
 const Server = @This();
 
 info: ServerInfo,
-clients: std.AutoHashMapUnmanaged(usize, ClientState) = .empty,
+clients: std.AutoHashMapUnmanaged(usize, *ClientState) = .empty,
 /// Map of subjects to a map of (client ID => SID)
 subscriptions: std.StringHashMapUnmanaged(std.AutoHashMapUnmanaged(usize, []const u8)) = .empty,
 
@@ -27,8 +27,10 @@ pub fn main(gpa: std.mem.Allocator, server_config: ServerInfo) !void {
 
     var id: usize = 0;
     while (true) : (id +%= 1) {
+        std.debug.print("in server loop\n", .{});
         if (server.clients.contains(id)) continue;
         const stream = try tcp_server.accept(io);
+        std.debug.print("accepted connection\n", .{});
         _ = io.concurrent(handleConnection, .{ &server, gpa, io, id, stream }) catch {
             std.debug.print("could not start concurrent handler for {d}\n", .{id});
             stream.close(io);
@@ -36,7 +38,7 @@ pub fn main(gpa: std.mem.Allocator, server_config: ServerInfo) !void {
     }
 }
 
-fn addClient(server: *Server, allocator: std.mem.Allocator, id: usize, client: ClientState) !void {
+fn addClient(server: *Server, allocator: std.mem.Allocator, id: usize, client: *ClientState) !void {
     // server.clients.lockPointers();
     try server.clients.put(allocator, id, client);
     // server.clients.unlockPointers();
@@ -71,7 +73,7 @@ fn handleConnection(
     const connect = (Message.next(connect_arena.allocator(), in) catch return).connect;
     var client_state: ClientState = try .init(io, allocator, id, connect, in, out);
 
-    try server.addClient(allocator, id, client_state);
+    try server.addClient(allocator, id, &client_state);
     defer server.removeClient(allocator, id);
 
     // defer {
@@ -113,12 +115,12 @@ fn publishMessage(server: *Server, io: std.Io, msg: Message.Pub) !void {
             const client_id = sub.key_ptr.*;
             const sid = sub.value_ptr.*;
 
-            var client = server.clients.getPtr(client_id) orelse {
+            const client = server.clients.getPtr(client_id) orelse {
                 std.debug.print("trying to publish to a client that no longer exists: {d}", .{client_id});
                 continue;
             };
 
-            _ = try client.send(io, .{ .msg = .{
+            _ = try client.*.send(io, .{ .msg = .{
                 .subject = msg.subject,
                 .sid = sid,
                 .reply_to = msg.reply_to,
@@ -141,9 +143,10 @@ pub fn processClient(server: *Server, gpa: std.mem.Allocator, io: std.Io, client
     defer client_state.deinit(gpa);
     std.debug.print("processing client: {d}\n", .{client_state.id});
 
-    while (true) {
-        std.debug.print("awaiting next message from client\n", .{});
-        switch (client_state.next(io)) {
+    std.debug.print("awaiting next message from client\n", .{});
+    while (client_state.next(io)) |msg| {
+        std.debug.print("message from client!: {any}\n", .{msg});
+        switch (msg) {
             .ping => {
                 std.debug.print("got a ping! sending a pong.\n", .{});
                 for (0..5) |_| {
@@ -156,22 +159,28 @@ pub fn processClient(server: *Server, gpa: std.mem.Allocator, io: std.Io, client
                     std.debug.print("could not pong to client {d}\n", .{client_state.id});
                 }
             },
-            .@"pub" => |msg| {
-                try server.publishMessage(io, msg);
-                if (client_state.connect.verbose) {
+            .@"pub" => |@"pub"| {
+                std.debug.print("pub: {any}\n", .{@"pub"});
+                try server.publishMessage(io, @"pub");
+                if (client_state.connect.connect.verbose) {
                     _ = try client_state.send(io, .@"+ok");
                 }
             },
-            .sub => |msg| {
-                try server.subscribe(gpa, client_state.id, msg);
+            .sub => |sub| {
+                try server.subscribe(gpa, client_state.id, sub);
+            },
+            .eos => {
+                client_state.io_group.wait(io);
+                break;
             },
-            else => |msg| {
-                std.debug.panic("Unimplemented message: {any}\n", .{msg});
+            else => |e| {
+                std.debug.panic("Unimplemented message: {any}\n", .{e});
             },
         }
 
         std.debug.print("processed message from client\n", .{});
-    }
+        std.debug.print("awaiting next message from client\n", .{});
+    } else |_| {}
 
     // while (!io.cancelRequested()) {
     //     if (client_state.send_queue.getOne(io)) |msg| {
diff --git a/src/server/message_parser.zig b/src/server/message_parser.zig
index e791d14..5669d2a 100644
--- a/src/server/message_parser.zig
+++ b/src/server/message_parser.zig
@@ -13,6 +13,7 @@ pub const MessageType = enum {
     pong,
     @"+ok",
     @"-err",
+    eos,
 
     fn parseMemEql(input: []const u8) ?MessageType {
         // if (std.mem.eql(u8, "INFO", input)) return .info;
@@ -44,6 +45,9 @@ pub const Message = union(MessageType) {
     pong,
     @"+ok": void,
     @"-err": void,
+    // Not an actual NATS message, but used to signal end of stream was reached in the input,
+    // and we should close the reader.
+    eos: void,
     pub const ServerInfo = struct {
         /// The unique identifier of the NATS server.
         server_id: []const u8,
@@ -75,7 +79,7 @@ pub const Message = union(MessageType) {
         allocator: std.heap.ArenaAllocator,
         connect: Connect,
 
-        pub fn deinit(self: *AllocatedConnect) void {
+        pub fn deinit(self: AllocatedConnect) void {
             self.allocator.deinit();
         }
     };
@@ -166,7 +170,14 @@ pub const Message = union(MessageType) {
                 try operation_string.appendBounded(byte);
                 try in.discardAll(1);
             } else break;
-        } else |err| return err;
+        } else |err| switch (err) {
+            error.EndOfStream => {
+                return .{ .eos = {} };
+            },
+            else => {
+                return err;
+            },
+        }
 
         const operation = parse(operation_string.items) orelse {
             return error.InvalidOperation;