spb/market-atlas

1import WebSocket from "ws";2import { backoffMs } from "./ratelimit.js";34export interface ManagedWebSocketOptions {5  /** Called on every (re)connect; send subscriptions here. */6  onOpen?: (ws: ManagedWebSocket) => void;7  onMessage: (data: string, ws: ManagedWebSocket) => void;8  onClose?: (code: number, reason: string) => void;9  onError?: (err: Error) => void;10  /** Consider the connection stale after this many ms without any frame (default 60 s). */11  staleAfterMs?: number;12  /** Ping payload sent when idle (WebSocket ping frame is used when null). */13  heartbeat?: { intervalMs: number; message?: string } | null;14  headers?: Record<string, string>;15  maxBackoffMs?: number;16  label?: string;17}1819type Listener = (event: "connected" | "reconnecting" | "closed" | "stale", detail?: Record<string, unknown>) => void;2021/**22 * Reconnecting WebSocket with heartbeat and stale detection.23 * One connection may carry many instruments — never open one socket per symbol.24 */25export class ManagedWebSocket {26  private ws: WebSocket | null = null;27  private attempt = 0;28  private closedByUser = false;29  private lastFrameAt = 0;30  private staleTimer: NodeJS.Timeout | null = null;31  private heartbeatTimer: NodeJS.Timeout | null = null;32  private reconnectTimer: NodeJS.Timeout | null = null;33  private listeners: Listener[] = [];34  reconnects = 0;35  connectedAt: number | null = null;3637  constructor(38    public readonly url: string,39    private opts: ManagedWebSocketOptions,40  ) {}4142  on(listener: Listener) {43    this.listeners.push(listener);44  }4546  private emit(...args: Parameters<Listener>) {47    for (const l of this.listeners) l(...args);48  }4950  connect(): void {51    this.closedByUser = false;52    this.open();53  }5455  private open() {56    if (this.closedByUser) return;57    const ws = new WebSocket(this.url, { headers: this.opts.headers, handshakeTimeout: 15_000 });58    this.ws = ws;59    ws.on("open", () => {60      this.attempt = 0;61      this.connectedAt = Date.now();62      this.lastFrameAt = Date.now();63      this.emit("connected");64      this.armStale();65      this.armHeartbeat();66      this.opts.onOpen?.(this);67    });68    ws.on("message", (data) => {69      this.lastFrameAt = Date.now();70      try {71        this.opts.onMessage(typeof data === "string" ? data : data.toString("utf8"), this);72      } catch (err) {73        this.opts.onError?.(err instanceof Error ? err : new Error(String(err)));74      }75    });76    ws.on("pong", () => (this.lastFrameAt = Date.now()));77    ws.on("error", (err) => this.opts.onError?.(err));78    ws.on("close", (code, reason) => {79      this.clearTimers();80      this.ws = null;81      this.opts.onClose?.(code, reason.toString());82      if (!this.closedByUser) this.scheduleReconnect();83      else this.emit("closed");84    });85  }8687  private scheduleReconnect() {88    if (this.reconnectTimer) return;89    const wait = backoffMs(this.attempt, 1000, this.opts.maxBackoffMs ?? 120_000);90    this.attempt++;91    this.reconnects++;92    this.emit("reconnecting", { attempt: this.attempt, waitMs: wait });93    this.reconnectTimer = setTimeout(() => {94      this.reconnectTimer = null;95      this.open();96    }, wait);97  }9899  private armStale() {100    const staleAfter = this.opts.staleAfterMs ?? 60_000;101    if (this.staleTimer) clearInterval(this.staleTimer);102    this.staleTimer = setInterval(() => {103      if (Date.now() - this.lastFrameAt > staleAfter) {104        this.emit("stale", { silentMs: Date.now() - this.lastFrameAt });105        this.ws?.terminate();106      }107    }, Math.max(5000, Math.floor(staleAfter / 3)));108  }109110  private armHeartbeat() {111    const hb = this.opts.heartbeat;112    if (hb === null) return;113    const intervalMs = hb?.intervalMs ?? 25_000;114    this.heartbeatTimer = setInterval(() => {115      if (this.ws?.readyState !== WebSocket.OPEN) return;116      if (hb?.message) this.ws.send(hb.message);117      else this.ws.ping();118    }, intervalMs);119  }120121  private clearTimers() {122    if (this.staleTimer) clearInterval(this.staleTimer);123    if (this.heartbeatTimer) clearInterval(this.heartbeatTimer);124    this.staleTimer = null;125    this.heartbeatTimer = null;126  }127128  send(data: string | object): boolean {129    if (this.ws?.readyState !== WebSocket.OPEN) return false;130    this.ws.send(typeof data === "string" ? data : JSON.stringify(data));131    return true;132  }133134  isOpen(): boolean {135    return this.ws?.readyState === WebSocket.OPEN;136  }137138  close(): void {139    this.closedByUser = true;140    if (this.reconnectTimer) clearTimeout(this.reconnectTimer);141    this.reconnectTimer = null;142    this.clearTimers();143    this.ws?.close();144    this.ws = null;145  }146}147