diff --git a/.changeset/world-postgres-listen-self-healing.md b/.changeset/world-postgres-listen-self-healing.md new file mode 100644 index 0000000000..d9b5b209ac --- /dev/null +++ b/.changeset/world-postgres-listen-self-healing.md @@ -0,0 +1,5 @@ +--- +'@workflow/world-postgres': patch +--- + +Make stream delivery durable across `LISTEN`/`NOTIFY` interruptions: the dedicated `pg.Client` now reconnects with bounded exponential backoff (250 ms → 30 s), and `readFromStream` runs a periodic re-query of `streams WHERE chunk_id > lastChunkId` as a polling safety net for chunks delivered while the LISTEN socket was down. The poll interval is configurable via `PostgresWorldConfig.streamPollIntervalMs` (default 5000 ms; set to 0 to disable). Tracks vercel/workflow#1855. diff --git a/packages/world-postgres/src/config.ts b/packages/world-postgres/src/config.ts index 42be065a2a..069a0a677a 100644 --- a/packages/world-postgres/src/config.ts +++ b/packages/world-postgres/src/config.ts @@ -1,5 +1,4 @@ import type { Pool } from 'pg'; -import type { ListenAdapter } from './listen-adapter.js'; type PgConnectionConfig = | { connectionString: string; maxPoolSize?: number; pool?: undefined } @@ -14,11 +13,11 @@ export type PostgresWorldConfig = PgConnectionConfig & { */ streamFlushIntervalMs?: number; /** - * Plug in an alternative LISTEN/NOTIFY transport for the streamer. - * Defaults to the bundled `pg`-backed adapter (with self-healing reconnect). - * See {@link ListenAdapter} for the contract; useful for swapping in - * `bun:sql`'s native listen (oven-sh/bun#29710) once it lands, or a - * different pub/sub backend entirely. + * How often (ms) `readFromStream` re-queries the `streams` table as a + * safety net for chunks delivered while the LISTEN client was reconnecting. + * Default is 5000. Set to 0 to disable polling entirely. + * + * See `CreateStreamerOptions.pollIntervalMs` for the full contract. */ - listenAdapter?: ListenAdapter; + streamPollIntervalMs?: number; }; diff --git a/packages/world-postgres/src/index.ts b/packages/world-postgres/src/index.ts index 6f1226e456..eed68cbade 100644 --- a/packages/world-postgres/src/index.ts +++ b/packages/world-postgres/src/index.ts @@ -54,7 +54,9 @@ export function createWorld( const drizzle = createClient(pool); const queue = createQueue(config, pool); const storage = createStorage(drizzle); - const streamer = createStreamer(pool, drizzle, config.listenAdapter); + const streamer = createStreamer(pool, drizzle, { + pollIntervalMs: config.streamPollIntervalMs, + }); return { specVersion: SPEC_VERSION_CURRENT, @@ -80,10 +82,4 @@ export function createWorld( // Re-export schema for users who want to extend or inspect the database schema export type { PostgresWorldConfig } from './config.js'; -export { - createBunSqlListenAdapter, - createPgListenAdapter, - type ListenAdapter, - type ListenSubscription, -} from './listen-adapter.js'; export * from './drizzle/schema.js'; diff --git a/packages/world-postgres/src/listen-adapter.ts b/packages/world-postgres/src/listen-adapter.ts deleted file mode 100644 index d912ee9e48..0000000000 --- a/packages/world-postgres/src/listen-adapter.ts +++ /dev/null @@ -1,203 +0,0 @@ -import { Client, type Pool } from 'pg'; - -/** - * A live subscription returned by a {@link ListenAdapter}. Calling `close` - * tears down the underlying connection (or pub/sub primitive) and stops - * delivering payloads. - */ -export interface ListenSubscription { - close(): Promise; -} - -/** - * Pluggable LISTEN/NOTIFY transport for the streamer. - * - * The default implementation ({@link createPgListenAdapter}) uses a dedicated - * `pg` client and the PostgreSQL `LISTEN`/`NOTIFY` protocol. Alternative - * adapters can plug in different runtimes (e.g. `bun:sql`'s native - * `sql.listen`, see oven-sh/bun#29710) or different transports (Redis pub/sub, - * NATS, etc.) without touching streamer/storage code. - * - * Adapters MUST: - * * resolve the `listen()` promise only after the subscription is live so - * callers may rely on "no missed events from this point on"; - * * survive transient connection drops (auto-reconnect) — events lost on - * the wire during a reconnect window are recovered by the streamer's - * polling fallback in `readFromStream`; - * * deliver each payload at-least-once. Duplicates are tolerated by the - * streamer (deduped via chunkId ordering). - */ -export interface ListenAdapter { - listen( - channel: string, - onPayload: (payload: string) => Promise | void - ): Promise; - notify(channel: string, payload: string): Promise; -} - -/** - * Default adapter backed by `pg` (`node-postgres`). - * - * Wraps the dedicated `Client` in a reconnect loop with exponential backoff - * (250 ms → 30 s cap). The initial connect must succeed (callers expect a - * live subscription before the promise resolves); subsequent reconnects are - * best-effort. Notifications fired while the dedicated client is reconnecting - * are lost on the wire — the polling fallback in the streamer's - * `readFromStream` picks them up from the database on its periodic tick, so - * end-to-end stream delivery stays correct. - * - * The dedicated `Client` is long-lived and will eventually be dropped by the - * server (idle TCP timeout, pgbouncer rotation, k8s CNI eviction). The - * unpatched `pg` behaviour does not reconnect, so a process running for more - * than a few hours stops receiving notifications and only a restart restores - * delivery (cf. brianc/node-postgres#967). - */ -export function createPgListenAdapter(pool: Pool): ListenAdapter { - const notify = async (channel: string, payload: string): Promise => { - await pool.query('SELECT pg_notify($1, $2)', [channel, payload]); - }; - - const listen = async ( - channel: string, - onPayload: (payload: string) => Promise | void - ): Promise => { - let client: Client | null = null; - let closed = false; - let reconnectAttempt = 0; - let reconnectTimer: ReturnType | null = null; - - const onNotification = (msg: { payload?: string | undefined }) => { - try { - const r = onPayload(msg.payload ?? ''); - if (r && typeof (r as Promise).catch === 'function') { - (r as Promise).catch(() => {}); - } - } catch { - // swallow handler errors - } - }; - - const detach = (c: Client | null) => { - if (!c) return; - try { - c.removeListener('notification', onNotification); - } catch { - // listener may already be detached - } - c.end().catch(() => {}); - }; - - const scheduleReconnect = () => { - if (closed || reconnectTimer) return; - const delay = Math.min(30_000, 250 * 2 ** reconnectAttempt); - reconnectAttempt++; - reconnectTimer = setTimeout(() => { - reconnectTimer = null; - if (closed) return; - connect().catch((err) => { - // eslint-disable-next-line no-console - console.warn( - '[world-postgres pg-listen] reconnect failed', - (err as Error)?.message ?? err - ); - scheduleReconnect(); - }); - }, delay); - }; - - const connect = async () => { - if (closed) return; - const next = new Client(pool.options); - next.on('error', (err) => { - // eslint-disable-next-line no-console - console.warn( - '[world-postgres pg-listen] client error', - (err as Error)?.message ?? err - ); - if (client === next) client = null; - detach(next); - scheduleReconnect(); - }); - next.on('end', () => { - if (closed) return; - if (client === next) client = null; - scheduleReconnect(); - }); - try { - await next.connect(); - await next.query(`LISTEN ${channel}`); - } catch (err) { - await next.end().catch(() => {}); - throw err; - } - next.on('notification', onNotification); - client = next; - reconnectAttempt = 0; - }; - - await connect(); - - return { - close: async () => { - closed = true; - if (reconnectTimer) { - clearTimeout(reconnectTimer); - reconnectTimer = null; - } - const c = client; - client = null; - if (!c) return; - try { - c.removeListener('notification', onNotification); - } catch { - // listener may already be detached - } - try { - await c.query(`UNLISTEN ${channel}`); - } finally { - await c.end().catch(() => {}); - } - }, - }; - }; - - return { listen, notify }; -} - -/** - * Stub adapter that defers to a future `bun:sql` native LISTEN/NOTIFY - * implementation (oven-sh/bun#29710). Until that PR lands, callers running - * on Bun should keep using {@link createPgListenAdapter} — `node-postgres` - * works fine on Bun's Node compat layer; only the LISTEN path was the - * historical pain point, and that is now self-healing in this package. - * - * Once `sql.listen` ships in Bun stable, the implementation becomes: - * - * ```ts - * import { SQL } from 'bun'; - * export function createBunSqlListenAdapter(sql: SQL): ListenAdapter { - * const notify = (channel: string, payload: string) => - * sql.notify(channel, payload); - * const listen = async (channel: string, onPayload) => { - * const { unlisten } = await sql.listen(channel, onPayload); - * return { close: () => unlisten() }; - * }; - * return { listen, notify }; - * } - * ``` - * - * Currently exported as a typed stub so consumers can compile-time-pick the - * adapter via `process.env` without conditional imports. - */ -export function createBunSqlListenAdapter(): ListenAdapter { - const notSupported = (): never => { - throw new Error( - '[world-postgres] bun:sql LISTEN/NOTIFY adapter is not yet available; ' + - 'tracked at oven-sh/bun#29710. Use createPgListenAdapter for now.' - ); - }; - return { - listen: () => Promise.reject(notSupported()), - notify: () => Promise.reject(notSupported()), - }; -} diff --git a/packages/world-postgres/src/queue.ts b/packages/world-postgres/src/queue.ts index bf2eca4bff..d432cb0d11 100644 --- a/packages/world-postgres/src/queue.ts +++ b/packages/world-postgres/src/queue.ts @@ -256,9 +256,10 @@ export function createQueue( `${baseUrl}/.well-known/workflow/v1/${pathname}`, { method: 'POST', + duplex: 'half', headers, body, - } + } as any ); const text = await response.text(); diff --git a/packages/world-postgres/src/streamer.ts b/packages/world-postgres/src/streamer.ts index 299ea44a2d..cf2d96620f 100644 --- a/packages/world-postgres/src/streamer.ts +++ b/packages/world-postgres/src/streamer.ts @@ -6,15 +6,10 @@ import type { StreamInfoResponse, } from '@workflow/world'; import { and, asc, eq, gt, sql } from 'drizzle-orm'; -import type { Pool } from 'pg'; +import { Client, type Pool } from 'pg'; import { monotonicFactory } from 'ulid'; import * as z from 'zod'; import { type Drizzle, Schema } from './drizzle/index.js'; -import { - createPgListenAdapter, - type ListenAdapter, - type ListenSubscription, -} from './listen-adapter.js'; import { Mutex } from './util.js'; const StreamPublishMessage = z.object({ @@ -49,35 +44,148 @@ class Rc { } /** - * Subscribe to a PostgreSQL NOTIFY channel. + * Subscribe to a PostgreSQL NOTIFY channel using a dedicated client created + * from the pool's connection options. `channel` must be a trusted identifier + * (interpolated into the LISTEN statement; `pg` does not parameterise + * identifiers). * - * Backwards-compatible thin wrapper around {@link createPgListenAdapter} — - * preserved as a named export because earlier versions of this package - * exposed it directly. New code should construct a {@link ListenAdapter} - * once (via `createPgListenAdapter` or a future `createBunSqlListenAdapter`) - * and pass it to {@link createStreamer}, which lets the streamer share a - * single dedicated subscriber across all stream consumers. + * The dedicated `Client` is long-lived and will eventually be dropped by the + * server (idle TCP timeout, pgbouncer rotation, k8s CNI eviction). Without + * reconnect handling, a process running for more than a few hours stops + * receiving notifications and only a restart restores delivery + * (cf. brianc/node-postgres#967). * - * `channel` must be a trusted identifier (it is interpolated into the - * `LISTEN` SQL — `pg` does not parameterise identifiers). + * This implementation wraps the client in a reconnect loop with bounded + * exponential backoff (250 ms → 30 s cap). The initial connect must succeed + * (callers expect a live subscription before the promise resolves); + * subsequent reconnects are best-effort. Notifications fired while the + * dedicated client is reconnecting are lost on the wire — the polling + * fallback in {@link createStreamer}'s `readFromStream` re-queries chunks + * from the database on its periodic tick, so end-to-end delivery stays + * correct even across LISTEN gaps. */ -export const listenChannel = ( +export const listenChannel = async ( pool: Pool, channel: string, onPayload: (payload: string) => Promise -): Promise => - createPgListenAdapter(pool).listen(channel, onPayload); +): Promise<{ close: () => Promise }> => { + let client: Client | null = null; + let closed = false; + let reconnectAttempt = 0; + let reconnectTimer: ReturnType | null = null; + + const onNotification = (msg: { payload?: string | undefined }) => { + onPayload(msg.payload ?? '').catch(() => {}); + }; + + const detach = (c: Client) => { + try { + c.removeListener('notification', onNotification); + } catch { + // listener may already be detached + } + c.end().catch(() => {}); + }; + + const scheduleReconnect = () => { + if (closed || reconnectTimer) return; + const delay = Math.min(30_000, 250 * 2 ** reconnectAttempt); + reconnectAttempt++; + reconnectTimer = setTimeout(() => { + reconnectTimer = null; + if (closed) return; + connect().catch((err) => { + console.warn( + '[world-postgres listenChannel] reconnect failed', + (err as Error)?.message ?? err + ); + scheduleReconnect(); + }); + }, delay); + }; + + const connect = async () => { + if (closed) return; + const next = new Client(pool.options); + next.on('error', (err) => { + console.warn( + '[world-postgres listenChannel] client error', + (err as Error)?.message ?? err + ); + if (client === next) client = null; + detach(next); + scheduleReconnect(); + }); + next.on('end', () => { + if (closed) return; + if (client === next) client = null; + scheduleReconnect(); + }); + try { + await next.connect(); + await next.query(`LISTEN ${channel}`); + } catch (err) { + await next.end().catch(() => {}); + throw err; + } + next.on('notification', onNotification); + client = next; + reconnectAttempt = 0; + }; + + await connect(); + + return { + close: async () => { + closed = true; + if (reconnectTimer) { + clearTimeout(reconnectTimer); + reconnectTimer = null; + } + const c = client; + client = null; + if (!c) return; + try { + c.removeListener('notification', onNotification); + } catch { + // listener may already be detached + } + try { + await c.query(`UNLISTEN ${channel}`); + } finally { + await c.end().catch(() => {}); + } + }, + }; +}; export type PostgresStreamer = Streamer & { /** Unlisten from the LISTEN subscription and release resources. */ close(): Promise; }; +export type CreateStreamerOptions = { + /** + * How often (ms) `readFromStream` re-queries the `streams` table for chunks + * past `lastChunkId` as a safety net for notifications dropped while the + * LISTEN client was reconnecting. The poll dedupes against in-band + * notifications via the existing `enqueue` ordering check, so it is safe + * to run alongside `LISTEN/NOTIFY`. + * + * Lower values reduce recovery latency after a LISTEN disconnect; higher + * values reduce baseline DB load (one extra `SELECT` per active reader per + * tick). Set to `0` to disable polling — only do this if you know the + * LISTEN connection cannot be interrupted (e.g. tests). Default: 5000. + */ + pollIntervalMs?: number; +}; + export function createStreamer( pool: Pool, drizzle: Drizzle, - listenAdapter: ListenAdapter = createPgListenAdapter(pool) + options: CreateStreamerOptions = {} ): PostgresStreamer { + const pollIntervalMs = options.pollIntervalMs ?? 5_000; const ulid = monotonicFactory(); const events = new EventEmitter<{ [key: `strm:${string}`]: [StreamChunkEvent]; @@ -99,7 +207,7 @@ export function createStreamer( const STREAM_TOPIC = 'workflow_event_chunk'; - const listenSubscription = listenAdapter.listen(STREAM_TOPIC, async (msg) => { + const listenSubscription = listenChannel(pool, STREAM_TOPIC, async (msg) => { const parsed = StreamPublishMessage.parse(JSON.parse(msg)); const key = `strm:${parsed.streamId}` as const; @@ -125,8 +233,9 @@ export function createStreamer( }); }); - const notifyStream = (payload: string) => - listenAdapter.notify(STREAM_TOPIC, payload); + const notifyStream = async (payload: string) => { + await pool.query('SELECT pg_notify($1, $2)', [STREAM_TOPIC, payload]); + }; // Helper to convert chunk to Buffer const toBuffer = (chunk: string | Uint8Array): Buffer => @@ -359,8 +468,8 @@ export function createStreamer( data: Uint8Array; eof: boolean; }) { - if (lastChunkId >= msg.id) { - // already sent or out of order + if (closed || lastChunkId >= msg.id) { + // already sent, out of order, or stream torn down return; } @@ -373,6 +482,7 @@ export function createStreamer( controller.enqueue(new Uint8Array(msg.data)); } if (msg.eof) { + closed = true; controller.close(); } lastChunkId = msg.id; @@ -416,38 +526,40 @@ export function createStreamer( buffer = null; // Polling fallback. NOTIFY is the fast path, but events are silently - // dropped while the dedicated `listenChannel` client is reconnecting - // (and would be dropped entirely on the unpatched implementation). - // A light periodic re-query of chunks past `lastChunkId` is the - // always-on safety net: every POLL_INTERVAL_MS it pulls any chunks + // dropped while the dedicated LISTEN client is reconnecting. A + // light periodic re-query of chunks past `lastChunkId` is the + // always-on safety net: every `pollIntervalMs` it pulls any chunks // the EventEmitter missed, deduped by the `enqueue` ordering check. - // Stops automatically on EOF (controller.close prevents further - // enqueues), on cancel, or on closed controller (try/catch in poll). - const POLL_INTERVAL_MS = 5_000; - const pollTimer = setInterval(async () => { + // Stops on EOF (enqueue sets `closed`), on cancel, or on pool error + // (try/catch keeps the timer alive for a future tick). + const runPoll = async () => { + const fresh = await drizzle + .select({ + id: streams.chunkId, + eof: streams.eof, + data: streams.chunkData, + }) + .from(streams) + .where( + and( + eq(streams.streamId, name), + gt(streams.chunkId, lastChunkId as `chnk_${string}`) + ) + ) + .orderBy(streams.chunkId); + for (const chunk of fresh) { + if (closed) return; + enqueue(chunk); + } + }; + + const tick = async () => { if (polling || closed) return; polling = true; try { - const fresh = await drizzle - .select({ - id: streams.chunkId, - eof: streams.eof, - data: streams.chunkData, - }) - .from(streams) - .where( - and( - eq(streams.streamId, name), - gt(streams.chunkId, lastChunkId as `chnk_${string}`) - ) - ) - .orderBy(streams.chunkId); - for (const chunk of fresh) { - enqueue(chunk); - } + await runPoll(); } catch (err) { // Best-effort. Logs only; the next tick retries. - // eslint-disable-next-line no-console console.warn( '[world-postgres readFromStream] poll failed', (err as Error)?.message ?? err @@ -455,10 +567,13 @@ export function createStreamer( } finally { polling = false; } - }, POLL_INTERVAL_MS); + }; + + const pollTimer = + pollIntervalMs > 0 ? setInterval(tick, pollIntervalMs) : null; cleanups.push(() => { closed = true; - clearInterval(pollTimer); + if (pollTimer) clearInterval(pollTimer); }); }, cancel() {