diff --git a/.server-changes/batch-item-streaming-rate-limit.md b/.server-changes/batch-item-streaming-rate-limit.md new file mode 100644 index 00000000000..f817195e9b1 --- /dev/null +++ b/.server-changes/batch-item-streaming-rate-limit.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: fix +--- + +Batch triggers no longer fail to start their runs when an environment is under heavy API load. If a batch still can't finish being created, `batchTriggerAndWait` now fails with an error instead of leaving the parent run waiting forever, and the batches page says so rather than reporting that it resumed. diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 2d28827f7d7..0bbc5d3870e 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -879,6 +879,20 @@ const EnvironmentSchema = z BATCH_RATE_LIMIT_MAX: z.coerce.number().int().default(1200), BATCH_RATE_LIMIT_REFILL_INTERVAL: z.string().default("10s"), BATCH_CONCURRENCY_LIMIT_DEFAULT: z.coerce.number().int().default(5), + /** + * How long a created batch may remain unsealed before the seal-timeout reaper + * aborts it and resumes any blocked parent with an error. Must exceed the SDK's + * worst-case stream-retry budget (maxAttempts x server request timeout). + * Doubles as the TTL of the phase 2 streaming grant, so the grant and the reaper + * always agree on how long a batch is allowed to be sealing. + */ + BATCH_SEAL_TIMEOUT_MS: z.coerce.number().int().positive().default(1_800_000), + /** + * Number of phase 2 (`POST /api/v3/batches/:id/items`) requests a created batch is + * granted, exempt from the general API rate limit. Sized above the SDK's stream + * maxAttempts so a batch admitted by the batch limiter can always finish streaming. + */ + BATCH_STREAM_GRANT_ATTEMPTS: z.coerce.number().int().positive().default(10), REALTIME_STREAM_VERSION: z.enum(["v1", "v2"]).default("v1"), REALTIME_STREAM_MAX_LENGTH: z.coerce.number().int().default(1000), diff --git a/apps/webapp/app/routes/resources.batches.$batchId.check-completion.ts b/apps/webapp/app/routes/resources.batches.$batchId.check-completion.ts index fffa0d09989..2d1c6bc702b 100644 --- a/apps/webapp/app/routes/resources.batches.$batchId.check-completion.ts +++ b/apps/webapp/app/routes/resources.batches.$batchId.check-completion.ts @@ -42,6 +42,20 @@ export const action: ActionFunction = async ({ request, params }) => { return redirectWithErrorMessage(safeRedirectUrl, request, "Batch not found"); } + const batch = await runStore.findBatchTaskRunById(ownedBatchRunId); + + if (!batch) { + return redirectWithErrorMessage(safeRedirectUrl, request, "Batch not found"); + } + + if (!batch.sealed) { + return redirectWithErrorMessage( + safeRedirectUrl, + request, + "This batch was never finished being created, so it can't be resumed. Please get in touch and we'll recover it for you." + ); + } + try { // v3 (engine V1) is retired; finalize the batch through the v2 completion path (no-op if not ready). await tryCompleteBatchV3(ownedBatchRunId, prisma, true); diff --git a/apps/webapp/app/runEngine/concerns/batchStreamGrants.server.ts b/apps/webapp/app/runEngine/concerns/batchStreamGrants.server.ts new file mode 100644 index 00000000000..17fc7324744 --- /dev/null +++ b/apps/webapp/app/runEngine/concerns/batchStreamGrants.server.ts @@ -0,0 +1,105 @@ +import { createRedisClient, type RedisClient, type RedisWithClusterOptions } from "~/redis.server"; +import { logger } from "~/services/logger.server"; + +export type BatchStreamGrantsOptions = { + redis: RedisWithClusterOptions; + /** How many phase 2 requests a created batch is allowed. */ + attempts: number; + /** How long the grant survives, matching how long a batch may legitimately be sealing. */ + ttlMs: number; +}; + +const KEY_PREFIX = "batch-stream-grant:"; + +/** + * Admission for phase 2 of the 2-phase batch API. + * + * Phase 1 (`POST /api/v3/batches`) already passes its own batch rate limiter, which fixes + * the batch's `expectedCount` and blocks the parent run on the batch's waitpoint. Phase 2 + * (`POST /api/v3/batches/:id/items`) is the only thing that can seal that batch, so having + * the general API limiter reject it strands the batch and the parent with it. + * + * Phase 1 therefore mints a bounded grant, and phase 2 spends it to bypass the general + * limiter. Admission stays a single decision made in phase 1, but the bypass is capped at + * `attempts` requests per batch rather than being unconditional. + */ +export class BatchStreamGrants { + private readonly redis: RedisClient; + + constructor(private readonly options: BatchStreamGrantsOptions) { + this.redis = createRedisClient("batchStreamGrants", options.redis); + this.#registerCommands(); + } + + /** + * Grant a newly created batch its phase 2 budget. Never throws: a batch that fails to get + * a grant still works, it just falls back to the general rate limiter for streaming. + */ + async mint(environmentId: string, batchId: string): Promise { + try { + await this.redis.set( + this.#key(environmentId, batchId), + this.options.attempts, + "PX", + this.options.ttlMs + ); + } catch (error) { + logger.warn("BatchStreamGrants: failed to mint grant", { + batchId, + error: error instanceof Error ? error.message : String(error), + }); + } + } + + /** + * Consume one phase 2 request from the batch's grant. + * + * Returns false when there is no grant, when the budget is spent, or when Redis is + * unreachable, so the caller falls back to the general rate limiter rather than opening + * an unbounded bypass. + */ + async spend(environmentId: string, batchId: string): Promise { + try { + // @ts-expect-error - Custom command defined via defineCommand + const remaining = (await this.redis.spendBatchStreamGrant( + this.#key(environmentId, batchId) + )) as number; + + return remaining >= 0; + } catch (error) { + logger.warn("BatchStreamGrants: failed to spend grant", { + batchId, + error: error instanceof Error ? error.message : String(error), + }); + + return false; + } + } + + async quit(): Promise { + await this.redis.quit(); + } + + /** + * Scoped to the environment as well as the batch, so a caller authenticated against a + * different environment can never spend this batch's grant even if they know its id. + */ + #key(environmentId: string, batchId: string): string { + return `${KEY_PREFIX}${environmentId}:${batchId}`; + } + + #registerCommands(): void { + this.redis.defineCommand("spendBatchStreamGrant", { + numberOfKeys: 1, + lua: ` +local remaining = tonumber(redis.call('GET', KEYS[1])) + +if not remaining or remaining <= 0 then + return -1 +end + +return redis.call('DECR', KEYS[1]) + `, + }); + } +} diff --git a/apps/webapp/app/runEngine/concerns/batchStreamGrantsInstance.server.ts b/apps/webapp/app/runEngine/concerns/batchStreamGrantsInstance.server.ts new file mode 100644 index 00000000000..1d79dc46674 --- /dev/null +++ b/apps/webapp/app/runEngine/concerns/batchStreamGrantsInstance.server.ts @@ -0,0 +1,20 @@ +import { env } from "~/env.server"; +import { singleton } from "~/utils/singleton"; +import { BatchStreamGrants } from "./batchStreamGrants.server"; + +export const batchStreamGrants = singleton( + "batchStreamGrants", + () => + new BatchStreamGrants({ + redis: { + port: env.RATE_LIMIT_REDIS_PORT, + host: env.RATE_LIMIT_REDIS_HOST, + username: env.RATE_LIMIT_REDIS_USERNAME, + password: env.RATE_LIMIT_REDIS_PASSWORD, + tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true", + clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1", + }, + attempts: env.BATCH_STREAM_GRANT_ATTEMPTS, + ttlMs: env.BATCH_SEAL_TIMEOUT_MS, + }) +); diff --git a/apps/webapp/app/runEngine/services/createBatch.server.ts b/apps/webapp/app/runEngine/services/createBatch.server.ts index 0095c48f2b5..0289e68e2c7 100644 --- a/apps/webapp/app/runEngine/services/createBatch.server.ts +++ b/apps/webapp/app/runEngine/services/createBatch.server.ts @@ -4,6 +4,8 @@ import { RunId } from "@trigger.dev/core/v3/isomorphic"; import { type BatchTaskRun, Prisma } from "@trigger.dev/database"; import { Evt } from "evt"; import { prisma, type PrismaClientOrTransaction } from "~/db.server"; +import { env } from "~/env.server"; +import { batchStreamGrants } from "../concerns/batchStreamGrantsInstance.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; @@ -120,8 +122,15 @@ export class CreateBatchService extends WithRunEngine { this.onBatchTaskRunCreated.post(batch); + await batchStreamGrants.mint(environment.id, friendlyId); + // Block parent run if this is a batchTriggerAndWait if (body.parentRunId && body.resumeParentOnCompletion) { + await this._engine.scheduleExpireBatch({ + batchId: batch.id, + availableAt: new Date(Date.now() + env.BATCH_SEAL_TIMEOUT_MS), + }); + await this._engine.blockRunWithCreatedBatch({ runId: RunId.fromFriendlyId(body.parentRunId), batchId: batch.id, diff --git a/apps/webapp/app/services/apiRateLimit.server.ts b/apps/webapp/app/services/apiRateLimit.server.ts index 837d3ee5895..1b8d8a3ed1f 100644 --- a/apps/webapp/app/services/apiRateLimit.server.ts +++ b/apps/webapp/app/services/apiRateLimit.server.ts @@ -1,8 +1,12 @@ +import { tryCatch } from "@trigger.dev/core/v3"; import { env } from "~/env.server"; +import { batchStreamGrants } from "~/runEngine/concerns/batchStreamGrantsInstance.server"; import { authenticateAuthorizationHeader } from "./apiAuth.server"; import { authorizationRateLimitMiddleware } from "./authorizationRateLimitMiddleware.server"; import type { Duration } from "./rateLimiter.server"; +const BATCH_STREAM_ITEMS_PATH = /^\/api\/v3\/batches\/([^/]+)\/items$/; + export const apiRateLimiter = authorizationRateLimitMiddleware({ redis: { port: env.RATE_LIMIT_REDIS_PORT, @@ -75,6 +79,32 @@ export const apiRateLimiter = authorizationRateLimitMiddleware({ /^\/api\/v2\/packets\//, /^\/api\/v1\/sessions\/[^/]+\/snapshot-url$/, ], + bypass: async (req) => { + const match = BATCH_STREAM_ITEMS_PATH.exec(req.path); + + if (!match) { + return false; + } + + const batchFriendlyId = match[1]; + const authorizationValue = req.headers.authorization; + + if (!batchFriendlyId || !authorizationValue) { + return false; + } + + const [authError, authenticated] = await tryCatch( + authenticateAuthorizationHeader(authorizationValue, { + allowPublicKey: true, + }) + ); + + if (authError || !authenticated || !authenticated.ok) { + return false; + } + + return batchStreamGrants.spend(authenticated.environment.id, batchFriendlyId); + }, log: { rejections: env.API_RATE_LIMIT_REJECTION_LOGS_ENABLED === "1", requests: env.API_RATE_LIMIT_REQUEST_LOGS_ENABLED === "1", diff --git a/apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts b/apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts index e6d4d3b1f5e..5fcd6eb5450 100644 --- a/apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts +++ b/apps/webapp/app/services/authorizationRateLimitMiddleware.server.ts @@ -5,7 +5,6 @@ import { Ratelimit } from "@upstash/ratelimit"; import type { Request as ExpressRequest, Response as ExpressResponse, NextFunction } from "express"; import { createHash } from "node:crypto"; import { z } from "zod"; -import { env } from "~/env.server"; import type { RedisWithClusterOptions } from "~/redis.server"; import { logger } from "./logger.server"; import type { Duration, Limiter } from "./rateLimiter.server"; @@ -56,10 +55,17 @@ export type RateLimiterConfig = z.infer; type LimitConfigOverrideFunction = (authorizationValue: string) => Promise; type Options = { - redis?: RedisWithClusterOptions; + redis: RedisWithClusterOptions; keyPrefix: string; pathMatchers: (RegExp | string)[]; pathWhiteList?: (RegExp | string)[]; + /** + * Escape hatch for requests that can only be admitted by consulting state, rather than by + * matching a path. Runs after the authorization header check, so an unauthenticated + * request is still rejected, and only skips the rate limit itself. Must not throw: a + * bypass that cannot decide should return false and let the limiter apply. + */ + bypass?: (req: ExpressRequest) => Promise; defaultLimiter: RateLimiterConfig; limiterConfigOverride?: LimitConfigOverrideFunction; limiterCache?: { @@ -151,6 +157,7 @@ export function authorizationRateLimitMiddleware({ defaultLimiter, pathMatchers, pathWhiteList = [], + bypass, log = { rejections: true, requests: true, @@ -176,16 +183,7 @@ export function authorizationRateLimitMiddleware({ }), }); - const redisClient = createRedisRateLimitClient( - redis ?? { - port: env.RATE_LIMIT_REDIS_PORT, - host: env.RATE_LIMIT_REDIS_HOST, - username: env.RATE_LIMIT_REDIS_USERNAME, - password: env.RATE_LIMIT_REDIS_PASSWORD, - tlsDisabled: env.RATE_LIMIT_REDIS_TLS_DISABLED === "true", - clusterMode: env.RATE_LIMIT_REDIS_CLUSTER_MODE_ENABLED === "1", - } - ); + const redisClient = createRedisRateLimitClient(redis); return async (req: ExpressRequest, res: ExpressResponse, next: NextFunction) => { if (log.requests) { @@ -247,6 +245,26 @@ export function authorizationRateLimitMiddleware({ ); } + if (bypass) { + let bypassed = false; + + try { + bypassed = await bypass(req); + } catch (error) { + logger.warn(`RateLimiter (${keyPrefix}): bypass threw, applying the limit`, { + path: req.path, + error: error instanceof Error ? error.message : String(error), + }); + } + + if (bypassed) { + if (log.requests) { + logger.info(`RateLimiter (${keyPrefix}): bypassed ${req.path}`); + } + return next(); + } + } + const hash = createHash("sha256"); hash.update(authorizationValue); const hashedAuthorizationValue = hash.digest("hex"); diff --git a/apps/webapp/app/v3/runOpsMigration/unblockRouteCatalog.ts b/apps/webapp/app/v3/runOpsMigration/unblockRouteCatalog.ts index 3296569ce0b..7d7f9d6f536 100644 --- a/apps/webapp/app/v3/runOpsMigration/unblockRouteCatalog.ts +++ b/apps/webapp/app/v3/runOpsMigration/unblockRouteCatalog.ts @@ -65,6 +65,12 @@ export const UNBLOCK_ROUTES: readonly UnblockRoute[] = [ site: BATCH_SYSTEM, symbol: "#tryCompleteBatch", }, + { + id: "batch.expireBatch", + kind: "RUN", + site: BATCH_SYSTEM, + symbol: "expireBatch", + }, { id: "ttl.expireRun", kind: "RUN", diff --git a/apps/webapp/test/authorizationRateLimitMiddleware.test.ts b/apps/webapp/test/authorizationRateLimitMiddleware.test.ts index 847327b4b7a..b6076cef0de 100644 --- a/apps/webapp/test/authorizationRateLimitMiddleware.test.ts +++ b/apps/webapp/test/authorizationRateLimitMiddleware.test.ts @@ -25,7 +25,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", redisTest("should allow requests within the rate limit", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test", defaultLimiter: { type: "tokenBucket", @@ -56,7 +56,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", redisTest("should reject requests without an Authorization header", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test", defaultLimiter: { type: "tokenBucket", @@ -80,7 +80,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", redisTest("should reject requests that exceed the rate limit", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test", defaultLimiter: { type: "tokenBucket", @@ -108,7 +108,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", redisTest("should not apply rate limiting to whitelisted paths", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test", defaultLimiter: { type: "tokenBucket", @@ -138,7 +138,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", "should apply different rate limits based on limiterConfigOverride", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test", defaultLimiter: { type: "tokenBucket", @@ -188,7 +188,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", // 1. Test different rate limit configurations redisTest("should enforce fixed window rate limiting", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test-fixed", defaultLimiter: { type: "fixedWindow", @@ -224,7 +224,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", redisTest("should enforce sliding window rate limiting", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test-sliding", defaultLimiter: { type: "slidingWindow", @@ -268,7 +268,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", // 2. Test edge cases around rate limit calculations redisTest("should handle token refill correctly", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test-refill", defaultLimiter: { type: "tokenBucket", @@ -309,7 +309,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", redisTest("should handle near-zero remaining tokens correctly", async ({ redisOptions }) => { const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test-near-zero", defaultLimiter: { type: "tokenBucket", @@ -357,7 +357,7 @@ describe.skipIf(process.env.GITHUB_ACTIONS)("authorizationRateLimitMiddleware", redisTest("should use cached limiter configurations", async ({ redisOptions }) => { let configOverrideCalls = 0; const rateLimitMiddleware = authorizationRateLimitMiddleware({ - redis: redisOptions, + redis: { ...redisOptions, tlsDisabled: true }, keyPrefix: "test-cache", defaultLimiter: { type: "tokenBucket", diff --git a/apps/webapp/test/authorizationRateLimitMiddlewareBypass.test.ts b/apps/webapp/test/authorizationRateLimitMiddlewareBypass.test.ts new file mode 100644 index 00000000000..b9d6c175d27 --- /dev/null +++ b/apps/webapp/test/authorizationRateLimitMiddlewareBypass.test.ts @@ -0,0 +1,92 @@ +import { redisTest } from "@internal/testcontainers"; +import { beforeEach, describe, expect, vi } from "vitest"; + +vi.setConfig({ testTimeout: 30_000 }); + +import type { Express } from "express"; +import express from "express"; +import request from "supertest"; +import { authorizationRateLimitMiddleware } from "../app/services/authorizationRateLimitMiddleware.server.js"; + +const exhaustedLimiter = { + type: "tokenBucket", + refillRate: 1, + interval: "1m", + maxTokens: 1, +} as const; + +describe("authorizationRateLimitMiddleware bypass", () => { + let app: Express; + + beforeEach(() => { + app = express(); + }); + + redisTest("lets a bypassed request through an exhausted limit", async ({ redisOptions }) => { + const rateLimitMiddleware = authorizationRateLimitMiddleware({ + redis: { ...redisOptions, tlsDisabled: true }, + keyPrefix: "test-bypass-allowed", + defaultLimiter: exhaustedLimiter, + pathMatchers: [/^\/api/], + bypass: async (req) => req.path === "/api/granted", + }); + + app.use(rateLimitMiddleware); + app.get("/api/granted", (req, res) => res.status(200).json({ message: "Granted" })); + app.get("/api/limited", (req, res) => res.status(200).json({ message: "Limited" })); + + await request(app).get("/api/limited").set("Authorization", "Bearer test-token"); + const limited = await request(app) + .get("/api/limited") + .set("Authorization", "Bearer test-token"); + expect(limited.status).toBe(429); + + const granted = await request(app) + .get("/api/granted") + .set("Authorization", "Bearer test-token"); + + expect(granted.status).toBe(200); + expect(granted.body).toEqual({ message: "Granted" }); + }); + + redisTest("falls back to the limiter when the bypass declines", async ({ redisOptions }) => { + const rateLimitMiddleware = authorizationRateLimitMiddleware({ + redis: { ...redisOptions, tlsDisabled: true }, + keyPrefix: "test-bypass-declined", + defaultLimiter: exhaustedLimiter, + pathMatchers: [/^\/api/], + bypass: async () => false, + }); + + app.use(rateLimitMiddleware); + app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" })); + + await request(app).get("/api/test").set("Authorization", "Bearer declined"); + const response = await request(app).get("/api/test").set("Authorization", "Bearer declined"); + + expect(response.status).toBe(429); + }); + + redisTest("does not let the bypass skip authentication", async ({ redisOptions }) => { + let bypassCalled = false; + + const rateLimitMiddleware = authorizationRateLimitMiddleware({ + redis: { ...redisOptions, tlsDisabled: true }, + keyPrefix: "test-bypass-unauthenticated", + defaultLimiter: exhaustedLimiter, + pathMatchers: [/^\/api/], + bypass: async () => { + bypassCalled = true; + return true; + }, + }); + + app.use(rateLimitMiddleware); + app.get("/api/test", (req, res) => res.status(200).json({ message: "Success" })); + + const response = await request(app).get("/api/test"); + + expect(response.status).toBe(401); + expect(bypassCalled).toBe(false); + }); +}); diff --git a/apps/webapp/test/batchStreamGrants.test.ts b/apps/webapp/test/batchStreamGrants.test.ts new file mode 100644 index 00000000000..82fbd50aa88 --- /dev/null +++ b/apps/webapp/test/batchStreamGrants.test.ts @@ -0,0 +1,102 @@ +import { redisTest } from "@internal/testcontainers"; +import { describe, expect, vi } from "vitest"; + +vi.setConfig({ testTimeout: 30_000 }); + +import { BatchStreamGrants } from "../app/runEngine/concerns/batchStreamGrants.server.js"; + +describe("BatchStreamGrants", () => { + const ENV = "env_1"; + + redisTest("spends exactly the granted number of attempts", async ({ redisOptions }) => { + const grants = new BatchStreamGrants({ + redis: { ...redisOptions, tlsDisabled: true }, + attempts: 3, + ttlMs: 60_000, + }); + + try { + await grants.mint(ENV, "batch_spend"); + + expect(await grants.spend(ENV, "batch_spend")).toBe(true); + expect(await grants.spend(ENV, "batch_spend")).toBe(true); + expect(await grants.spend(ENV, "batch_spend")).toBe(true); + expect(await grants.spend(ENV, "batch_spend")).toBe(false); + expect(await grants.spend(ENV, "batch_spend")).toBe(false); + } finally { + await grants.quit(); + } + }); + + redisTest("declines a batch that was never granted", async ({ redisOptions }) => { + const grants = new BatchStreamGrants({ + redis: { ...redisOptions, tlsDisabled: true }, + attempts: 5, + ttlMs: 60_000, + }); + + try { + expect(await grants.spend(ENV, "batch_never_minted")).toBe(false); + } finally { + await grants.quit(); + } + }); + + redisTest("keeps grants isolated per batch", async ({ redisOptions }) => { + const grants = new BatchStreamGrants({ + redis: { ...redisOptions, tlsDisabled: true }, + attempts: 1, + ttlMs: 60_000, + }); + + try { + await grants.mint(ENV, "batch_a"); + await grants.mint(ENV, "batch_b"); + + expect(await grants.spend(ENV, "batch_a")).toBe(true); + expect(await grants.spend(ENV, "batch_a")).toBe(false); + expect(await grants.spend(ENV, "batch_b")).toBe(true); + } finally { + await grants.quit(); + } + }); + + redisTest("expires the grant so it cannot outlive the seal window", async ({ redisOptions }) => { + const grants = new BatchStreamGrants({ + redis: { ...redisOptions, tlsDisabled: true }, + attempts: 100_000, + ttlMs: 150, + }); + + try { + await grants.mint(ENV, "batch_expiring"); + expect(await grants.spend(ENV, "batch_expiring")).toBe(true); + + await vi.waitFor( + async () => { + expect(await grants.spend(ENV, "batch_expiring")).toBe(false); + }, + { timeout: 5_000, interval: 50 } + ); + } finally { + await grants.quit(); + } + }); + + redisTest("another environment cannot spend this batch's grant", async ({ redisOptions }) => { + const grants = new BatchStreamGrants({ + redis: { ...redisOptions, tlsDisabled: true }, + attempts: 1, + ttlMs: 60_000, + }); + + try { + await grants.mint(ENV, "batch_scoped"); + + expect(await grants.spend("env_intruder", "batch_scoped")).toBe(false); + expect(await grants.spend(ENV, "batch_scoped")).toBe(true); + } finally { + await grants.quit(); + } + }); +}); diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 3c1f4330a0f..808c6c68fef 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -285,6 +285,9 @@ export class RunEngine { tryCompleteBatch: async ({ payload }) => { await this.batchSystem.performCompleteBatch({ batchId: payload.batchId }); }, + expireBatch: async ({ payload }) => { + await this.batchSystem.expireBatch({ batchId: payload.batchId }); + }, continueRunIfUnblocked: async ({ payload }) => { await this.waitpointSystem.continueRunIfUnblocked({ runId: payload.runId, @@ -1777,6 +1780,28 @@ export class RunEngine { return this.batchSystem.scheduleCompleteBatch({ batchId }); } + /** + * Terminally fail a batch whose phase 2 item stream never sealed it, completing the + * parent's batchTriggerAndWait waitpoint with an error so the parent resumes. + */ + async expireBatch({ batchId }: { batchId: string }): Promise { + return this.batchSystem.expireBatch({ batchId }); + } + + /** + * Schedule the seal-timeout reaper. Only worth scheduling for a batch that blocks a + * parent, since a fire-and-forget batch has nothing to strand. + */ + async scheduleExpireBatch({ + batchId, + availableAt, + }: { + batchId: string; + availableAt: Date; + }): Promise { + return this.batchSystem.scheduleExpireBatch({ batchId, availableAt }); + } + // ============================================================================ // BatchQueue methods (DRR-based batch processing) // ============================================================================ diff --git a/internal-packages/run-engine/src/engine/systems/batchSystem.test.ts b/internal-packages/run-engine/src/engine/systems/batchSystem.test.ts index 39a0f965e13..defd36f6147 100644 --- a/internal-packages/run-engine/src/engine/systems/batchSystem.test.ts +++ b/internal-packages/run-engine/src/engine/systems/batchSystem.test.ts @@ -88,6 +88,14 @@ class CountingPostgresRunStore extends PostgresRunStore { return super.updateBatchTaskRun(args, tx); } + override async updateManyBatchTaskRun( + args: Prisma.BatchTaskRunUpdateManyArgs, + tx?: any + ): Promise { + this.batchUpdates++; + return super.updateManyBatchTaskRun(args, tx); + } + override async findWaitpoint( args: any, client?: any diff --git a/internal-packages/run-engine/src/engine/systems/batchSystem.ts b/internal-packages/run-engine/src/engine/systems/batchSystem.ts index 7ca78703dd8..6d86a4519bf 100644 --- a/internal-packages/run-engine/src/engine/systems/batchSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/batchSystem.ts @@ -1,4 +1,5 @@ -import { startSpan } from "@internal/tracing"; +import { startSpan, type Counter } from "@internal/tracing"; +import type { TaskRunError } from "@trigger.dev/core/v3/schemas"; import { isFinalRunStatus } from "../statuses.js"; import type { SystemResources } from "./systems.js"; import type { WaitpointSystem } from "./waitpointSystem.js"; @@ -11,10 +12,15 @@ export type BatchSystemOptions = { export class BatchSystem { private readonly $: SystemResources; private readonly waitpointSystem: WaitpointSystem; + private readonly expirationsCounter: Counter; constructor(private readonly options: BatchSystemOptions) { this.$ = options.resources; this.waitpointSystem = options.waitpointSystem; + this.expirationsCounter = this.$.meter.createCounter("batch_system.expirations", { + description: "Seal-timeout reaper runs, by outcome", + unit: "batches", + }); } public async scheduleCompleteBatch({ batchId }: { batchId: string }): Promise { @@ -31,6 +37,127 @@ export class BatchSystem { await this.#tryCompleteBatch({ batchId }); } + public async scheduleExpireBatch({ + batchId, + availableAt, + }: { + batchId: string; + availableAt: Date; + }): Promise { + await this.$.worker.enqueue({ + id: `expireBatch:${batchId}`, + job: "expireBatch", + payload: { batchId }, + availableAt, + }); + } + + /** + * Terminally fail a batch whose phase 2 item stream never sealed it, completing the + * parent's batchTriggerAndWait waitpoint with an error so the parent resumes with a + * failure instead of hanging forever. + * + * A batch is only sealed once every item has streamed, but the parent is blocked on the + * batch's waitpoint from the moment the batch is created. Nothing else completes that + * waitpoint, so without this a stream that never finishes strands the parent permanently. + * + * Idempotent and race-safe: a batch that sealed in the meantime is left alone. + * + * Resumable too. The abort and the waitpoint completion are two writes, so a crash between + * them would otherwise leave the parent blocked with no way back: the retry would see a + * non-PENDING batch and bail. An already-ABORTED batch therefore falls through to complete + * its waitpoint again, which {@link WaitpointSystem.completeWaitpoint} treats as a no-op. + */ + public async expireBatch({ batchId }: { batchId: string }): Promise { + return startSpan(this.$.tracer, "expireBatch", async (span) => { + span.setAttribute("batchId", batchId); + + const batch = await this.$.runStore.findBatchTaskRunById(batchId, undefined, this.$.prisma); + + if (!batch) { + this.$.logger.debug("expireBatch: batch doesn't exist", { batchId }); + this.expirationsCounter.add(1, { outcome: "missing" }); + return; + } + + if (batch.sealed || (batch.status !== "PENDING" && batch.status !== "ABORTED")) { + this.$.logger.debug("expireBatch: batch sealed or already finished, nothing to do", { + batchId, + status: batch.status, + sealed: batch.sealed, + }); + this.expirationsCounter.add(1, { outcome: "already_settled" }); + return; + } + + if (batch.status === "PENDING") { + const aborted = await this.$.runStore.updateManyBatchTaskRun( + { + where: { id: batchId, sealed: false, status: "PENDING" }, + data: { + status: "ABORTED", + completedAt: new Date(), + processingCompletedAt: new Date(), + }, + }, + this.$.prisma + ); + + if (aborted.count === 0) { + const current = await this.$.runStore.findBatchTaskRunById( + batchId, + undefined, + this.$.prisma + ); + + if (!current || current.sealed || current.status !== "ABORTED") { + this.$.logger.debug("expireBatch: lost the race to a seal, no-op", { batchId }); + this.expirationsCounter.add(1, { outcome: "lost_race" }); + return; + } + } + } + + const waitpoint = await this.$.runStore.findWaitpoint( + { + where: { completedByBatchId: batchId }, + }, + this.$.prisma + ); + + if (!waitpoint) { + this.$.logger.debug("expireBatch: no waitpoint, nothing was blocked on this batch", { + batchId, + }); + this.expirationsCounter.add(1, { outcome: "aborted_no_parent" }); + return; + } + + const enqueuedCount = (batch.successfulRunCount ?? 0) + (batch.failedRunCount ?? 0); + const error: TaskRunError = { + type: "STRING_ERROR", + raw: + `Batch ${batch.friendlyId} was never fully created: only ${enqueuedCount} of ` + + `${batch.expectedCount} items were received before it timed out, so it can never ` + + `complete. batchTriggerAndWait failed rather than waiting forever.`, + }; + + await this.waitpointSystem.completeWaitpoint({ + id: waitpoint.id, + output: { value: JSON.stringify(error), isError: true }, + }); + + this.expirationsCounter.add(1, { outcome: "aborted_parent_resumed" }); + + this.$.logger.warn("expireBatch: aborted an unsealed batch and resumed its parent", { + batchId, + waitpointId: waitpoint.id, + enqueuedCount, + expectedCount: batch.expectedCount, + }); + }); + } + /** * Checks to see if all runs for a BatchTaskRun are completed, if they are then update the status. * This isn't used operationally, but it's used for the Batches dashboard page. @@ -89,21 +216,22 @@ export class BatchSystem { if (runs.every((r) => isFinalRunStatus(r.status))) { this.$.logger.debug("#tryCompleteBatch: All runs are completed", { batchId }); - await this.$.runStore.updateBatchTaskRun( + + const completed = await this.$.runStore.updateManyBatchTaskRun( { - where: { - id: batchId, - }, - data: { - status: "COMPLETED", - }, - select: { - id: true, - }, + where: { id: batchId, status: { notIn: ["ABORTED", "COMPLETED"] } }, + data: { status: "COMPLETED" }, }, this.$.prisma ); + if (completed.count === 0) { + this.$.logger.debug("#tryCompleteBatch: batch already reached a terminal status", { + batchId, + }); + return; + } + //get waitpoint (if there is one) const waitpoint = await this.$.runStore.findWaitpoint( { diff --git a/internal-packages/run-engine/src/engine/tests/batchTwoPhase.test.ts b/internal-packages/run-engine/src/engine/tests/batchTwoPhase.test.ts index 8471c07844b..106a4c6b73b 100644 --- a/internal-packages/run-engine/src/engine/tests/batchTwoPhase.test.ts +++ b/internal-packages/run-engine/src/engine/tests/batchTwoPhase.test.ts @@ -625,4 +625,324 @@ describe("RunEngine 2-Phase Batch API", () => { } } ); + + describe("seal-timeout reaper", () => { + type TestPrisma = Parameters[0]; + type TestEnvironment = Awaited>; + type TestRedisOptions = ConstructorParameters[0]["runLock"]["redis"]; + + function createEngine(prisma: TestPrisma, redisOptions: TestRedisOptions) { + const engine = new RunEngine({ + prisma, + worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 20 }, + queue: { + redis: redisOptions, + masterQueueConsumersDisabled: true, + processWorkerQueueDebounceMs: 50, + }, + runLock: { redis: redisOptions }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 }, + }, + baseCostInCents: 0.0001, + }, + batchQueue: { + redis: redisOptions, + consumerCount: 2, + consumerIntervalMs: 50, + drr: { quantum: 10, maxDeficit: 100 }, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + + engine.setBatchProcessItemCallback(async () => ({ + success: true, + runId: "should-not-be-called", + })); + engine.setBatchCompletionCallback(async () => {}); + + return engine; + } + + /** + * Phase 1 without phase 2: a created, unsealed batch with its parent blocked on the + * batch waitpoint. This is the exact state a stream that never completes leaves behind. + */ + async function setupUnsealedBatchBlockingParent( + engine: RunEngine, + prisma: TestPrisma, + authenticatedEnvironment: TestEnvironment, + expectedCount = 2 + ) { + const parentTask = "parent-task"; + await setupBackgroundWorker(engine, authenticatedEnvironment, [parentTask]); + + const { id: batchId, friendlyId: batchFriendlyId } = BatchId.generate(); + await prisma.batchTaskRun.create({ + data: { + id: batchId, + friendlyId: batchFriendlyId, + runtimeEnvironmentId: authenticatedEnvironment.id, + status: "PENDING", + runCount: expectedCount, + expectedCount, + sealed: false, + batchVersion: "runengine:v2", + }, + }); + + const parentRun = await engine.trigger( + { + friendlyId: generateFriendlyId("run"), + environment: authenticatedEnvironment, + taskIdentifier: parentTask, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t_parent", + spanId: "s_parent", + workerQueue: "main", + queue: `task/${parentTask}`, + isTest: false, + tags: [], + }, + prisma + ); + + await vi.waitFor( + async () => { + const dequeued = await engine.dequeueFromWorkerQueue({ + consumerId: "test_12345", + workerQueue: "main", + }); + expect(dequeued.length).toBe(1); + }, + { timeout: 15_000, interval: 100 } + ); + + const initialExecutionData = await engine.getRunExecutionData({ runId: parentRun.id }); + assertNonNullable(initialExecutionData); + await engine.startRunAttempt({ + runId: parentRun.id, + snapshotId: initialExecutionData.snapshot.id, + }); + + await engine.blockRunWithCreatedBatch({ + runId: parentRun.id, + batchId, + environmentId: authenticatedEnvironment.id, + projectId: authenticatedEnvironment.projectId, + organizationId: authenticatedEnvironment.organizationId, + }); + + await engine.initializeBatch({ + batchId, + friendlyId: batchFriendlyId, + environmentId: authenticatedEnvironment.id, + environmentType: authenticatedEnvironment.type, + organizationId: authenticatedEnvironment.organizationId, + projectId: authenticatedEnvironment.projectId, + runCount: expectedCount, + parentRunId: parentRun.id, + resumeParentOnCompletion: true, + }); + + const afterBlocked = await engine.getRunExecutionData({ runId: parentRun.id }); + assertNonNullable(afterBlocked); + expect(afterBlocked.snapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + return { batchId, batchFriendlyId, parentRun }; + } + + containerTest( + "aborts an unsealed batch and resumes the parent with an error", + async ({ prisma, redisOptions }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = createEngine(prisma, redisOptions); + + try { + const { batchId, parentRun } = await setupUnsealedBatchBlockingParent( + engine, + prisma, + authenticatedEnvironment + ); + + await engine.expireBatch({ batchId }); + + const abortedBatch = await prisma.batchTaskRun.findFirst({ where: { id: batchId } }); + expect(abortedBatch?.status).toBe("ABORTED"); + expect(abortedBatch?.completedAt).not.toBeNull(); + + const waitpoint = await prisma.waitpoint.findFirst({ + where: { completedByBatchId: batchId }, + }); + assertNonNullable(waitpoint); + expect(waitpoint.status).toBe("COMPLETED"); + expect(waitpoint.outputIsError).toBe(true); + + await vi.waitFor( + async () => { + const waitpoints = await prisma.taskRunWaitpoint.findMany({ + where: { taskRunId: parentRun.id }, + }); + expect(waitpoints.length).toBe(0); + }, + { timeout: 15_000 } + ); + } finally { + await engine.quit(); + } + } + ); + + containerTest( + "leaves a batch alone when the stream sealed it first", + async ({ prisma, redisOptions }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = createEngine(prisma, redisOptions); + + try { + const { batchId } = await setupUnsealedBatchBlockingParent( + engine, + prisma, + authenticatedEnvironment + ); + + await prisma.batchTaskRun.update({ + where: { id: batchId }, + data: { sealed: true, sealedAt: new Date() }, + }); + + await engine.expireBatch({ batchId }); + + const batch = await prisma.batchTaskRun.findFirst({ where: { id: batchId } }); + expect(batch?.status).toBe("PENDING"); + + const waitpoint = await prisma.waitpoint.findFirst({ + where: { completedByBatchId: batchId }, + }); + assertNonNullable(waitpoint); + expect(waitpoint.status).toBe("PENDING"); + } finally { + await engine.quit(); + } + } + ); + + containerTest("is idempotent when it runs twice", async ({ prisma, redisOptions }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = createEngine(prisma, redisOptions); + + try { + const { batchId } = await setupUnsealedBatchBlockingParent( + engine, + prisma, + authenticatedEnvironment + ); + + await engine.expireBatch({ batchId }); + const afterFirst = await prisma.batchTaskRun.findFirst({ where: { id: batchId } }); + + await engine.expireBatch({ batchId }); + const afterSecond = await prisma.batchTaskRun.findFirst({ where: { id: batchId } }); + + expect(afterSecond?.status).toBe("ABORTED"); + expect(afterSecond?.completedAt?.getTime()).toBe(afterFirst?.completedAt?.getTime()); + + const waitpoints = await prisma.waitpoint.findMany({ + where: { completedByBatchId: batchId }, + }); + expect(waitpoints.length).toBe(1); + expect(waitpoints[0]?.status).toBe("COMPLETED"); + } finally { + await engine.quit(); + } + }); + + containerTest( + "resumes the parent when a previous run aborted the batch but died before the waitpoint", + async ({ prisma, redisOptions }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = createEngine(prisma, redisOptions); + + try { + const { batchId, parentRun } = await setupUnsealedBatchBlockingParent( + engine, + prisma, + authenticatedEnvironment + ); + + await prisma.batchTaskRun.update({ + where: { id: batchId }, + data: { + status: "ABORTED", + completedAt: new Date(), + processingCompletedAt: new Date(), + }, + }); + + const beforeRetry = await prisma.waitpoint.findFirst({ + where: { completedByBatchId: batchId }, + }); + assertNonNullable(beforeRetry); + expect(beforeRetry.status).toBe("PENDING"); + + await engine.expireBatch({ batchId }); + + const waitpoint = await prisma.waitpoint.findFirst({ + where: { completedByBatchId: batchId }, + }); + assertNonNullable(waitpoint); + expect(waitpoint.status).toBe("COMPLETED"); + expect(waitpoint.outputIsError).toBe(true); + + await vi.waitFor( + async () => { + const waitpoints = await prisma.taskRunWaitpoint.findMany({ + where: { taskRunId: parentRun.id }, + }); + expect(waitpoints.length).toBe(0); + }, + { timeout: 15_000 } + ); + } finally { + await engine.quit(); + } + } + ); + + containerTest( + "aborts a fire-and-forget batch without a waitpoint to resolve", + async ({ prisma, redisOptions }) => { + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const engine = createEngine(prisma, redisOptions); + + try { + const { id: batchId, friendlyId } = BatchId.generate(); + await prisma.batchTaskRun.create({ + data: { + id: batchId, + friendlyId, + runtimeEnvironmentId: authenticatedEnvironment.id, + status: "PENDING", + runCount: 1, + expectedCount: 1, + sealed: false, + batchVersion: "runengine:v2", + }, + }); + + await engine.expireBatch({ batchId }); + + const batch = await prisma.batchTaskRun.findFirst({ where: { id: batchId } }); + expect(batch?.status).toBe("ABORTED"); + } finally { + await engine.quit(); + } + } + ); + }); }); diff --git a/internal-packages/run-engine/src/engine/workerCatalog.ts b/internal-packages/run-engine/src/engine/workerCatalog.ts index de54c1ece00..8f49d7d6068 100644 --- a/internal-packages/run-engine/src/engine/workerCatalog.ts +++ b/internal-packages/run-engine/src/engine/workerCatalog.ts @@ -58,6 +58,12 @@ export const workerCatalog = { }), visibilityTimeoutMs: 30_000, }, + expireBatch: { + schema: z.object({ + batchId: z.string(), + }), + visibilityTimeoutMs: 30_000, + }, continueRunIfUnblocked: { schema: z.object({ runId: z.string(),