Cleanup for deinitializing

2 files changed, 40 insertions, 7 deletions

diff --git a/src/server/client.zig b/src/server/client.zig
index 5c8d0cd..ff7e5c1 100644
--- a/src/server/client.zig
+++ b/src/server/client.zig
@@ -29,7 +29,20 @@ pub fn start(self: *Client, io: std.Io, alloc: std.mem.Allocator, queue: *std.Io
     while (true) {
         const len = try queue.get(io, &msgs, 1);
         std.debug.assert(len <= msgs.len);
-        for (msgs[0..len]) |msg| {
+        for (0..len) |i| {
+            const msg = msgs[i];
+            defer switch (msg) {
+                .msg => |m| m.deinit(alloc),
+                else => {},
+            };
+            errdefer {
+                for (msgs[i + 1 .. len]) |mg| switch (mg) {
+                    .msg => |m| {
+                        m.deinit(alloc);
+                    },
+                    else => {},
+                };
+            }
             switch (msg) {
                 .@"+ok" => {
                     _ = try self.to_client.write("+OK\r\n");
@@ -43,7 +56,6 @@ pub fn start(self: *Client, io: std.Io, alloc: std.mem.Allocator, queue: *std.Io
                     _ = try self.to_client.write("\r\n");
                 },
                 .msg => |m| {
-                    defer m.deinit(alloc);
                     try self.to_client.print(
                         "MSG {s} {s} {s} {d}\r\n{s}\r\n",
                         .{
diff --git a/src/server/main.zig b/src/server/main.zig
index dab0f0a..f3728ea 100644
--- a/src/server/main.zig
+++ b/src/server/main.zig
@@ -25,6 +25,18 @@ fn handleSigInt(sig: std.os.linux.SIG) callconv(.c) void {
     keep_running.store(false, .monotonic);
 }
 
+pub fn deinit(server: *Server, io: std.Io, alloc: std.mem.Allocator) void {
+    server.subs_lock.lockUncancelable(io);
+    defer server.subs_lock.unlock(io);
+    for (server.subscriptions.items) |sub| {
+        alloc.free(sub.sid);
+        alloc.free(sub.subject);
+    }
+    server.subscriptions.shrinkAndFree(alloc, 0);
+
+    server.clients.clearAndFree(alloc);
+}
+
 pub fn main(alloc: std.mem.Allocator, server_config: ServerInfo) !void {
     // Configure the signal action
     const act = std.posix.Sigaction{
@@ -42,14 +54,15 @@ pub fn main(alloc: std.mem.Allocator, server_config: ServerInfo) !void {
         defer _ = dba.deinit();
         const gpa = if (builtin.mode == .Debug or builtin.mode == .ReleaseSafe) dba.allocator() else alloc;
 
-        var server: Server = .{
-            .info = server_config,
-        };
-
         var threaded: std.Io.Threaded = .init(gpa, .{});
         defer threaded.deinit();
         const io = threaded.io();
 
+        var server: Server = .{
+            .info = server_config,
+        };
+        defer server.deinit(io, gpa);
+
         var server_task = try io.concurrent(start, .{ &server, io, gpa });
         defer server_task.cancel(io) catch {};
 
@@ -136,7 +149,15 @@ fn handleConnection(
 
     var qbuf: [8]Message = undefined;
     var queue: std.Io.Queue(Message) = .init(&qbuf);
-    defer queue.close(io);
+    defer {
+        queue.close(io);
+        while (queue.getOne(io)) |msg| {
+            switch (msg) {
+                .msg => |m| m.deinit(server_allocator),
+                else => {},
+            }
+        } else |_| {}
+    }
 
     var client_task = try io.concurrent(Client.start, .{ &client, io, server_allocator, &queue });
     defer client_task.cancel(io) catch {};