From cab3f89c29887c0e39d920f48e02a033b23aa79e Mon Sep 17 00:00:00 2001 From: Roscoe Rubin-Rottenberg Date: Mon, 2 Feb 2026 19:57:48 -0500 Subject: [PATCH] refactor: use ws for keepalive --- common/util.ts | 6 + deno.lock | 32 +--- sync/deno.json | 2 +- sync/firehose/index.ts | 2 +- xrpc-server/deno.json | 5 +- xrpc-server/stream/websocket-keepalive.ts | 223 +++++++--------------- 6 files changed, 90 insertions(+), 180 deletions(-) diff --git a/common/util.ts b/common/util.ts index 224e30b..813eece 100644 --- a/common/util.ts +++ b/common/util.ts @@ -181,3 +181,9 @@ export const mapDefined = ( } return output; }; + +export const isErrnoException = ( + err: unknown, +): err is NodeJS.ErrnoException => { + return !!err && (err as Record)["code"] !== undefined; +}; diff --git a/deno.lock b/deno.lock index ea41748..dbc7f6f 100644 --- a/deno.lock +++ b/deno.lock @@ -35,7 +35,7 @@ "jsr:@ts-morph/ts-morph@26": "26.0.0", "jsr:@zod/zod@^4.1.11": "4.1.13", "jsr:@zod/zod@^4.1.13": "4.1.13", - "npm:@atproto/crypto@*": "0.4.4", + "npm:@atproto/crypto@*": "0.1.0", "npm:@did-plc/lib@^0.0.4": "0.0.4", "npm:@did-plc/server@^0.0.1": "0.0.1_express@4.21.2", "npm:@ipld/dag-cbor@^9.2.5": "9.2.5", @@ -45,7 +45,8 @@ "npm:multiformats@^13.4.1": "13.4.1", "npm:p-queue@^8.1.1": "8.1.1", "npm:prettier@^3.6.2": "3.6.2", - "npm:rate-limiter-flexible@9": "9.0.0" + "npm:rate-limiter-flexible@9": "9.0.0", + "npm:ws@^8.18.0": "8.18.3" }, "jsr": { "@cliffy/ansi@1.0.0-rc.8": { @@ -221,19 +222,11 @@ "uint8arrays" ] }, - "@atproto/crypto@0.4.4": { - "integrity": "sha512-Yq9+crJ7WQl7sxStVpHgie5Z51R05etaK9DLWYG/7bR5T4bhdcIgF6IfklLShtZwLYdVVj+K15s0BqW9a8PSDA==", - "dependencies": [ - "@noble/curves", - "@noble/hashes", - "uint8arrays" - ] - }, "@did-plc/lib@0.0.4": { "integrity": "sha512-Omeawq3b8G/c/5CtkTtzovSOnWuvIuCI4GTJNrt1AmCskwEQV7zbX5d6km1mjJNbE0gHuQPTVqZxLVqetNbfwA==", "dependencies": [ "@atproto/common@0.1.1", - "@atproto/crypto@0.1.0", + "@atproto/crypto", "@ipld/dag-cbor@7.0.3", "axios", "multiformats@9.9.0", @@ -245,7 +238,7 @@ "integrity": "sha512-GtxxHcOrOQ6fNI1ufq3Zqjc2PtWqPZOdsuzlwtxiH9XibUGwDkb0GmaBHyU5GiOxOKZEW1GspZ8mreBA6XOlTQ==", "dependencies": [ "@atproto/common@0.1.0", - "@atproto/crypto@0.1.0", + "@atproto/crypto", "@did-plc/lib", "axios", "cors", @@ -273,15 +266,6 @@ "multiformats@13.4.1" ] }, - "@noble/curves@1.9.7": { - "integrity": "sha512-gbKGcRUYIjA3/zCCNaWDciTMFI0dCkvou3TL8Zmy5Nc7sJ47a0jtOeZoTaMxkuqRo9cRhjOdZJXegxYE5FN/xw==", - "dependencies": [ - "@noble/hashes" - ] - }, - "@noble/hashes@1.8.0": { - "integrity": "sha512-jCs9ldd7NwzpgXDIf6P3+NrHh9/sD6CQdxHyjQI+h/6rDNo88ypBxxz45UDuZHz9r3tNz7N/VInSVoVdtXEI4A==" - }, "@noble/secp256k1@1.7.2": { "integrity": "sha512-/qzwYl5eFLH8OWIecQWM31qld2g1NfjgylK+TNhqtaUKP37Nm+Y+z30Fjhw0Ct8p9yCQEm2N3W/AckdIb3SMcQ==" }, @@ -1094,6 +1078,9 @@ "vary@1.1.2": { "integrity": "sha512-BNGbWLfd0eUPabhkXUVm0j8uuvREyTh5ovRa/dyow/BqAbZJyC+5fU+IzQOzmAKzYqYRAISoRhdQr3eIZ/PXqg==" }, + "ws@8.18.3": { + "integrity": "sha512-PEIGCY5tSlUt50cqyMXfCzX+oOPqN0vuGqWzbcJ2xvnkzkq46oOpz7dQaTDBdfICb4N14+GARUDw2XV2N4tvzg==" + }, "xtend@4.0.2": { "integrity": "sha512-LKYU1iAXJXUgAXn9URjiu+MWhyUXHsvfp7mcuYm9dSUKK0/CjtrUwFAxD82/mCWbtLsGjFIad0wIsod4zrTAEQ==" }, @@ -1193,7 +1180,8 @@ "npm:get-port@^7.1.0", "npm:key-encoder@^2.0.3", "npm:multiformats@^13.4.1", - "npm:rate-limiter-flexible@9" + "npm:rate-limiter-flexible@9", + "npm:ws@^8.18.0" ] } } diff --git a/sync/deno.json b/sync/deno.json index 7c100a8..b9c64f6 100644 --- a/sync/deno.json +++ b/sync/deno.json @@ -1,6 +1,6 @@ { "name": "@atp/sync", - "version": "0.1.0-alpha.5", + "version": "0.1.0-alpha.6", "exports": "./mod.ts", "license": "MIT", "imports": { diff --git a/sync/firehose/index.ts b/sync/firehose/index.ts index ce79561..00e3db8 100644 --- a/sync/firehose/index.ts +++ b/sync/firehose/index.ts @@ -250,7 +250,7 @@ export class Firehose { if (!parsed) { continue; } - await this.opts.runner.trackEvent( + this.opts.runner.trackEvent( parsed.did, parsed.seq, async () => { diff --git a/xrpc-server/deno.json b/xrpc-server/deno.json index 1b5fa7d..a7b2523 100644 --- a/xrpc-server/deno.json +++ b/xrpc-server/deno.json @@ -1,6 +1,6 @@ { "name": "@atp/xrpc-server", - "version": "0.1.0-alpha.6", + "version": "0.1.0-alpha.7", "exports": "./mod.ts", "license": "MIT", "imports": { @@ -11,7 +11,8 @@ "multiformats": "npm:multiformats@^13.4.1", "zod": "jsr:@zod/zod@^4.1.13", "hono": "jsr:@hono/hono@^4.10.8", - "rate-limiter-flexible": "npm:rate-limiter-flexible@^9.0.0" + "rate-limiter-flexible": "npm:rate-limiter-flexible@^9.0.0", + "ws": "npm:ws@^8.18.0" }, "test": { "permissions": { diff --git a/xrpc-server/stream/websocket-keepalive.ts b/xrpc-server/stream/websocket-keepalive.ts index 00aa7e1..fadcfea 100644 --- a/xrpc-server/stream/websocket-keepalive.ts +++ b/xrpc-server/stream/websocket-keepalive.ts @@ -1,41 +1,17 @@ -// websocket-keepalive.ts -// Runtime-agnostic (Deno / Workers / Bun / Browser) - -import { SECOND, wait } from "@atp/common"; +import { type ClientOptions, WebSocket, createWebSocketStream } from "ws"; +import { isErrnoException, SECOND, wait } from "@atp/common"; import { CloseCode, DisconnectError } from "./types.ts"; -import { iterateBinary } from "./stream.ts"; -/** - * Options for a {@link WebSocketKeepAlive} instance - * - * @prop getUrl Method to get the current URL of the websocket endpoint - * @prop maxReconnectSeconds Maximum time a request can take to reconnect - * @prop signal Abort signal to send when aborting connection - * - * @prop heartbeatIntervalMs Interval to send provided heartbeatPayload on, - * @prop heartbeatPayload Method to create payload to send for heartbeat - * @prop isPong If provided, we mark alive only when it returns true for a message - * if omitted, *any* message is considered proof of life - * - * @prop onReconnectError Reconnect hook - * - * @prop createSocket Socket factory override (lets you use custom client if needed) - * @prop protocols Override value for accepted protocols - */ -export type KeepAliveOptions = { +export type KeepAliveOptions = ClientOptions & { getUrl: () => Promise; maxReconnectSeconds?: number; signal?: AbortSignal; - - heartbeatIntervalMs?: number; // default 10 * SECOND - heartbeatPayload?: () => string | ArrayBuffer | Uint8Array | Blob; - isPong?: (data: unknown) => boolean; - - // Reconnect hook - onReconnectError?: (error: unknown, n: number, initialSetup: boolean) => void; - - createSocket?: (url: string, protocols?: string | string[]) => WebSocket; - protocols?: string | string[]; + heartbeatIntervalMs?: number; + onReconnectError?: ( + error: unknown, + n: number, + initialSetup: boolean, + ) => void; }; export class WebSocketKeepAlive { @@ -43,15 +19,10 @@ export class WebSocketKeepAlive { public initialSetup = true; public reconnects: number | null = null; - /** - * Creates a new WebSocketKeepAlive instance. - * @param opts Configuration options for keepalive, heartbeat, reconnect, and socket creation. - */ constructor(public opts: KeepAliveOptions) {} async *[Symbol.asyncIterator](): AsyncGenerator { const maxReconnectMs = 1000 * (this.opts.maxReconnectSeconds ?? 64); - while (true) { if (this.reconnects !== null) { const duration = this.initialSetup @@ -59,141 +30,86 @@ export class WebSocketKeepAlive { : backoffMs(this.reconnects++, maxReconnectMs); await wait(duration); } - const url = await this.opts.getUrl(); - - // Create a web-standard WebSocket (or a custom one if provided). - const ws = this.opts.createSocket?.(url, this.opts.protocols) ?? - new WebSocket(url, this.opts.protocols); - this.ws = ws; - + this.ws = new WebSocket(url, this.opts); const ac = new AbortController(); if (this.opts.signal) { forwardSignal(this.opts.signal, ac); } - - // Track liveness (application-level heartbeat) - this.startHeartbeat(ws, ac); - - // When the socket opens, reset backoff. - ws.addEventListener( - "open", - () => { - this.initialSetup = false; - this.reconnects = 0; - }, - { once: true }, - ); - - // Distinguish abnormal close → treat as reconnectable error - ws.addEventListener( - "close", - (ev) => { - if (ev.code === CloseCode.Abnormal) { - ac.abort( - new AbnormalCloseError( - `Abnormal ws close: ${String(ev.reason || "")}`, - ), - ); - } - }, - { once: true }, - ); + this.ws.once("open", () => { + this.initialSetup = false; + this.reconnects = 0; + if (this.ws) { + this.startHeartbeat(this.ws); + } + }); + this.ws.once("close", (code: number, reason: Uint8Array) => { + if (code === CloseCode.Abnormal) { + ac.abort( + new AbnormalCloseError( + `Abnormal ws close: ${new TextDecoder().decode(reason)}`, + ), + ); + } + }); try { - // Iterate incoming binary chunks - for await (const chunk of iterateBinary(ws)) { + const wsStream = createWebSocketStream(this.ws, { + signal: ac.signal, + readableObjectMode: true, + }); + for await (const chunk of wsStream) { yield chunk; } - } catch (error) { - // Normalize Abort into same shape your old code expected. - const err = (error as Error)?.name === "AbortError" - ? (error as Error).cause ?? error - : error; - + } catch (_err) { + const err = + isErrnoException(_err) && _err.code === "ABORT_ERR" + ? _err.cause + : _err; if (err instanceof DisconnectError) { - // We cleanly end the connection - ws?.close(err.wsCode); + this.ws?.close(err.wsCode); break; } - - // Close if not already closing - ws.close(); - + this.ws?.close(); if (isReconnectable(err)) { - this.reconnects ??= 0; // Never reconnect when null + this.reconnects ??= 0; this.opts.onReconnectError?.(err, this.reconnects, this.initialSetup); - continue; // loop to reconnect + continue; } else { throw err; } } - - // Other side ended stream cleanly; stop iterating. break; } } - /** Application-level heartbeat (web standard). - * - * In Node's `ws` you used `ping`/`pong`. Those do not exist in web sockets. - * Here we: - * - periodically send `heartbeatPayload()` if provided - * - consider the connection "alive" when: - * * `isPong(ev.data)` returns true (if provided), OR - * * *any* message is received (fallback) - * - if no proof of life for one interval, we close the socket (which triggers reconnect) - */ - private startHeartbeat(ws: WebSocket, ac: AbortController) { - const intervalMs = this.opts.heartbeatIntervalMs ?? 10 * SECOND; - + startHeartbeat(ws: WebSocket) { let isAlive = true; - let timer: number | null = null; - - const onMessage = (ev: MessageEvent) => { - // If a custom pong detector exists, use it; otherwise any message counts. - if (!this.opts.isPong || this.opts.isPong(ev.data)) { - isAlive = true; - } - }; + let heartbeatInterval: ReturnType | null = null; - const tick = () => { + const checkAlive = () => { if (!isAlive) { - // No pong/traffic since last tick → consider dead and close. - ws.close(1000); - // Abort the iterator with a recognizable shape like before. - const domErr = new DOMException("Aborted", "AbortError"); - domErr.cause = new DisconnectError( - CloseCode.Abnormal, - "HeartbeatTimeout", - ); - ac.abort(domErr); - return; + return ws.terminate(); } isAlive = false; - - const payload = this.opts.heartbeatPayload?.(); - if (payload !== undefined) { - ws.send(payload); - } + ws.ping(); }; - // Prime one cycle and schedule subsequent ones - tick(); - timer = setInterval(tick, intervalMs) as unknown as number; - - ws.addEventListener("message", onMessage); - ws.addEventListener( - "close", - () => { - if (timer !== null) { - clearInterval(timer); - timer = null; - } - ws.removeEventListener("message", onMessage); - }, - { once: true }, + checkAlive(); + heartbeatInterval = setInterval( + checkAlive, + this.opts.heartbeatIntervalMs ?? 10 * SECOND, ); + + ws.on("pong", () => { + isAlive = true; + }); + ws.once("close", () => { + if (heartbeatInterval) { + clearInterval(heartbeatInterval); + heartbeatInterval = null; + } + }); } } @@ -204,11 +120,10 @@ class AbnormalCloseError extends Error { } function isReconnectable(err: unknown): boolean { - // Network-ish errors are reconnectable. Keep your previous codes. - if (!err || typeof err !== "object") return false; - const e = err as { name?: unknown; code?: unknown }; - if (typeof e.name !== "string") return false; - return typeof e.code === "string" && networkErrorCodes.includes(e.code); + if (isErrnoException(err) && typeof err.code === "string") { + return networkErrorCodes.includes(err.code); + } + return false; } const networkErrorCodes = [ @@ -219,12 +134,11 @@ const networkErrorCodes = [ "EPIPE", "ETIMEDOUT", "ECANCELED", - "ABORT_ERR", // surface our aborts as reconnectable if you want ]; function backoffMs(n: number, maxMs: number) { - const baseSec = Math.pow(2, n); // 1, 2, 4, ... - const randSec = Math.random() - 0.5; // jitter [-0.5, +0.5] + const baseSec = Math.pow(2, n); + const randSec = Math.random() - 0.5; const ms = 1000 * (baseSec + randSec); return Math.min(ms, maxMs); } @@ -232,8 +146,9 @@ function backoffMs(n: number, maxMs: number) { function forwardSignal(signal: AbortSignal, ac: AbortController) { if (signal.aborted) { return ac.abort(signal.reason); + } else { + signal.addEventListener("abort", () => ac.abort(signal.reason), { + signal: ac.signal, + }); } - const onAbort = () => ac.abort(signal.reason); - // Use AbortSignal.any? Not universally available; just add/remove. - signal.addEventListener("abort", onAbort, { signal: ac.signal }); } -- 2.51.2