From 9bcd399079ea3af5873d06c6fac747f34f26bcdd Mon Sep 17 00:00:00 2001 From: Remon Oldenbeuving Date: Thu, 6 Aug 2026 17:05:42 +0200 Subject: [PATCH] diag(ai-gateway): measure whether pool slots are leaked The instrumentation from #5098 established that pool acquisition dominates /api/internal/usage/record: idle == 0 on 96-98% of sampled requests with up to 245 queued against max 10, while PostgreSQL reports zero active queries and Supavisor zero waiting clients. Connections are checked out without doing work, and there are two explanations that call for opposite fixes. Either Fluid compute packs high concurrency onto few instances and ten connections is simply too few, or slots are leaked and never return. drizzle-orm 0.45.2 issues BEGIN outside its try/finally, so a failed BEGIN never reaches release() and burns that slot for the life of the process; ten of those kill an instance permanently. If that is what is happening, raising max postpones and hides it, and max is a single module-level constant shared by every route in both Vercel projects whose comment records that raising it previously exhausted Supabase's connection limit. The arithmetic is what makes this worth measuring rather than assuming: at ~47 req/s with Frankfurt-local statements of 2-10ms, roughly 5 connections are needed, not 245 queued. The discriminator is the low-water mark of checked-out connections. A healthy pool returns to zero checked out whenever the instance goes quiet, and idleTimeoutMillis of 5s then closes the idle clients. A leaked client is never returned to _idle, so it counts as checked out forever and the low-water mark cannot fall below the number of leaks. A low-water mark rather than an instantaneous reading because this pool is shared with every other route on the instance: one quiet moment resets it, so unrelated traffic cannot inflate it. Adds, all read-only counters updated on pg-pool events, with no queries and no timers: - acquires/releases/connects/removes, plus outstanding as a cross-check that should track pool.totalCount - pool.idleCount - min_checked_out and max_checked_out - begin_failures, counted where the retry loop already inspects the error. Matching the message is safe only for BEGIN, which takes no parameters and so carries no prompt text or client IP - ms_since_last_request and checked_out_at_entry, read before the request acquires anything Emission now also triggers on the first request after 5s of instance quiet. Those requests are fast, so the duration threshold would never surface them, yet they are the only moment a leaked slot is distinguishable from a busy one. Counters do not attach under NODE_ENV=test, matching how drizzle.ts already treats the test pool. --- .../app/api/internal/usage/record/route.ts | 9 +- apps/web/src/lib/ai-gateway/processUsage.ts | 7 + .../usage-record-diagnostics.test.ts | 41 ++++- .../ai-gateway/usage-record-diagnostics.ts | 47 ++++- apps/web/src/lib/db-pool-leak-probe.test.ts | 102 ++++++++++ apps/web/src/lib/db-pool-leak-probe.ts | 174 ++++++++++++++++++ 6 files changed, 373 insertions(+), 7 deletions(-) create mode 100644 apps/web/src/lib/db-pool-leak-probe.test.ts create mode 100644 apps/web/src/lib/db-pool-leak-probe.ts diff --git a/apps/web/src/app/api/internal/usage/record/route.ts b/apps/web/src/app/api/internal/usage/record/route.ts index e94b416d8e..ae53dc8f1b 100644 --- a/apps/web/src/app/api/internal/usage/record/route.ts +++ b/apps/web/src/app/api/internal/usage/record/route.ts @@ -14,6 +14,7 @@ import { import { createPhaseTimer, emitUsageRecordTiming, + noteRequestStart, readPoolGauges, shouldEmitUsageRecordTiming, } from '@/lib/ai-gateway/usage-record-diagnostics'; @@ -48,7 +49,11 @@ export async function POST(request: NextRequest): Promise { // pairing them with the in-process pool gauges distinguishes waiting for a pool // connection from waiting for PostgreSQL from a stalled event loop. const timer = createPhaseTimer(); + // Read before this request acquires anything: on an instance that has been quiet + // longer than `idleTimeoutMillis`, anything still checked out belongs to nobody. + const msSinceLastRequest = noteRequestStart(); const poolBefore = readPoolGauges(); + const checkedOutAtEntry = poolBefore.total - poolBefore.idle; let poolWaitingPeak = poolBefore.waiting; const samplePool = () => { const gauges = readPoolGauges(); @@ -57,7 +62,7 @@ export async function POST(request: NextRequest): Promise { }; const reportTiming = (usageId: string, outcome: UsageRecordResponse['status']) => { const totalMs = timer.totalMs(); - if (!shouldEmitUsageRecordTiming(totalMs)) return; + if (!shouldEmitUsageRecordTiming(totalMs, msSinceLastRequest)) return; emitUsageRecordTiming({ usageId, outcome, @@ -66,6 +71,8 @@ export async function POST(request: NextRequest): Promise { poolBefore, poolAfter: samplePool(), poolWaitingPeak, + msSinceLastRequest, + checkedOutAtEntry, }); }; diff --git a/apps/web/src/lib/ai-gateway/processUsage.ts b/apps/web/src/lib/ai-gateway/processUsage.ts index 5100c6ac95..0ef2c50eb1 100644 --- a/apps/web/src/lib/ai-gateway/processUsage.ts +++ b/apps/web/src/lib/ai-gateway/processUsage.ts @@ -6,6 +6,7 @@ import { isUsageRowConflict, stackFramesUnderHeader, } from './usage-record-diagnostics'; +import { isTransactionBeginFailure, recordTransactionBeginFailure } from '@/lib/db-pool-leak-probe'; import type { MicrodollarUsage } from '@kilocode/db/schema'; import { microdollar_usage } from '@kilocode/db/schema'; import { createTimer } from '@/lib/timer'; @@ -689,6 +690,12 @@ export async function insertUsageRecord( // retrying — `id` is fixed for the delivery — so stop immediately and // let the outer handler recover the committed row's identity. Retrying // it burned three attempts and rebuilt the statement each time. + // Count a failed `BEGIN` before deciding what to do with it. Under + // drizzle 0.45.2 the `BEGIN` is issued outside the `try`/`finally`, so + // this failure never reaches `release()` and burns a pool slot for the + // life of the process. The count is what makes the leak measurable + // against the pool's low-water mark. + if (isTransactionBeginFailure(error)) recordTransactionBeginFailure(); if (attempt >= 2 || isUsageRowConflict(error)) throw error; // Never log the raw error: its message is the interpolated statement, // roughly 30KB including prompt prefixes and the client IP. diff --git a/apps/web/src/lib/ai-gateway/usage-record-diagnostics.test.ts b/apps/web/src/lib/ai-gateway/usage-record-diagnostics.test.ts index 95a51e5975..8d4d774691 100644 --- a/apps/web/src/lib/ai-gateway/usage-record-diagnostics.test.ts +++ b/apps/web/src/lib/ai-gateway/usage-record-diagnostics.test.ts @@ -2,15 +2,17 @@ jest.mock('@/lib/drizzle', () => ({ pool: { totalCount: 4, idleCount: 1, waitingCount: 0, options: { max: 10 } }, })); -import { describe, expect, test } from '@jest/globals'; +import { beforeEach, describe, expect, test } from '@jest/globals'; import { createPhaseTimer, describeDatabaseError, isUsageRowConflict, + noteRequestStart, readPoolGauges, shouldEmitUsageRecordTiming, stackFramesUnderHeader, + __resetRequestClockForTest, } from './usage-record-diagnostics'; /** Shape of the drizzle wrapper: a message carrying the statement, plus a cause. */ @@ -254,12 +256,41 @@ describe('readPoolGauges', () => { describe('shouldEmitUsageRecordTiming', () => { test('always emits at or above the slow threshold', () => { - expect(shouldEmitUsageRecordTiming(1_000, () => 1)).toBe(true); - expect(shouldEmitUsageRecordTiming(50_000, () => 1)).toBe(true); + expect(shouldEmitUsageRecordTiming(1_000, null, () => 1)).toBe(true); + expect(shouldEmitUsageRecordTiming(50_000, null, () => 1)).toBe(true); }); test('samples a small fraction of fast requests for a baseline', () => { - expect(shouldEmitUsageRecordTiming(5, () => 0.005)).toBe(true); - expect(shouldEmitUsageRecordTiming(5, () => 0.5)).toBe(false); + expect(shouldEmitUsageRecordTiming(5, null, () => 0.005)).toBe(true); + expect(shouldEmitUsageRecordTiming(5, null, () => 0.5)).toBe(false); + }); + + // The leak observation: a request landing on a quiet instance is fast, so the + // duration threshold would never surface it, yet it is the only moment when a + // leaked slot is distinguishable from a busy one. + test('always emits the first request after the instance has been quiet', () => { + expect(shouldEmitUsageRecordTiming(5, 5_000, () => 1)).toBe(true); + expect(shouldEmitUsageRecordTiming(5, 60_000, () => 1)).toBe(true); + }); + + test('does not treat a busy instance as an observation point', () => { + expect(shouldEmitUsageRecordTiming(5, 4_999, () => 1)).toBe(false); + expect(shouldEmitUsageRecordTiming(5, 0, () => 1)).toBe(false); + }); +}); + +describe('noteRequestStart', () => { + beforeEach(() => { + __resetRequestClockForTest(); + }); + + test('reports null for the first request on an instance', () => { + expect(noteRequestStart(1_000)).toBeNull(); + }); + + test('reports the gap since the previous request', () => { + noteRequestStart(1_000); + expect(noteRequestStart(7_500)).toBe(6_500); + expect(noteRequestStart(7_600)).toBe(100); }); }); diff --git a/apps/web/src/lib/ai-gateway/usage-record-diagnostics.ts b/apps/web/src/lib/ai-gateway/usage-record-diagnostics.ts index d1de25d674..2992d7808f 100644 --- a/apps/web/src/lib/ai-gateway/usage-record-diagnostics.ts +++ b/apps/web/src/lib/ai-gateway/usage-record-diagnostics.ts @@ -2,6 +2,7 @@ import 'server-only'; import { monitorEventLoopDelay } from 'node:perf_hooks'; import { pool } from '@/lib/drizzle'; +import { readPoolLeakStats } from '@/lib/db-pool-leak-probe'; /** * Diagnostics for the Frankfurt-local usage write. @@ -95,13 +96,54 @@ export type UsageRecordTiming = { poolAfter: PoolGauges; /** Highest `waitingCount` sampled during the request. */ poolWaitingPeak: number; + /** + * Milliseconds since the previous request on this instance, `null` for the + * first. Above `QUIET_INSTANCE_MS` the pool should have drained, so a non-zero + * `checked_out` in the same line is a leaked slot. + */ + msSinceLastRequest: number | null; + /** + * Checked-out connections observed at handler entry, before this request + * acquired anything. On a quiet instance this should be zero. + */ + checkedOutAtEntry: number; }; +/** + * A request arriving after this much instance quiet is treated as an observation + * point regardless of how fast it was. `idleTimeoutMillis` is 5s, so by then a + * healthy pool has closed its idle clients and `checked_out` should be zero — + * which is precisely when a leaked slot is visible. These requests are fast, so + * the duration threshold would never surface them. + */ +const QUIET_INSTANCE_MS = 5_000; + export function shouldEmitUsageRecordTiming( totalMs: number, + msSinceLastRequest: number | null = null, random: () => number = Math.random ): boolean { - return totalMs >= SLOW_EMIT_THRESHOLD_MS || random() < BASELINE_SAMPLE_RATE; + if (totalMs >= SLOW_EMIT_THRESHOLD_MS) return true; + if (msSinceLastRequest !== null && msSinceLastRequest >= QUIET_INSTANCE_MS) return true; + return random() < BASELINE_SAMPLE_RATE; +} + +/** + * Milliseconds since the previous request on this instance, or `null` for the + * first one. Tracked here rather than in the route so the notion of "quiet" and + * its emission threshold stay together. + */ +let lastRequestAtMs: number | null = null; + +export function noteRequestStart(now: number = Date.now()): number | null { + const since = lastRequestAtMs === null ? null : now - lastRequestAtMs; + lastRequestAtMs = now; + return since; +} + +/** Test-only reset so quiet-period assertions do not depend on suite ordering. */ +export function __resetRequestClockForTest(): void { + lastRequestAtMs = null; } /** @@ -120,6 +162,9 @@ export function emitUsageRecordTiming(timing: UsageRecordTiming): void { pool_waiting_peak: timing.poolWaitingPeak, pool_max: poolMax(), event_loop_lag: eventLoopLagMs(), + ms_since_last_request: timing.msSinceLastRequest, + checked_out_at_entry: timing.checkedOutAtEntry, + pool_leak: readPoolLeakStats(), }) ); } diff --git a/apps/web/src/lib/db-pool-leak-probe.test.ts b/apps/web/src/lib/db-pool-leak-probe.test.ts new file mode 100644 index 0000000000..244aa70899 --- /dev/null +++ b/apps/web/src/lib/db-pool-leak-probe.test.ts @@ -0,0 +1,102 @@ +// Counters attach to the real pool only outside the test environment, so the pool +// is stubbed here and the gauges are driven directly. The stub is built inside the +// factory because `jest.mock` is hoisted above this file's own declarations, so +// referencing an outer const here throws "Cannot access before initialization". +jest.mock('@/lib/drizzle', () => ({ + pool: { totalCount: 0, idleCount: 0, waitingCount: 0, on: jest.fn() }, +})); + +import { beforeEach, describe, expect, test } from '@jest/globals'; + +import { + isTransactionBeginFailure, + readPoolLeakStats, + recordTransactionBeginFailure, + __resetPoolLeakCountersForTest, +} from './db-pool-leak-probe'; + +const { pool: mockPool } = jest.requireMock<{ + pool: { totalCount: number; idleCount: number }; +}>('@/lib/drizzle'); + +function drizzleError(message: string, cause?: unknown) { + const error = new Error(message); + error.name = 'DrizzleQueryError'; + if (cause !== undefined) (error as unknown as { cause: unknown }).cause = cause; + return error; +} + +beforeEach(() => { + __resetPoolLeakCountersForTest(); + mockPool.totalCount = 0; + mockPool.idleCount = 0; +}); + +describe('isTransactionBeginFailure', () => { + // BEGIN takes no parameters, which is what makes matching the message safe here + // and unsafe for every other statement in this path. + test.each([ + 'Failed query: begin', + 'Failed query: begin\nparams: ', + 'Failed query: BEGIN', + ' Failed query: begin isolation level serializable', + ])('recognises %j', message => { + expect(isTransactionBeginFailure(drizzleError(message))).toBe(true); + }); + + test('recognises a begin failure nested in the cause chain', () => { + expect( + isTransactionBeginFailure(drizzleError('wrapper', drizzleError('Failed query: begin'))) + ).toBe(true); + }); + + // These must not be counted as leaks: the statement failed inside the + // try/finally, so drizzle released the client. + test.each([ + 'Failed query: WITH microdollar_usage_ins AS (...) params: 3f5826e7,You are Kilo', + 'Failed query: commit', + 'Failed query: rollback', + 'Failed query: select 1', + 'Connection terminated unexpectedly', + ])('does not match %j', message => { + expect(isTransactionBeginFailure(drizzleError(message))).toBe(false); + }); + + test('does not match a statement that merely mentions begin', () => { + expect(isTransactionBeginFailure(drizzleError('Failed query: select * from beginnings'))).toBe( + false + ); + }); + + test('does not throw on null or a primitive', () => { + expect(isTransactionBeginFailure(null)).toBe(false); + expect(isTransactionBeginFailure('nope')).toBe(false); + }); +}); + +describe('readPoolLeakStats', () => { + test('derives checked_out from the pool gauges', () => { + mockPool.totalCount = 10; + mockPool.idleCount = 3; + expect(readPoolLeakStats().checked_out).toBe(7); + }); + + test('reports min_checked_out as null until a pool event is seen', () => { + expect(readPoolLeakStats().min_checked_out).toBeNull(); + }); + + test('counts begin failures, the per-leak unit', () => { + recordTransactionBeginFailure(); + recordTransactionBeginFailure(); + expect(readPoolLeakStats().begin_failures).toBe(2); + }); + + test('reports outstanding as acquires minus releases', () => { + const stats = readPoolLeakStats(); + expect(stats.outstanding).toBe(stats.acquires - stats.releases); + }); + + test('exposes uptime so a low-water mark can be read against instance age', () => { + expect(readPoolLeakStats().instance_uptime_ms).toBeGreaterThanOrEqual(0); + }); +}); diff --git a/apps/web/src/lib/db-pool-leak-probe.ts b/apps/web/src/lib/db-pool-leak-probe.ts new file mode 100644 index 0000000000..8ac549fede --- /dev/null +++ b/apps/web/src/lib/db-pool-leak-probe.ts @@ -0,0 +1,174 @@ +import 'server-only'; + +import { pool } from '@/lib/drizzle'; + +/** + * Read-only probe that decides whether pooled connections are being leaked. + * + * `/api/internal/usage/record` shows the in-process pool pinned at `idle: 0` with + * up to 245 requests queued against `max: 10`, while PostgreSQL reports zero + * active queries and Supavisor reports zero waiting clients. Connections are + * therefore checked out without doing work, and there are two very different + * explanations: + * + * 1. Vercel Fluid compute packs high concurrency onto few instances, so ten + * connections are genuinely in use and the cap is simply too low. Fix: more + * connections. + * 2. Slots are leaked and never come back. drizzle-orm 0.45.2 issues `BEGIN` + * *outside* its `try`/`finally`, so a `BEGIN` that fails never reaches the + * `release()` in the `finally` and burns that slot for the life of the + * process. Ten of those kill an instance permanently. Fix: release on a failed + * `BEGIN`. Raising `max` would only postpone and hide it. + * + * Guessing wrong is expensive in opposite directions, so measure instead. Note + * `max` in `drizzle.ts` is a single module-level constant shared by every route in + * both Vercel projects, and its comment records that raising it previously + * exhausted Supabase's connection limit across ~2,200 instances. + * + * The discriminator is the **low-water mark** of checked-out connections. A + * healthy pool returns to zero checked out whenever the instance goes quiet, and + * `idleTimeoutMillis` of 5s then closes the idle clients. A leaked client is never + * returned to `_idle`, so it is counted as checked out forever and the low-water + * mark can never fall back below the number of leaks. + * + * A low-water mark is deliberately chosen over an instantaneous reading because + * this pool is shared with every other route on the instance: a single quiet + * moment resets it, so concurrent unrelated traffic cannot inflate it. If it + * climbs over an instance's lifetime and tracks `beginFailures`, explanation 2 is + * confirmed with a number. + * + * Everything here is counters updated on pool events. No queries, no timers. + */ + +type PoolLeakCounters = { + acquires: number; + releases: number; + connects: number; + removes: number; + /** + * Failures of the `BEGIN` statement itself. Under the drizzle bug each one is + * exactly one permanently leaked slot, which is what makes this directly + * comparable to `minCheckedOut`. + */ + beginFailures: number; + maxCheckedOut: number; + /** Low-water mark of checked-out connections. The leak discriminator. */ + minCheckedOut: number; + startedAtMs: number; +}; + +const counters: PoolLeakCounters = { + acquires: 0, + releases: 0, + connects: 0, + removes: 0, + beginFailures: 0, + maxCheckedOut: 0, + minCheckedOut: Number.POSITIVE_INFINITY, + startedAtMs: Date.now(), +}; + +function checkedOut(): number { + return pool.totalCount - pool.idleCount; +} + +function sampleCheckedOut(): void { + const current = checkedOut(); + if (current > counters.maxCheckedOut) counters.maxCheckedOut = current; + if (current < counters.minCheckedOut) counters.minCheckedOut = current; +} + +// Attaching in the test environment would leave the counters at the mercy of +// whatever the shared Jest pool does, and `drizzle.ts` already treats the test +// pool as a special case for the same reason. +if (process.env.NODE_ENV !== 'test') { + pool.on('acquire', () => { + counters.acquires++; + sampleCheckedOut(); + }); + pool.on('release', () => { + counters.releases++; + sampleCheckedOut(); + }); + pool.on('connect', () => { + counters.connects++; + }); + pool.on('remove', () => { + counters.removes++; + sampleCheckedOut(); + }); +} + +export function recordTransactionBeginFailure(): void { + counters.beginFailures++; +} + +/** + * True when the failing statement was the transaction's own `BEGIN`. + * + * Safe to match on the message, unlike every other statement in this path: + * `BEGIN` takes no parameters, so the drizzle message carries no prompt text, + * client IP or other request data. + */ +export function isTransactionBeginFailure(error: unknown): boolean { + let current: unknown = error; + for (let depth = 0; depth < 5 && current !== null && typeof current === 'object'; depth++) { + const candidate = current as { message?: unknown; cause?: unknown }; + if ( + typeof candidate.message === 'string' && + /^failed query:\s*begin\b/i.test(candidate.message.trim()) + ) { + return true; + } + current = candidate.cause; + } + return false; +} + +export type PoolLeakReading = { + acquires: number; + releases: number; + connects: number; + removes: number; + begin_failures: number; + checked_out: number; + max_checked_out: number; + /** + * `null` until the first pool event. Above zero after a quiet period is the + * leak signature; compare against `begin_failures`. + */ + min_checked_out: number | null; + /** + * `acquires - releases`. Should equal `checked_out`; a persistent divergence + * means the accounting itself is wrong and the rest should not be trusted. + */ + outstanding: number; + instance_uptime_ms: number; +}; + +export function readPoolLeakStats(): PoolLeakReading { + return { + acquires: counters.acquires, + releases: counters.releases, + connects: counters.connects, + removes: counters.removes, + begin_failures: counters.beginFailures, + checked_out: checkedOut(), + max_checked_out: counters.maxCheckedOut, + min_checked_out: Number.isFinite(counters.minCheckedOut) ? counters.minCheckedOut : null, + outstanding: counters.acquires - counters.releases, + instance_uptime_ms: Date.now() - counters.startedAtMs, + }; +} + +/** Test-only reset so counter assertions do not depend on suite ordering. */ +export function __resetPoolLeakCountersForTest(): void { + counters.acquires = 0; + counters.releases = 0; + counters.connects = 0; + counters.removes = 0; + counters.beginFailures = 0; + counters.maxCheckedOut = 0; + counters.minCheckedOut = Number.POSITIVE_INFINITY; + counters.startedAtMs = Date.now(); +}