made some progress on subscriptions

3 files changed, 124 insertions, 153 deletions

diff --git a/src/server/client.zig b/src/server/client.zig
index 458bbbb..32a53b7 100644
--- a/src/server/client.zig
+++ b/src/server/client.zig
@@ -66,12 +66,14 @@ pub const ClientState = struct {
                 },
             }
         }
+        self.task.cancel(io);
     }
 
     pub fn deinit(self: *ClientState, io: std.Io, allocator: std.mem.Allocator) void {
         self.task.cancel(io);
         self.connect.deinit();
-        allocator.destroy(self.recv_queue);
+        _ = allocator;
+        // allocator.destroy(self.recv_queue);
     }
 
     /// Return true if the value was put in the clients buffer to process, else false.
@@ -95,14 +97,11 @@ fn writeOk(out: *std.Io.Writer) !void {
 }
 
 fn writePong(out: *std.Io.Writer) !void {
-    std.debug.print("out pointer: {*}\n", .{out});
-    std.debug.print("writing pong\n", .{});
     _ = try out.write("PONG\r\n");
     try out.flush();
 }
 
 pub fn writeInfo(out: *std.Io.Writer, info: Message.ServerInfo) !void {
-    std.debug.print("writing info: {any}\n", .{info});
     _ = try out.write("INFO ");
     try std.json.Stringify.value(info, .{}, out);
     _ = try out.write("\r\n");
diff --git a/src/server/main.zig b/src/server/main.zig
index 2ebf96d..e1f9891 100644
--- a/src/server/main.zig
+++ b/src/server/main.zig
@@ -53,11 +53,18 @@ fn removeClient(server: *Server, allocator: std.mem.Allocator, id: usize) void {
 
 fn handleConnection(
     server: *Server,
-    allocator: std.mem.Allocator,
+    server_allocator: std.mem.Allocator,
     io: std.Io,
     id: usize,
     stream: std.Io.net.Stream,
 ) !void {
+    _ = server_allocator;
+    var client_allocator: std.heap.DebugAllocator(.{}) = .init;
+    defer {
+        std.debug.print("deinitializing debug allocator\n", .{});
+        _ = client_allocator.deinit();
+    }
+    const allocator = client_allocator.allocator();
     defer stream.close(io);
     var w_buffer: [1024]u8 = undefined;
     var writer = stream.writer(io, &w_buffer);
@@ -79,155 +86,46 @@ fn handleConnection(
     try server.addClient(allocator, id, &client_state);
     defer server.removeClient(allocator, id);
 
-    // defer {
-    //     server.clients.lockPointers();
-    //     server.clients.remove(allocator, id);
-    //     server.clients.unlockPointers();
-    //     server.subscriptions.lockPointers();
-    //     var sub_iter = server.subscriptions.iterator();
-    //     var to_free: std.ArrayList(usize) = .empty;
-    //     defer to_free.deinit(allocator);
-    //     while (sub_iter.next()) |sub| {
-    //         while (std.simd.firstIndexOfValue(sub.value_ptr.*, id)) |i| {
-    //             sub.value_ptr.*.orderedRemove(i);
-    //         }
-    //         if (sub.value_ptr.items.len == 0) {
-    //             to_free.append(allocator, sub.index);
-    //         }
-    //     }
-    //     server.subscriptions.orderedRemoveAtMany(allocator, to_free.items);
-    //     server.subscriptions.unlockPointers();
-    // }
-
-    {
-        defer std.debug.print("done processing client??\n", .{});
-        std.debug.print("processing client: {d}\n", .{client_state.id});
-
-        std.debug.print("awaiting next message from client\n", .{});
-        while (client_state.next(allocator)) |msg| {
-            std.debug.print("message from client!: {any}\n", .{msg});
-            switch (msg) {
-                .ping => {
-                    std.debug.print("got a ping! sending a pong.\n", .{});
-
-                    std.debug.print("recv queue in server loop: {*}\n", .{&client_state.recv_queue});
-                    // @import("./client.zig").writePong(out) catch return;
-                    for (0..5) |_| {
-                        if (try client_state.send(io, .pong)) {
-                            std.debug.print("sent pong\n", .{});
-                            break;
-                        }
-                        std.debug.print("trying to send a pong again.\n", .{});
-                    } else {
-                        std.debug.print("could not pong to client {d}\n", .{client_state.id});
+    while (client_state.next(allocator)) |msg| {
+        std.debug.print("message from client: {any}\n", .{msg});
+        switch (msg) {
+            .ping => {
+                // Respond to ping with pong.
+                for (0..5) |_| {
+                    if (try client_state.send(io, .pong)) {
+                        break;
                     }
-                },
-                .@"pub" => |@"pub"| {
-                    std.debug.print("pub: {any}\n", .{@"pub"});
-                    try server.publishMessage(io, @"pub");
-                    if (client_state.connect.connect.verbose) {
-                        std.debug.print("server IS sending +ok\n", .{});
-                        _ = try client_state.send(io, .@"+ok");
-                    } else {
-                        std.debug.print("server NOT sending +ok\n", .{});
-                    }
-                },
-                .sub => |sub| {
-                    try server.subscribe(allocator, client_state.id, sub);
-                },
-                .eos => {
-                    break;
-                },
-                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| {
-        //         switch (msg) {
-        //             // Respond to ping with pong.
-        //             .ping => {
-        //                 try client_state.recv_queue.putOne(io, .pong);
-        //             },
-        //             .@"pub" => |p| {
-        //                 std.debug.print("subs (in pub): {any}\n", .{server.subscriptions});
-        //                 std.debug.print("subs size: {d}\n", .{server.subscriptions.size});
-        //                 std.debug.print("subs subjects:\n", .{});
-        //                 var key_iter = server.subscriptions.keyIterator();
-        //                 while (key_iter.next()) |k| {
-        //                     std.debug.print("- {s}\n", .{k.*});
-        //                 } else std.debug.print("<none>\n", .{});
-        //                 std.debug.print("pub subject: '{s}'\n", .{p.subject});
-        //                 std.debug.print("pub: {any}\n", .{p});
-        //                 errdefer _ = client_state.recv_queue.put(io, &.{.@"-err"}, 1) catch {};
-        //                 // Just publishing to exact matches right now.
-        //                 // TODO: Publish to pattern matching subjects.
-        //                 if (server.subscriptions.get(p.subject)) |subs| {
-        //                     var subs_iter = subs.iterator();
-        //                     while (subs_iter.next()) |sub| {
-        //                         var client = server.clients.get(sub.key_ptr.*) orelse std.debug.panic("Trying to pub to a client that doesn't exist!\n", .{});
-        //                         std.debug.print("{d} is pubbing to {d}\n", .{ client_state.id, client.id });
-        //                         if ((try client.recv_queue.put(
-        //                             io,
-        //                             &.{
-        //                                 .{
-        //                                     .msg = .{
-        //                                         .subject = p.subject,
-        //                                         .sid = sub.value_ptr.*,
-        //                                         .reply_to = p.reply_to,
-        //                                         .payload = p.payload,
-        //                                     },
-        //                                 },
-        //                             },
-        //                             0,
-        //                         )) > 0) {
-        //                             std.debug.print("published message!\n", .{});
-        //                         } else {
-        //                             std.debug.print("skipped publishing for some reason\n", .{});
-        //                         }
-        //                     }
-        //                     try client_state.recv_queue.putOne(io, .@"+ok");
-        //                 } else {
-        //                     std.debug.print("no subs with subject\n", .{});
-        //                 }
-        //             },
-        //             .sub => |s| {
-        //                 var subscribers = try server.subscriptions.getOrPut(gpa, s.subject);
-        //                 if (!subscribers.found_existing) {
-        //                     subscribers.value_ptr.* = .empty;
-        //                 }
-        //                 try subscribers.value_ptr.*.put(gpa, client_state.id, s.sid);
-
-        //                 std.debug.print("subs: {any}\n", .{server.subscriptions});
-        //             },
-        //             .info => |info| {
-        //                 std.debug.panic("got an info message? : {any}\n", .{info});
-        //             },
-        //             else => |m| {
-        //                 std.debug.panic("Unimplemented: {any}\n", .{m});
-        //             },
-        //         }
-        //     } else |err| return err;
-        // }
+                } else {}
+            },
+            .@"pub" => |@"pub"| {
+                std.debug.print("pub: {any}\n", .{@"pub"});
+                try server.publishMessage(io, @"pub");
+                if (client_state.connect.connect.verbose) {
+                    std.debug.print("server IS sending +ok\n", .{});
+                    _ = try client_state.send(io, .@"+ok");
+                } else {
+                    std.debug.print("server NOT sending +ok\n", .{});
+                }
+            },
+            .sub => |sub| {
+                try server.subscribe(allocator, client_state.id, sub);
+            },
+            .unsub => |unsub| {
+                try server.unsubscribe(client_state.id, unsub);
+            },
+            .eos => {
+                break;
+            },
+            else => |e| {
+                std.debug.panic("Unimplemented message: {any}\n", .{e});
+            },
+        }
 
-        // while (true) {
-        //     switch (next_message) {
-        //         .connect => |connect| {
-        //             std.debug.panic("Connection message after already connected: {any}\n", .{connect});
-        //         },
-        //         .ping => try writePong(out),
-        //         .@"pub" => try writeOk(out),
-        //         else => |msg| std.debug.panic("Message type not implemented: {any}\n", .{msg}),
-        //     }
-        // }
-    }
+        std.debug.print("processed message from client\n", .{});
+        std.debug.print("awaiting next message from client\n", .{});
+    } else |_| {}
 
-    client_state.task.await(io);
+    // client_state.task.await(io);
 }
 
 // // Result is owned by the caller
@@ -267,6 +165,25 @@ fn subscribe(server: *Server, gpa: std.mem.Allocator, id: usize, msg: Message.Su
     try server.subscriptions.put(gpa, msg.subject, subs_for_subject);
 }
 
+fn unsubscribe(server: *Server, id: usize, msg: Message.Unsub) !void {
+    // Get the subscription in subscriptions by looping over all the subjects,
+    // and getting the SID for that subject for the current client ID.
+    // If the SID matches, remove the kv for the client ID from subscriptions for that subject.
+    // If the value for that subject is empty, remove the subject.
+    var subscriptions_iter = server.subscriptions.iterator();
+    while (subscriptions_iter.next()) |*subs_for_sub| {
+        if (subs_for_sub.value_ptr.get(id)) |client_sub| {
+            if (std.mem.eql(u8, client_sub, msg.sid)) {
+                _ = subs_for_sub.value_ptr.*.remove(id);
+                if (subs_for_sub.value_ptr.count() == 0) {
+                    _ = server.subscriptions.remove(subs_for_sub.key_ptr.*);
+                }
+                break;
+            }
+        }
+    }
+}
+
 pub fn createId() []const u8 {
     return "SERVERID";
 }
diff --git a/src/server/message_parser.zig b/src/server/message_parser.zig
index 5669d2a..0ff3dd4 100644
--- a/src/server/message_parser.zig
+++ b/src/server/message_parser.zig
@@ -38,7 +38,7 @@ pub const Message = union(MessageType) {
     @"pub": Pub,
     hpub: void,
     sub: Sub,
-    unsub: void,
+    unsub: Unsub,
     msg: Msg,
     hmsg: void,
     ping,
@@ -117,6 +117,12 @@ pub const Message = union(MessageType) {
         /// A unique alphanumeric subscription ID, generated by the client.
         sid: []const u8,
     };
+    pub const Unsub = struct {
+        /// The unique alphanumeric subscription ID of the subject to unsubscribe from.
+        sid: []const u8,
+        /// A number of messages to wait for before automatically unsubscribing.
+        max_msgs: ?usize = null,
+    };
     pub const Msg = struct {
         subject: []const u8,
         sid: []const u8,
@@ -225,7 +231,7 @@ pub const Message = union(MessageType) {
                         }
                     } else return error.InvalidStream;
 
-                    break :blk try std.fmt.parseUnsigned(u64, byte_count_list.items, 10);
+                    break :blk try std.fmt.parseUnsigned(usize, byte_count_list.items, 10);
                 };
 
                 const payload = blk: {
@@ -280,6 +286,55 @@ pub const Message = union(MessageType) {
                     },
                 };
             },
+            .unsub => {
+                std.debug.assert(std.ascii.isWhitespace(try in.takeByte()));
+                // Parse byte count
+                const sid = blk: {
+                    var acc: std.ArrayList(u8) = try .initCapacity(alloc, 8);
+                    while (in.peekByte() catch null) |byte| {
+                        if (std.ascii.isWhitespace(byte)) break;
+                        try acc.append(alloc, byte);
+                        in.toss(1);
+                    } else return error.InvalidStream;
+                    break :blk try acc.toOwnedSlice(alloc);
+                };
+
+                if ((try in.peekByte()) == '\r') {
+                    std.debug.assert(std.mem.eql(u8, try in.take(2), "\r\n"));
+                    return .{
+                        .unsub = .{
+                            .sid = sid,
+                        },
+                    };
+                } else if (std.ascii.isWhitespace(try in.peekByte())) {
+                    in.toss(1);
+                    const max_msgs = blk: {
+                        var max_msgs_list: std.ArrayList(u8) = try .initCapacity(alloc, 64);
+                        while (in.takeByte() catch null) |byte| {
+                            if (std.ascii.isWhitespace(byte)) {
+                                std.debug.assert(byte == '\r');
+                                std.debug.assert(try in.takeByte() == '\n');
+                                break;
+                            }
+
+                            if (std.ascii.isDigit(byte)) {
+                                try max_msgs_list.append(alloc, byte);
+                            } else {
+                                return error.InvalidStream;
+                            }
+                        } else return error.InvalidStream;
+
+                        break :blk try std.fmt.parseUnsigned(usize, max_msgs_list.items, 10);
+                    };
+
+                    return .{
+                        .unsub = .{
+                            .sid = sid,
+                            .max_msgs = max_msgs,
+                        },
+                    };
+                } else return error.InvalidStream;
+            },
             else => |msg| std.debug.panic("Not implemented: {}\n", .{msg}),
         }
     }