From bd9829f6842f0c989389aa4ce9784ab6e3cb4ee5 Mon Sep 17 00:00:00 2001 From: Robby Zambito Date: Sat, 3 Jan 2026 05:33:13 +0000 Subject: Organize things Making it easier to use the server as a library --- src/server/Client.zig | 221 +++++++++++++++++++++++++++++++++ src/server/Server.zig | 305 ++++++++++++++++++++++++++++++++++++++++++++++ src/server/client.zig | 221 --------------------------------- src/server/main.zig | 329 ++++---------------------------------------------- src/server/test.zig | 3 - 5 files changed, 548 insertions(+), 531 deletions(-) create mode 100644 src/server/Client.zig create mode 100644 src/server/Server.zig delete mode 100644 src/server/client.zig delete mode 100644 src/server/test.zig (limited to 'src/server') diff --git a/src/server/Client.zig b/src/server/Client.zig new file mode 100644 index 0000000..2ce3c38 --- /dev/null +++ b/src/server/Client.zig @@ -0,0 +1,221 @@ +const Message = @import("message_parser.zig").Message; +const std = @import("std"); + +const Client = @This(); + +connect: ?Message.Connect, + +// Messages for this client to receive. +recv_queue: *std.Io.Queue(Message), + +from_client: *std.Io.Reader, +to_client: *std.Io.Writer, + +pub fn init( + connect: ?Message.Connect, + recv_queue: *std.Io.Queue(Message), + in: *std.Io.Reader, + out: *std.Io.Writer, +) Client { + return .{ + .connect = connect, + .recv_queue = recv_queue, + .from_client = in, + .to_client = out, + }; +} + +pub fn deinit(self: *Client, alloc: std.mem.Allocator) void { + if (self.connect) |c| { + c.deinit(alloc); + } + self.* = undefined; +} + +pub fn start(self: *Client, io: std.Io, alloc: std.mem.Allocator) !void { + var msgs: [8]Message = undefined; + + while (true) { + const len = try self.recv_queue.get(io, &msgs, 1); + std.debug.assert(len <= msgs.len); + for (0..len) |i| { + const msg = msgs[i]; + defer switch (msg) { + .msg => |m| m.deinit(alloc), + .hmsg => |h| h.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"); + }, + .pong => { + _ = try self.to_client.write("PONG\r\n"); + }, + .info => |info| { + _ = try self.to_client.write("INFO "); + try std.json.Stringify.value(info, .{}, self.to_client); + _ = try self.to_client.write("\r\n"); + }, + .msg => |m| { + @branchHint(.likely); + try self.to_client.print( + "MSG {s} {s} {s} {d}\r\n{s}\r\n", + .{ + m.subject, + m.sid, + m.reply_to orelse "", + m.payload.len, + m.payload, + }, + ); + }, + .hmsg => |hmsg| { + @branchHint(.likely); + try self.to_client.print("HMSG {s} {s} {s} {d} {d}\r\n{s}\r\n", .{ + hmsg.msg.subject, + hmsg.msg.sid, + hmsg.msg.reply_to orelse "", + hmsg.header_bytes, + hmsg.msg.payload.len, + hmsg.msg.payload, + }); + }, + .@"-err" => |s| { + _ = try self.to_client.print("-ERR '{s}'\r\n", .{s}); + }, + else => |m| { + std.debug.panic("unimplemented write: {any}\n", .{m}); + }, + } + } + try self.to_client.flush(); + } +} + +pub fn send(self: *Client, io: std.Io, msg: Message) !void { + try self.recv_queue.putOne(io, msg); +} + +pub fn next(self: *Client, allocator: std.mem.Allocator) !Message { + return Message.next(allocator, self.from_client); +} + +test { + const io = std.testing.io; + const gpa = std.testing.allocator; + + var from_client: std.Io.Reader = .fixed( + "CONNECT {\"verbose\":false,\"pedantic\":false,\"tls_required\":false,\"name\":\"NATS CLI Version v0.2.4\",\"lang\":\"go\",\"version\":\"1.43.0\",\"protocol\":1,\"echo\":true,\"headers\":true,\"no_responders\":true}\r\n" ++ + "PING\r\n", + ); + var from_client_buf: [1024]Message = undefined; + var from_client_queue: std.Io.Queue(Message) = .init(&from_client_buf); + + { + // Simulate stream + while (Message.next(gpa, &from_client)) |msg| { + try from_client_queue.putOne(io, msg); + } else |err| switch (err) { + error.EndOfStream => from_client_queue.close(io), + else => return err, + } + + while (from_client_queue.getOne(io)) |msg| { + switch (msg) { + .connect => |*c| { + std.debug.print("Message: {any}\n", .{msg}); + c.deinit(gpa); + }, + else => { + std.debug.print("Message: {any}\n", .{msg}); + }, + } + } else |_| {} + } + + from_client_queue = .init(&from_client_buf); + // Reset the reader to process it again. + from_client.seek = 0; + + // { + // const SemiClient = struct { + // q: std.Io.Queue(Message), + + // fn parseClientInput(self: *@This(), ioh: std.Io, in: *std.Io.Reader) void { + // defer std.debug.print("done parse\n", .{}); + // while (Message.next(gpa, in)) |msg| { + // self.q.putOne(ioh, msg) catch return; + // } else |_| {} + // } + + // fn next(self: *@This(), ioh: std.Io) !Message { + // return self.q.getOne(ioh); + // } + + // fn printAll(self: *@This(), ioh: std.Io) void { + // defer std.debug.print("done print\n", .{}); + // while (self.next(ioh)) |*msg| { + // std.debug.print("Client msg: {any}\n", .{msg}); + // switch (msg.*) { + // .connect => |c| { + // c.deinit(gpa); + // }, + // else => {}, + // } + // } else |_| {} + // } + // }; + + // var c: SemiClient = .{ .q = from_client_queue }; + // var group: std.Io.Group = .init; + // defer group.wait(io); + + // group.concurrent(io, SemiClient.printAll, .{ &c, io }) catch { + // @panic("could not start printAll\n"); + // }; + + // group.concurrent(io, SemiClient.parseClientInput, .{ &c, io, &from_client }) catch { + // @panic("could not start printAll\n"); + // }; + // } + + //////// + + // const connect = (Message.next(gpa, &from_client) catch unreachable).connect; + + // var to_client_alloc: std.Io.Writer.Allocating = .init(gpa); + // defer to_client_alloc.deinit(); + // var to_client = to_client_alloc.writer; + + // var client: ClientState = try .init(io, gpa, 0, connect, &from_client, &to_client); + // defer client.deinit(gpa); + + // { + // var get_next = io.concurrent(ClientState.next, .{ &client, io }) catch unreachable; + // defer if (get_next.cancel(io)) |_| {} else |_| @panic("fail"); + + // var timeout = io.concurrent(std.Io.sleep, .{ io, .fromMilliseconds(1000), .awake }) catch unreachable; + // defer timeout.cancel(io) catch {}; + + // switch (try io.select(.{ + // .get_next = &get_next, + // .timeout = &timeout, + // })) { + // .get_next => |next| { + // std.debug.print("next is {any}\n", .{next}); + // try std.testing.expect((next catch |err| return err) == .ping); + // }, + // .timeout => { + // std.debug.print("reached timeout\n", .{}); + // return error.TestUnexpectedResult; + // }, + // } + // } +} diff --git a/src/server/Server.zig b/src/server/Server.zig new file mode 100644 index 0000000..ebbac19 --- /dev/null +++ b/src/server/Server.zig @@ -0,0 +1,305 @@ +const std = @import("std"); +const Allocator = std.mem.Allocator; +const ArrayList = std.ArrayList; +const AutoHashMapUnmanaged = std.AutoHashMapUnmanaged; + +const Io = std.Io; +const Group = Io.Group; +const IpAddress = std.Io.net.IpAddress; +const Mutex = Io.Mutex; +const Queue = Io.Queue; +const Stream = std.Io.net.Stream; + +const message_parser = @import("./message_parser.zig"); +pub const MessageType = message_parser.MessageType; +pub const Message = message_parser.Message; +const ServerInfo = Message.ServerInfo; +pub const Client = @import("./Client.zig"); +const Server = @This(); + +pub const Subscription = struct { + subject: []const u8, + client_id: usize, + sid: []const u8, + + fn deinit(self: Subscription, alloc: Allocator) void { + alloc.free(self.subject); + alloc.free(self.sid); + } +}; + +const eql = std.mem.eql; +const log = std.log; +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); + } + server.subscriptions.deinit(alloc); + server.clients.deinit(alloc); +} + +pub fn start(server: *Server, io: Io, gpa: Allocator) !void { + var tcp_server = try IpAddress.listen(try IpAddress.parse( + server.info.host, + server.info.port, + ), io, .{ .reuse_address = true }); + defer tcp_server.deinit(io); + log.debug("Server headers: {s}", .{if (server.info.headers) "true" else "false"}); + log.debug("Server max payload: {d}", .{server.info.max_payload}); + log.info("Server ID: {s}", .{server.info.server_id}); + log.info("Server name: {s}", .{server.info.server_name}); + log.info("Server listening on {s}:{d}", .{ server.info.host, server.info.port }); + + var client_group: Group = .init; + defer client_group.cancel(io); + + var id: usize = 0; + while (true) : (id +%= 1) { + if (server.clients.contains(id)) continue; + log.debug("Accepting next client", .{}); + const stream = try tcp_server.accept(io); + log.debug("Accepted connection {d}", .{id}); + _ = client_group.concurrent(io, handleConnectionInfallible, .{ server, gpa, io, id, stream }) catch { + log.err("Could not start concurrent handler for {d}", .{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 { + 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); + } + } + } +} + +fn handleConnectionInfallible(server: *Server, server_allocator: Allocator, io: Io, id: usize, stream: Stream) void { + handleConnection(server, server_allocator, io, id, stream) catch |err| { + log.err("Failed processing client {d}: {any}", .{ id, err }); + }; +} + +fn handleConnection(server: *Server, server_allocator: Allocator, io: Io, id: usize, stream: Stream) !void { + defer stream.close(io); + + // TODO: use a client allocator for things that should only live for as long as the client? + // I had this before, but it seemed to have made lifetimes harder to track. + // Messages made sense to parse using a client allocator, but that makes it hard to free + // messages when done processing them (usually outside the client process, ie: publish). + + // Set up client writer + // TODO: how many bytes can fit in a network write syscall? cat /proc/sys/net/core/wmem_max + var w_buffer: [1024 * 16]u8 = undefined; + var writer = stream.writer(io, &w_buffer); + const out = &writer.interface; + + // Set up client reader + // TODO: how many bytes can fit in a network read syscall? cat /proc/sys/net/core/rmem_max + var r_buffer: [1024 * 16]u8 = undefined; + var reader = stream.reader(io, &r_buffer); + const in = &reader.interface; + + // Set up buffer queue + var qbuf: [8]Message = undefined; + var queue: Queue(Message) = .init(&qbuf); + defer { + queue.close(io); + while (queue.getOne(io)) |msg| { + switch (msg) { + .msg => |m| m.deinit(server_allocator), + .hmsg => |h| h.deinit(server_allocator), + else => {}, + } + } else |_| {} + } + + // Create client + var client: Client = .init(null, &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 + try queue.putOne(io, .{ .info = server.info }); + + var client_task = try io.concurrent(Client.start, .{ &client, io, server_allocator }); + defer client_task.cancel(io) catch {}; + + // Messages are owned by the server after they are received from the client + while (client.next(server_allocator)) |msg| { + switch (msg) { + .ping => { + // Respond to ping with pong. + try client.send(io, .pong); + }, + .@"pub", .hpub => { + defer switch (msg) { + .@"pub" => |pb| pb.deinit(server_allocator), + .hpub => |hp| hp.deinit(server_allocator), + else => unreachable, + }; + try server.publishMessage(io, server_allocator, &client, msg); + }, + .sub => |sub| { + defer sub.deinit(server_allocator); + try server.subscribe(io, server_allocator, id, sub); + }, + .unsub => |unsub| { + defer unsub.deinit(server_allocator); + try server.unsubscribe(io, server_allocator, id, unsub); + }, + .connect => |connect| { + if (client.connect) |*current| { + current.deinit(server_allocator); + } + client.connect = connect; + }, + else => |e| { + panic("Unimplemented message: {any}\n", .{e}); + }, + } + } else |err| switch (err) { + error.EndOfStream => { + log.debug("Client {d} disconnected", .{id}); + }, + else => { + return err; + }, + } +} + +fn subjectMatches(sub_subject: []const u8, pub_subject: []const u8) bool { + // TODO: assert that sub_subject and pub_subject are valid. + var sub_iter = std.mem.splitScalar(u8, sub_subject, '.'); + var pub_iter = std.mem.splitScalar(u8, pub_subject, '.'); + + while (sub_iter.next()) |st| { + const pt = pub_iter.next() orelse return false; + + if (eql(u8, st, ">")) return true; + + if (!eql(u8, st, "*") and !eql(u8, st, pt)) { + return false; + } + } + + return pub_iter.next() == null; +} + +test subjectMatches { + const expect = std.testing.expect; + try expect(subjectMatches("foo", "foo")); + try expect(!subjectMatches("foo", "bar")); + + try expect(subjectMatches("foo.*", "foo.bar")); + try expect(!subjectMatches("foo.*", "foo")); + try expect(!subjectMatches("foo.>", "foo")); + + // the wildcard subscriptions foo.*.quux and foo.> both match foo.bar.quux, but only the latter matches foo.bar.baz. + try expect(subjectMatches("foo.*.quux", "foo.bar.quux")); + try expect(subjectMatches("foo.>", "foo.bar.quux")); + try expect(!subjectMatches("foo.*.quux", "foo.bar.baz")); + try expect(subjectMatches("foo.>", "foo.bar.baz")); +} + +fn publishMessage(server: *Server, io: Io, alloc: Allocator, source_client: *Client, msg: Message) !void { + errdefer { + if (source_client.connect) |c| { + if (c.verbose) { + source_client.send(io, .{ .@"-err" = "Slow Consumer" }) catch {}; + } + } + } + const subject = switch (msg) { + .@"pub" => |pb| pb.subject, + .hpub => |hp| hp.@"pub".subject, + else => unreachable, + }; + try server.subs_lock.lock(io); + defer server.subs_lock.unlock(io); + for (server.subscriptions.items) |subscription| { + if (subjectMatches(subscription.subject, subject)) { + const client = server.clients.get(subscription.client_id) orelse { + log.debug("Trying to publish to a client that no longer exists: {d}\n", .{subscription.client_id}); + continue; + }; + + switch (msg) { + .@"pub" => |pb| client.send(io, .{ + .msg = try pb.toMsg(alloc, subscription.sid), + }) catch |err| switch (err) { + error.Canceled => return err, + else => {}, + }, + .hpub => |hp| client.send(io, .{ .hmsg = try hp.toHMsg( + alloc, + subscription.sid, + ) }) catch |err| switch (err) { + error.Canceled => return err, + else => {}, + }, + else => unreachable, + } + } + } + if (source_client.connect) |c| { + if (c.verbose) { + source_client.send(io, .@"+ok") catch {}; + } + } +} + +fn subscribe(server: *Server, io: Io, gpa: Allocator, id: usize, msg: Message.Sub) !void { + try server.subs_lock.lock(io); + defer server.subs_lock.unlock(io); + const subject = try gpa.dupe(u8, msg.subject); + errdefer gpa.free(subject); + const sid = try gpa.dupe(u8, msg.sid); + errdefer gpa.free(sid); + try server.subscriptions.append(gpa, .{ + .subject = subject, + .client_id = id, + .sid = sid, + }); +} + +fn unsubscribe(server: *Server, io: Io, gpa: 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) |from_end| { + const i = len - from_end - 1; + const sub = server.subscriptions.items[i]; + if (sub.client_id == id and eql(u8, sub.sid, msg.sid)) { + sub.deinit(gpa); + _ = server.subscriptions.swapRemove(i); + } + } +} + +pub const default_id = "server-id-123"; +pub const default_name = "Zits Server"; diff --git a/src/server/client.zig b/src/server/client.zig deleted file mode 100644 index 2ce3c38..0000000 --- a/src/server/client.zig +++ /dev/null @@ -1,221 +0,0 @@ -const Message = @import("message_parser.zig").Message; -const std = @import("std"); - -const Client = @This(); - -connect: ?Message.Connect, - -// Messages for this client to receive. -recv_queue: *std.Io.Queue(Message), - -from_client: *std.Io.Reader, -to_client: *std.Io.Writer, - -pub fn init( - connect: ?Message.Connect, - recv_queue: *std.Io.Queue(Message), - in: *std.Io.Reader, - out: *std.Io.Writer, -) Client { - return .{ - .connect = connect, - .recv_queue = recv_queue, - .from_client = in, - .to_client = out, - }; -} - -pub fn deinit(self: *Client, alloc: std.mem.Allocator) void { - if (self.connect) |c| { - c.deinit(alloc); - } - self.* = undefined; -} - -pub fn start(self: *Client, io: std.Io, alloc: std.mem.Allocator) !void { - var msgs: [8]Message = undefined; - - while (true) { - const len = try self.recv_queue.get(io, &msgs, 1); - std.debug.assert(len <= msgs.len); - for (0..len) |i| { - const msg = msgs[i]; - defer switch (msg) { - .msg => |m| m.deinit(alloc), - .hmsg => |h| h.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"); - }, - .pong => { - _ = try self.to_client.write("PONG\r\n"); - }, - .info => |info| { - _ = try self.to_client.write("INFO "); - try std.json.Stringify.value(info, .{}, self.to_client); - _ = try self.to_client.write("\r\n"); - }, - .msg => |m| { - @branchHint(.likely); - try self.to_client.print( - "MSG {s} {s} {s} {d}\r\n{s}\r\n", - .{ - m.subject, - m.sid, - m.reply_to orelse "", - m.payload.len, - m.payload, - }, - ); - }, - .hmsg => |hmsg| { - @branchHint(.likely); - try self.to_client.print("HMSG {s} {s} {s} {d} {d}\r\n{s}\r\n", .{ - hmsg.msg.subject, - hmsg.msg.sid, - hmsg.msg.reply_to orelse "", - hmsg.header_bytes, - hmsg.msg.payload.len, - hmsg.msg.payload, - }); - }, - .@"-err" => |s| { - _ = try self.to_client.print("-ERR '{s}'\r\n", .{s}); - }, - else => |m| { - std.debug.panic("unimplemented write: {any}\n", .{m}); - }, - } - } - try self.to_client.flush(); - } -} - -pub fn send(self: *Client, io: std.Io, msg: Message) !void { - try self.recv_queue.putOne(io, msg); -} - -pub fn next(self: *Client, allocator: std.mem.Allocator) !Message { - return Message.next(allocator, self.from_client); -} - -test { - const io = std.testing.io; - const gpa = std.testing.allocator; - - var from_client: std.Io.Reader = .fixed( - "CONNECT {\"verbose\":false,\"pedantic\":false,\"tls_required\":false,\"name\":\"NATS CLI Version v0.2.4\",\"lang\":\"go\",\"version\":\"1.43.0\",\"protocol\":1,\"echo\":true,\"headers\":true,\"no_responders\":true}\r\n" ++ - "PING\r\n", - ); - var from_client_buf: [1024]Message = undefined; - var from_client_queue: std.Io.Queue(Message) = .init(&from_client_buf); - - { - // Simulate stream - while (Message.next(gpa, &from_client)) |msg| { - try from_client_queue.putOne(io, msg); - } else |err| switch (err) { - error.EndOfStream => from_client_queue.close(io), - else => return err, - } - - while (from_client_queue.getOne(io)) |msg| { - switch (msg) { - .connect => |*c| { - std.debug.print("Message: {any}\n", .{msg}); - c.deinit(gpa); - }, - else => { - std.debug.print("Message: {any}\n", .{msg}); - }, - } - } else |_| {} - } - - from_client_queue = .init(&from_client_buf); - // Reset the reader to process it again. - from_client.seek = 0; - - // { - // const SemiClient = struct { - // q: std.Io.Queue(Message), - - // fn parseClientInput(self: *@This(), ioh: std.Io, in: *std.Io.Reader) void { - // defer std.debug.print("done parse\n", .{}); - // while (Message.next(gpa, in)) |msg| { - // self.q.putOne(ioh, msg) catch return; - // } else |_| {} - // } - - // fn next(self: *@This(), ioh: std.Io) !Message { - // return self.q.getOne(ioh); - // } - - // fn printAll(self: *@This(), ioh: std.Io) void { - // defer std.debug.print("done print\n", .{}); - // while (self.next(ioh)) |*msg| { - // std.debug.print("Client msg: {any}\n", .{msg}); - // switch (msg.*) { - // .connect => |c| { - // c.deinit(gpa); - // }, - // else => {}, - // } - // } else |_| {} - // } - // }; - - // var c: SemiClient = .{ .q = from_client_queue }; - // var group: std.Io.Group = .init; - // defer group.wait(io); - - // group.concurrent(io, SemiClient.printAll, .{ &c, io }) catch { - // @panic("could not start printAll\n"); - // }; - - // group.concurrent(io, SemiClient.parseClientInput, .{ &c, io, &from_client }) catch { - // @panic("could not start printAll\n"); - // }; - // } - - //////// - - // const connect = (Message.next(gpa, &from_client) catch unreachable).connect; - - // var to_client_alloc: std.Io.Writer.Allocating = .init(gpa); - // defer to_client_alloc.deinit(); - // var to_client = to_client_alloc.writer; - - // var client: ClientState = try .init(io, gpa, 0, connect, &from_client, &to_client); - // defer client.deinit(gpa); - - // { - // var get_next = io.concurrent(ClientState.next, .{ &client, io }) catch unreachable; - // defer if (get_next.cancel(io)) |_| {} else |_| @panic("fail"); - - // var timeout = io.concurrent(std.Io.sleep, .{ io, .fromMilliseconds(1000), .awake }) catch unreachable; - // defer timeout.cancel(io) catch {}; - - // switch (try io.select(.{ - // .get_next = &get_next, - // .timeout = &timeout, - // })) { - // .get_next => |next| { - // std.debug.print("next is {any}\n", .{next}); - // try std.testing.expect((next catch |err| return err) == .ping); - // }, - // .timeout => { - // std.debug.print("reached timeout\n", .{}); - // return error.TestUnexpectedResult; - // }, - // } - // } -} diff --git a/src/server/main.zig b/src/server/main.zig index 4e47c10..1aaf572 100644 --- a/src/server/main.zig +++ b/src/server/main.zig @@ -1,48 +1,32 @@ const std = @import("std"); -const builtin = @import("builtin"); -const Message = @import("./message_parser.zig").Message; -pub const ServerInfo = Message.ServerInfo; +const Allocator = std.mem.Allocator; +const AtomicValue = std.atomic.Value; +const DebugAllocator = std.heap.DebugAllocator; +const Sigaction = std.posix.Sigaction; -const Client = @import("./client.zig"); -const Server = @This(); +const Io = std.Io; +const Threaded = Io.Threaded; -const Subscription = struct { - subject: []const u8, - client_id: usize, - sid: []const u8, +const builtin = @import("builtin"); - fn deinit(self: Subscription, alloc: std.mem.Allocator) void { - alloc.free(self.subject); - alloc.free(self.sid); - } -}; +const zits = @import("zits"); +const Message = zits.Server.Message; +const ServerInfo = Message.ServerInfo; -info: ServerInfo, -clients: std.AutoHashMapUnmanaged(usize, *Client) = .empty, +const Server = zits.Server; -subs_lock: std.Io.Mutex = .init, -subscriptions: std.ArrayList(Subscription) = .empty, +const safe_build = builtin.mode == .Debug or builtin.mode == .ReleaseSafe; -var keep_running = std.atomic.Value(bool).init(true); +var keep_running = AtomicValue(bool).init(true); fn handleSigInt(sig: std.os.linux.SIG) callconv(.c) void { _ = sig; 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| { - sub.deinit(alloc); - } - server.subscriptions.deinit(alloc); - server.clients.deinit(alloc); -} - -pub fn main(alloc: std.mem.Allocator, server_config: ServerInfo) !void { +pub fn main(outer_alloc: Allocator, server_config: ServerInfo) !void { // Configure the signal action - const act = std.posix.Sigaction{ + const act = Sigaction{ .handler = .{ .handler = handleSigInt }, .mask = std.posix.sigemptyset(), .flags = 0, @@ -52,23 +36,24 @@ pub fn main(alloc: std.mem.Allocator, server_config: ServerInfo) !void { std.posix.sigaction(std.posix.SIG.INT, &act, null); { - var dba: std.heap.DebugAllocator(.{}) = .init; - dba.backing_allocator = alloc; + var dba: DebugAllocator(.{}) = .init; + dba.backing_allocator = outer_alloc; defer _ = dba.deinit(); - const gpa = if (builtin.mode == .Debug or builtin.mode == .ReleaseSafe) dba.allocator() else alloc; + const alloc = if (safe_build) dba.allocator() else outer_alloc; - var threaded: std.Io.Threaded = .init(gpa, .{}); + var threaded: Threaded = .init(alloc, .{}); defer threaded.deinit(); const io = threaded.io(); var server: Server = .{ .info = server_config, }; - defer server.deinit(io, gpa); + defer server.deinit(io, alloc); - var server_task = try io.concurrent(start, .{ &server, io, gpa }); + var server_task = try io.concurrent(Server.start, .{ &server, io, alloc }); defer server_task.cancel(io) catch {}; + // Block until Ctrl+C while (keep_running.load(.monotonic)) { try io.sleep(.fromMilliseconds(1), .awake); } @@ -79,273 +64,3 @@ pub fn main(alloc: std.mem.Allocator, server_config: ServerInfo) !void { } std.log.info("Goodbye", .{}); } - -pub fn start(server: *Server, io: std.Io, gpa: std.mem.Allocator) !void { - var tcp_server = try std.Io.net.IpAddress.listen(try std.Io.net.IpAddress.parse( - server.info.host, - server.info.port, - ), io, .{ .reuse_address = true }); - defer tcp_server.deinit(io); - std.log.debug("Server headers: {s}", .{if (server.info.headers) "true" else "false"}); - std.log.debug("Server max payload: {d}", .{server.info.max_payload}); - std.log.info("Server ID: {s}", .{server.info.server_id}); - std.log.info("Server name: {s}", .{server.info.server_name}); - std.log.info("Server listening on {s}:{d}", .{ server.info.host, server.info.port }); - - var client_group: std.Io.Group = .init; - defer client_group.cancel(io); - - var id: usize = 0; - while (true) : (id +%= 1) { - if (server.clients.contains(id)) continue; - std.log.debug("Accepting next client", .{}); - const stream = try tcp_server.accept(io); - std.log.debug("Accepted connection {d}", .{id}); - _ = client_group.concurrent(io, handleConnectionInfallible, .{ server, gpa, io, id, stream }) catch { - std.log.err("Could not start concurrent handler for {d}", .{id}); - stream.close(io); - }; - } -} - -fn addClient(server: *Server, allocator: std.mem.Allocator, id: usize, client: *Client) !void { - try server.clients.put(allocator, id, client); -} - -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); - 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); - } - } - } -} - -fn handleConnectionInfallible( - server: *Server, - server_allocator: std.mem.Allocator, - io: std.Io, - id: usize, - stream: std.Io.net.Stream, -) void { - handleConnection(server, server_allocator, io, id, stream) catch |err| { - std.log.err("Failed processing client {d}: {any}", .{ id, err }); - }; -} - -fn handleConnection( - server: *Server, - server_allocator: std.mem.Allocator, - io: std.Io, - id: usize, - stream: std.Io.net.Stream, -) !void { - defer stream.close(io); - - //var client_allocator: std.heap.DebugAllocator(.{}) = .init; - //client_allocator.backing_allocator = server_allocator; - //defer _ = client_allocator.deinit(); - //const allocator = if (builtin.mode == .Debug or builtin.mode == .ReleaseSafe) client_allocator.allocator() else server_allocator; - - // Set up client writer - var w_buffer: [1024 * 16]u8 = undefined; - var writer = stream.writer(io, &w_buffer); - const out = &writer.interface; - - // Set up client reader - var r_buffer: [1024 * 16]u8 = undefined; - var reader = stream.reader(io, &r_buffer); - const in = &reader.interface; - - // Set up buffer queue - var qbuf: [8]Message = undefined; - var queue: std.Io.Queue(Message) = .init(&qbuf); - defer { - queue.close(io); - while (queue.getOne(io)) |msg| { - switch (msg) { - .msg => |m| m.deinit(server_allocator), - .hmsg => |h| h.deinit(server_allocator), - else => {}, - } - } else |_| {} - } - - // Create client - var client: Client = .init(null, &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 - try queue.putOne(io, .{ .info = server.info }); - - var client_task = try io.concurrent(Client.start, .{ &client, io, server_allocator }); - defer client_task.cancel(io) catch {}; - - // Messages are owned by the server after they are received from the client - while (client.next(server_allocator)) |msg| { - switch (msg) { - .ping => { - // Respond to ping with pong. - try client.send(io, .pong); - }, - .@"pub", .hpub => { - defer switch (msg) { - .@"pub" => |pb| pb.deinit(server_allocator), - .hpub => |hp| hp.deinit(server_allocator), - else => unreachable, - }; - try server.publishMessage(io, server_allocator, &client, msg); - }, - .sub => |sub| { - defer sub.deinit(server_allocator); - try server.subscribe(io, server_allocator, id, sub); - }, - .unsub => |unsub| { - defer unsub.deinit(server_allocator); - try server.unsubscribe(io, server_allocator, id, unsub); - }, - .connect => |connect| { - if (client.connect) |*current| { - current.deinit(server_allocator); - } - client.connect = connect; - }, - else => |e| { - std.debug.panic("Unimplemented message: {any}\n", .{e}); - }, - } - } else |err| switch (err) { - error.EndOfStream => { - std.log.debug("Client {d} disconnected", .{id}); - }, - else => { - return err; - }, - } -} - -fn subjectMatches(sub_subject: []const u8, pub_subject: []const u8) bool { - // TODO: assert that sub_subject and pub_subject are valid. - var sub_iter = std.mem.splitScalar(u8, sub_subject, '.'); - var pub_iter = std.mem.splitScalar(u8, pub_subject, '.'); - - while (sub_iter.next()) |st| { - const pt = pub_iter.next() orelse return false; - - if (std.mem.eql(u8, st, ">")) return true; - - if (!std.mem.eql(u8, st, "*") and !std.mem.eql(u8, st, pt)) { - return false; - } - } - - return pub_iter.next() == null; -} - -test subjectMatches { - try std.testing.expect(subjectMatches("foo", "foo")); - try std.testing.expect(!subjectMatches("foo", "bar")); - - try std.testing.expect(subjectMatches("foo.*", "foo.bar")); - try std.testing.expect(!subjectMatches("foo.*", "foo")); - try std.testing.expect(!subjectMatches("foo.>", "foo")); - - // the wildcard subscriptions foo.*.quux and foo.> both match foo.bar.quux, but only the latter matches foo.bar.baz. - try std.testing.expect(subjectMatches("foo.*.quux", "foo.bar.quux")); - try std.testing.expect(subjectMatches("foo.>", "foo.bar.quux")); - try std.testing.expect(!subjectMatches("foo.*.quux", "foo.bar.baz")); - try std.testing.expect(subjectMatches("foo.>", "foo.bar.baz")); -} - -fn publishMessage(server: *Server, io: std.Io, alloc: std.mem.Allocator, source_client: *Client, msg: Message) !void { - errdefer { - if (source_client.connect) |c| { - if (c.verbose) { - source_client.send(io, .{ .@"-err" = "Slow Consumer" }) catch {}; - } - } - } - const subject = switch (msg) { - .@"pub" => |pb| pb.subject, - .hpub => |hp| hp.@"pub".subject, - else => unreachable, - }; - try server.subs_lock.lock(io); - defer server.subs_lock.unlock(io); - for (server.subscriptions.items) |subscription| { - if (subjectMatches(subscription.subject, subject)) { - const client = server.clients.get(subscription.client_id) orelse { - std.debug.print("trying to publish to a client that no longer exists: {d}\n", .{subscription.client_id}); - continue; - }; - - switch (msg) { - .@"pub" => |pb| client.send(io, .{ - .msg = try pb.toMsg(alloc, subscription.sid), - }) catch |err| switch (err) { - error.Canceled => return err, - else => {}, - }, - .hpub => |hp| client.send(io, .{ .hmsg = try hp.toHMsg( - alloc, - subscription.sid, - ) }) catch |err| switch (err) { - error.Canceled => return err, - else => {}, - }, - else => unreachable, - } - } - } - if (source_client.connect) |c| { - if (c.verbose) { - source_client.send(io, .@"+ok") catch {}; - } - } -} - -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); - const subject = try gpa.dupe(u8, msg.subject); - errdefer gpa.free(subject); - const sid = try gpa.dupe(u8, msg.sid); - errdefer gpa.free(sid); - try server.subscriptions.append(gpa, .{ - .subject = subject, - .client_id = id, - .sid = sid, - }); -} - -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) |from_end| { - const i = len - from_end - 1; - const sub = server.subscriptions.items[i]; - if (sub.client_id == id and std.mem.eql(u8, sub.sid, msg.sid)) { - sub.deinit(gpa); - _ = server.subscriptions.swapRemove(i); - } - } -} - -pub fn createId() []const u8 { - return "SERVERID"; -} - -pub fn createName() []const u8 { - return "SERVERNAME"; -} diff --git a/src/server/test.zig b/src/server/test.zig deleted file mode 100644 index e89e354..0000000 --- a/src/server/test.zig +++ /dev/null @@ -1,3 +0,0 @@ -const std = @import("std"); -const Server = @import("./main.zig"); -const Client = @import("./client.zig").ClientState; -- cgit