more graceful shutdown handling

This commit is contained in:
Emil Lerch 2026-06-29 20:20:17 -07:00
parent 279659188f
commit eadba4696c
Signed by: lobo
GPG key ID: A7B62D657EF764F8

View file

@ -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"};