diff --git a/src/internal/streaming/firehose.zig b/src/internal/streaming/firehose.zig index a462c29..b23996e 100644 --- a/src/internal/streaming/firehose.zig +++ b/src/internal/streaming/firehose.zig @@ -708,15 +708,23 @@ pub const FirehoseClient = struct { /// declared, raw websocket frames are delivered INSTEAD of decoded /// events (the caller owns decoding; cursor tracking still works by /// peeking the seq). for consumers that decode on their own threads. - /// blocks forever — reconnects with exponential backoff on disconnect. - /// rotates through hosts on each reconnect attempt. + /// optional: fn shouldStop(*@TypeOf(handler)) bool — checked before + /// every frame delivery and around every reconnect. returning true + /// makes subscribe return cleanly (bounded consumers: bootstrap + /// capture, tests, graceful shutdown). without it, an interrupted + /// read is indistinguishable from a connection error and the + /// reconnect loop absorbs it forever; pair a stop flag with + /// future.cancel to unblock an idle read. + /// blocks until shouldStop (forever without it) — reconnects with + /// exponential backoff on disconnect. rotates through hosts on each + /// reconnect attempt. pub fn subscribe(self: *FirehoseClient, handler: anytype) Io.Cancelable!void { var backoff: u64 = 1; var host_index: usize = 0; const max_backoff: u64 = 60; var prev_host_index: usize = 0; - while (true) { + while (!stopRequested(handler)) { const host = self.options.hosts[host_index % self.options.hosts.len]; const effective_index = host_index % self.options.hosts.len; @@ -728,12 +736,16 @@ pub const FirehoseClient = struct { log.info("connecting to host {d}/{d}: {s}", .{ effective_index + 1, self.options.hosts.len, host }); self.connectAndRead(host, handler) catch |err| { + // a stop-interrupted read surfaces as a connection error; + // don't report it as one + if (stopRequested(handler)) return; if (comptime @hasDecl(@TypeOf(handler.*), "onError")) { handler.onError(err); } else { log.err("firehose error: {s}, reconnecting in {d}s...", .{ @errorName(err), backoff }); } }; + if (stopRequested(handler)) return; prev_host_index = effective_index; host_index += 1; @@ -786,6 +798,14 @@ pub const FirehoseClient = struct { } }; +/// true iff the handler declares shouldStop and it currently returns true +fn stopRequested(handler: anytype) bool { + if (comptime @hasDecl(@TypeOf(handler.*), "shouldStop")) { + return handler.shouldStop(); + } + return false; +} + fn WsHandler(comptime H: type) type { return struct { allocator: Allocator, @@ -795,6 +815,10 @@ fn WsHandler(comptime H: type) type { const Self = @This(); pub fn serverMessage(self: *Self, data: []const u8) !void { + if (comptime @hasDecl(H, "shouldStop")) { + // breaks the read loop; subscribe sees the stop and returns + if (self.handler.shouldStop()) return error.SubscriptionStopped; + } if (comptime @hasDecl(H, "onRawFrame")) { // track the reconnect cursor with a cheap peek (header + // payload map; no CAR hydration — that's the caller's job) @@ -1346,3 +1370,47 @@ test "peekSeq extracts seq without full decode" { try std.testing.expectEqual(@as(?i64, 42), peekSeq(std.testing.allocator, &frame)); try std.testing.expectEqual(@as(?i64, null), peekSeq(std.testing.allocator, "junk")); } + +test "shouldStop hook: subscribe returns cleanly instead of reconnecting" { + var threaded: Io.Threaded = .init(std.testing.allocator, .{}); + defer threaded.deinit(); + const io = threaded.io(); + + const StopHandler = struct { + stop: bool = false, + errors: usize = 0, + fn onEvent(_: *@This(), _: Event) void {} + fn shouldStop(self: *@This()) bool { + return self.stop; + } + fn onError(self: *@This(), _: anyerror) void { + // first failed connect (port 9, nothing listening) requests stop + self.errors += 1; + self.stop = true; + } + }; + + var client = FirehoseClient.init(io, std.testing.allocator, .{ + .hosts = &.{"ws://127.0.0.1:9"}, + }); + defer client.deinit(); + + // stop set before the first connect: returns without dialing + var pre_stopped: StopHandler = .{ .stop = true }; + try client.subscribe(&pre_stopped); + try std.testing.expectEqual(@as(usize, 0), pre_stopped.errors); + + // stop set from onError: returns after the first failed connect, + // without sleeping into the reconnect backoff + var stops_on_error: StopHandler = .{}; + try client.subscribe(&stops_on_error); + try std.testing.expectEqual(@as(usize, 1), stops_on_error.errors); +} + +test "handlers without shouldStop never stop-request" { + const Plain = struct { + fn onEvent(_: *@This(), _: Event) void {} + }; + var h: Plain = .{}; + try std.testing.expect(!stopRequested(&h)); +}