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(); +}