diff --git a/src/lib/adapters/cloudflare-runtime.ts b/src/lib/adapters/cloudflare-runtime.ts index 047d61150..397c39d03 100644 --- a/src/lib/adapters/cloudflare-runtime.ts +++ b/src/lib/adapters/cloudflare-runtime.ts @@ -117,6 +117,7 @@ type CloudflareRuntimeEnv = Record & { type CloudflareRuntimeContext = { cache: Map; cacheStorage?: CloudflareCacheStorage; + cleanups: Set<() => Promise | void>; env?: CloudflareRuntimeEnv; request?: CloudflareRequestContext; scheduleTask?: CloudflareTaskScheduler; @@ -177,6 +178,61 @@ function getCurrentCloudflareRuntimeEnv() { return globalForCloudflareRuntime.__lifeUstcCloudflareRuntimeEnv; } +async function cleanupCloudflareRuntimeContext( + context: CloudflareRuntimeContext, +) { + const cleanupResults = await Promise.allSettled( + [...context.cleanups].map((cleanup) => Promise.resolve().then(cleanup)), + ); + context.cache.clear(); + context.cleanups.clear(); + const failures = cleanupResults + .filter( + (cleanupResult): cleanupResult is PromiseRejectedResult => + cleanupResult.status === "rejected", + ) + .map((cleanupResult) => cleanupResult.reason); + if (failures.length === 1) throw failures[0]; + if (failures.length > 1) { + throw new AggregateError(failures, "Cloudflare runtime cleanup failed"); + } +} + +function responseWithRuntimeCleanup( + response: Response, + cleanup: () => Promise, +) { + if (!response.body) return response; + const reader = response.body.getReader(); + const body = new ReadableStream( + { + async pull(controller) { + try { + const chunk = await reader.read(); + if (!chunk.done) { + controller.enqueue(chunk.value); + return; + } + await cleanup(); + controller.close(); + } catch (error) { + await cleanup().catch(() => undefined); + controller.error(error); + } + }, + async cancel(reason) { + try { + await reader.cancel(reason); + } finally { + await cleanup(); + } + }, + }, + { highWaterMark: 0 }, + ); + return new Response(body, response); +} + export function runWithCloudflareRuntimeEnv( env: unknown, callback: () => T | Promise, @@ -195,20 +251,43 @@ export function runWithCloudflareRuntimeEnv( const context: CloudflareRuntimeContext = { cache: new Map(), cacheStorage: normalizeCloudflareCacheStorage(), + cleanups: new Set(), env: normalizeCloudflareRuntimeEnv(env), scheduleTask: normalizeCloudflareTaskScheduler(executionContext), tracing, }; return cloudflareRuntimeStorage.run(context, async () => { + let cleanupPromise: Promise | undefined; + const cleanup = () => { + cleanupPromise ??= cleanupCloudflareRuntimeContext(context); + return cleanupPromise; + }; + let result: T; try { - return await callback(); - } finally { - context.cache.clear(); + result = await callback(); + } catch (error) { + await cleanup().catch(() => undefined); + throw error; + } + if ( + result instanceof Response && + result.body && + context.cleanups.size > 0 + ) { + return responseWithRuntimeCleanup(result, cleanup) as T; } + await cleanup(); + return result; }); } +export function registerCloudflareRuntimeCleanup( + cleanup: () => Promise | void, +) { + cloudflareRuntimeStorage.getStore()?.cleanups.add(cleanup); +} + export function runCloudflareTraceSpan( name: string, attributes: Record, diff --git a/src/lib/db/auth-prisma.ts b/src/lib/db/auth-prisma.ts index 77f84201e..2889467aa 100644 --- a/src/lib/db/auth-prisma.ts +++ b/src/lib/db/auth-prisma.ts @@ -3,6 +3,7 @@ import type { PrismaClient } from "@/generated/prisma/client"; import { getCloudflareRuntimeContext, hasCloudflareRuntimeEnv, + registerCloudflareRuntimeCleanup, } from "@/lib/adapters/cloudflare-runtime"; import { createBasePrisma, logPrismaQuery } from "@/lib/db/prisma-query-events"; import { shouldEnablePrismaQueryLogging } from "@/lib/db/prisma-query-logging"; @@ -31,6 +32,7 @@ function getBaseAuthPrisma() { | undefined; if (cached) return cached; const client = createAuthPrismaClient(); + registerCloudflareRuntimeCleanup(() => client.$disconnect()); cache.set(cloudflareAuthPrismaCacheKey, client); return client; } diff --git a/src/lib/db/maintenance-prisma.ts b/src/lib/db/maintenance-prisma.ts index 557c881da..085010d39 100644 --- a/src/lib/db/maintenance-prisma.ts +++ b/src/lib/db/maintenance-prisma.ts @@ -3,6 +3,7 @@ import type { PrismaClient } from "@/generated/prisma/client"; import { getCloudflareRuntimeContext, hasCloudflareRuntimeEnv, + registerCloudflareRuntimeCleanup, } from "@/lib/adapters/cloudflare-runtime"; import { createBasePrisma, logPrismaQuery } from "@/lib/db/prisma-query-events"; import { shouldEnablePrismaQueryLogging } from "@/lib/db/prisma-query-logging"; @@ -33,6 +34,7 @@ function getBaseMaintenancePrisma() { | undefined; if (cached) return cached; const client = createMaintenancePrismaClient(); + registerCloudflareRuntimeCleanup(() => client.$disconnect()); cache.set(cloudflareMaintenancePrismaCacheKey, client); return client; } diff --git a/src/lib/db/prisma-adapter.ts b/src/lib/db/prisma-adapter.ts index f11bc89e7..4bcca3524 100644 --- a/src/lib/db/prisma-adapter.ts +++ b/src/lib/db/prisma-adapter.ts @@ -79,14 +79,13 @@ export function createPrismaAdapter( { connectionString: resolvedConnectionString, // On Workers every request builds a fresh pool (pg sockets cannot be - // reused across requests) and each new connection pays full SCRAM - // deriveBits CPU. Concurrent queries (Promise.all, RLS tx + session - // lookup) otherwise open up to pg's default of 10 connections per - // request; cap the pool so a single request can never open more than 3. + // reused across requests), which the runtime context disconnects before + // the request completes. Concurrent queries (Promise.all, RLS tx + + // session lookup) otherwise open up to pg's default of 10 connections; + // cap the pool so a single request can never open more than 3. max: 3, - // Idle connections from a finished request are never reusable, so close - // them quickly instead of holding sockets (and server slots) for pg's - // 10s default. + // Keep a short idle timeout as a safety net for clients created outside + // the managed request context. idleTimeoutMillis: 5_000, }, { diff --git a/src/lib/db/prisma.ts b/src/lib/db/prisma.ts index 8a86746d1..103b3f499 100644 --- a/src/lib/db/prisma.ts +++ b/src/lib/db/prisma.ts @@ -3,6 +3,7 @@ import type { PrismaClient } from "@/generated/prisma/client"; import { getCloudflareRuntimeContext, hasCloudflareRuntimeEnv, + registerCloudflareRuntimeCleanup, } from "@/lib/adapters/cloudflare-runtime"; import { localizedNamesExtension } from "@/lib/db/prisma-localized-names"; import { createBasePrisma, logPrismaQuery } from "@/lib/db/prisma-query-events"; @@ -67,11 +68,16 @@ function getBasePrisma() { if (hasCloudflareRuntimeEnv()) { const cache = getCloudflarePrismaCache(); if (cache) { - cache.base ??= createPrismaClient(cache); + if (!cache.base) { + cache.base = createPrismaClient(cache); + registerCloudflareRuntimeCleanup(() => cache.base?.$disconnect()); + } return cache.base; } - return createPrismaClient(); + const client = createPrismaClient(); + registerCloudflareRuntimeCleanup(() => client.$disconnect()); + return client; } const cached = globalForPrisma.prisma ?? basePrisma; diff --git a/tests/unit/auth-prisma-boundary.test.ts b/tests/unit/auth-prisma-boundary.test.ts index 967a95554..624086280 100644 --- a/tests/unit/auth-prisma-boundary.test.ts +++ b/tests/unit/auth-prisma-boundary.test.ts @@ -1,4 +1,4 @@ -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import { getCloudflareAuthHyperdriveConnectionString, runWithCloudflareRuntimeEnv, @@ -6,10 +6,16 @@ import { const { appClient, createBasePrismaMock, firstClient, secondClient } = vi.hoisted(() => ({ - appClient: { user: { boundary: "app" } }, + appClient: { $disconnect: vi.fn(), user: { boundary: "app" } }, createBasePrismaMock: vi.fn(), - firstClient: { user: { boundary: "first-auth" } }, - secondClient: { user: { boundary: "second-auth" } }, + firstClient: { + $disconnect: vi.fn(), + user: { boundary: "first-auth" }, + }, + secondClient: { + $disconnect: vi.fn(), + user: { boundary: "second-auth" }, + }, })); vi.mock("@/lib/db/prisma-query-events", () => ({ @@ -22,6 +28,12 @@ vi.mock("@/lib/db/prisma-query-logging", () => ({ })); describe("auth Prisma boundary", () => { + beforeEach(() => { + appClient.$disconnect.mockClear(); + firstClient.$disconnect.mockClear(); + secondClient.$disconnect.mockClear(); + }); + it("keeps overlapping Cloudflare requests on their own auth clients", async () => { createBasePrismaMock.mockReset().mockImplementation((_url, database) => { if (database !== "auth") return appClient; @@ -65,6 +77,8 @@ describe("auth Prisma boundary", () => { expect(createBasePrismaMock).toHaveBeenCalledTimes(2); expect(createBasePrismaMock).toHaveBeenCalledWith(undefined, "auth"); + expect(firstClient.$disconnect).toHaveBeenCalledOnce(); + expect(secondClient.$disconnect).toHaveBeenCalledOnce(); }); it("keeps app and auth clients distinct inside one request", async () => { @@ -94,5 +108,7 @@ describe("auth Prisma boundary", () => { expect(createBasePrismaMock).toHaveBeenCalledTimes(2); expect(createBasePrismaMock).toHaveBeenCalledWith(); expect(createBasePrismaMock).toHaveBeenCalledWith(undefined, "auth"); + expect(appClient.$disconnect).toHaveBeenCalledOnce(); + expect(firstClient.$disconnect).toHaveBeenCalledOnce(); }); }); diff --git a/tests/unit/cloudflare-runtime-tracing.test.ts b/tests/unit/cloudflare-runtime-tracing.test.ts index fd3019717..79168ab21 100644 --- a/tests/unit/cloudflare-runtime-tracing.test.ts +++ b/tests/unit/cloudflare-runtime-tracing.test.ts @@ -2,6 +2,7 @@ import { afterEach, describe, expect, it, vi } from "vitest"; import { getCloudflareNamedCache, getCloudflareRuntimeTaskScheduler, + registerCloudflareRuntimeCleanup, runCloudflareTraceSpan, runWithCloudflareRuntimeEnv, } from "@/lib/adapters/cloudflare-runtime"; @@ -148,4 +149,58 @@ describe("Cloudflare runtime tracing", () => { expect(getCloudflareNamedCache("outside-request")).toBeUndefined(); expect(getCloudflareRuntimeTaskScheduler()).toBeUndefined(); }); + + it("awaits request-scoped cleanup before resolving", async () => { + const events: string[] = []; + + await runWithCloudflareRuntimeEnv({}, async () => { + registerCloudflareRuntimeCleanup(async () => { + await Promise.resolve(); + events.push("cleanup"); + }); + events.push("callback"); + }); + + expect(events).toEqual(["callback", "cleanup"]); + }); + + it("defers cleanup until a response body finishes streaming", async () => { + const cleanup = vi.fn(); + + const response = await runWithCloudflareRuntimeEnv({}, () => { + registerCloudflareRuntimeCleanup(cleanup); + return new Response("streamed"); + }); + + expect(cleanup).not.toHaveBeenCalled(); + await expect(response.text()).resolves.toBe("streamed"); + expect(cleanup).toHaveBeenCalledOnce(); + }); + + it("cleans up when a response body is canceled", async () => { + const cleanup = vi.fn(); + const cancel = vi.fn(); + + const response = await runWithCloudflareRuntimeEnv({}, () => { + registerCloudflareRuntimeCleanup(cleanup); + return new Response(new ReadableStream({ cancel })); + }); + + await response.body?.cancel("client disconnected"); + expect(cancel).toHaveBeenCalledWith("client disconnected"); + expect(cleanup).toHaveBeenCalledOnce(); + }); + + it("preserves callback errors when cleanup also fails", async () => { + const callbackFailure = new Error("callback failed"); + + await expect( + runWithCloudflareRuntimeEnv({}, () => { + registerCloudflareRuntimeCleanup(() => { + throw new Error("cleanup failed"); + }); + throw callbackFailure; + }), + ).rejects.toBe(callbackFailure); + }); }); diff --git a/tests/unit/prisma-rls-context.test.ts b/tests/unit/prisma-rls-context.test.ts index a470a25ff..449810468 100644 --- a/tests/unit/prisma-rls-context.test.ts +++ b/tests/unit/prisma-rls-context.test.ts @@ -17,6 +17,7 @@ const { baseClient, extendedClient, todoFindManyMock } = vi.hoisted(() => { }; return { baseClient: { + $disconnect: vi.fn(), $extends: vi.fn(() => extended), }, extendedClient: extended, @@ -92,6 +93,7 @@ describe("localized Prisma clients in RLS context", () => { ); }); }); + expect(baseClient.$disconnect).toHaveBeenCalledOnce(); }); it("blocks saved localized clients, delegates, and methods inside RLS context", async () => {