feat(relay): expose preloaded PostgreSQL statement statistics (#20712)

cloud/apps/relay/src/database-postgres-timeout.test.ts

@@ -1,4 +1,5 @@
 import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'
+import { POSTGRES_STATEMENT_STATS_MIGRATION } from './postgres-statement-stats.js'
 
 const fakes = vi.hoisted(() => ({
   configs: [] as Array<Record<string, unknown>>,
@@ -119,11 +120,16 @@ describe('PostgreSQL relay deadlines', () => {
     })
 
     expect(ddl.length).toBeGreaterThan(0)
+    expect(ddl).toContain(POSTGRES_STATEMENT_STATS_MIGRATION)
     // Statements can open with a leading `--` rationale comment.
     const body = (statement: string): string =>
       statement.replace(/^(?:\s*--[^\n]*\n)*\s*/, '')
     expect(
-      ddl.every((statement) => /^(?:CREATE|ALTER TABLE)\b/i.test(body(statement)))
+      ddl.every(
+        (statement) =>
+          statement === POSTGRES_STATEMENT_STATS_MIGRATION ||
+          /^(?:CREATE|ALTER TABLE)\b/i.test(body(statement))
+      )
     ).toBe(true)
     // The backfill is DML, so it stays on the deadline-bearing serving pool.
     expect(ddl.some((statement) => statement.includes('INSERT INTO'))).toBe(false)

cloud/apps/relay/src/database.ts

@@ -10,6 +10,7 @@ import {
   type PostgresPoolPressureCounts
 } from './postgres-pool-pressure.js'
 import { applyPostgresSchema } from './postgres-schema-startup.js'
+import { POSTGRES_STATEMENT_STATS_MIGRATION } from './postgres-statement-stats.js'
 import {
   CellInventoryHoldSamples,
   emptyCellInventoryHoldCounts,
@@ -619,6 +620,7 @@ CREATE INDEX IF NOT EXISTS relay_audit_events_at ON relay_audit_events(at);
 // auto-named; the replacement is named, so both statements are no-ops on a
 // database the current schema created and neither can drop the other.
 export const POSTGRES_SCHEMA_MIGRATIONS = [
+  POSTGRES_STATEMENT_STATS_MIGRATION,
   `ALTER TABLE relay_region_decisions ADD COLUMN IF NOT EXISTS last_considered_at BIGINT NOT NULL DEFAULT 0`,
   `ALTER TABLE relay_region_decisions ADD COLUMN IF NOT EXISTS cohort_bucket BIGINT NOT NULL DEFAULT 0`,
   `ALTER TABLE relay_region_rehome_attempts

cloud/apps/relay/src/postgres-statement-stats-postgres.test.ts

@@ -0,0 +1,121 @@
+import { randomUUID } from 'node:crypto'
+import pg from 'pg'
+import { afterAll, beforeAll, describe, expect, it } from 'vitest'
+import { openRelayDatabase } from './database.js'
+import { POSTGRES_STATEMENT_STATS_MIGRATION } from './postgres-statement-stats.js'
+
+const databaseUrl = process.env.ORCA_RELAY_TEST_POSTGRES_URL
+const describePostgres = databaseUrl ? describe : describe.skip
+
+describePostgres('optional PostgreSQL statement statistics', () => {
+  let admin: pg.Client
+  let preloaded: boolean
+  const databases: string[] = []
+  const roles: string[] = []
+
+  beforeAll(async () => {
+    admin = new pg.Client({ connectionString: databaseUrl })
+    await admin.connect()
+    const result = await admin.query<{ loaded: boolean }>(
+      `SELECT 'pg_stat_statements' = ANY(string_to_array(
+         replace(current_setting('shared_preload_libraries'), ' ', ''), ','
+       )) AS loaded`
+    )
+    preloaded = result.rows[0]!.loaded
+  })
+
+  afterAll(async () => {
+    for (const database of databases) await admin.query(`DROP DATABASE IF EXISTS ${database}`)
+    for (const role of roles) await admin.query(`DROP ROLE IF EXISTS ${role}`)
+    await admin.end()
+  })
+
+  async function freshDatabase(): Promise<string> {
+    const name = `relay_stats_${randomUUID().replaceAll('-', '')}`
+    await admin.query(`CREATE DATABASE ${name}`)
+    databases.push(name)
+    const url = new URL(databaseUrl!)
+    url.pathname = `/${name}`
+    return url.toString()
+  }
+
+  async function connect(url: string): Promise<pg.Client> {
+    const client = new pg.Client({ connectionString: url, statement_timeout: 2_000 })
+    await client.connect()
+    return client
+  }
+
+  async function installed(client: pg.Client): Promise<boolean> {
+    const result = await client.query<{ present: boolean }>(
+      `SELECT EXISTS (SELECT 1 FROM pg_extension WHERE extname = 'pg_stat_statements') AS present`
+    )
+    return result.rows[0]!.present
+  }
+
+  it('exposes an existing collector idempotently, and skips servers without one', async () => {
+    const url = await freshDatabase()
+    const database = await openRelayDatabase({ databaseUrl: url, dataDir: '' })
+    await database.close()
+    const client = await connect(url)
+    try {
+      expect(await installed(client)).toBe(preloaded)
+      if (preloaded) {
+        const before = await client.query('SELECT stats_reset FROM public.pg_stat_statements_info')
+        await client.query(POSTGRES_STATEMENT_STATS_MIGRATION)
+        const after = await client.query('SELECT stats_reset FROM public.pg_stat_statements_info')
+        expect(after.rows).toEqual(before.rows)
+        await client.query('SELECT calls, wal_bytes, shared_blks_dirtied FROM public.pg_stat_statements LIMIT 1')
+      } else {
+        await client.query(POSTGRES_STATEMENT_STATS_MIGRATION)
+        expect(await installed(client)).toBe(false)
+      }
+    } finally {
+      await client.end()
+    }
+  })
+
+  it.each([false, true])('tolerates missing extension privileges (read settings: %s)', async (readSettings) => {
+    const client = await connect(await freshDatabase())
+    const role = `relay_stats_role_${randomUUID().replaceAll('-', '')}`
+    await admin.query(`CREATE ROLE ${role}`)
+    roles.push(role)
+    if (readSettings) await admin.query(`GRANT pg_read_all_settings TO ${role}`)
+    try {
+      await client.query(`SET ROLE ${role}`)
+      await client.query(POSTGRES_STATEMENT_STATS_MIGRATION)
+      expect(await installed(client)).toBe(false)
+      expect((await client.query<{ value: number }>('SELECT 42 AS value')).rows[0]!.value).toBe(42)
+    } finally {
+      await client.end()
+    }
+  })
+
+  it('serializes concurrent catalog creation across directors', async () => {
+    const url = await freshDatabase()
+    const clients = await Promise.all(Array.from({ length: 5 }, async () => await connect(url)))
+    try {
+      await Promise.all(clients.map(async (client) => await client.query(POSTGRES_STATEMENT_STATS_MIGRATION)))
+      expect(await installed(clients[0]!)).toBe(preloaded)
+    } finally {
+      await Promise.all(clients.map(async (client) => await client.end()))
+    }
+  })
+
+  it('yields to an in-progress installer instead of blocking startup', async () => {
+    const url = await freshDatabase()
+    const owner = await connect(url)
+    const contender = await connect(url)
+    try {
+      await owner.query('BEGIN')
+      await owner.query(`SELECT pg_advisory_xact_lock(hashtext('orca-relay'), hashtext('statement-stats'))`)
+      await contender.query(POSTGRES_STATEMENT_STATS_MIGRATION)
+      expect(await installed(contender)).toBe(false)
+      await owner.query('COMMIT')
+      await contender.query(POSTGRES_STATEMENT_STATS_MIGRATION)
+      expect(await installed(contender)).toBe(preloaded)
+    } finally {
+      await owner.end()
+      await contender.end()
+    }
+  })
+})

cloud/apps/relay/src/postgres-statement-stats.ts

@@ -0,0 +1,28 @@
+// Expose an already-running collector; never preload a module or require elevated runtime privileges.
+export const POSTGRES_STATEMENT_STATS_MIGRATION = `
+DO $relay_statement_stats$
+BEGIN
+  IF NOT EXISTS (
+    SELECT 1 FROM pg_catalog.pg_settings
+    WHERE name = 'shared_preload_libraries'
+      AND 'pg_stat_statements' = ANY(string_to_array(replace(setting, ' ', ''), ','))
+  ) OR EXISTS (
+    SELECT 1 FROM pg_catalog.pg_extension WHERE extname = 'pg_stat_statements'
+  ) OR NOT EXISTS (
+    SELECT 1 FROM pg_catalog.pg_available_extensions WHERE name = 'pg_stat_statements'
+  ) THEN
+    RETURN;
+  END IF;
+
+  IF NOT pg_try_advisory_xact_lock(hashtext('orca-relay'), hashtext('statement-stats')) THEN
+    RETURN;
+  END IF;
+
+  BEGIN
+    CREATE EXTENSION IF NOT EXISTS pg_stat_statements WITH SCHEMA public;
+  EXCEPTION WHEN insufficient_privilege THEN
+    RAISE WARNING 'orca_relay_statement_stats_unavailable: insufficient privilege';
+  END;
+END
+$relay_statement_stats$;
+`

cloud/docs/orca-relay-operations.md

@@ -2,6 +2,21 @@
 
 This runbook applies to the stable Cloud Run director and the production-shaped GCE cells in both environments. It does not authorize a full Terraform apply: staging and production contain unrelated drift, so inspect a saved targeted plan and its destroy count before every apply.
 
+## PostgreSQL statement statistics
+
+Relay schema startup exposes `pg_stat_statements` when the server already preloads
+that collector and the schema identity can install its extension. Servers without
+the collector or the required privileges continue normally. Installation does not
+change preload settings, reset collected counters, or require a database restart;
+concurrent startups yield to one installer. An existing extension is left in place.
+
+For SQL incidents, inspect bounded aggregates of `calls`, `total_exec_time`,
+`shared_blks_read`, `shared_blks_dirtied`, and `wal_bytes`, scoped to the relay
+database and identified query IDs. Compare counter deltas over the same interval
+as fleet runtime metrics; retain the statistics reset timestamp. Do not export
+query text, identities, credentials, or invoke `pg_stat_statements_reset()` during
+an investigation. Treat an unavailable view as missing evidence, not zero work.
+
 The relay is automatically active for entitled signed-in desktops. There is no rollout flag, cohort, or user toggle. The emergency product kill switch is the auth plane refusing relay-token exchange; use cell drains only to move or terminate existing data-plane work.
 
 ## Safety rules