switch to uuids for clients

1 files changed, 70 insertions, 54 deletions

diff --git a/src/Server.zig b/src/Server.zig
index d183593..c45db19 100644
--- a/src/Server.zig
+++ b/src/Server.zig
@@ -27,7 +27,7 @@ const builtin = @import("builtin");
 
 const Subscription = struct {
     subject: []const u8,
-    client_id: usize,
+    client_id: u128,
     sid: []const u8,
     queue_group: ?[]const u8,
     queue_lock: *Mutex,
@@ -38,6 +38,14 @@ const Subscription = struct {
         alloc.free(self.sid);
         if (self.queue_group) |g| alloc.free(g);
     }
+
+    fn send(self: *Subscription, io: Io, bytes: []const []const u8) !void {
+        try self.queue_lock.lock(io);
+        defer self.queue_lock.unlock(io);
+        for (bytes) |chunk| {
+            try self.queue.putAll(io, chunk);
+        }
+    }
 };
 
 const eql = std.mem.eql;
@@ -45,20 +53,18 @@ const log = std.log.scoped(.zits);
 const panic = std.debug.panic;
 
 info: ServerInfo,
-clients: AutoHashMapUnmanaged(usize, *Client) = .empty,
 
 subs_lock: Mutex = .init,
 subscriptions: ArrayList(Subscription) = .empty,
 
 pub fn deinit(server: *Server, io: Io, alloc: Allocator) void {
     server.subs_lock.lockUncancelable(io);
-    defer server.subs_lock.unlock(io);
     for (server.subscriptions.items) |sub| {
         sub.deinit(alloc);
     }
     // TODO drain subscription queues
     server.subscriptions.deinit(alloc);
-    server.clients.deinit(alloc);
+    server.subs_lock.unlock(io);
 }
 
 pub fn start(server: *Server, io: Io, gpa: Allocator) !void {
@@ -79,43 +85,40 @@ pub fn start(server: *Server, io: Io, gpa: Allocator) !void {
     const read_buffer_size, const write_buffer_size = getBufferSizes(io);
     log.debug("read buf: {d} write buf: {d}", .{ read_buffer_size, write_buffer_size });
 
-    var id: usize = 0;
-    while (true) : (id +%= 1) {
-        if (server.clients.contains(id)) continue;
+    const rand_io: std.Random.IoSource = .{ .io = io };
+    const rand: std.Random = rand_io.interface();
+
+    var id = rand.int(u128);
+    while (true) : (id = rand.int(u128)) {
         log.debug("Accepting next client", .{});
         const stream = try tcp_server.accept(io);
-        log.debug("Accepted connection {d}", .{id});
+        log.debug("Accepted connection {s}", .{idToStr(id)});
         _ = client_group.concurrent(io, handleConnectionInfallible, .{
             server,
             gpa,
             io,
+            rand,
             id,
             stream,
             read_buffer_size,
             write_buffer_size,
         }) catch {
-            log.err("Could not start concurrent handler for {d}", .{id});
+            log.err("Could not start concurrent handler for {s}", .{idToStr(id)});
             stream.close(io);
         };
     }
 }
 
-fn addClient(server: *Server, allocator: Allocator, id: usize, client: *Client) !void {
-    try server.clients.put(allocator, id, client);
-}
-
-fn removeClient(server: *Server, io: Io, allocator: Allocator, id: usize) void {
+fn removeClient(server: *Server, io: Io, allocator: Allocator, id: u128) void {
     server.subs_lock.lockUncancelable(io);
     defer server.subs_lock.unlock(io);
-    if (server.clients.remove(id)) {
-        const len = server.subscriptions.items.len;
-        for (0..len) |from_end| {
-            const i = len - from_end - 1;
-            const sub = server.subscriptions.items[i];
-            if (sub.client_id == id) {
-                sub.deinit(allocator);
-                _ = server.subscriptions.swapRemove(i);
-            }
+    const len = server.subscriptions.items.len;
+    for (0..len) |from_end| {
+        const i = len - from_end - 1;
+        const sub = server.subscriptions.items[i];
+        if (sub.client_id == id) {
+            sub.deinit(allocator);
+            _ = server.subscriptions.swapRemove(i);
         }
     }
 }
@@ -124,15 +127,16 @@ fn handleConnectionInfallible(
     server: *Server,
     server_allocator: Allocator,
     io: Io,
-    id: usize,
+    rand: std.Random,
+    id: u128,
     stream: Stream,
     r_buf_size: usize,
     w_buf_size: usize,
 ) !void {
-    handleConnection(server, server_allocator, io, id, stream, r_buf_size, w_buf_size) catch |err| switch (err) {
+    handleConnection(server, server_allocator, io, rand, id, stream, r_buf_size, w_buf_size) catch |err| switch (err) {
         error.Canceled => return error.Canceled,
-        error.ClientDisconnected => log.debug("Client {d} disconnected", .{id}),
-        else => log.err("Failed processing client {d}: {t}", .{ id, err }),
+        error.ClientDisconnected => log.debug("Client {s} disconnected", .{idToStr(id)}),
+        else => log.err("Failed processing client {s}: {t}", .{ idToStr(id), err }),
     };
 }
 
@@ -140,7 +144,8 @@ fn handleConnection(
     server: *Server,
     server_allocator: Allocator,
     io: Io,
-    id: usize,
+    rand: std.Random,
+    id: u128,
     stream: Stream,
     r_buf_size: usize,
     w_buf_size: usize,
@@ -176,8 +181,6 @@ fn handleConnection(
     // Create client
     var client: Client = .init(null, &recv_queue, in, out);
     defer client.deinit(server_allocator);
-
-    try server.addClient(server_allocator, id, &client);
     defer server.removeClient(io, server_allocator, id);
 
     // Do initial handshake with client
@@ -201,7 +204,7 @@ fn handleConnection(
             .PUB => {
                 @branchHint(.likely);
                 // log.debug("received a pub msg", .{});
-                server.publishMessage(io, server_allocator, &client, .@"pub") catch |err| switch (err) {
+                server.publishMessage(io, rand, server_allocator, &client, .@"pub") catch |err| switch (err) {
                     error.WriteFailed => return writer.err.?,
                     error.ReadFailed => return reader.err.?,
                     error.EndOfStream => return error.ClientDisconnected,
@@ -210,7 +213,7 @@ fn handleConnection(
             },
             .HPUB => {
                 @branchHint(.likely);
-                server.publishMessage(io, server_allocator, &client, .hpub) catch |err| switch (err) {
+                server.publishMessage(io, rand, server_allocator, &client, .hpub) catch |err| switch (err) {
                     error.WriteFailed => return writer.err.?,
                     error.ReadFailed => return reader.err.?,
                     error.EndOfStream => return error.ClientDisconnected,
@@ -225,7 +228,7 @@ fn handleConnection(
                 };
             },
             .UNSUB => {
-                server.unsubscribe(io, server_allocator, client, id) catch |err| switch (err) {
+                server.unsubscribe(io, server_allocator, client.from_client, id) catch |err| switch (err) {
                     error.ReadFailed => return reader.err.?,
                     error.EndOfStream => return error.ClientDisconnected,
                     else => |e| return e,
@@ -289,6 +292,7 @@ test subjectMatches {
 fn publishMessage(
     server: *Server,
     io: Io,
+    rand: std.Random,
     alloc: Allocator,
     source_client: *Client,
     comptime pub_or_hpub: enum { @"pub", hpub },
@@ -314,17 +318,16 @@ fn publishMessage(
 
     try server.subs_lock.lock(io);
     defer server.subs_lock.unlock(io);
+
     var published_queue_groups: ArrayList([]const u8) = .empty;
     defer published_queue_groups.deinit(alloc);
-    var published_queue_sub_idxs: ArrayList(usize) = .empty;
-    defer published_queue_sub_idxs.deinit(alloc);
 
     var line_writer_allocating: std.Io.Writer.Allocating = .init(alloc);
     defer line_writer_allocating.deinit();
     var line_writer = &line_writer_allocating.writer;
 
     subs: for (0..server.subscriptions.items.len) |i| {
-        const subscription = server.subscriptions.items[i];
+        var subscription = server.subscriptions.items[i];
         if (subjectMatches(subscription.subject, msg.subject)) {
             if (subscription.queue_group) |sg| {
                 for (published_queue_groups.items) |g| {
@@ -334,9 +337,6 @@ fn publishMessage(
                 }
                 // Don't republish to the same queue
                 try published_queue_groups.append(alloc, sg);
-                // Move this index to the end of the subscription list,
-                // to prioritize other subscriptions in the queue next time.
-                try published_queue_sub_idxs.append(alloc, i);
             }
 
             line_writer_allocating.clearRetainingCapacity();
@@ -357,17 +357,11 @@ fn publishMessage(
             }
             try line_writer.print("{d}\r\n", .{msg.payload.len - 2});
 
-            try subscription.queue_lock.lock(io);
-            defer subscription.queue_lock.unlock(io);
-            try subscription.queue.putAll(io, line_writer.buffered());
-            try subscription.queue.putAll(io, msg.payload);
+            try subscription.send(io, &.{ line_writer.buffered(), msg.payload });
         }
     }
 
-    for (0..published_queue_sub_idxs.items.len) |from_end| {
-        const i = published_queue_sub_idxs.items.len - from_end - 1;
-        server.subscriptions.appendAssumeCapacity(server.subscriptions.orderedRemove(i));
-    }
+    rand.shuffle(Subscription, server.subscriptions.items);
 }
 
 fn subscribe(
@@ -375,8 +369,7 @@ fn subscribe(
     io: Io,
     gpa: Allocator,
     client: *Client,
-    id: usize,
-    // msg: Message.Sub,
+    id: u128,
 ) !void {
     const msg = try parse.sub(client.from_client);
     try server.subs_lock.lock(io);
@@ -395,17 +388,17 @@ fn subscribe(
         .queue_lock = &client.recv_queue_write_lock,
         .queue = client.recv_queue,
     });
-    log.debug("Client {d} subscribed to {s}", .{ id, msg.subject });
+    log.debug("Client {s} subscribed to {s}", .{ idToStr(id), msg.subject });
 }
 
 fn unsubscribe(
     server: *Server,
     io: Io,
     gpa: Allocator,
-    client: Client,
-    id: usize,
+    senders_reader: *std.Io.Reader,
+    id: u128,
 ) !void {
-    const msg = try parse.unsub(client.from_client);
+    const msg = try parse.unsub(senders_reader);
     try server.subs_lock.lock(io);
     defer server.subs_lock.unlock(io);
     const len = server.subscriptions.items.len;
@@ -413,7 +406,7 @@ fn unsubscribe(
         const i = len - from_end - 1;
         const sub = server.subscriptions.items[i];
         if (sub.client_id == id and eql(u8, sub.sid, msg.sid)) {
-            log.debug("Client {d} unsubscribed from {s}", .{ id, server.subscriptions.items[i].subject });
+            log.debug("Client {s} unsubscribed from {s}", .{ idToStr(id), server.subscriptions.items[i].subject });
             sub.deinit(gpa);
             _ = server.subscriptions.swapRemove(i);
             break;
@@ -453,5 +446,28 @@ fn readBufferSize(io: Io, dir: anytype, filename: []const u8, buf: []u8, default
     };
 }
 
+const uuid_len = 36;
+fn idToStr(in: u128) [uuid_len]u8 {
+    // Extract segments using bit shifting and casting
+    const part1: u32 = @intCast(in >> 96);
+    const part2: u16 = @intCast((in >> 80) & 0xFFFF);
+    const part3: u16 = @intCast((in >> 64) & 0xFFFF);
+    const part4: u16 = @intCast((in >> 48) & 0xFFFF);
+    const part5: u64 = @intCast(in & 0xFFFFFFFFFFFF);
+
+    var res: [uuid_len]u8 = undefined;
+
+    // bufPrint returns a slice of the buffer; we ignore it as we return the whole array
+    _ = std.fmt.bufPrint(&res, "{x:0>8}-{x:0>4}-{x:0>4}-{x:0>4}-{x:0>12}", .{
+        part1,
+        part2,
+        part3,
+        part4,
+        part5,
+    }) catch unreachable; // unreachable because the buffer size is guaranteed
+
+    return res;
+}
+
 pub const default_id = "server-id-123";
 pub const default_name = "Zits Server";