diff options
| -rw-r--r-- | src/main.zig | 27 | ||||
| -rw-r--r-- | src/server/client.zig | 58 | ||||
| -rw-r--r-- | src/server/main.zig | 274 | ||||
| -rw-r--r-- | src/server/message_parser.zig | 154 |
4 files changed, 369 insertions, 144 deletions
diff --git a/src/main.zig b/src/main.zig index 367cb62..b0588eb 100644 --- a/src/main.zig +++ b/src/main.zig @@ -51,7 +51,7 @@ pub fn main() !void { .server_id = zits.Server.createId(), .server_name = zits.Server.createName(), .version = "zits-master", - .max_payload = 99999, + .max_payload = 1048576, }; if (serve_matches.getSingleValue("port")) |port| { info.port = std.fmt.parseUnsigned(@TypeOf(info.port), port, 10) catch |err| std.process.fatal("Could not parse port ({s}): {}\n", .{ port, err }); @@ -70,31 +70,6 @@ pub fn main() !void { // _ = iter; // _ = main_args; -// var threaded: std.Io.Threaded = .init(gpa); -// defer threaded.deinit(); -// const io = threaded.io(); - -// const info: ServerInfo = .{ -// .server_id = "NBEK5DBBB4ZO5LTBGPXACZSB2QUTODC6GGN5NLOSPIGSRFWJID4XU52C", -// .server_name = "bar", -// .version = "2.11.8", -// .go = "go1.24.6", -// .headers = true, -// .max_payload = 1048576, -// }; - -// var server = try std.Io.net.IpAddress.listen( -// .{ -// .ip4 = .{ -// .bytes = .{ 0, 0, 0, 0 }, -// .port = info.port, -// }, -// }, -// io, -// .{}, -// ); -// defer server.deinit(io); - // var group: std.Io.Group = .init; // defer group.wait(io); // for (0..5) |_| { diff --git a/src/server/client.zig b/src/server/client.zig index bc02611..037a4fa 100644 --- a/src/server/client.zig +++ b/src/server/client.zig @@ -1,17 +1,51 @@ const Message = @import("message_parser.zig").Message; const std = @import("std"); -const ClientState = struct { - id: u32, - /// Used to back `connect` strings. - string_buffer: [4096]u8, +pub const ClientState = struct { + id: usize, connect: Message.Connect, - send_queue: std.Io.Queue(Message) = blk: { - var send_queue_buffer: [1024]Message = undefined; - break :blk .init(&send_queue_buffer); - }, - recv_queue: std.Io.Queue(Message) = blk: { - var recv_queue_buffer: [1024]Message = undefined; - break :blk .init(&recv_queue_buffer); - }, + /// Messages that this client is trying to send. + send_queue: std.Io.Queue(Message), + send_queue_buffer: [1024]Message = undefined, + /// Messages that this client should receive. + recv_queue: std.Io.Queue(Message), + recv_queue_buffer: [1024]Message = undefined, + + /// Mapping of the subjects to the IDs. + subscription_ids: std.StringHashMapUnmanaged([]const u8) = .empty, + + pub fn init(id: usize, connect: Message.Connect) ClientState { + var res: ClientState = .{ + .id = id, + .connect = connect, + .send_queue = undefined, + .recv_queue = undefined, + }; + res.send_queue = .init(&res.send_queue_buffer); + res.recv_queue = .init(&res.recv_queue_buffer); + return res; + } + + pub fn deinit(self: *ClientState, alloc: std.mem.Allocator) void { + self.subscription_ids.clearAndFree(alloc); + } + + pub fn subscribe(self: *ClientState, gpa: std.mem.Allocator, sub: Message.Sub) !void { + self.subscription_ids.lockPointers(); + defer self.subscription_ids.unlockPointers(); + try self.subscription_ids.putNoClobber(gpa, sub.subject, sub.sid); + } + + pub fn unsubscribe(self: *ClientState, gpa: std.mem.Allocator, sid: []const u8) void { + _ = gpa; + self.subscription_ids.lockPointers(); + defer self.subscription_ids.unlockPointers(); + var iter = self.subscription_ids.iterator(); + while (iter.next()) |entry| { + if (std.mem.eql(u8, sid, entry.value_ptr.*)) { + self.subscription_ids.swapRemove(entry.value_ptr.*); + break; + } + } else unreachable; // Assert that the SID is in the subscriptions. + } }; diff --git a/src/server/main.zig b/src/server/main.zig index b26f4b0..6d82e2d 100644 --- a/src/server/main.zig +++ b/src/server/main.zig @@ -1,47 +1,39 @@ const std = @import("std"); -const Message = @import("./message_parser.zig"); - -const ClientState = @import("./client.zig"); - -pub const ServerInfo = struct { - /// The unique identifier of the NATS server. - server_id: []const u8, - /// The name of the NATS server. - server_name: []const u8, - /// The version of NATS. - version: []const u8, - /// The version of golang the NATS server was built with. - go: []const u8 = "0.0.0", - /// The IP address used to start the NATS server, - /// by default this will be 0.0.0.0 and can be - /// configured with -client_advertise host:port. - host: []const u8 = "0.0.0.0", - /// The port number the NATS server is configured - /// to listen on. - port: u16 = 4222, - /// Whether the server supports headers. - headers: bool = false, - /// Maximum payload size, in bytes, that the server - /// will accept from the client. - max_payload: u64, - /// An integer indicating the protocol version of - /// the server. The server version 1.2.0 sets this - /// to 1 to indicate that it supports the "Echo" - /// feature. - proto: u32 = 1, -}; - -server_info: ServerInfo, -clients: std.AutoHashMapUnmanaged(u64, ClientState) = .empty, -/// Map of subjects to client IDs that are subscribed to that subject. -subscriptions: std.StringHashMapUnmanaged(std.ArrayList(u64)), - -pub fn main(gpa: std.mem.Allocator, main_args: anytype) !void { - _ = gpa; - _ = main_args; +const Message = @import("./message_parser.zig").Message; +pub const ServerInfo = Message.ServerInfo; + +const ClientState = @import("./client.zig").ClientState; +const Server = @This(); + +info: ServerInfo, +clients: std.AutoHashMapUnmanaged(usize, ClientState) = .empty, +/// Map of subjects to a map of (client ID => SID) +subscriptions: std.StringHashMapUnmanaged(std.AutoHashMapUnmanaged(usize, []const u8)) = .empty, + +pub fn main(gpa: std.mem.Allocator, server_config: ServerInfo) !void { + var server: Server = .{ + .info = server_config, + }; + + var threaded: std.Io.Threaded = .init(gpa); + defer threaded.deinit(); + const io = threaded.io(); + + var tcp_server = try std.Io.net.IpAddress.listen(try std.Io.net.IpAddress.parse( + server.info.host, + server.info.port, + ), io, .{}); + defer tcp_server.deinit(io); + + var id: usize = 0; + while (true) : (id +%= 1) { + if (server.clients.contains(id)) continue; + const stream = try tcp_server.accept(io); + _ = io.async(handleConnection, .{ gpa, io, id, stream, &server }); + } } -fn handleConnection(allocator: std.mem.Allocator, io: std.Io, stream: std.Io.net.Stream, info: ServerInfo) void { +fn handleConnection(allocator: std.mem.Allocator, io: std.Io, id: usize, stream: std.Io.net.Stream, server: *Server) !void { defer stream.close(io); var w_buffer: [1024]u8 = undefined; var writer = stream.writer(io, &w_buffer); @@ -51,42 +43,169 @@ fn handleConnection(allocator: std.mem.Allocator, io: std.Io, stream: std.Io.net var reader = stream.reader(io, &r_buffer); const in = &reader.interface; - processClient(allocator, in, out, info) catch |err| { + writeInfo(out, server.info) catch return; + var connect_arena: std.heap.ArenaAllocator = .init(allocator); + defer connect_arena.deinit(); + const connect = (Message.next(connect_arena.allocator(), in) catch return).connect; + var client_state: ClientState = .init(id, connect); + + // server.clients.lockPointers(); + try server.clients.put(allocator, id, client_state); + // server.clients.unlockPointers(); + + // 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(); + // } + + processClient(allocator, io, in, out, server, &client_state) catch |err| { std.debug.panic("Error processing client: {}\n", .{err}); }; } -fn processClient(gpa: std.mem.Allocator, in: *std.Io.Reader, out: *std.Io.Writer, info: ServerInfo) !void { - try writeInfo(out, info); +fn processClient(gpa: std.mem.Allocator, io: std.Io, in: *std.Io.Reader, out: *std.Io.Writer, server: *Server, client_state: *ClientState) !void { + var parse_task = io.async(parseMessages, .{ gpa, io, in, client_state }); + defer if (parse_task.cancel(io)) {} else |err| { + std.debug.print("Error canceling parse_task for {d}: {any}", .{ client_state.id, err }); + }; + var write_task = io.async(writeMessages, .{ io, out, server.*, client_state }); + defer if (write_task.cancel(io)) {} else |err| { + std.debug.print("Error canceling write_task for {d}: {any}", .{ client_state.id, err }); + }; - var client_state_arena: std.heap.ArenaAllocator = .init(gpa); - defer client_state_arena.deinit(); - const client_state = (try Message.next(client_state_arena.allocator(), in)).connect; - _ = client_state; + while (!io.cancelRequested()) { + if (client_state.send_queue.getOne(io)) |msg| { + switch (msg) { + // TODO: REMOVE + .not_real => {}, + // 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); - var message_parsing_arena: std.heap.ArenaAllocator = .init(gpa); - defer message_parsing_arena.deinit(); - const message_parsing_allocator = message_parsing_arena.allocator(); - while (true) { - defer _ = message_parsing_arena.reset(.retain_capacity); - const next_message = Message.next(message_parsing_allocator, in) catch |err| { - switch (err) { - error.EndOfStream => { - break; + std.debug.print("subs: {any}\n", .{server.subscriptions}); }, - else => { - return err; + .info => |info| { + std.debug.panic("got an info message? : {any}\n", .{info}); + }, + else => |m| { + std.debug.panic("Unimplemented: {any}\n", .{m}); }, } - }; - 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}), - } + } else |err| return err; + } + + // 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}), + // } + // } +} + +fn parseMessages(gpa: std.mem.Allocator, io: std.Io, in: *std.Io.Reader, client_state: *ClientState) !void { + // var message_parsing_arena: std.heap.ArenaAllocator = .init(); + // defer message_parsing_arena.deinit(); + // const message_parsing_allocator = message_parsing_arena.allocator(); + + while (!io.cancelRequested()) { + // defer _ = message_parsing_arena.reset(.retain_capacity); + const next_message = // Message.next(message_parsing_allocator, in) + Message.next(gpa, in) catch |err| { + switch (err) { + error.EndOfStream => { + break; + }, + else => { + return err; + }, + } + }; + std.debug.print("received message from client {d}: {any}\n'{s}'\n", .{ client_state.id, next_message, in.buffered() }); + try client_state.send_queue.putOne(io, next_message); + } +} + +fn writeMessages(io: std.Io, out: *std.Io.Writer, server: Server, client_state: *ClientState) !void { + while (true) { + std.debug.print("in writeMessage loop for {d}\n", .{client_state.id}); + if (client_state.recv_queue.getOne(io)) |msg| { + std.debug.print("attempting to write message for {d}: {any}\n", .{ client_state.id, msg }); + switch (msg) { + .@"+ok" => try writeOk(out), + .pong => try writePong(out), + .info => try writeInfo(out, server.info), + .msg => |m| try writeMsg(out, m), + else => std.debug.panic("unimplemented write", .{}), + } + } else |err| return err; } } @@ -107,6 +226,21 @@ fn writeInfo(out: *std.Io.Writer, info: ServerInfo) !void { try out.flush(); } +fn writeMsg(out: *std.Io.Writer, msg: Message.Msg) !void { + std.debug.print("PRINTING MESSAGE\n\n\n\n", .{}); + try out.print( + "MSG {s} {s} {s} {d}\r\n{s}\r\n", + .{ + msg.subject, + msg.sid, + msg.reply_to orelse "", + msg.payload.len, + msg.payload, + }, + ); + try out.flush(); +} + pub fn createId() []const u8 { return "SERVERID"; } @@ -114,3 +248,7 @@ pub fn createId() []const u8 { pub fn createName() []const u8 { return "SERVERNAME"; } + +// test "handle pub" { +// const io = std.testing.io; +// } diff --git a/src/server/message_parser.zig b/src/server/message_parser.zig index 75e13d2..1a6b213 100644 --- a/src/server/message_parser.zig +++ b/src/server/message_parser.zig @@ -32,20 +32,48 @@ pub const MessageType = enum { }; pub const Message = union(enum) { - info: void, - + /// TODO: REMOVE + not_real: void, + info: ServerInfo, connect: Connect, @"pub": Pub, hpub: void, - sub: void, + sub: Sub, unsub: void, - msg: void, + msg: Msg, hmsg: void, ping, pong, @"+ok": void, @"-err": void, - const Connect = struct { + pub const ServerInfo = struct { + /// The unique identifier of the NATS server. + server_id: []const u8, + /// The name of the NATS server. + server_name: []const u8, + /// The version of NATS. + version: []const u8, + /// The version of golang the NATS server was built with. + go: []const u8 = "0.0.0", + /// The IP address used to start the NATS server, + /// by default this will be 0.0.0.0 and can be + /// configured with -client_advertise host:port. + host: []const u8 = "0.0.0.0", + /// The port number the NATS server is configured + /// to listen on. + port: u16 = 4222, + /// Whether the server supports headers. + headers: bool = false, + /// Maximum payload size, in bytes, that the server + /// will accept from the client. + max_payload: u64, + /// An integer indicating the protocol version of + /// the server. The server version 1.2.0 sets this + /// to 1 to indicate that it supports the "Echo" + /// feature. + proto: u32 = 1, + }; + pub const Connect = struct { verbose: bool = false, pedantic: bool = false, tls_required: bool = false, @@ -63,9 +91,26 @@ pub const Message = union(enum) { headers: ?bool = null, nkey: ?[]const u8 = null, }; - const Pub = struct { + pub const Pub = struct { + /// The destination subject to publish to. subject: []const u8, + /// 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, + }; + pub const Sub = struct { + /// The subject name to subscribe to. + subject: []const u8, + /// If specified, the subscriber will join this queue group. + queue_group: ?[]const u8, + /// A unique alphanumeric subscription ID, generated by the client. + sid: []const u8, + }; + pub const Msg = struct { + subject: []const u8, + sid: []const u8, + reply_to: ?[]const u8, payload: []const u8, }; @@ -132,32 +177,7 @@ pub const Message = union(enum) { try in.discardAll(1); // throw away space // Parse subject - const subject: []const u8 = blk: { - // TODO: should be ARENA allocator - var subject_list: std.ArrayList(u8) = try .initCapacity(alloc, 1024); - - // Handle the first character - { - const byte = try in.takeByte(); - if (byte == '.' or std.ascii.isWhitespace(byte)) - return error.InvalidSubject; - - try subject_list.append(alloc, byte); - } - - while (in.takeByte() catch null) |byte| { - if (std.ascii.isWhitespace(byte)) break; - if (std.ascii.isAscii(byte)) { - if (byte == '.') { - const next_byte = try in.peekByte(); - if (next_byte == '.' or std.ascii.isWhitespace(next_byte)) - return error.InvalidSubject; - } - try subject_list.append(alloc, byte); - } - } else return error.InvalidStream; - break :blk subject_list.items; - }; + const subject: []const u8 = try readSubject(alloc, in); // Parse byte count const byte_count = blk: { @@ -186,12 +206,12 @@ pub const Message = union(enum) { break :blk bytes; }; - std.debug.print("buffer: '{s}'\n", .{in.buffered()}); - // return std.debug.panic("not implemented", .{}); - return .{ .@"pub" = .{ - .subject = subject, - .payload = payload, - } }; + return .{ + .@"pub" = .{ + .subject = subject, + .payload = payload, + }, + }; }, .ping => { std.debug.assert(std.mem.eql(u8, try in.take(2), "\r\n")); @@ -201,11 +221,69 @@ pub const Message = union(enum) { std.debug.assert(std.mem.eql(u8, try in.take(2), "\r\n")); return .pong; }, + .sub => { + std.debug.assert(std.ascii.isWhitespace(try in.takeByte())); + const subject = try readSubject(alloc, in); + const second = blk: { + // Drop whitespace + while (in.peekByte() catch null) |byte| { + if (std.ascii.isWhitespace(byte)) { + in.toss(1); + } else break; + } else return error.InvalidStream; + + var acc: std.ArrayList(u8) = try .initCapacity(alloc, 32); + while (in.takeByte() catch null) |byte| { + if (std.ascii.isWhitespace(byte)) break; + try acc.append(alloc, byte); + } else return error.InvalidStream; + + break :blk try acc.toOwnedSlice(alloc); + }; + const queue_group = if ((try in.peekByte()) != '\r') second else null; + const sid = if (queue_group) |_| try in.takeDelimiterExclusive('\r') else second; + std.debug.assert(std.mem.eql(u8, try in.take(2), "\r\n")); + return .{ + .sub = .{ + .subject = subject, + .queue_group = queue_group, + .sid = sid, + }, + }; + }, else => |msg| std.debug.panic("Not implemented: {}\n", .{msg}), } } }; +fn readSubject(alloc: std.mem.Allocator, in: *std.Io.Reader) ![]const u8 { + // TODO: should be ARENA allocator + var subject_list: std.ArrayList(u8) = try .initCapacity(alloc, 1024); + + // Handle the first character + { + const byte = try in.takeByte(); + std.debug.assert(!std.ascii.isWhitespace(byte)); + if (byte == '.') + return error.InvalidSubject; + + try subject_list.append(alloc, byte); + } + + while (in.takeByte() catch null) |byte| { + if (std.ascii.isWhitespace(byte)) break; + if (std.ascii.isAscii(byte)) { + if (byte == '.') { + const next_byte = try in.peekByte(); + if (next_byte == '.' or std.ascii.isWhitespace(next_byte)) + return error.InvalidSubject; + } + try subject_list.append(alloc, byte); + } + } else return error.InvalidStream; + return subject_list.toOwnedSlice(alloc); +} + fn parseJsonMessage(T: type, alloc: std.mem.Allocator, in: *std.Io.Reader) !T { var reader: std.json.Reader = .init(alloc, in); return std.json.innerParse(T, alloc, &reader, .{ |