Reorganized things

1 files changed, 66 insertions, 67 deletions

diff --git a/src/server/main.zig b/src/server/main.zig
index a78bab2..e5f6e59 100644
--- a/src/server/main.zig
+++ b/src/server/main.zig
@@ -14,9 +14,11 @@ const Subscription = struct {
 info: ServerInfo,
 clients: std.AutoHashMapUnmanaged(usize, *ClientState) = .empty,
 
-subs_lock: std.Thread.Mutex = .{},
+subs_lock: std.Io.Mutex = .init,
 subscriptions: std.ArrayList(Subscription) = .empty,
 
+msg_queue: std.Io.Queue(Message.Pub),
+
 var keep_running = std.atomic.Value(bool).init(true);
 
 fn handleSigInt(sig: std.os.linux.SIG) callconv(.c) void {
@@ -35,14 +37,22 @@ pub fn main(gpa: std.mem.Allocator, server_config: ServerInfo) !void {
     // // Register the handler for SIGINT (Ctrl+C)
     // std.posix.sigaction(std.posix.SIG.INT, &act, null);
 
+    // 64 mb buffer for messages
+    const queue_buf = try gpa.alloc(Message.Pub, 1024 * 1024);
+    defer gpa.free(queue_buf);
+
     var server: Server = .{
         .info = server_config,
+        .msg_queue = .init(queue_buf),
     };
 
-    var threaded: std.Io.Threaded = .init(gpa);
+    var threaded: std.Io.Threaded = .init(gpa, .{});
     defer threaded.deinit();
     const io = threaded.io();
 
+    var msgProcess = try io.concurrent(processMsgs, .{ &server, io, gpa });
+    defer msgProcess.cancel(io) catch {};
+
     var tcp_server = try std.Io.net.IpAddress.listen(try std.Io.net.IpAddress.parse(
         server.info.host,
         server.info.port,
@@ -65,15 +75,37 @@ pub fn main(gpa: std.mem.Allocator, server_config: ServerInfo) !void {
     std.debug.print("Exiting gracefully\n", .{});
 }
 
+fn processMsgs(server: *Server, io: std.Io, alloc: std.mem.Allocator) !void {
+    while (true) {
+        const msg = try server.msg_queue.getOne(io);
+        defer msg.deinit(alloc);
+
+        for (server.subscriptions.items) |subscription| {
+            if (subjectMatches(subscription.subject, msg.subject)) {
+                const client = server.clients.get(subscription.client_id) orelse {
+                    std.debug.print("trying to publish to a client that no longer exists: {d}", .{subscription.client_id});
+                    continue;
+                };
+                try client.send(io, .{ .msg = .{
+                    .subject = msg.subject,
+                    .sid = subscription.sid,
+                    .reply_to = msg.reply_to,
+                    .payload = msg.payload,
+                } });
+            }
+        }
+    }
+}
+
 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();
 }
 
-fn removeClient(server: *Server, allocator: std.mem.Allocator, id: usize) void {
-    server.subs_lock.lock();
-    defer server.subs_lock.unlock();
+fn removeClient(server: *Server, io: std.Io, allocator: std.mem.Allocator, id: usize) void {
+    server.subs_lock.lockUncancelable(io);
+    defer server.subs_lock.unlock(io);
     _ = server.clients.remove(id);
     const len = server.subscriptions.items.len;
     for (0..len) |i| {
@@ -95,72 +127,49 @@ fn handleConnection(
 ) !void {
     var client_allocator: std.heap.DebugAllocator(.{}) = .init;
     client_allocator.backing_allocator = server_allocator;
-    defer {
-        std.debug.print("deinitializing debug allocator\n", .{});
-        _ = client_allocator.deinit();
-    }
+    defer _ = client_allocator.deinit();
+
     const allocator = client_allocator.allocator();
     defer stream.close(io);
     var w_buffer: [4096]u8 = undefined;
     var writer = stream.writer(io, &w_buffer);
     const out = &writer.interface;
 
-    std.debug.print("out pointer in client handler: {*}\n", .{out});
-
     var r_buffer: [8192]u8 = undefined;
     var reader = stream.reader(io, &r_buffer);
     const in = &reader.interface;
 
-    @import("./client.zig").writeInfo(out, server.info) catch return;
+    var client_state: ClientState = .init(null, in, out);
+    try client_state.send(io, .{ .info = server.info });
+
     var connect_arena: std.heap.ArenaAllocator = .init(allocator);
     defer connect_arena.deinit();
-    const connect = (Message.next(connect_arena.allocator(), in) catch return).connect;
-    var client_state: ClientState = try .init(connect, allocator, in, out);
-    try client_state.start(io);
-    defer client_state.deinit(io, allocator);
+    client_state.connect = (Message.next(connect_arena.allocator(), in) catch return).connect;
 
     try server.addClient(server_allocator, id, &client_state);
-    defer server.removeClient(server_allocator, id);
+    defer server.removeClient(io, server_allocator, id);
 
-    var msg_arena: std.heap.ArenaAllocator = .init(allocator);
-    defer msg_arena.deinit();
-    const msg_allocator = msg_arena.allocator();
-
-    while (client_state.next(msg_allocator)) |msg| {
-        defer _ = msg_arena.reset(.retain_capacity);
+    // Messages are owned by the server after they are received from the client
+    while (client_state.next(server_allocator)) |msg| {
         switch (msg) {
             .ping => {
                 // Respond to ping with pong.
                 try client_state.send(io, .pong);
             },
             .@"pub" => |pb| {
-                defer {
-                    msg_allocator.free(pb.payload);
-                    msg_allocator.free(pb.subject);
-                    if (pb.reply_to) |r| {
-                        msg_allocator.free(r);
-                    }
-                }
+                // Do not free pb, server.publishMessage takes ownership.
                 try server.publishMessage(io, pb);
-                if (client_state.connect.connect.verbose) {
-                    try client_state.send(io, .@"+ok");
+                if (client_state.connect) |c| {
+                    if (c.verbose) {
+                        try client_state.send(io, .@"+ok");
+                    }
                 }
             },
             .sub => |sub| {
-                defer {
-                    msg_allocator.free(sub.subject);
-                    msg_allocator.free(sub.sid);
-                    if (sub.queue_group) |q| {
-                        msg_allocator.free(q);
-                    }
-                }
-                try server.subscribe(server_allocator, id, sub);
+                try server.subscribe(io, server_allocator, id, sub);
             },
             .unsub => |unsub| {
-                defer {
-                    msg_allocator.free(unsub.sid);
-                }
-                try server.unsubscribe(server_allocator, id, unsub);
+                try server.unsubscribe(io, server_allocator, id, unsub);
             },
             else => |e| {
                 std.debug.panic("Unimplemented message: {any}\n", .{e});
@@ -186,36 +195,26 @@ fn subjectMatches(expected: []const u8, actual: []const u8) bool {
 }
 
 fn publishMessage(server: *Server, io: std.Io, msg: Message.Pub) !void {
-    for (server.subscriptions.items) |subscription| {
-        if (subjectMatches(subscription.subject, msg.subject)) {
-            const client = server.clients.get(subscription.client_id) orelse {
-                std.debug.print("trying to publish to a client that no longer exists: {d}", .{subscription.client_id});
-                continue;
-            };
-            try client.send(io, .{ .msg = .{
-                .subject = msg.subject,
-                .sid = subscription.sid,
-                .reply_to = msg.reply_to,
-                .payload = msg.payload,
-            } });
-        }
-    }
+    try server.msg_queue.putOne(io, .{
+        .payload = msg.payload,
+        .reply_to = msg.reply_to,
+        .subject = msg.subject,
+    });
 }
 
-fn subscribe(server: *Server, gpa: std.mem.Allocator, id: usize, msg: Message.Sub) !void {
-    std.debug.print("Recieved SUBSCRIBE message: {any}\n\n", .{msg});
-    server.subs_lock.lock();
-    defer server.subs_lock.unlock();
+fn subscribe(server: *Server, io: std.Io, gpa: std.mem.Allocator, id: usize, msg: Message.Sub) !void {
+    try server.subs_lock.lock(io);
+    defer server.subs_lock.unlock(io);
     try server.subscriptions.append(gpa, .{
-        .subject = try gpa.dupe(u8, msg.subject),
+        .subject = msg.subject,
         .client_id = id,
-        .sid = try gpa.dupe(u8, msg.sid),
+        .sid = msg.sid,
     });
 }
 
-fn unsubscribe(server: *Server, gpa: std.mem.Allocator, id: usize, msg: Message.Unsub) !void {
-    server.subs_lock.lock();
-    defer server.subs_lock.unlock();
+fn unsubscribe(server: *Server, io: std.Io, gpa: std.mem.Allocator, id: usize, msg: Message.Unsub) !void {
+    try server.subs_lock.lock(io);
+    defer server.subs_lock.unlock(io);
     const len = server.subscriptions.items.len;
     for (0..len) |i| {
         const sub = server.subscriptions.items[len - i - 1];