From eadba4696c51ef82604bf11956ee06675bd3156f Mon Sep 17 00:00:00 2001 From: Emil Lerch Date: Mon, 29 Jun 2026 20:20:17 -0700 Subject: [PATCH] more graceful shutdown handling --- src/net/LiveStream.zig | 125 ++++++++++++++++++++++++----------------- 1 file changed, 73 insertions(+), 52 deletions(-) diff --git a/src/net/LiveStream.zig b/src/net/LiveStream.zig index 8038887..8a380da 100644 --- a/src/net/LiveStream.zig +++ b/src/net/LiveStream.zig @@ -27,22 +27,28 @@ //! //! ## Why a raw background thread (not std.Io tasks) //! -//! The read loop is a single long-lived blocking loop (`readLoop`), -//! so it runs on its own OS thread that coexists with the app's -//! `Io.Threaded` event loop. The pinned `websocket.zig` client is -//! itself Io-aware (connect, TLS, reads, writes go through `std.Io`, -//! which is safe to call from this thread); we thread `io` in for it. +//! The read loop is a single long-lived loop on its own OS thread that +//! coexists with the app's `Io.Threaded` event loop. The pinned +//! `websocket.zig` client is itself Io-aware (connect, TLS, reads, +//! writes go through `std.Io`, which is safe to call from this thread); +//! we thread `io` in for it. //! //! ## Threading and shutdown //! -//! `prices` and `client` are guarded by `mutex` (a `std.atomic.Mutex` -//! spinlock, since the read loop runs outside the Io task model and -//! critical sections are single hashmap ops held for microseconds). -//! Shutdown: `stop()` sets `should_stop` and calls `client.close()` -//! from the caller's thread, which unblocks the read thread's -//! `readLoop`, after which we `join`. The blocking connect (DNS + TCP -//! + TLS) runs WITHOUT holding `mutex`, so a quit during a slow connect -//! isn't blocked. +//! `prices` is guarded by `mutex` (a `std.atomic.Mutex` spinlock, since +//! the read loop runs outside the Io task model and critical sections +//! are single hashmap ops held for microseconds). +//! +//! Shutdown is COOPERATIVE: the read loop polls with a short read +//! timeout (`read_poll_ms`) and re-checks `should_stop` between reads, +//! so `stop()` just sets `should_stop` and `join`s - the read thread +//! returns on its OWN thread and tears the client down there. We do NOT +//! close the socket from another thread: closing an Io-aware TLS stream +//! while a read is in flight crashes the std TLS reader (the read +//! thread faults dereferencing the half-torn-down client). The only +//! cost is up to one `read_poll_ms` of latency on quit, plus (if a quit +//! lands during the blocking connect/handshake) that operation's own +//! timeout. const std = @import("std"); const websocket = @import("websocket"); @@ -65,6 +71,10 @@ const max_ticker_len = 24; const handshake_timeout_ms = 10_000; const max_message_size = 512 * 1024; const read_buffer_size = 32 * 1024; +/// Read-poll timeout (ms). The read loop wakes at least this often to +/// re-check `should_stop`, bounding quit latency while streaming. Short +/// enough to feel instant, long enough to be free when idle. +const read_poll_ms: u32 = 250; const reconnect_backoff_ms = 2_000; /// Yahoo pricing protobufs are ~60-90 bytes decoded; 1 KiB is ample. /// A frame whose decoded size exceeds this is skipped (returns null). @@ -110,16 +120,13 @@ transport: Transport, api_key: ?[]u8 = null, /// Owned, duped tickers to subscribe. Set in `start`, freed in `stop`. tickers: [][]u8 = &.{}, -/// Guards `prices` and `client`. Spinlock (not an Io mutex) because the -/// read loop is a raw OS thread outside the Io task model; critical -/// sections are one hashmap op or a short pointer read. +/// Guards `prices`. Spinlock (not an Io mutex) because the read loop is +/// a raw OS thread outside the Io task model; critical sections are one +/// hashmap op or a short pointer read. mutex: std.atomic.Mutex = .unlocked, /// Latest price per ticker, keyed by UPPERCASED ticker (owned dup). /// Written by the read thread, read via `snapshotInto`. prices: std.StringHashMap(f64), -/// The live client, published so `stop()` can close it cross-thread. -/// Heap-owned by the read thread for the duration of one connection. -client: ?*websocket.Client = null, thread: ?std.Thread = null, should_stop: std.atomic.Value(bool) = .init(false), @@ -173,16 +180,13 @@ pub fn start(self: *LiveStream, symbols: []const []const u8) !void { self.thread = try std.Thread.spawn(.{}, threadMain, .{self}); } -/// Signal the read thread to stop, close the connection to unblock it, -/// join, and free the subscribed-ticker list. Safe to call when not -/// running. +/// Signal the read thread to stop and join it, then free the +/// subscribed-ticker list. Cooperative: the read loop notices +/// `should_stop` within one `read_poll_ms` and returns on its own +/// thread (we never touch the socket/client from here - see the +/// shutdown note at the top). Safe to call when not running. pub fn stop(self: *LiveStream) void { self.should_stop.store(true, .release); - // Close the socket so the read thread's blocking readLoop returns. - // close() is thread-safe per the library. - self.lock(); - if (self.client) |c| c.close(.{}) catch |err| log.debug("close on stop: {t}", .{err}); - self.unlock(); if (self.thread) |t| { t.join(); @@ -226,11 +230,11 @@ pub fn snapshotInto(self: *LiveStream, symbols: []const []const u8, out: *std.St } } -/// Read-thread handler hook (called by `websocket.Client.readLoop`). -/// Parses one frame and updates `prices`. Never propagates an error (a -/// bad frame must not tear down the stream); failures are swallowed so -/// the connection stays up. -pub fn serverMessage(self: *LiveStream, data: []const u8) !void { +/// Parse one frame and update `prices`. Never fails: a bad/unparseable +/// frame is simply ignored (returns without touching `prices`) so the +/// connection stays up. Returns `void` precisely because there's no +/// error worth propagating to the read loop. +pub fn serverMessage(self: *LiveStream, data: []const u8) void { const tick = switch (self.transport) { .yahoo => parseYahooMessage(self.allocator, data), } orelse return; @@ -275,8 +279,6 @@ fn transportConfig(self: *LiveStream) TransportConfig { fn connectAndRun(self: *LiveStream) !void { const cfg = self.transportConfig(); - // Blocking connect (DNS + TCP + TLS) happens BEFORE publishing the - // client, so a concurrent stop() isn't blocked waiting on it. const cptr = try self.allocator.create(websocket.Client); cptr.* = websocket.Client.init(self.io, self.allocator, .{ .host = cfg.host, @@ -288,29 +290,48 @@ fn connectAndRun(self: *LiveStream) !void { self.allocator.destroy(cptr); return err; }; - // From here cptr owns a live connection; always tear down + unpublish. + // The client is created, used, and torn down entirely on THIS (read) + // thread - no other thread ever touches it. That's what makes the + // cooperative shutdown safe (see the shutdown note at the top): on + // `should_stop`, `readUntilStopped` returns and this defer closes + + // frees the client here, not from the quitting thread. defer { - self.lock(); - self.client = null; - self.unlock(); cptr.deinit(); self.allocator.destroy(cptr); } - // Publish so stop() can close the socket to unblock readLoop. - self.lock(); - self.client = cptr; - self.unlock(); - try cptr.handshake(cfg.path, .{ .timeout_ms = handshake_timeout_ms, .headers = cfg.headers, }); try self.sendSubscribe(cptr); - // Blocks until the connection closes (or stop() closes it). The - // library auto-handles ping/pong and close frames; our - // serverMessage handles the text data frames. - try cptr.readLoop(self); + try self.readUntilStopped(cptr); +} + +/// Cooperative read loop. Polls with a `read_poll_ms` timeout so a +/// no-data wait returns `null` (WouldBlock) and we re-check +/// `should_stop` between reads - letting the thread exit on its own +/// when asked to stop. Replaces websocket.zig's `readLoop`, which +/// blocks indefinitely (read timeout 0) and trips `unreachable` on a +/// timeout, so it can only be unblocked by a cross-thread close - the +/// exact thing that crashes the Io-aware TLS reader. We replicate its +/// frame handling: text/binary -> parser; ping -> pong; close -> done. +fn readUntilStopped(self: *LiveStream, c: *websocket.Client) !void { + try c.readTimeout(read_poll_ms); + while (!self.should_stop.load(.acquire)) { + const message = c.read() catch |err| switch (err) { + error.Closed => return, + else => return err, + } orelse continue; // timeout / no data yet -> re-check should_stop + defer c.done(message); + switch (message.type) { + .text, .binary => self.serverMessage(message.data), + // @constCast is safe: data points into our own read buffer. + .ping => c.writeFrame(.pong, @constCast(message.data)) catch |err| log.debug("pong: {t}", .{err}), + .close => return, + .pong => {}, + } + } } fn sendSubscribe(self: *LiveStream, c: *websocket.Client) !void { @@ -593,14 +614,14 @@ test "LiveStream: serverMessage stores prices; snapshotInto matches case-insensi const s = try LiveStream.create(std.testing.allocator, std.testing.io, .yahoo, null); defer s.destroy(); - try s.serverMessage( + s.serverMessage( \\{"type":"pricing","message":"CgRBQVBMFUgBjUMY4JKAz+JnKgNOTVMwCDgBRQisH79InK3UKmUAj+K/sAFQ2AEE"} ); - try s.serverMessage( + s.serverMessage( \\{"type":"pricing","message":"CgNTUFkVUhg5RBiwooDP4mcqA1BDWDAUOAFF9v3HP0jQ+bonZYA9NkGwAVDYAQQ="} ); // A non-pricing frame is ignored (early return, no entry, no crash). - try s.serverMessage("not a frame"); + s.serverMessage("not a frame"); var out = std.StringHashMap(f64).init(std.testing.allocator); defer out.deinit(); @@ -619,8 +640,8 @@ test "LiveStream: serverMessage overwrites an existing ticker (found_existing pa const frame = \\{"type":"pricing","message":"CgRBQVBMFUgBjUMY4JKAz+JnKgNOTVMwCDgBRQisH79InK3UKmUAj+K/sAFQ2AEE"} ; - try s.serverMessage(frame); - try s.serverMessage(frame); // second time reuses the existing key + s.serverMessage(frame); + s.serverMessage(frame); // second time reuses the existing key var out = std.StringHashMap(f64).init(std.testing.allocator); defer out.deinit(); const syms = [_][]const u8{"AAPL"};