orca/cloud/apps/relay/src/database/postgres-lost-reply-postgres.test.ts

import { createServer, connect, type Server, type Socket } from 'node:net'
import { afterEach, describe, expect, it } from 'vitest'
import {
  isRelayDatabaseTransientError,
  openRelayDatabase,
  type RelayDatabase
} from './database.js'

const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
const describePostgres = databaseUrl ? describe : describe.skip

const STATEMENT_TIMEOUT_MS = 500
// A short margin keeps the blackhole case fast; production's is 10 s.
const READ_TIMEOUT_MARGIN_MS = 2_000
const PROBE_TABLE = 'relay_lost_reply_probe'

// Drop server replies after a statement naming the trigger reaches PostgreSQL.
async function blackholeProxy(target: URL, trigger: string) {
  const sockets = new Set<Socket>()
  const afterTrigger: Buffer[] = []
  const server: Server = createServer((client) => {
    const upstream = connect(Number(target.port || 5432), target.hostname)
    sockets.add(client)
    sockets.add(upstream)
    let swallowing = false
    client.on('data', (chunk: Buffer) => {
      if (swallowing) afterTrigger.push(chunk)
      else if (chunk.includes(trigger)) swallowing = true
      upstream.write(chunk)
    })
    upstream.on('data', (chunk: Buffer) => {
      if (!swallowing) {
        client.write(chunk)
      }
    })
    for (const socket of [client, upstream]) {
      socket.on('error', () => undefined)
      socket.on('close', () => {
        client.destroy()
        upstream.destroy()
      })
    }
  })
  await new Promise<void>((resolve) => server.listen(0, '127.0.0.1', resolve))
  const address = server.address()
  if (typeof address !== 'object' || address === null) throw new Error('proxy not listening')
  const url = new URL(target)
  url.hostname = '127.0.0.1'
  url.port = String(address.port)
  return {
    url: url.toString(),
    sentAfterTrigger: () => Buffer.concat(afterTrigger).toString('latin1'),
    close: async () => {
      for (const socket of sockets) socket.destroy()
      await new Promise<void>((resolve) => server.close(() => resolve()))
    }
  }
}

describePostgres('a lost PostgreSQL reply against a real server', () => {
  const cleanups: Array<() => Promise<void>> = []
  afterEach(async () => {
    for (const cleanup of cleanups.splice(0).reverse()) await cleanup()
  })

  async function open(
    url: string,
    options: {
      statementTimeoutMs?: number
    } = {}
  ): Promise<RelayDatabase> {
    const database = await openRelayDatabase({
      databaseUrl: url,
      dataDir: '',
      appliesPostgresSchema: false,
      statementTimeoutMs: options.statementTimeoutMs ?? STATEMENT_TIMEOUT_MS,
      readTimeoutMarginMs: READ_TIMEOUT_MARGIN_MS
    })
    cleanups.push(async () => await database.close())
    return database
  }

  it('rejects a lost transaction reply without sending more SQL and recovers', async () => {
    const direct = await open(databaseUrl!, { statementTimeoutMs: 5_000 })
    await direct.query(`CREATE TABLE IF NOT EXISTS ${PROBE_TABLE} (id TEXT PRIMARY KEY)`)
    await direct.query(`INSERT INTO ${PROBE_TABLE} (id) VALUES ('c1') ON CONFLICT DO NOTHING`)
    cleanups.push(async () => {
      await direct.query(`DROP TABLE IF EXISTS ${PROBE_TABLE}`)
    })
    const proxy = await blackholeProxy(new URL(databaseUrl!), `FROM ${PROBE_TABLE}`)
    cleanups.push(proxy.close)
    const database = await open(proxy.url)

    const failure = await database
      .transaction(
        async (transaction) =>
          await transaction.queryLocked(`SELECT id FROM ${PROBE_TABLE} WHERE id = ?`, ['c1'], {
            lockTimeoutMs: 200
          })
      )
      .catch((error: unknown) => error)

    expect(failure).toBeInstanceOf(Error)
    expect(isRelayDatabaseTransientError(failure)).toBe(true)
    // pg's Terminate ('X') may follow as the client is destroyed; no statement may.
    const after = proxy.sentAfterTrigger()
    for (const statement of ['lock_timeout', 'ROLLBACK', 'COMMIT', 'SELECT']) {
      expect(after).not.toContain(statement)
    }
    // The next checkout is a fresh connection, so it is not stuck behind the lost reply.
    await expect(database.query('SELECT 1 AS one')).resolves.toEqual([{ one: 1 }])
    // The row lock went with the destroyed session.
    await expect(
      direct.transaction(
        async (transaction) =>
          await transaction.queryLocked(`SELECT id FROM ${PROBE_TABLE} WHERE id = ?`, ['c1'], {
            failIfUnavailable: true
          })
      )
    ).resolves.toEqual([{ id: 'c1' }])
  }, 20_000)
})