diff --git a/.server-changes/transaction-resilience-during-db-interruptions.md b/.server-changes/transaction-resilience-during-db-interruptions.md new file mode 100644 index 0000000000..c05bfd5f22 --- /dev/null +++ b/.server-changes/transaction-resilience-during-db-interruptions.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: improvement +--- + +Triggering tasks is now more resilient to brief, transient service interruptions, so short stalls are less likely to surface as errors. diff --git a/apps/webapp/app/db.server.ts b/apps/webapp/app/db.server.ts index a06321e167..e049631398 100644 --- a/apps/webapp/app/db.server.ts +++ b/apps/webapp/app/db.server.ts @@ -32,6 +32,13 @@ import { import { computeRunOpsSplitReadEnabled } from "./v3/runOpsMigration/runOpsSplitReadGate"; import { assertControlPlaneCoresidencyAdvisory } from "./v3/runOpsMigration/controlPlaneCoresidencySentinel.server"; import { DATASOURCE_CONTEXT_KEY, startActiveSpan } from "./v3/tracer.server"; +import { + controlPlaneTransactionResilience, + registerTransactionResilience, + resilienceForClient, + runOpsLegacyTransactionResilience, + runOpsTransactionResilience, +} from "./v3/transactionResilience.server"; import type { Span } from "@opentelemetry/api"; import { context, trace } from "@opentelemetry/api"; import { queryPerformanceMonitor } from "./utils/queryPerformanceMonitor.server"; @@ -59,6 +66,18 @@ function logTransactionPrismaError(error: Prisma.PrismaClientKnownRequestError) }); } +function withTransactionDefaults( + client: PrismaClientOrTransaction, + options?: PrismaTransactionOptions +): PrismaTransactionOptions { + const resilience = resilienceForClient(client as object); + return { + maxWait: resilience.maxWait, + ...options, + startRetry: options?.startRetry ?? resilience.startRetry, + }; +} + export async function $transaction( prisma: PrismaClientOrTransaction, name: string, @@ -93,35 +112,41 @@ async function $transactionInner( options?: PrismaTransactionOptions ): Promise { if (typeof fnOrName === "string") { + const effectiveOptions = withTransactionDefaults(prisma, options); return await startActiveSpan(fnOrName, async (span) => { span.setAttribute("$transaction", true); - if (options?.isolationLevel) { - span.setAttribute("isolation_level", options.isolationLevel); + if (effectiveOptions.isolationLevel) { + span.setAttribute("isolation_level", effectiveOptions.isolationLevel); } - if (options?.timeout) { - span.setAttribute("timeout", options.timeout); + if (effectiveOptions.timeout) { + span.setAttribute("timeout", effectiveOptions.timeout); } - if (options?.maxWait) { - span.setAttribute("max_wait", options.maxWait); + if (effectiveOptions.maxWait) { + span.setAttribute("max_wait", effectiveOptions.maxWait); } - if (options?.swallowPrismaErrors) { - span.setAttribute("swallow_prisma_errors", options.swallowPrismaErrors); + if (effectiveOptions.swallowPrismaErrors) { + span.setAttribute("swallow_prisma_errors", effectiveOptions.swallowPrismaErrors); } const fn = fnOrOptions as (prisma: PrismaTransactionClient, span: Span) => Promise; - return transac(prisma, (client) => fn(client, span), logTransactionPrismaError, options); + return transac( + prisma, + (client) => fn(client, span), + logTransactionPrismaError, + effectiveOptions + ); }); } else { return transac( prisma, fnOrName, logTransactionPrismaError, - typeof fnOrOptions === "function" ? undefined : fnOrOptions + withTransactionDefaults(prisma, typeof fnOrOptions === "function" ? undefined : fnOrOptions) ); } } @@ -180,7 +205,10 @@ function captureInfraErrorsRunOps(client: RunOpsPrismaClient): RunOpsPrismaClien } export const prisma = singleton("prisma", () => - captureInfrastructureErrors(tagDatasource("control-plane-writer", getClient())) + registerTransactionResilience( + captureInfrastructureErrors(tagDatasource("control-plane-writer", getClient())), + controlPlaneTransactionResilience + ) ); export const $replica: PrismaReplicaClient = singleton("replica", () => { @@ -309,15 +337,18 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => { { controlPlane: { writer: prisma, replica: $replica }, buildNewWriter: (url, clientType) => - captureInfraErrorsRunOps( - tagDatasourceRunOps( - "run-ops-writer", - buildRunOpsWriterClient({ - url, - clientType, - useDriverAdapter: env.RUN_OPS_DATABASE_WRITER_DRIVER_ADAPTER === "1", - }) - ) + registerTransactionResilience( + captureInfraErrorsRunOps( + tagDatasourceRunOps( + "run-ops-writer", + buildRunOpsWriterClient({ + url, + clientType, + useDriverAdapter: env.RUN_OPS_DATABASE_WRITER_DRIVER_ADAPTER === "1", + }) + ) + ), + runOpsTransactionResilience ), // Brand the run-ops replica (only built for a real replica URL) so routed replica reads stay // off the primary. When no replica URL is set, selectRunOpsTopology reuses the writer here — @@ -338,17 +369,20 @@ const runOpsTopology: RunOpsTopology = singleton("runOpsTopology", () => { // Legacy client shares the exact control-plane wrapper stack (the legacy DB carries the full // control-plane schema); markReadReplicaClient only on a real replica URL, as with the NEW replica. buildLegacyWriter: (url, clientType) => - captureInfrastructureErrors( - tagDatasource( - "legacy-run-ops-writer", - buildWriterClient({ - url, - clientType, - poolTimeout: env.RUN_OPS_LEGACY_DATABASE_WRITER_POOL_TIMEOUT, - connectTimeout: env.RUN_OPS_LEGACY_DATABASE_WRITER_CONNECTION_TIMEOUT, - useDriverAdapter: env.RUN_OPS_LEGACY_DATABASE_WRITER_DRIVER_ADAPTER === "1", - }) - ) + registerTransactionResilience( + captureInfrastructureErrors( + tagDatasource( + "legacy-run-ops-writer", + buildWriterClient({ + url, + clientType, + poolTimeout: env.RUN_OPS_LEGACY_DATABASE_WRITER_POOL_TIMEOUT, + connectTimeout: env.RUN_OPS_LEGACY_DATABASE_WRITER_CONNECTION_TIMEOUT, + useDriverAdapter: env.RUN_OPS_LEGACY_DATABASE_WRITER_DRIVER_ADAPTER === "1", + }) + ) + ), + runOpsLegacyTransactionResilience ), buildLegacyReplica: (url, clientType) => markReadReplicaClient( diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 3ce98dd521..6c058c3e13 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -114,6 +114,26 @@ const OptionalIntEnv = z.preprocess( z.coerce.number().int().optional() ); +/** Optional boolean env var; blank/whitespace/unset normalises to undefined (so it falls back). */ +const OptionalBoolEnv = z.preprocess((v) => { + if (typeof v !== "string" || v.trim() === "") return undefined; + return ["true", "1"].includes(v.toLowerCase().trim()); +}, z.boolean().optional()); + +/** Boolean env var with a default where blank/whitespace falls back to the default instead of parsing as false. */ +const BoolEnvWithDefault = (defaultValue: boolean) => + z.preprocess((v) => { + if (typeof v !== "string" || v.trim() === "") return undefined; + return ["true", "1"].includes(v.toLowerCase().trim()); + }, z.boolean().default(defaultValue)); + +/** Int env var with a default where a blank/whitespace value falls back to the default instead of coercing to 0. */ +const IntEnvWithDefault = (defaultValue: number) => + z.preprocess( + (v) => (typeof v === "string" && v.trim() === "" ? undefined : v), + z.coerce.number().int().default(defaultValue) + ); + /** * Optional int env var for a limit that can be switched off. Blank, whitespace and `0` all mean * "no limit" and normalise to undefined; anything else that is set must be greater than zero. @@ -142,6 +162,27 @@ const EnvironmentSchema = z DATABASE_WRITER_CONNECTION_TIMEOUT: OptionalIntEnv, DATABASE_READ_REPLICA_POOL_TIMEOUT: OptionalIntEnv, DATABASE_READ_REPLICA_CONNECTION_TIMEOUT: OptionalIntEnv, + DATABASE_TRANSACTION_MAX_WAIT_MS: IntEnvWithDefault(10000), + DATABASE_TRANSACTION_START_RETRY_ENABLED: BoolEnvWithDefault(true), + DATABASE_TRANSACTION_START_RETRY_MAX_ATTEMPTS: IntEnvWithDefault(3), + DATABASE_TRANSACTION_START_RETRY_BACKOFF_MIN_MS: IntEnvWithDefault(50), + DATABASE_TRANSACTION_START_RETRY_BACKOFF_MAX_MS: IntEnvWithDefault(250), + DATABASE_TRANSACTION_START_RETRY_BUDGET_PER_SEC: IntEnvWithDefault(50), + DATABASE_TRANSACTION_START_RETRY_BUDGET_BURST: IntEnvWithDefault(100), + RUN_OPS_DATABASE_TRANSACTION_MAX_WAIT_MS: OptionalIntEnv, + RUN_OPS_DATABASE_TRANSACTION_START_RETRY_ENABLED: OptionalBoolEnv, + RUN_OPS_DATABASE_TRANSACTION_START_RETRY_MAX_ATTEMPTS: OptionalIntEnv, + RUN_OPS_DATABASE_TRANSACTION_START_RETRY_BACKOFF_MIN_MS: OptionalIntEnv, + RUN_OPS_DATABASE_TRANSACTION_START_RETRY_BACKOFF_MAX_MS: OptionalIntEnv, + RUN_OPS_DATABASE_TRANSACTION_START_RETRY_BUDGET_PER_SEC: OptionalIntEnv, + RUN_OPS_DATABASE_TRANSACTION_START_RETRY_BUDGET_BURST: OptionalIntEnv, + RUN_OPS_LEGACY_DATABASE_TRANSACTION_MAX_WAIT_MS: OptionalIntEnv, + RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_ENABLED: OptionalBoolEnv, + RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_MAX_ATTEMPTS: OptionalIntEnv, + RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_BACKOFF_MIN_MS: OptionalIntEnv, + RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_BACKOFF_MAX_MS: OptionalIntEnv, + RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_BUDGET_PER_SEC: OptionalIntEnv, + RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_BUDGET_BURST: OptionalIntEnv, // Dashboard-agent conversation store. Cloud points this at a dedicated // database; when unset it falls back to DATABASE_URL (OSS), where // the tables live in the isolated `trigger_dashboard_agent` schema. diff --git a/apps/webapp/app/v3/runStore.server.ts b/apps/webapp/app/v3/runStore.server.ts index bf1642b5f4..9ccf84b511 100644 --- a/apps/webapp/app/v3/runStore.server.ts +++ b/apps/webapp/app/v3/runStore.server.ts @@ -12,6 +12,10 @@ import { } from "~/db.server"; import { env } from "~/env.server"; import { singleton } from "~/utils/singleton"; +import { + resilienceForClient, + type TransactionResilienceConfig, +} from "./transactionResilience.server"; type BuildRunStoreDeps = { /** Boot constant: true only when both run-ops DBs are configured and the split flag is on. */ @@ -27,6 +31,10 @@ type BuildRunStoreDeps = { singleReplica: PrismaReplicaClient; /** Residency classifier; defaults to ownerEngine inside RoutingRunStore. */ classify?: (id: string) => Residency; + /** Per-pool transaction-resilience configs threaded into the store(s) this builds (IoC). */ + singleResilience?: TransactionResilienceConfig; + newResilience?: TransactionResilienceConfig; + legacyResilience?: TransactionResilienceConfig; }; /** @@ -46,6 +54,8 @@ export function buildRunStore(deps: BuildRunStoreDeps): RunStore { return new PostgresRunStore({ prisma: deps.singleWriter, readOnlyPrisma: deps.singleReplica, + maxWait: deps.singleResilience?.maxWait, + transactionStartRetry: deps.singleResilience?.startRetry, }); } @@ -59,10 +69,14 @@ export function buildRunStore(deps: BuildRunStoreDeps): RunStore { prisma: deps.newWriter, readOnlyPrisma: deps.newReplica, schemaVariant: "dedicated", + maxWait: deps.newResilience?.maxWait, + transactionStartRetry: deps.newResilience?.startRetry, }); const legacyStore = new PostgresRunStore({ prisma: deps.legacyWriter, readOnlyPrisma: deps.legacyReplica, + maxWait: deps.legacyResilience?.maxWait, + transactionStartRetry: deps.legacyResilience?.startRetry, }); return new RoutingRunStore({ @@ -110,6 +124,7 @@ export const runStore: RunStore = singleton("RunStore", () => { splitEnabled: false, singleWriter: prisma, singleReplica: $replica, + singleResilience: resilienceForClient(prisma), }); } return buildRunStore({ @@ -117,5 +132,8 @@ export const runStore: RunStore = singleton("RunStore", () => { ...handles, singleWriter: prisma, singleReplica: $replica, + singleResilience: resilienceForClient(prisma), + newResilience: resilienceForClient(handles.newWriter), + legacyResilience: resilienceForClient(handles.legacyWriter), }); }); diff --git a/apps/webapp/app/v3/transactionResilience.server.ts b/apps/webapp/app/v3/transactionResilience.server.ts new file mode 100644 index 0000000000..ae5678c987 --- /dev/null +++ b/apps/webapp/app/v3/transactionResilience.server.ts @@ -0,0 +1,100 @@ +import { TokenBucketRetryBudget, type TransactionStartRetryConfig } from "@trigger.dev/database"; +import { env } from "~/env.server"; +import { logger } from "~/services/logger.server"; + +/** + * Resolved transaction-resilience config for one writer pool. Each pool gets its own + * {@link TransactionStartRetryConfig} (with its OWN token bucket, so a storm on one pool cannot + * drain another's retry budget) plus the `maxWait` applied when that pool opens a transaction. + * Env is read here at the app boundary (IoC); the library never reads env. + * + * Kept out of `db.server` on purpose: `db.server` is mocked wholesale by ~150 tests, and a new + * export there breaks every mock that does not list it. Both `db.server` and `runStore.server` + * import these from here instead. + */ +export type TransactionResilienceConfig = { + maxWait: number; + startRetry: TransactionStartRetryConfig; +}; + +function resolveTransactionResilience( + pool: "control-plane" | "run-ops" | "run-ops-legacy", + overrides: { + maxWaitMs?: number; + enabled?: boolean; + maxAttempts?: number; + backoffMinMs?: number; + backoffMaxMs?: number; + budgetPerSec?: number; + budgetBurst?: number; + } +): TransactionResilienceConfig { + const budgetPerSec = + overrides.budgetPerSec ?? env.DATABASE_TRANSACTION_START_RETRY_BUDGET_PER_SEC; + const budgetBurst = overrides.budgetBurst ?? env.DATABASE_TRANSACTION_START_RETRY_BUDGET_BURST; + return { + maxWait: Math.max(0, overrides.maxWaitMs ?? env.DATABASE_TRANSACTION_MAX_WAIT_MS), + startRetry: { + options: { + enabled: overrides.enabled ?? env.DATABASE_TRANSACTION_START_RETRY_ENABLED, + maxAttempts: overrides.maxAttempts ?? env.DATABASE_TRANSACTION_START_RETRY_MAX_ATTEMPTS, + backoffMinMs: overrides.backoffMinMs ?? env.DATABASE_TRANSACTION_START_RETRY_BACKOFF_MIN_MS, + backoffMaxMs: overrides.backoffMaxMs ?? env.DATABASE_TRANSACTION_START_RETRY_BACKOFF_MAX_MS, + }, + budget: new TokenBucketRetryBudget({ ratePerSec: budgetPerSec, burst: budgetBurst }), + onRetry: ({ attempt, delayMs }) => + logger.warn("retrying transaction start after acquisition failure", { + pool, + attempt, + delayMs, + }), + }, + }; +} + +export const controlPlaneTransactionResilience = resolveTransactionResilience("control-plane", {}); + +export const runOpsTransactionResilience = resolveTransactionResilience("run-ops", { + maxWaitMs: env.RUN_OPS_DATABASE_TRANSACTION_MAX_WAIT_MS, + enabled: env.RUN_OPS_DATABASE_TRANSACTION_START_RETRY_ENABLED, + maxAttempts: env.RUN_OPS_DATABASE_TRANSACTION_START_RETRY_MAX_ATTEMPTS, + backoffMinMs: env.RUN_OPS_DATABASE_TRANSACTION_START_RETRY_BACKOFF_MIN_MS, + backoffMaxMs: env.RUN_OPS_DATABASE_TRANSACTION_START_RETRY_BACKOFF_MAX_MS, + budgetPerSec: env.RUN_OPS_DATABASE_TRANSACTION_START_RETRY_BUDGET_PER_SEC, + budgetBurst: env.RUN_OPS_DATABASE_TRANSACTION_START_RETRY_BUDGET_BURST, +}); + +export const runOpsLegacyTransactionResilience = resolveTransactionResilience("run-ops-legacy", { + maxWaitMs: env.RUN_OPS_LEGACY_DATABASE_TRANSACTION_MAX_WAIT_MS, + enabled: env.RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_ENABLED, + maxAttempts: env.RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_MAX_ATTEMPTS, + backoffMinMs: env.RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_BACKOFF_MIN_MS, + backoffMaxMs: env.RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_BACKOFF_MAX_MS, + budgetPerSec: env.RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_BUDGET_PER_SEC, + budgetBurst: env.RUN_OPS_LEGACY_DATABASE_TRANSACTION_START_RETRY_BUDGET_BURST, +}); + +const transactionResilienceByClient = new WeakMap(); + +/** + * Associate a writer client with its pool's resilience config. Returns the client for inline use at + * construction. Kept here (not in db.server) so nothing new lands on db.server's wholesale-mocked + * export surface. + */ +export function registerTransactionResilience( + client: T, + resilience: TransactionResilienceConfig +): T { + transactionResilienceByClient.set(client, resilience); + return client; +} + +/** + * The resilience config registered for a writer client, or the control-plane config as a safe + * fallback. Derives resilience from the ACTUAL client identity rather than an assumed routing role, + * so run-ops clients aliased onto the control-plane pool (split flag off) correctly get the + * control-plane config instead of a run-ops override. + */ +export function resilienceForClient(client: object): TransactionResilienceConfig { + return transactionResilienceByClient.get(client) ?? controlPlaneTransactionResilience; +} diff --git a/internal-packages/database/src/transaction.test.ts b/internal-packages/database/src/transaction.test.ts new file mode 100644 index 0000000000..f29297757d --- /dev/null +++ b/internal-packages/database/src/transaction.test.ts @@ -0,0 +1,340 @@ +import { describe, expect, it } from "vitest"; +import { + $transaction, + isTransactionAcquisitionError, + TokenBucketRetryBudget, + UNLIMITED_RETRY_BUDGET, + withTransactionStartRetry, + type TransactionStartRetryConfig, +} from "./transaction"; + +function acquisitionError() { + return { code: "P2028", message: "Unable to start a transaction in the given time." }; +} + +function poolTimeoutError() { + return { + code: "P2024", + message: "Timed out fetching a new connection from the connection pool.", + }; +} + +function inTxP2028() { + return { code: "P2028", message: "Transaction API error: Transaction already closed." }; +} + +const noSleep = (_ms: number) => Promise.resolve(); + +function config( + overrides?: Partial +): TransactionStartRetryConfig { + return { + options: { + enabled: true, + maxAttempts: 2, + backoffMinMs: 50, + backoffMaxMs: 250, + ...overrides, + }, + sleep: noSleep, + random: () => 0, + }; +} + +function counter() { + let calls = 0; + return { + get calls() { + return calls; + }, + tick() { + calls += 1; + }, + }; +} + +describe("isTransactionAcquisitionError", () => { + it("is true only for P2028 raised at acquisition", () => { + expect(isTransactionAcquisitionError(acquisitionError())).toBe(true); + }); + + it("is false for P2024 (pool exhaustion — retrying makes it worse)", () => { + expect(isTransactionAcquisitionError(poolTimeoutError())).toBe(false); + }); + + it("is false for a P2028 raised inside a running transaction", () => { + expect(isTransactionAcquisitionError(inTxP2028())).toBe(false); + }); + + it("is true for the driver-adapter acquire timeout (pre-BEGIN, no code)", () => { + expect( + isTransactionAcquisitionError({ message: "Timeout exceeded when trying to connect" }) + ).toBe(true); + expect( + isTransactionAcquisitionError(new Error("timeout exceeded when trying to connect")) + ).toBe(true); + }); + + it("is false for non-Prisma errors", () => { + expect(isTransactionAcquisitionError(new Error("boom"))).toBe(false); + expect(isTransactionAcquisitionError(undefined)).toBe(false); + }); +}); + +describe("withTransactionStartRetry", () => { + it("runs once on success", async () => { + const c = counter(); + const run = async () => { + c.tick(); + return "ok"; + }; + await expect(withTransactionStartRetry(run, config())).resolves.toBe("ok"); + expect(c.calls).toBe(1); + }); + + it("retries an acquisition failure then succeeds", async () => { + const c = counter(); + const run = async () => { + c.tick(); + if (c.calls === 1) throw acquisitionError(); + return "ok"; + }; + await expect(withTransactionStartRetry(run, config())).resolves.toBe("ok"); + expect(c.calls).toBe(2); + }); + + it("does NOT retry P2024", async () => { + const c = counter(); + const err = poolTimeoutError(); + const run = async () => { + c.tick(); + throw err; + }; + await expect(withTransactionStartRetry(run, config())).rejects.toBe(err); + expect(c.calls).toBe(1); + }); + + it("stops after maxAttempts total attempts", async () => { + const c = counter(); + const run = async () => { + c.tick(); + throw acquisitionError(); + }; + await expect(withTransactionStartRetry(run, config({ maxAttempts: 3 }))).rejects.toMatchObject({ + code: "P2028", + }); + expect(c.calls).toBe(3); + }); + + it("runs once when disabled", async () => { + const c = counter(); + const run = async () => { + c.tick(); + throw acquisitionError(); + }; + await expect(withTransactionStartRetry(run, config({ enabled: false }))).rejects.toMatchObject({ + code: "P2028", + }); + expect(c.calls).toBe(1); + }); + + it("does not retry when the budget is exhausted", async () => { + const c = counter(); + const budgetChecks = counter(); + const run = async () => { + c.tick(); + throw acquisitionError(); + }; + const budget = { + tryConsume() { + budgetChecks.tick(); + return false; + }, + }; + await expect( + withTransactionStartRetry(run, { ...config({ maxAttempts: 5 }), budget }) + ).rejects.toMatchObject({ code: "P2028" }); + expect(c.calls).toBe(1); + expect(budgetChecks.calls).toBe(1); + }); + + it("does not retry when canRetry() is false (callback already entered)", async () => { + const c = counter(); + const run = async () => { + c.tick(); + throw acquisitionError(); + }; + await expect( + withTransactionStartRetry(run, config({ maxAttempts: 5 }), () => false) + ).rejects.toMatchObject({ code: "P2028" }); + expect(c.calls).toBe(1); + }); + + it("sleeps a jittered delay within [min, max]", async () => { + const c = counter(); + const delays: number[] = []; + const run = async () => { + c.tick(); + if (c.calls === 1) throw acquisitionError(); + return "ok"; + }; + await withTransactionStartRetry(run, { + options: { enabled: true, maxAttempts: 2, backoffMinMs: 50, backoffMaxMs: 250 }, + sleep: (ms) => { + delays.push(ms); + return Promise.resolve(); + }, + random: () => 0.5, + }); + expect(delays).toEqual([150]); + }); +}); + +describe("TokenBucketRetryBudget", () => { + it("allows up to burst then denies", () => { + const budget = new TokenBucketRetryBudget({ ratePerSec: 0, burst: 2, now: () => 1000 }); + expect(budget.tryConsume()).toBe(true); + expect(budget.tryConsume()).toBe(true); + expect(budget.tryConsume()).toBe(false); + }); + + it("refills over time", () => { + let now = 1000; + const budget = new TokenBucketRetryBudget({ ratePerSec: 10, burst: 1, now: () => now }); + expect(budget.tryConsume()).toBe(true); + expect(budget.tryConsume()).toBe(false); + now += 100; + expect(budget.tryConsume()).toBe(true); + }); +}); + +describe("$transaction startRetry wiring", () => { + function fakeClient(behavior: (call: number) => Promise) { + const c = counter(); + return { + client: { + $transaction: (fn: (tx: unknown) => Promise, _options?: unknown) => { + c.tick(); + return behavior(c.calls).then(() => fn({})); + }, + } as any, + get calls() { + return c.calls; + }, + }; + } + + it("retries a transaction start that fails with an acquisition error", async () => { + const fake = fakeClient((call) => + call === 1 ? Promise.reject(acquisitionError()) : Promise.resolve() + ); + const result = await $transaction( + fake.client, + async () => "done", + () => {}, + { startRetry: config() } + ); + expect(result).toBe("done"); + expect(fake.calls).toBe(2); + }); + + it("does not retry without a startRetry config", async () => { + const err = acquisitionError(); + const fake = fakeClient(() => Promise.reject(err)); + let captured: unknown; + await expect( + $transaction( + fake.client, + async () => "x", + (e) => { + captured = e; + }, + {} + ) + ).rejects.toBe(err); + expect(fake.calls).toBe(1); + expect(captured).toBe(err); + }); + + it("passes maxWait through to prisma.$transaction options", async () => { + let seenOptions: unknown; + const client = { + $transaction: (fn: (tx: unknown) => Promise, options?: unknown) => { + seenOptions = options; + return fn({}); + }, + } as any; + await $transaction( + client, + async () => "x", + () => {}, + { maxWait: 10000 } + ); + expect(seenOptions).toEqual({ maxWait: 10000 }); + }); + + it("does not let maxRetries retry an acquisition error while startRetry is active", async () => { + const fake = fakeClient(() => Promise.reject(acquisitionError())); + await expect( + $transaction( + fake.client, + async () => "x", + () => {}, + { startRetry: config({ maxAttempts: 2 }), maxRetries: 3 } + ) + ).rejects.toMatchObject({ code: "P2028" }); + expect(fake.calls).toBe(2); + }); + + it("falls back to maxRetries for acquisition errors when startRetry is disabled", async () => { + const fake = fakeClient(() => Promise.reject(acquisitionError())); + await expect( + $transaction( + fake.client, + async () => "x", + () => {}, + { startRetry: config({ enabled: false }), maxRetries: 3 } + ) + ).rejects.toMatchObject({ code: "P2028" }); + expect(fake.calls).toBe(4); + }); + + it("still lets maxRetries retry a serialization error (P2034)", async () => { + const serializationError = { code: "P2034", message: "write conflict / deadlock" }; + const fake = fakeClient(() => Promise.reject(serializationError)); + await expect( + $transaction( + fake.client, + async () => "x", + () => {}, + { maxRetries: 2 } + ) + ).rejects.toMatchObject({ code: "P2034" }); + expect(fake.calls).toBe(3); + }); + + it("does not retry an acquisition-shaped error thrown after the callback entered", async () => { + let calls = 0; + const prisma = { + $transaction: (fn: (tx: unknown) => Promise) => { + calls += 1; + return fn({}).then(() => { + throw acquisitionError(); + }); + }, + } as any; + await expect( + $transaction( + prisma, + async () => "x", + () => {}, + { startRetry: config() } + ) + ).rejects.toMatchObject({ code: "P2028" }); + expect(calls).toBe(1); + }); + + it("UNLIMITED_RETRY_BUDGET always consumes", () => { + expect(UNLIMITED_RETRY_BUDGET.tryConsume()).toBe(true); + }); +}); diff --git a/internal-packages/database/src/transaction.ts b/internal-packages/database/src/transaction.ts index 806ae60403..a11afe5d1e 100644 --- a/internal-packages/database/src/transaction.ts +++ b/internal-packages/database/src/transaction.ts @@ -59,6 +59,148 @@ export function isPrismaRaceConditionError(error: unknown): boolean { return error.code === "P2025"; } +const TRANSACTION_ACQUISITION_MESSAGE = /Unable to start a transaction in the given time/i; + +/** + * True for a connection-acquisition failure raised BEFORE any SQL ran, so there + * is nothing to undo and retrying it is safe. Two shapes: + * - the default engine's P2028 "Unable to start a transaction in the given time" + * (couldn't borrow a connection within `maxWait`), and + * - the pg driver adapter's "timeout exceeded when trying to connect" (the pg + * Pool couldn't hand out a connection within `connectionTimeoutMillis`). + * Deliberately narrower than {@link isPrismaRetriableError}: it excludes P2024 + * (pool exhausted) and P2028s raised from inside a running transaction. Callers + * additionally gate retries on the transaction callback not having entered, so a + * same-shaped error from a nested transaction never re-runs a side-effectful body. + */ +export function isTransactionAcquisitionError(error: unknown): boolean { + const message = (error as { message?: unknown })?.message; + + if (isPrismaKnownError(error) && error.code === "P2028") { + return typeof message === "string" && TRANSACTION_ACQUISITION_MESSAGE.test(message); + } + + return typeof message === "string" && ADAPTER_ACQUIRE_TIMEOUT.test(message); +} + +/** Retry tuning for transaction-start (P2028-at-acquisition) failures. */ +export type TransactionStartRetryOptions = { + /** Kill switch. When false, {@link withTransactionStartRetry} runs the thunk once. */ + enabled: boolean; + /** Total attempts including the first. `1` (or less) disables retrying. */ + maxAttempts: number; + /** Lower bound of the jittered backoff between attempts, in ms. */ + backoffMinMs: number; + /** Upper bound of the jittered backoff between attempts, in ms. */ + backoffMaxMs: number; +}; + +/** + * A shared rate limiter so a mass freeze (every request failing tx-start at + * once) cannot amplify into a retry storm. One instance is shared across all + * retrying call sites; `tryConsume` returns false when the budget is spent. + */ +export interface RetryBudget { + tryConsume(): boolean; +} + +/** Budget that never denies — the default when a caller supplies no budget. */ +export const UNLIMITED_RETRY_BUDGET: RetryBudget = { tryConsume: () => true }; + +/** + * Token-bucket {@link RetryBudget}: refills at `ratePerSec` tokens/second up to + * `burst`, one token per retry. `now` is injectable for tests. + */ +export class TokenBucketRetryBudget implements RetryBudget { + #tokens: number; + readonly #capacity: number; + readonly #refillPerMs: number; + readonly #now: () => number; + #lastRefillAt: number; + + constructor(options: { ratePerSec: number; burst: number; now?: () => number }) { + this.#capacity = Math.max(0, options.burst); + this.#refillPerMs = Math.max(0, options.ratePerSec) / 1000; + this.#now = options.now ?? Date.now; + this.#tokens = this.#capacity; + this.#lastRefillAt = this.#now(); + } + + tryConsume(): boolean { + const now = this.#now(); + const elapsed = Math.max(0, now - this.#lastRefillAt); + if (elapsed > 0) { + this.#tokens = Math.min(this.#capacity, this.#tokens + elapsed * this.#refillPerMs); + this.#lastRefillAt = now; + } + if (this.#tokens >= 1) { + this.#tokens -= 1; + return true; + } + return false; + } +} + +export type TransactionStartRetryConfig = { + options: TransactionStartRetryOptions; + /** Shared budget; defaults to {@link UNLIMITED_RETRY_BUDGET} when omitted. */ + budget?: RetryBudget; + /** Injectable for tests. */ + sleep?: (ms: number) => Promise; + /** Injectable for tests; returns [0, 1). */ + random?: () => number; + /** Observability hook fired just before each backoff sleep. */ + onRetry?: (info: { attempt: number; delayMs: number; error: unknown }) => void; +}; + +function defaultSleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +/** + * Runs `run` and, ONLY when it throws a P2028-at-acquisition error + * ({@link isTransactionAcquisitionError}), retries it up to + * `options.maxAttempts` total attempts with jittered backoff, gated by the + * shared `budget`. Any other error (and an exhausted budget) rethrows + * immediately. When retrying is disabled the thunk runs exactly once. + */ +export async function withTransactionStartRetry( + run: () => Promise, + config?: TransactionStartRetryConfig, + canRetry?: () => boolean +): Promise { + if (!config || !config.options.enabled || config.options.maxAttempts <= 1) { + return run(); + } + + const { maxAttempts, backoffMinMs, backoffMaxMs } = config.options; + const budget = config.budget ?? UNLIMITED_RETRY_BUDGET; + const sleep = config.sleep ?? defaultSleep; + const random = config.random ?? Math.random; + + let attempt = 1; + while (true) { + try { + return await run(); + } catch (error) { + if ( + attempt >= maxAttempts || + (canRetry !== undefined && !canRetry()) || + !isTransactionAcquisitionError(error) || + !budget.tryConsume() + ) { + throw error; + } + const low = Math.max(0, Math.min(backoffMinMs, backoffMaxMs)); + const high = Math.max(low, backoffMaxMs); + const delayMs = Math.round(low + random() * (high - low)); + config.onRetry?.({ attempt, delayMs, error }); + await sleep(delayMs); + attempt += 1; + } + } +} + export type PrismaTransactionOptions = { /** The maximum amount of time (in ms) Prisma Client will wait to acquire a transaction from the database. The default value is 2000ms. */ maxWait?: number; @@ -77,6 +219,14 @@ export type PrismaTransactionOptions = { * See https://www.prisma.io/docs/orm/prisma-client/queries/transactions#transaction-timing-issues */ maxRetries?: number; + + /** + * When set, retry a transaction START that fails with P2028-at-acquisition + * (see {@link withTransactionStartRetry}). Config is threaded in from the app + * boundary (IoC); this library never reads env. Only wraps the top-level + * attempt, so it never nests with the `maxRetries` serialization retry. + */ + startRetry?: TransactionStartRetryConfig; }; export async function $transaction( @@ -90,11 +240,25 @@ export async function $transaction( return fn(prisma); } + const startRetry = attempt === 0 ? options?.startRetry : undefined; + const startRetryActive = + !!startRetry && startRetry.options.enabled && startRetry.options.maxAttempts > 1; + + let entered = false; try { - return await (prisma as PrismaClient).$transaction(fn, options); + return await withTransactionStartRetry( + () => + (prisma as PrismaClient).$transaction((tx) => { + entered = true; + return fn(tx as PrismaTransactionClient); + }, options), + startRetry, + () => !entered + ); } catch (error) { if ( isPrismaRetriableError(error) && + !(startRetryActive && isTransactionAcquisitionError(error)) && typeof options?.maxRetries === "number" && attempt < options.maxRetries ) { diff --git a/internal-packages/run-store/src/PostgresRunStore.ts b/internal-packages/run-store/src/PostgresRunStore.ts index b3d24c5a65..d11d51fb2c 100644 --- a/internal-packages/run-store/src/PostgresRunStore.ts +++ b/internal-packages/run-store/src/PostgresRunStore.ts @@ -1,4 +1,4 @@ -import { Prisma, boundedIn } from "@trigger.dev/database"; +import { Prisma, boundedIn, withTransactionStartRetry } from "@trigger.dev/database"; import type { BatchTaskRun, BatchTaskRunItemStatus, @@ -6,6 +6,7 @@ import type { PrismaClientOrTransaction, TaskRun, TaskRunStatus, + TransactionStartRetryConfig, WaitpointTag, } from "@trigger.dev/database"; import type { @@ -109,6 +110,13 @@ export type PostgresRunStoreOptions = { readOnlyPrisma: RunOpsCapableClient; /** Defaults to `"legacy"` so existing callers/tests are unaffected. */ schemaVariant?: RunStoreSchemaVariant; + /** + * `maxWait` (ms) applied when THIS store opens its own transaction. Threaded from the app + * boundary (IoC). Undefined leaves Prisma's own default (2000ms), preserving test behavior. + */ + maxWait?: number; + /** Env-driven P2028-at-acquisition retry config, threaded from the app boundary (IoC). */ + transactionStartRetry?: TransactionStartRetryConfig; }; // A caller sub-select for a relation: `{ select?, include? }` or `true` for a bare `key: true`. @@ -629,6 +637,8 @@ export class PostgresRunStore implements RunStore { private readonly prisma: RunOpsCapableClient; private readonly readOnlyPrisma: RunOpsCapableClient; private readonly schemaVariant: RunStoreSchemaVariant; + private readonly maxWait?: number; + private readonly transactionStartRetry?: TransactionStartRetryConfig; constructor(options: PostgresRunStoreOptions) { // Normalize foreign (run-ops-generation) Prisma known-request-errors to the control-plane @@ -637,6 +647,8 @@ export class PostgresRunStore implements RunStore { this.prisma = wrapRunOpsClientForErrorNormalization(options.prisma); this.readOnlyPrisma = wrapRunOpsClientForErrorNormalization(options.readOnlyPrisma); this.schemaVariant = options.schemaVariant ?? "legacy"; + this.maxWait = options.maxWait; + this.transactionStartRetry = options.transactionStartRetry; } // The writer handle in read-client form, so the routing layer can honor a caller-passed client @@ -654,8 +666,18 @@ export class PostgresRunStore implements RunStore { _runId: string | undefined, fn: (store: RunStore, tx: PrismaClientOrTransaction) => Promise ): Promise { - return (this.prisma as RunOpsTransactionalClient).$transaction((tx) => - fn(this, tx as unknown as PrismaClientOrTransaction) + let entered = false; + return withTransactionStartRetry( + () => + (this.prisma as RunOpsTransactionalClient).$transaction( + (tx) => { + entered = true; + return fn(this, tx as unknown as PrismaClientOrTransaction); + }, + { maxWait: this.maxWait } + ), + this.transactionStartRetry, + () => !entered ); } @@ -674,9 +696,16 @@ export class PostgresRunStore implements RunStore { if (alreadyInTransaction) { return fn(tx); } - return (this.prisma as RunOpsTransactionalClient).$transaction( - (t) => fn(t as unknown as PrismaClientOrTransaction), - options + const txOptions = { maxWait: this.maxWait, ...options }; + let entered = false; + return withTransactionStartRetry( + () => + (this.prisma as RunOpsTransactionalClient).$transaction((t) => { + entered = true; + return fn(t as unknown as PrismaClientOrTransaction); + }, txOptions), + this.transactionStartRetry, + () => !entered ); }