Fix cloud dev-stack degradation under sustained MCP load

The cloud e2e stack degrades after a few minutes of sustained suite load and eventually hangs, which is the CI cascade flake: MCP-heavy shards (toolkits-mcp et al) run, then unrelated specs die with Playwright TimeoutErrors because requests hang rather than erroring. Root cause is postgres connection churn against the dev PGlite server, not SSE frame buffering or OTel span growth. Two compounding problems: - The dev-db PGlite socket server (apps/cloud/scripts/dev-db.ts) used the library default maxConnections: 1. The cloud worker opens a fresh postgres pool per request (the MCP auth seam rebuilds one on every /mcp request to respect workerd's per-request I/O rule), so under concurrent load the second-and-later connects were answered with "Too many connections" and an immediate socket close. postgres.js then reconnected in a tight loop, piling up thousands of half-closed sockets, starving real queries and driving request latency into the tens of seconds until the stack hung. PGlite runs queries serially anyway (its QueryQueueManager uses runExclusive), so allowing many connections is safe: they queue instead of storming. Raise the cap (overridable via DEV_DB_MAX_CONNECTIONS). - The postgres pool finalizer (apps/cloud/src/db/db.ts) was fire-and-forget: Effect.runFork(sql.end({ timeout: 0 })) returned before the socket was torn down, so closing sockets accumulated faster than the server reaped them. Await a graceful end() inside the acquireRelease scope so the number of live-plus-closing sockets stays bounded to what is actually in flight. Extracted as closePostgres() and covered by a unit test. Measured on a booted cloud stack (3 concurrent MCP clients, 1s think-time): baseline (maxConnections=1): ~18 sessions in 135s, avg 15s/session, p95 up to 50s, latency climbing, stack stalling, thousands of "Too many connections". fixed: 347 sessions in 135s, avg ~125ms/session, p95 ~150ms, flat, zero "Too many connections".

Rhys Sullivan committed 4 months ago

Commit df86ec51 parent dce15d53 files · +169 −7Browse files

Unified · Split

3 files

apps/cloud/scripts/dev-db.ts +18 −0
apps/cloud/src/db/db.close.test.ts +109 −0
apps/cloud/src/db/db.ts +42 −7

apps/cloud/scripts/dev-db.ts+18−0

@@ -82,10 +82,28 @@ const db = await PGlite.create(DB_PATH); · 81 unmodified lines
82	82	console.log(`[dev-db] Running migrations from ${MIGRATIONS_FOLDER}`);
83	83	await migrate(drizzle(db), { migrationsFolder: MIGRATIONS_FOLDER });
84	84
	85	+// `PGLiteSocketServer` defaults to `maxConnections: 1` and answers every extra
	86	+// concurrent connection with "Too many connections" + an immediate socket
	87	+// close. (pglite-socket 0.1.4's published index.d.ts documents "default: 100",
	88	+// but the shipped runtime JS is `maxConnections ?? 1`, verified in the shipped
	89	+// chunk, so the runtime default really is 1.) The cloud worker opens a fresh
	90	+// postgres pool per request (the MCP auth seam rebuilds one on EVERY `/mcp`
	91	+// request, see apps/cloud/src/mcp/auth.ts), so under concurrent load, exactly
	92	+// what the e2e suite generates against one shared dev stack, the
	93	+// second-and-later connects were rejected, and postgres.js reconnected in a
	94	+// tight loop against the closed socket. That reconnect storm piled up
	95	+// thousands of half-closed sockets, starved real queries, drove request
	96	+// latency into the tens of seconds, and eventually hung the stack: the CI e2e
	97	+// "cloud dev stack degrades after minutes of sustained load" cascade flake.
	98	+// PGlite still runs queries serially (its internal QueryQueueManager executes
	99	+// each under `runExclusive`), so allowing many connections is safe: they queue
	100	+// instead of being rejected. Raise the cap so concurrent requests wait their
	101	+// turn instead of storming.
85	102	const server = new PGLiteSocketServer({
86	103	db,
87	104	port: PORT,
88	105	host: "127.0.0.1",
	106	+ maxConnections: Number(process.env.DEV_DB_MAX_CONNECTIONS ?? 1000),
89	107	});
90	108
91	109	await server.start();

apps/cloud/src/db/db.close.test.ts+109−0

@@ -0,0 +1,109 @@
	1	+// Contract: the postgres pool finalizer must AWAIT the connection teardown.
	2	+//
	3	+// The cloud MCP auth seam (`makeMcpOrganizationAuthServices`) builds a fresh
	4	+// postgres pool on EVERY `/mcp` request and closes it in its `acquireRelease`
	5	+// finalizer. That finalizer used to be fire-and-forget (`Effect.runFork(
	6	+// sql.end({ timeout: 0 }))`): it returned before the socket was torn down, so
	7	+// under sustained MCP load closed-but-unreaped sockets piled up against the
	8	+// dev PGlite server (effectively single-connection) faster than it reaped
	9	+// them. New connects queued behind the backlog, request latency climbed into
	10	+// the tens of seconds, and the e2e cloud dev stack hung after a few minutes:
	11	+// the CI cascade flake.
	12	+//
	13	+// These tests characterize the contract the old fire-and-forget close violated
	14	+// (`closePostgres` itself is new alongside them): it must (a) call `sql.end`
	15	+// with a NON-zero drain window (a clean Terminate, not an abandon) and (b)
	16	+// return an Effect that does not complete until `sql.end` has resolved, even
	17	+// when the teardown takes real wall-clock time. All asserted with a fake `sql`
	18	+// whose `end()` completion is observable, so the tests are fast and need no
	19	+// live database.
	20	+
	21	+import { describe, expect, it } from "@effect/vitest";
	22	+import { Effect, Exit } from "effect";
	23	+
	24	+import { POSTGRES_END_TIMEOUT_SECONDS, closePostgres } from "./db";
	25	+
	26	+describe("closePostgres", () => {
	27	+ it.effect("passes a non-zero drain window to sql.end (clean Terminate, not abandon)", () =>
	28	+ Effect.gen(function* () {
	29	+ let received: { timeout?: number } \| undefined;
	30	+ const fakeSql = {
	31	+ end: (options?: { timeout?: number }) => {
	32	+ received = options;
	33	+ return Promise.resolve();
	34	+ },
	35	+ };
	36	+
	37	+ yield* closePostgres(fakeSql);
	38	+
	39	+ expect(received?.timeout).toBe(POSTGRES_END_TIMEOUT_SECONDS);
	40	+ expect(POSTGRES_END_TIMEOUT_SECONDS).toBeGreaterThan(0);
	41	+ }),
	42	+ );
	43	+
	44	+ it.effect("does not complete until sql.end resolves (awaits the teardown)", () =>
	45	+ Effect.gen(function* () {
	46	+ // `end` records an ordering marker only after an async tick. If
	47	+ // `closePostgres` awaits it, the "close completed" marker lands AFTER the
	48	+ // "end resolved" marker. The old fire-and-forget close returned before
	49	+ // `end` ran, so "close completed" would land FIRST.
	50	+ const order: string[] = [];
	51	+ const fakeSql = {
	52	+ end: () =>
	53	+ // Defer resolution across a microtask so a non-awaiting close would
	54	+ // observably finish before this runs.
	55	+ Promise.resolve()
	56	+ .then(() => Promise.resolve())
	57	+ .then(() => {
	58	+ order.push("end-resolved");
	59	+ }),
	60	+ };
	61	+
	62	+ yield* closePostgres(fakeSql);
	63	+ order.push("close-completed");
	64	+
	65	+ // Awaiting the teardown means end resolved strictly before close returned.
	66	+ expect(order).toEqual(["end-resolved", "close-completed"]);
	67	+ }),
	68	+ );
	69	+
	70	+ it.effect("awaits a teardown that takes real wall-clock time (bounded by the ceiling)", () =>
	71	+ Effect.gen(function* () {
	72	+ // `end` resolves only after a real timer delay, not just a microtask.
	73	+ // This pins the "we await to completion, the timeout is only a ceiling"
	74	+ // contract: closePostgres must stay suspended across the delay rather
	75	+ // than resolving early.
	76	+ const order: string[] = [];
	77	+ const fakeSql = {
	78	+ end: () =>
	79	+ new Promise<void>((resolve) => {
	80	+ setTimeout(() => {
	81	+ order.push("end-resolved");
	82	+ resolve();
	83	+ }, 75);
	84	+ }),
	85	+ };
	86	+
	87	+ const startedAt = Date.now();
	88	+ yield* closePostgres(fakeSql);
	89	+ order.push("close-completed");
	90	+
	91	+ expect(order).toEqual(["end-resolved", "close-completed"]);
	92	+ // Slack of a few ms: platform timers may fire marginally early.
	93	+ expect(Date.now() - startedAt).toBeGreaterThanOrEqual(70);
	94	+ }),
	95	+ );
	96	+
	97	+ it.effect("swallows sql.end failures (a teardown error must not fail the request scope)", () =>
	98	+ Effect.gen(function* () {
	99	+ const fakeSql = {
	100	+ // A rejected teardown (connection already gone) must not surface as a
	101	+ // scope failure.
	102	+ // oxlint-disable-next-line executor/no-promise-reject -- test fake: model `sql.end` (a raw postgres.js promise) rejecting
	103	+ end: () => new Promise<void>((_resolve, reject) => reject("connection already gone")),
	104	+ };
	105	+ const exit = yield* Effect.exit(closePostgres(fakeSql));
	106	+ expect(Exit.isSuccess(exit)).toBe(true);
	107	+ }),
	108	+ );
	109	+});

apps/cloud/src/db/db.ts+42−7

@@ -67,18 +67,53 @@ const makeSql = (): Sql => · 66 unmodified lines
67	67	onnotice: () => undefined,
68	68	});
69	69
	70	+/**
	71	+ * Graceful drain window (seconds) handed to `postgres`'s `sql.end`. Small but
	72	+ * non-zero: `timeout: 0` closes immediately without waiting for a clean
	73	+ * Terminate, which leaves the socket half-closed for the server to reap on its
	74	+ * own schedule. A short drain lets postgres.js finish the wire teardown.
	75	+ *
	76	+ * This is a `Promise.race` CEILING, not a fixed wait: postgres.js races
	77	+ * `end({ timeout })` against the actual teardown, and an idle connection's
	78	+ * `end()` calls `terminate()` immediately and resolves as soon as the socket
	79	+ * closes (sub-millisecond). Awaiting it in the request scope therefore adds no
	80	+ * meaningful latency on the common path; the 5s only bounds how long a
	81	+ * connection still mid-query can hold the scope open.
	82	+ */
	83	+export const POSTGRES_END_TIMEOUT_SECONDS = 5;
	84	+
	85	+/**
	86	+ * Close a postgres pool and AWAIT its teardown.
	87	+ *
	88	+ * The `DbService` layers ({@link DbService.Live}, {@link makeDbLayer}) run this
	89	+ * as their `acquireRelease` finalizer. The MCP auth seam
	90	+ * (`makeMcpOrganizationAuthServices`) builds a FRESH pool on EVERY `/mcp`
	91	+ * request, so this finalizer runs per request under sustained load.
	92	+ *
	93	+ * It used to be fire-and-forget (`Effect.runFork(sql.end({ timeout: 0 }))`),
	94	+ * which returned before the connection was actually torn down. The abandoned
	95	+ * sockets piled up against the dev PGlite server (effectively single-connection)
	96	+ * faster than it reaped them; new connects then queued behind the backlog, so
	97	+ * request latency climbed into the tens of seconds and the stack eventually
	98	+ * hung — the CI e2e "cloud dev stack degrades after minutes of sustained load"
	99	+ * cascade. Awaiting the close bounds the number of live-plus-closing sockets to
	100	+ * what is actually in flight. It runs inside the request's own Effect scope, so
	101	+ * it respects workerd's per-request I/O rule.
	102	+ */
	103	+export const closePostgres = (sql: Pick<Sql, "end">): Effect.Effect<void> =>
	104	+ Effect.ignore(
	105	+ Effect.tryPromise({
	106	+ try: () => sql.end({ timeout: POSTGRES_END_TIMEOUT_SECONDS }),
	107	+ catch: (cause) => cause,
	108	+ }),
	109	+ );
	110	+
70	111	const makePostgresResource = (): DbResource => {
71	112	const sql = makeSql();
72	113	return {
73	114	sql,
74	115	db: drizzle(sql, { schema: combinedSchema }) as DrizzleDb,
75		− close: () =>
76		− Effect.ignore(
77		− Effect.tryPromise({
78		− try: () => sql.end({ timeout: 0 }),
79		− catch: (cause) => cause,
80		− }),
81		− ),
	116	+ close: () => closePostgres(sql),
82	117	};
83	118	};
84	119