From 4bca7260afa768988348a9b1a3e63b583da8aac3 Mon Sep 17 00:00:00 2001 From: DualFroz Date: Fri, 11 Sep 2026 15:12:40 +0200 Subject: [PATCH] feat: add handleSubscribe and watchForDisconnect to server --- src/server.zig | 158 +++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 158 insertions(+) diff --git a/src/server.zig b/src/server.zig index 216d55b..56f63b2 100644 --- a/src/server.zig +++ b/src/server.zig @@ -1383,3 +1383,161 @@ pub const Server = struct { }; defer querymod.freeRows(self.allocator, rows); + if (args.groupField.len == 0) { + const values = aggregatemod.runGlobal(self.allocator, rows, args.specs) catch |err| { + return switch (err) { + error.NonNumericField, error.Overflow => writeErr(writer, .bad_request, @errorName(err)), + else => writeErr(writer, .internal, @errorName(err)), + }; + }; + defer aggregatemod.freeGlobalResult(self.allocator, values); + const matches = if (havingPred) |hp| + aggregatemod.globalMatchesHaving(self.allocator, args.specs, values, hp) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + } + else + true; + const out = protocol.encodeAggGlobalResult(self.allocator, if (matches) values else &.{}) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer self.allocator.free(out); + try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), out); + } else { + const unpaged = aggregatemod.run(self.allocator, rows, args.groupField, args.specs) catch |err| { + return switch (err) { + error.NonNumericField, error.Overflow => writeErr(writer, .bad_request, @errorName(err)), + else => writeErr(writer, .internal, @errorName(err)), + }; + }; + const filtered = if (havingPred) |hp| + aggregatemod.filterHaving(self.allocator, unpaged, args.groupField, args.specs, hp) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + } + else + unpaged; + if (args.orderBy.len > 0) { + aggregatemod.sortGroupsBy(self.allocator, filtered, args.groupField, args.specs, args.orderBy) catch |err| { + protocol.freeAggGroupRows(self.allocator, filtered); + return switch (err) { + error.UnknownOrderField => writeErr(writer, .bad_request, @errorName(err)), + else => writeErr(writer, .internal, @errorName(err)), + }; + }; + } + const page = querymod.Page{ .offset = args.offset, .limit = args.limit }; + const groups = aggregatemod.applyPageGroups(self.allocator, filtered, page) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer protocol.freeAggGroupRows(self.allocator, groups); + const out = protocol.encodeAggGroupRows(self.allocator, groups) catch |err| { + return writeErr(writer, .internal, @errorName(err)); + }; + defer self.allocator.free(out); + try protocol.writeFrame(writer, @intFromEnum(protocol.Status.ok), out); + } + } + + /// Executes a `table_window` request: opens the named table, filters + /// rows with the same `expr.zig` grammar every other `table_*` opcode + /// uses (empty text selects every row), then hands the filtered rows to + /// `window.compute` for the single `args.spec` annotation, along with + /// `args.frame` (ROWS or RANGE - see `window.zig`'s `FrameMode` doc + /// comment) - unlike `table_aggregate`, this NEVER collapses rows, so + /// `results[i]` stays aligned 1:1 with `rows[i]` regardless of paging (a + /// `NonNumericField`/`Overflow`/`MissingField`/`RangeNotApplicable` from + /// `window.zig` is a client mistake, `bad_request`, same treatment as + /// `table_aggregate`). + /// `args.offset`/`args.limit` then window which rows get PRINTED, walked + /// in `window.orderedIndices`' partition-then-order sequence (the same + /// display order local `table window` uses, not raw scan order) rather + /// than slicing `rows`/`results` directly - the exact bug local `table + /// window` hit and fixed (see `window.zig`'s `orderedIndices` doc + /// comment) before this network path existed, so it is built paginated + /// the correct way from the start instead of repeating that mistake. + /// Each surviving row is paired with its already-computed value into a + /// `protocol.WindowRow` and sent back via `encodeWindowRows` - the + /// client prints them in the order received, no further sorting needed. + /// A `subscribe` request never returns in the normal sense: once the + /// connection is registered in `self.pubsub`, this function owns the + /// rest of the connection's lifetime (both directions of the socket) + /// until the client disconnects, then propagates `error.SubscriptionEnded` + /// so `handleConnection`'s `catch return` unwinds cleanly - the same + /// "any error just ends the connection" contract every other opcode + /// already relies on for a plain disconnect mid-frame. + /// + /// Design: a `subscribe` connection becomes receive-only after the ack. + /// A dedicated watcher thread does exactly one blocking read attempt on + /// `stream` - whatever it gets back (a byte the client sent, or an + /// error/EOF from a disconnect) closes the subscription, since there is + /// no way to both keep reading further commands AND push messages on + /// one TCP connection without either a second concurrent writer (a + /// write race against the push loop below) or re-entering this same + /// read/dispatch loop (a read race against the watcher). Unlike a + /// protocol like Redis's, which multiplexes further commands + /// (UNSUBSCRIBE, PING) onto the same subscribed connection, this + /// connection itself never reads another request - but `unsubscribe` + /// (see `handleUnsubscribe`) reaches the SAME `sub` from a completely + /// different connection, keyed by the id this function's ack carries + /// (`sub.id`), so dropping a subset of channels still does not require + /// closing the connection outright; only dropping every remaining + /// channel does (`pubsub.Hub.unsubscribe`'s doc comment). Once the push loop ends (for any reason - the + /// watcher closed `sub`, or a write here failed because the peer is + /// gone), `std.posix.shutdown` unblocks the watcher's still-pending read + /// (rather than relying on `conn.stream.close()` racing a live read on + /// the same fd from another thread, which is a well-known way to + /// recycle a file descriptor into an unrelated, concurrently-accepted + /// connection) before joining it - so by the time this returns, no + /// other thread still touches `stream` or `sub`. + fn handleSubscribe(self: *Server, payload: []const u8, stream: std.net.Stream) !void { + const args = protocol.decodeSubscribe(self.allocator, payload) catch { + return writeErr(stream.writer(), .bad_request, "malformed subscribe payload"); + }; + defer args.deinit(self.allocator); + if (args.channels.len == 0) { + return writeErr(stream.writer(), .bad_request, "subscribe needs at least one channel"); + } + + const channelNames = try self.allocator.alloc([]const u8, args.channels.len); + defer self.allocator.free(channelNames); + for (args.channels, 0..) |c, i| channelNames[i] = c; + + const sub = self.pubsub.subscribe(channelNames, pubsubmod.defaultCapacity) catch { + return writeErr(stream.writer(), .internal, "failed to register subscription"); + }; + + const ack = try protocol.encodeSubscribeAck(self.allocator, sub.id, @intCast(channelNames.len)); + defer self.allocator.free(ack); + protocol.writeFrame(stream.writer(), @intFromEnum(protocol.Status.ok), ack) catch { + self.pubsub.closeSubscriber(sub); + self.pubsub.destroySubscriber(sub); + return error.SubscriptionEnded; + }; + + const watcher = try std.Thread.spawn(.{}, watchForDisconnect, .{ &self.pubsub, sub, stream }); + + while (sub.waitMessage()) |msg| { + defer msg.deinit(self.allocator); + const framePayload = protocol.encodePubSubMessage(self.allocator, msg.channel, msg.payload) catch break; + defer self.allocator.free(framePayload); + protocol.writeFrame(stream.writer(), @intFromEnum(protocol.Status.ok), framePayload) catch break; + } + + std.posix.shutdown(stream.handle, .both) catch {}; + watcher.join(); + self.pubsub.closeSubscriber(sub); + self.pubsub.destroySubscriber(sub); + return error.SubscriptionEnded; + } + + /// The watcher half of `handleSubscribe`: blocks on exactly one read + /// attempt, then closes `sub` regardless of the outcome (data, EOF, or + /// an error - all three mean "stop pushing to this connection"). Runs + /// on its own thread so `handleSubscribe`'s push loop can block + /// independently on `sub.waitMessage()` - see that function's doc + /// comment for why this is a single read, not a loop. + fn watchForDisconnect(hub: *pubsubmod.Hub, sub: *pubsubmod.Subscriber, stream: std.net.Stream) void { + var buf: [1]u8 = undefined; + _ = stream.read(&buf) catch {}; + hub.closeSubscriber(sub); + } +