more graceful shutdown handling

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`),
+//! The read loop is a single long-lived loop on its own OS thread that
-//! so it runs on its own OS thread that coexists with the app's
+//! coexists with the app's `Io.Threaded` event loop. The pinned
-//! `Io.Threaded` event loop. The pinned `websocket.zig` client is
+//! `websocket.zig` client is itself Io-aware (connect, TLS, reads,
-//! itself Io-aware (connect, TLS, reads, writes go through `std.Io`,
+//! writes go through `std.Io`, which is safe to call from this thread);
-//! which is safe to call from this thread); we thread `io` in for it.
+//! we thread `io` in for it.
 //!
 //! ## Threading and shutdown
 //!
-//! `prices` and `client` are guarded by `mutex` (a `std.atomic.Mutex`
+//! `prices` is guarded by `mutex` (a `std.atomic.Mutex` spinlock, since
-//! spinlock, since the read loop runs outside the Io task model and
+//! the read loop runs outside the Io task model and critical sections
-//! critical sections are single hashmap ops held for microseconds).
+//! 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
+//! Shutdown is COOPERATIVE: the read loop polls with a short read
-//! `readLoop`, after which we `join`. The blocking connect (DNS + TCP
+//! timeout (`read_poll_ms`) and re-checks `should_stop` between reads,
-//! + TLS) runs WITHOUT holding `mutex`, so a quit during a slow connect
+//! so `stop()` just sets `should_stop` and `join`s - the read thread
-//! isn't blocked.
+//! 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
+/// Guards `prices`. Spinlock (not an Io mutex) because the read loop is
-/// read loop is a raw OS thread outside the Io task model; critical
+/// a raw OS thread outside the Io task model; critical sections are one
-/// sections are one hashmap op or a short pointer read.
+/// 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,
+/// Signal the read thread to stop and join it, then free the
-/// join, and free the subscribed-ticker list. Safe to call when not
+/// subscribed-ticker list. Cooperative: the read loop notices
-/// running.
+/// `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`).
+/// Parse one frame and update `prices`. Never fails: a bad/unparseable
-/// Parses one frame and updates `prices`. Never propagates an error (a
+/// frame is simply ignored (returns without touching `prices`) so the
-/// bad frame must not tear down the stream); failures are swallowed so
+/// connection stays up. Returns `void` precisely because there's no
-/// the connection stays up.
+/// error worth propagating to the read loop.
-pub fn serverMessage(self: *LiveStream, data: []const u8) !void {
+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
+    try self.readUntilStopped(cptr);
-    // library auto-handles ping/pong and close frames; our
+}
-    // serverMessage handles the text data frames.
+
-    try cptr.readLoop(self);
+/// 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);
+    s.serverMessage(frame);
-    try s.serverMessage(frame); // second time reuses the existing key
+    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"};