diff --git a/.changeset/world-postgres-listen-reconnect-test.md b/.changeset/world-postgres-listen-reconnect-test.md new file mode 100644 index 0000000000..84ab68539a --- /dev/null +++ b/.changeset/world-postgres-listen-reconnect-test.md @@ -0,0 +1,4 @@ +--- +--- + +Make the world-postgres LISTEN reconnect test retry its reconnect the way consumers do. diff --git a/packages/world-postgres/test/invoke.test.ts b/packages/world-postgres/test/invoke.test.ts index d86a335d16..5ec9e37d5d 100644 --- a/packages/world-postgres/test/invoke.test.ts +++ b/packages/world-postgres/test/invoke.test.ts @@ -15,7 +15,7 @@ import { encode } from 'cbor-x'; import { drizzle } from 'drizzle-orm/node-postgres'; import { migrate } from 'drizzle-orm/node-postgres/migrator'; import { makeWorkerUtils } from 'graphile-worker'; -import { Pool } from 'pg'; +import { Client, Pool } from 'pg'; import { ulid } from 'ulid'; import { afterAll, beforeAll, describe, expect, it, vi } from 'vitest'; import { registerStepFunction } from '../../core/dist/private.js'; @@ -34,6 +34,7 @@ import { import { createWorld } from '../src/index.js'; import { createInvocationNotifications, + INVOCATION_FALLBACK_MS, INVOCATION_INPUT_TOPIC, INVOCATION_RESULT_TOPIC, invocationNotificationKey, @@ -1148,55 +1149,97 @@ describe.skipIf(process.platform === 'win32')( } }); - it('reconnects a terminated real LISTEN client and receives later notifications', async () => { + it('reconnects a terminated real LISTEN client, retrying a failed attempt, and receives later notifications', async () => { const name = `notify-${randomUUID()}`; const listenerPool = new Pool({ connectionString: container.getConnectionUri(), application_name: name, }); + // The listener's first reconnect fails, as a connect timeout under load + // does. Nothing retries it but a later wait(). + const connect = Client.prototype.connect; + let listenerConnects = 0; + const connectSpy = vi + .spyOn(Client.prototype, 'connect') + .mockImplementation(function (this: Client, ...args: never[]) { + const { connectionParameters } = this as unknown as { + connectionParameters: { application_name?: string }; + }; + if ( + connectionParameters.application_name === `${name}:invocations` && + ++listenerConnects === 2 + ) { + return Promise.reject( + new Error('Connection terminated due to connection timeout') + ); + } + return Reflect.apply(connect, this, args); + } as typeof connect); + const warn = vi.spyOn(console, 'warn'); + // Logged right before the reconnected listener wakes its watches. + const restorations = () => + warn.mock.calls.filter(([message]) => + String(message).includes('Invocation notifications restored') + ).length; const observer = createInvocationNotifications(listenerPool); const key = invocationNotificationKey('reconnect'); const watch = observer.watch(INVOCATION_INPUT_TOPIC, key); const signal = new AbortController().signal; + // A backend shows as idle before its LISTEN has run. const listenerPid = async () => ( await pool.query<{ pid: number }>( - 'SELECT pid FROM pg_stat_activity WHERE application_name = $1 AND state = $2', - [`${name}:invocations`, 'idle'] + "SELECT pid FROM pg_stat_activity WHERE application_name = $1 AND state = 'idle' AND query LIKE 'LISTEN %'", + [`${name}:invocations`] ) ).rows[0]?.pid; - try { + // Read, then wait, the way consumers do: each wait() past the backoff + // retries the connection, so a failed attempt cannot strand the watch. + const poll = (read: () => Promise, ready: (value: T) => boolean) => + until(async () => { + const value = await read(); + if (!ready(value)) + await watch.wait(watch.revision, INVOCATION_FALLBACK_MS, signal); + return value; + }, ready); + const notified = async () => { + const before = watch.revision; + await pool.query('SELECT pg_notify($1, $2)', [ + INVOCATION_INPUT_TOPIC, + key, + ]); await until( async () => watch.revision, - (revision) => revision > 0 + (revision) => revision > before ); - const firstPid = await until(listenerPid, (pid) => pid !== undefined); + }; + try { + const firstPid = await poll(listenerPid, (pid) => pid !== undefined); + // The server lists the LISTEN before the client has read its reply, + // and a failed first connect leaves a restore log still to come. A + // notification follows that reply on the same socket, so once it is + // seen the first subscription has settled, wake and log included. + await notified(); const beforeDisconnect = watch.revision; + const beforeRestore = restorations(); await pool.query('SELECT pg_terminate_backend($1)', [firstPid]); + // Consumers reread on disconnect, then fall back to timed reads. await until( async () => watch.revision, (revision) => revision > beforeDisconnect ); - // Normal consumers reread on disconnect, then use the slow fallback - // while reconnect is backed off, then wait/read again. - await watch.wait(watch.revision, 1_000, signal); - const reconnected = watch.wait(watch.revision, 5_000, signal); - const secondPid = await until( - listenerPid, - (pid) => pid !== undefined && pid !== firstPid + await poll( + async () => restorations(), + (count) => count > beforeRestore ); - await reconnected; + expect(listenerConnects).toBeGreaterThanOrEqual(3); + const secondPid = await listenerPid(); + expect(secondPid).toBeDefined(); expect(secondPid).not.toBe(firstPid); - const beforeNotify = watch.revision; - await pool.query('SELECT pg_notify($1, $2)', [ - INVOCATION_INPUT_TOPIC, - key, - ]); - await until( - async () => watch.revision, - (revision) => revision > beforeNotify - ); + await notified(); } finally { + connectSpy.mockRestore(); + warn.mockRestore(); watch.dispose(); await observer.close(); await listenerPool.end();