Merge pull request #851 from unix/codex/fix-listen-heartbeat

fix: recover silent LISTEN/NOTIFY connection failures

Author: timgit

src/db.ts

@@ -3,6 +3,13 @@ import pg from 'pg'
 import assert from 'node:assert'
 import type * as types from './types.ts'
 
+// Keep silent network failures below the default 30-second notify polling backstop: in the
+// worst case a failure happens immediately after a successful check, then takes one interval,
+// one query timeout, and the existing first reconnect backoff (1s) to restore LISTEN.
+const LISTEN_HEARTBEAT_INTERVAL_MS = 10000
+const LISTEN_HEARTBEAT_TIMEOUT_MS = 5000
+const LISTEN_KEEP_ALIVE_INITIAL_DELAY_MS = 10000
+
 class Db extends EventEmitter implements types.IDatabase, types.EventsMixin {
   private pool!: pg.Pool
   private config: types.DatabaseOptions
@@ -57,8 +64,9 @@ class Db extends EventEmitter implements types.IDatabase, types.EventsMixin {
 
   // Opens a dedicated, session-pinned connection for LISTEN/NOTIFY. A separate pg.Client
   // (not a pooled connection) is used so the listener never depletes the query pool and so
-  // reconnection is self-contained. On any drop the client reconnects with capped backoff
-  // and re-runs LISTEN, then calls onReconnect so the caller can recover missed messages.
+  // reconnection is self-contained. TCP keepalive plus a same-session heartbeat detect silent
+  // drops and lost subscriptions. The client reconnects with capped backoff, re-runs LISTEN,
+  // then calls onReconnect so the caller can recover missed messages.
   async listen (
     channel: string,
     onNotification: (payload: string) => void,
@@ -69,14 +77,23 @@ class Db extends EventEmitter implements types.IDatabase, types.EventsMixin {
     let closed = false
     let client: pg.Client | null = null
     let reconnectTimer: ReturnType<typeof setTimeout> | null = null
+    let heartbeatTimer: ReturnType<typeof setTimeout> | null = null
     let attempt = 0
+    const heartbeatInterval = this.config.__test__listenHeartbeatIntervalMs ?? LISTEN_HEARTBEAT_INTERVAL_MS
+    const heartbeatTimeout = this.config.__test__listenHeartbeatTimeoutMs ?? LISTEN_HEARTBEAT_TIMEOUT_MS
     // Only self-heal once the listener has been established at least once. If the INITIAL connect
     // fails, the rejection propagates to the caller (Notifier.start), which falls back to
     // polling-only and discards this subscription's close handle — so a reconnect scheduled from
     // the client 'error' handler would be an untracked connection nothing can close, keeping the
     // event loop alive and delivering notifications into a stopped manager.
     let established = false
 
+    const clearHeartbeat = () => {
+      if (!heartbeatTimer) return
+      clearTimeout(heartbeatTimer)
+      heartbeatTimer = null
+    }
+
     const scheduleReconnect = () => {
       if (closed || reconnectTimer) return
       const backoff = Math.min(30000, 1000 * 2 ** Math.min(attempt, 5))
@@ -87,21 +104,72 @@ class Db extends EventEmitter implements types.IDatabase, types.EventsMixin {
       }, backoff)
     }
 
+    const disconnect = (target: pg.Client, error: Error) => {
+      if (closed || client !== target) return
+
+      clearHeartbeat()
+      client = null
+      target.removeAllListeners()
+      target.end().catch(() => {})
+      this.emit('error', error)
+      if (established) scheduleReconnect()
+    }
+
+    const scheduleHeartbeat = (target: pg.Client) => {
+      if (closed || client !== target) return
+      heartbeatTimer = setTimeout(() => {
+        heartbeatTimer = null
+        heartbeat(target).catch(error => disconnect(target, error))
+      }, heartbeatInterval)
+    }
+
+    const heartbeat = async (target: pg.Client) => {
+      if (closed || client !== target) return
+
+      let timeout: ReturnType<typeof setTimeout> | null = null
+      const query = target.query(
+        `SELECT EXISTS (
+           SELECT 1
+             FROM pg_listening_channels() AS active(channel)
+            WHERE channel = $1
+         ) AS listening`,
+        [channel]
+      )
+      query.catch(() => {})
+
+      try {
+        const result = await Promise.race([
+          query,
+          new Promise<never>((resolve, reject) => {
+            timeout = setTimeout(() => reject(new Error('LISTEN/NOTIFY heartbeat timed out')), heartbeatTimeout)
+          })
+        ])
+
+        if (!result.rows[0]?.listening) {
+          throw new Error('LISTEN/NOTIFY channel registration was lost')
+        }
+      } finally {
+        if (timeout) clearTimeout(timeout)
+      }
+
+      scheduleHeartbeat(target)
+    }
+
     const connect = async () => {
       if (closed) return
 
-      const next = new pg.Client(this.config)
+      const next = new pg.Client({
+        ...this.config,
+        keepAlive: true,
+        keepAliveInitialDelayMillis: LISTEN_KEEP_ALIVE_INITIAL_DELAY_MS
+      })
 
       next.on('error', error => {
-        this.emit('error', error)
-        if (!closed) {
-          next.removeAllListeners()
-          next.end().catch(() => {})
-          if (client === next) client = null
-          if (established) scheduleReconnect()
-        }
+        disconnect(next, error)
       })
 
+      next.on('end', () => disconnect(next, new Error('LISTEN/NOTIFY connection ended')))
+
       next.on('notification', msg => {
         if (msg.payload !== undefined) onNotification(msg.payload)
       })
@@ -125,6 +193,7 @@ class Db extends EventEmitter implements types.IDatabase, types.EventsMixin {
 
       attempt = 0
       established = true
+      scheduleHeartbeat(next)
       onReconnect()
     }
 
@@ -137,6 +206,7 @@ class Db extends EventEmitter implements types.IDatabase, types.EventsMixin {
           clearTimeout(reconnectTimer)
           reconnectTimer = null
         }
+        clearHeartbeat()
         if (client) {
           client.removeAllListeners()
           await client.end().catch(() => {})

src/types.ts

@@ -47,6 +47,10 @@ export interface DatabaseOptions {
   connectionTimeoutMillis?: number;
   /** @internal */
   debug?: boolean;
+  /** @internal */
+  __test__listenHeartbeatIntervalMs?: number;
+  /** @internal */
+  __test__listenHeartbeatTimeoutMs?: number;
 }
 
 export interface SchedulingOptions {

test/dbListenTest.ts

@@ -1,6 +1,9 @@
-import { it, expect } from 'vitest'
+import pg from 'pg'
+import { it, expect, vi } from 'vitest'
+import Db from '../src/db.ts'
 import * as helper from './testHelper.ts'
 import { delay } from '../src/tools.ts'
+import HalfOpenProxy from './halfOpenProxy.ts'
 
 // Exercises the low-level LISTEN/NOTIFY connection lifecycle on Db directly: the
 // dedicated session-pinned client, capped-backoff reconnection after a dropped
@@ -18,6 +21,110 @@ async function terminateListener (db: any, channel: string): Promise<void> {
 // pg_terminate_backend, and capped-backoff reconnection. None of that exists for embedded
 // single-connection PGlite, so skip the whole file there (PGlite's listen is covered via fromPglite).
 helper.describePglite('db listen/notify', function () {
+  it('detects a silent half-open listener and restores the subscription', async function () {
+    const config = helper.getConfig()
+    if (!config.host || !config.port) throw new Error('Postgres host and port are required')
+
+    const proxy = new HalfOpenProxy(config.host, config.port)
+    const proxyPort = await proxy.start()
+    const db = new Db({
+      ...config,
+      application_name: 'pgboss_half_open_test',
+      host: '127.0.0.1',
+      port: proxyPort,
+      __test__listenHeartbeatIntervalMs: 100,
+      __test__listenHeartbeatTimeoutMs: 100
+    })
+    const channel = 'pgboss_db_half_open_test'
+    const payloads: string[] = []
+    const errors: Error[] = []
+    let reconnects = 0
+    let handle: Awaited<ReturnType<Db['listen']>> | undefined
+
+    db.on('error', error => errors.push(error))
+    await db.open()
+
+    try {
+      const generation = proxy.listenerGeneration
+      handle = await db.listen(channel, payload => payloads.push(payload), () => { reconnects++ })
+      const listener = await proxy.waitForListenerAfter(generation)
+      expect(reconnects).toBe(1)
+
+      proxy.blackhole(listener)
+
+      for (let i = 0; i < 40; i++) {
+        if (reconnects >= 2) break
+        await delay(100)
+      }
+
+      expect(reconnects).toBe(2)
+      expect(errors.some(error => error.message === 'LISTEN/NOTIFY heartbeat timed out')).toBe(true)
+
+      await db.executeSql(`NOTIFY "${channel}", 'recovered'`)
+      for (let i = 0; i < 30; i++) {
+        if (payloads.length) break
+        await delay(100)
+      }
+      expect(payloads).toContain('recovered')
+    } finally {
+      await handle?.close()
+      await db.close()
+      await proxy.close()
+    }
+  })
+
+  it('reconnects when the heartbeat finds the channel registration missing', async function () {
+    const db = new Db({
+      ...helper.getConfig(),
+      application_name: 'pgboss_lost_subscription_test',
+      __test__listenHeartbeatIntervalMs: 100,
+      __test__listenHeartbeatTimeoutMs: 100
+    })
+    const channel = 'pgboss_db_lost_subscription_test'
+    const payloads: string[] = []
+    const errors: Error[] = []
+    let reconnects = 0
+    let reportMissing = true
+    let handle: Awaited<ReturnType<Db['listen']>> | undefined
+    const originalQuery = pg.Client.prototype.query
+    const querySpy = vi.spyOn(pg.Client.prototype, 'query').mockImplementation(function (this: pg.Client, ...args: any[]) {
+      const [text] = args
+      if (reportMissing && typeof text === 'string' && text.includes('FROM pg_listening_channels()')) {
+        reportMissing = false
+        return Promise.resolve({ rows: [{ listening: false }] }) as any
+      }
+
+      return originalQuery.apply(this, args as any)
+    })
+
+    db.on('error', error => errors.push(error))
+
+    try {
+      await db.open()
+      handle = await db.listen(channel, payload => payloads.push(payload), () => { reconnects++ })
+      expect(reconnects).toBe(1)
+
+      for (let i = 0; i < 30; i++) {
+        if (reconnects >= 2) break
+        await delay(100)
+      }
+
+      expect(reconnects).toBe(2)
+      expect(errors.some(error => error.message === 'LISTEN/NOTIFY channel registration was lost')).toBe(true)
+
+      await db.executeSql(`NOTIFY "${channel}", 'recovered'`)
+      for (let i = 0; i < 30; i++) {
+        if (payloads.length) break
+        await delay(100)
+      }
+      expect(payloads).toContain('recovered')
+    } finally {
+      querySpy.mockRestore()
+      await handle?.close()
+      await db.close()
+    }
+  })
+
   it('reconnects after the listen connection drops and delivers later notifications', async function () {
     const db = await helper.getDb()
     const channel = 'pgboss_db_reconnect_test'