Use memorypool for hot messageshot-message-pool

1 files changed, 127 insertions, 58 deletions

diff --git a/src/server/message_parser.zig b/src/server/message_parser.zig
index 1e7527d..109174f 100644
--- a/src/server/message_parser.zig
+++ b/src/server/message_parser.zig
@@ -6,6 +6,7 @@ const StaticStringMap = std.StaticStringMap;
 
 const Io = std.Io;
 const AllocatingWriter = Io.Writer.Allocating;
+const Mutex = Io.Mutex;
 const Reader = Io.Reader;
 
 const ascii = std.ascii;
@@ -22,11 +23,11 @@ pub const MessageType = @typeInfo(Message).@"union".tag_type.?;
 pub const Message = union(enum) {
     info: ServerInfo,
     connect: Connect,
-    @"pub": Pub,
+    @"pub": *Pub,
     hpub: HPub,
     sub: Sub,
     unsub: Unsub,
-    msg: Msg,
+    msg: *Msg,
     hmsg: HMsg,
     ping,
     pong,
@@ -77,7 +78,7 @@ pub const Message = union(enum) {
         headers: ?bool = null,
         nkey: ?[]const u8 = null,
 
-        pub fn deinit(self: Connect, alloc: Allocator) void {
+        pub fn deinit(self: *Connect, alloc: Allocator) void {
             if (self.auth_token) |a| alloc.free(a);
             if (self.user) |u| alloc.free(u);
             if (self.pass) |p| alloc.free(p);
@@ -87,6 +88,7 @@ pub const Message = union(enum) {
             if (self.sig) |s| alloc.free(s);
             if (self.jwt) |j| alloc.free(j);
             if (self.nkey) |n| alloc.free(n);
+            self.* = undefined;
         }
 
         pub fn dupe(self: Connect, alloc: Allocator) !Connect {
@@ -118,46 +120,58 @@ pub const Message = union(enum) {
         /// The reply subject that subscribers can use to send a response back to the publisher/requestor.
         reply_to: ?[]const u8 = null,
         /// The message payload data.
-        payload: []const u8,
+        payload_len: usize,
+        payload: [128]u8,
+        payload_extra: []const u8 = &.{},
 
-        pub fn deinit(self: Pub, alloc: Allocator) void {
+        pub fn deinit(self: *Pub, alloc: Allocator) void {
             alloc.free(self.subject);
-            alloc.free(self.payload);
+            alloc.free(self.payload_extra);
             if (self.reply_to) |r| alloc.free(r);
+            self.* = undefined;
         }
 
-        pub fn toMsg(self: Pub, alloc: Allocator, sid: []const u8) !Msg {
+        pub fn toMsg(
+            self: Pub,
+            alloc: Allocator,
+            hot_msg_manager: *HotMessageManager,
+            sid: []const u8,
+        ) !*Msg {
             const res: Msg = .{
                 .subject = self.subject,
                 .sid = sid,
                 .reply_to = self.reply_to,
+                .payload_len = self.payload_len,
                 .payload = self.payload,
+                .payload_extra = self.payload_extra,
             };
-            return res.dupe(alloc);
+            return res.dupe(alloc, hot_msg_manager);
         }
     };
     pub const HPub = struct {
         header_bytes: usize,
-        @"pub": Pub,
+        @"pub": *Pub,
 
-        pub fn deinit(self: HPub, alloc: Allocator) void {
+        pub fn deinit(self: *HPub, alloc: Allocator) void {
             self.@"pub".deinit(alloc);
+            self.* = undefined;
         }
 
-        pub fn toHMsg(self: HPub, alloc: Allocator, sid: []const u8) !HMsg {
+        pub fn toHMsg(self: HPub, alloc: Allocator, hot_msg_manager: *HotMessageManager, sid: []const u8) !HMsg {
             return .{
                 .header_bytes = self.header_bytes,
-                .msg = try self.@"pub".toMsg(alloc, sid),
+                .msg = try self.@"pub".toMsg(alloc, hot_msg_manager, sid),
             };
         }
     };
 
     pub const HMsg = struct {
         header_bytes: usize,
-        msg: Msg,
+        msg: *Msg,
 
-        pub fn deinit(self: HMsg, alloc: Allocator) void {
+        pub fn deinit(self: *HMsg, alloc: Allocator) void {
             self.msg.deinit(alloc);
+            self.* = undefined;
         }
 
         pub fn dupe(self: HMsg, alloc: Allocator) !HMsg {
@@ -174,10 +188,11 @@ pub const Message = union(enum) {
         /// A unique alphanumeric subscription ID, generated by the client.
         sid: []const u8,
 
-        pub fn deinit(self: Sub, alloc: Allocator) void {
+        pub fn deinit(self: *Sub, alloc: Allocator) void {
             alloc.free(self.subject);
             alloc.free(self.sid);
             if (self.queue_group) |q| alloc.free(q);
+            self.* = undefined;
         }
     };
     pub const Unsub = struct {
@@ -186,33 +201,39 @@ pub const Message = union(enum) {
         /// A number of messages to wait for before automatically unsubscribing.
         max_msgs: ?usize = null,
 
-        pub fn deinit(self: Unsub, alloc: Allocator) void {
+        pub fn deinit(self: *Unsub, alloc: Allocator) void {
             alloc.free(self.sid);
+            self.* = undefined;
         }
     };
     pub const Msg = struct {
         subject: []const u8,
         sid: []const u8,
         reply_to: ?[]const u8,
-        payload: []const u8,
+        payload_len: usize,
+        payload: [128]u8,
+        payload_extra: []const u8,
 
-        pub fn deinit(self: Msg, alloc: Allocator) void {
+        pub fn deinit(self: *Msg, alloc: Allocator) void {
             alloc.free(self.subject);
             alloc.free(self.sid);
             if (self.reply_to) |r| alloc.free(r);
-            alloc.free(self.payload);
+            alloc.free(self.payload_extra);
+            self.* = undefined;
         }
 
-        pub fn dupe(self: Msg, alloc: Allocator) !Msg {
-            var res: Msg = undefined;
+        pub fn dupe(self: Msg, alloc: Allocator, hot_msg_manager: *HotMessageManager) !*Msg {
+            var res = try hot_msg_manager.createMsg(alloc);
             res.subject = try alloc.dupe(u8, self.subject);
             errdefer alloc.free(res.subject);
             res.sid = try alloc.dupe(u8, self.sid);
             errdefer alloc.free(res.sid);
             res.reply_to = if (self.reply_to) |r| try alloc.dupe(u8, r) else null;
             errdefer if (res.reply_to) |r| alloc.free(r);
-            res.payload = try alloc.dupe(u8, self.payload);
-            errdefer alloc.free(res.payload);
+            res.payload_len = self.payload_len;
+            @memcpy(&res.payload, &self.payload);
+            res.payload_extra = try alloc.dupe(u8, self.payload_extra);
+            errdefer alloc.free(res.payload_extra);
             return res;
         }
     };
@@ -240,7 +261,11 @@ pub const Message = union(enum) {
     pub const parse = parseStaticStringMap;
 
     /// An error should be handled by cleaning up this connection.
-    pub fn next(alloc: Allocator, in: *Reader) !Message {
+    pub fn next(
+        alloc: Allocator,
+        in: *Reader,
+        hot_msg_manager: *HotMessageManager,
+    ) !Message {
         var operation_string: ArrayList(u8) = blk: {
             comptime var buf_len = 0;
             comptime {
@@ -299,11 +324,11 @@ pub const Message = union(enum) {
             },
             .@"pub" => {
                 @branchHint(.likely);
-                return parsePub(alloc, in);
+                return parsePub(alloc, in, hot_msg_manager);
             },
             .hpub => {
                 @branchHint(.likely);
-                return parseHPub(alloc, in);
+                return parseHPub(alloc, in, hot_msg_manager);
             },
             .ping => {
                 try expectStreamBytes(in, "\r\n");
@@ -590,12 +615,15 @@ test parseUnsub {
     }
 }
 
-fn parsePub(alloc: Allocator, in: *Reader) !Message {
+fn parsePub(alloc: Allocator, in: *Reader, hot_msg_manager: *HotMessageManager) !Message {
     try in.discardAll(1); // throw away space
 
+    const res = try hot_msg_manager.createPub(alloc);
+    errdefer hot_msg_manager.destroyPub(res);
+
     // Parse subject
-    const subject: []const u8 = try readSubject(alloc, in, .@"pub");
-    errdefer alloc.free(subject);
+    res.*.subject = try readSubject(alloc, in, .@"pub");
+    errdefer alloc.free(res.*.subject);
 
     const States = enum {
         before_second,
@@ -609,8 +637,6 @@ fn parsePub(alloc: Allocator, in: *Reader) !Message {
     defer second.deinit(alloc);
     var third: ?ArrayList(u8) = null;
     defer if (third) |*t| t.deinit(alloc);
-    var payload: AllocatingWriter = .init(alloc);
-    errdefer payload.deinit();
 
     sw: switch (@as(States, .before_second)) {
         .before_second => {
@@ -658,25 +684,26 @@ fn parsePub(alloc: Allocator, in: *Reader) !Message {
         },
     }
 
-    const reply_to: ?[]const u8, const bytes: usize =
+    res.*.reply_to, res.*.payload_len =
         if (third) |t| .{
-            try alloc.dupe(u8, second.items),
+            try second.toOwnedSlice(alloc),
             try parseUnsigned(usize, t.items, 10),
         } else .{
             null,
             try parseUnsigned(usize, second.items, 10),
         };
 
-    try in.streamExact(&payload.writer, bytes);
+    try in.readSliceAll(res.*.payload[0..@min(res.*.payload.len, res.*.payload_len)]);
+    {
+        const remaining_payload_bytes = res.*.payload_len - res.*.payload.len;
+        var payload: AllocatingWriter = .init(alloc);
+        errdefer payload.deinit();
+        try in.streamExact(&payload.writer, remaining_payload_bytes);
+        res.*.payload_extra = try payload.toOwnedSlice();
+    }
     try expectStreamBytes(in, "\r\n");
 
-    return .{
-        .@"pub" = .{
-            .subject = subject,
-            .payload = try payload.toOwnedSlice(),
-            .reply_to = reply_to,
-        },
-    };
+    return .{ .@"pub" = res };
 }
 
 test parsePub {
@@ -736,12 +763,15 @@ test parsePub {
     }
 }
 
-fn parseHPub(alloc: Allocator, in: *Reader) !Message {
+fn parseHPub(alloc: Allocator, in: *Reader, hot_msg_manager: *HotMessageManager) !Message {
     try in.discardAll(1); // throw away space
 
+    const nested_pub = try hot_msg_manager.createPub(alloc);
+    errdefer hot_msg_manager.destroyPub(nested_pub);
+
     // Parse subject
-    const subject: []const u8 = try readSubject(alloc, in, .@"pub");
-    errdefer alloc.free(subject);
+    nested_pub.*.subject = try readSubject(alloc, in, .@"pub");
+    errdefer alloc.free(nested_pub.*.subject);
 
     const States = enum {
         before_second,
@@ -759,8 +789,6 @@ fn parseHPub(alloc: Allocator, in: *Reader) !Message {
     defer third.deinit(alloc);
     var fourth: ?ArrayList(u8) = null;
     defer if (fourth) |*f| f.deinit(alloc);
-    var payload: AllocatingWriter = .init(alloc);
-    errdefer payload.deinit();
 
     sw: switch (@as(States, .before_second)) {
         .before_second => {
@@ -828,7 +856,12 @@ fn parseHPub(alloc: Allocator, in: *Reader) !Message {
         },
     }
 
-    const reply_to: ?[]const u8, const header_bytes: usize, const total_bytes: usize =
+    var res: Message.HPub = .{
+        .header_bytes = 0,
+        .@"pub" = nested_pub,
+    };
+
+    res.@"pub".*.reply_to, res.header_bytes, res.@"pub".*.payload_len =
         if (fourth) |f| .{
             try alloc.dupe(u8, second.items),
             try parseUnsigned(usize, third.items, 10),
@@ -839,19 +872,17 @@ fn parseHPub(alloc: Allocator, in: *Reader) !Message {
             try parseUnsigned(usize, third.items, 10),
         };
 
-    try in.streamExact(&payload.writer, total_bytes);
+    try in.readSliceAll(res.@"pub".*.payload[0..@min(res.@"pub".*.payload.len, res.@"pub".*.payload_len)]);
+    {
+        const remaining_payload_bytes = res.@"pub".*.payload_len - res.@"pub".*.payload.len;
+        var payload: AllocatingWriter = .init(alloc);
+        errdefer payload.deinit();
+        try in.streamExact(&payload.writer, remaining_payload_bytes);
+        res.@"pub".*.payload_extra = try payload.toOwnedSlice();
+    }
     try expectStreamBytes(in, "\r\n");
 
-    return .{
-        .hpub = .{
-            .header_bytes = header_bytes,
-            .@"pub" = .{
-                .subject = subject,
-                .payload = try payload.toOwnedSlice(),
-                .reply_to = reply_to,
-            },
-        },
-    };
+    return .{ .hpub = res };
 }
 
 test parseHPub {
@@ -976,6 +1007,44 @@ inline fn expectStreamBytes(reader: *Reader, expected: []const u8) !void {
     }
 }
 
+pub const HotMessageManager = struct {
+    pub_pool_lock: Mutex = .init,
+    pub_pool: std.heap.MemoryPool(Message.Pub) = .empty,
+    msg_pool_lock: Mutex = .init,
+    msg_pool: std.heap.MemoryPool(Message.Msg) = .empty,
+    // used for locking
+    io: Io,
+
+    pub fn deinit(self: *HotMessageManager, alloc: Allocator) void {
+        self.pub_pool.deinit(alloc);
+        self.msg_pool.deinit(alloc);
+    }
+
+    pub fn createPub(self: *HotMessageManager, alloc: Allocator) !*Message.Pub {
+        try self.pub_pool_lock.lock(self.io);
+        defer self.pub_pool_lock.unlock(self.io);
+        return self.pub_pool.create(alloc);
+    }
+
+    pub fn destroyPub(self: *HotMessageManager, msg: *Message.Pub) void {
+        self.pub_pool_lock.lockUncancelable(self.io);
+        defer self.pub_pool_lock.unlock(self.io);
+        self.pub_pool.destroy(msg);
+    }
+
+    pub fn createMsg(self: *HotMessageManager, alloc: Allocator) !*Message.Msg {
+        try self.msg_pool_lock.lock(self.io);
+        defer self.msg_pool_lock.unlock(self.io);
+        return self.msg_pool.create(alloc);
+    }
+
+    pub fn destroyMsg(self: *HotMessageManager, msg: *Message.Msg) void {
+        self.msg_pool_lock.lockUncancelable(self.io);
+        defer self.msg_pool_lock.unlock(self.io);
+        self.msg_pool.destroy(msg);
+    }
+};
+
 test "parsing a stream" {
     const alloc = std.testing.allocator;
     const expectEqualDeep = std.testing.expectEqualDeep;