diff --git a/.changeset/queue-concurrency-overrides.md b/.changeset/queue-concurrency-overrides.md new file mode 100644 index 00000000000..65579b0d81a --- /dev/null +++ b/.changeset/queue-concurrency-overrides.md @@ -0,0 +1,18 @@ +--- +"@trigger.dev/sdk": patch +"@trigger.dev/core": patch +--- + +Adjust queue concurrency at runtime, per key and in total. `queues.overrideConcurrencyLimit` accepts a `concurrencyKey` to raise or lower one key's limit without touching the rest of the queue, and the new `queues.overrideTotalConcurrencyLimit` and `queues.resetTotalConcurrencyLimit` adjust the cap across all keys. + +```ts +import { queues } from "@trigger.dev/sdk"; + +await queues.overrideConcurrencyLimit("my-queue", 20, { concurrencyKey: "tenant-123" }); +await queues.resetConcurrencyLimit("my-queue", { concurrencyKey: "tenant-123" }); + +await queues.overrideTotalConcurrencyLimit("my-queue", 100); +await queues.resetTotalConcurrencyLimit("my-queue"); +``` + +Overrides survive deploys and reset back to the declared configuration. Enforcement happens server-side on servers with total concurrency limits enabled. diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.key.override.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.key.override.ts new file mode 100644 index 00000000000..9a72155d4fc --- /dev/null +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.key.override.ts @@ -0,0 +1,103 @@ +import { json } from "@remix-run/server-runtime"; +import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { z } from "zod"; +import { toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencySystem } from "~/v3/services/concurrencySystemInstance.server"; + +const BodySchema = z.object({ + type: RetrieveQueueType.default("id"), + concurrencyKey: z.string().min(1).max(128), + concurrencyLimit: z.number().int().min(0).max(100000), +}); + +const route = createActionApiRoute( + { + body: BodySchema, + params: z.object({ + queueParam: z.string().transform((val) => val.replace(/%2F/g, "/")), + }), + authorization: { + action: "write", + resource: () => ({ type: "queues" }), + }, + }, + async ({ params, body, authentication }) => { + const input: RetrieveQueueParam = + body.type === "id" + ? params.queueParam + : { + type: body.type, + name: decodeURIComponent(params.queueParam).replace(/%2F/g, "/"), + }; + + return concurrencySystem.queues + .overrideConcurrencyKeyLimit( + authentication.environment, + input, + body.concurrencyKey, + body.concurrencyLimit + ) + .match( + (queue) => { + return json( + toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + running: queue.running, + queued: queue.queued, + concurrencyLimit: queue.concurrencyLimit, + concurrencyLimitBase: queue.concurrencyLimitBase, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, + concurrencyLimitOverriddenBy: null, + paused: queue.paused, + }), + { status: 200 } + ); + }, + (error) => { + switch (error.type) { + case "queue_not_found": { + return json({ error: "Queue not found" }, { status: 404 }); + } + case "invalid_override": + case "concurrency_limit_exceeds_maximum": + case "too_many_key_overrides": { + return json({ error: error.message }, { status: 400 }); + } + case "queue_update_failed": { + return json( + { error: "Failed to update queue concurrency key limit" }, + { status: 500 } + ); + } + case "sync_queue_concurrency_to_engine_failed": { + return json({ error: "Failed to sync the concurrency key limit" }, { status: 500 }); + } + case "get_queue_stats_failed": { + return json({ error: "Failed to read queue stats" }, { status: 500 }); + } + case "other": { + return json( + { error: "Failed to update queue concurrency key limit" }, + { + status: 500, + } + ); + } + default: { + return json( + { error: "Failed to update queue concurrency key limit" }, + { + status: 500, + } + ); + } + } + } + ); + } +); + +export const { action } = route; diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.key.reset.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.key.reset.ts new file mode 100644 index 00000000000..3024ac84b72 --- /dev/null +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.key.reset.ts @@ -0,0 +1,95 @@ +import { json } from "@remix-run/server-runtime"; +import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { z } from "zod"; +import { toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencySystem } from "~/v3/services/concurrencySystemInstance.server"; + +const BodySchema = z.object({ + type: RetrieveQueueType.default("id"), + concurrencyKey: z.string().min(1).max(128), +}); + +const route = createActionApiRoute( + { + body: BodySchema, + params: z.object({ + queueParam: z.string().transform((val) => val.replace(/%2F/g, "/")), + }), + authorization: { + action: "write", + resource: () => ({ type: "queues" }), + }, + }, + async ({ params, body, authentication }) => { + const input: RetrieveQueueParam = + body.type === "id" + ? params.queueParam + : { + type: body.type, + name: decodeURIComponent(params.queueParam).replace(/%2F/g, "/"), + }; + + return concurrencySystem.queues + .resetConcurrencyKeyLimit(authentication.environment, input, body.concurrencyKey) + .match( + (queue) => { + return json( + toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + running: queue.running, + queued: queue.queued, + concurrencyLimit: queue.concurrencyLimit, + concurrencyLimitBase: queue.concurrencyLimitBase, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, + concurrencyLimitOverriddenBy: null, + paused: queue.paused, + }), + { status: 200 } + ); + }, + (error) => { + switch (error.type) { + case "queue_not_found": { + return json({ error: "Queue not found" }, { status: 404 }); + } + case "queue_not_overridden": { + return json( + { error: "This concurrency key does not have an override" }, + { status: 400 } + ); + } + case "queue_update_failed": { + return json({ error: "Failed to reset the concurrency key limit" }, { status: 500 }); + } + case "sync_queue_concurrency_to_engine_failed": { + return json({ error: "Failed to sync the concurrency key limit" }, { status: 500 }); + } + case "get_queue_stats_failed": { + return json({ error: "Failed to read queue stats" }, { status: 500 }); + } + case "other": { + return json( + { error: "Failed to reset the concurrency key limit" }, + { + status: 500, + } + ); + } + default: { + return json( + { error: "Failed to reset the concurrency key limit" }, + { + status: 500, + } + ); + } + } + } + ); + } +); + +export const { action } = route; diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.total.override.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.total.override.ts new file mode 100644 index 00000000000..b0c34ae1379 --- /dev/null +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.total.override.ts @@ -0,0 +1,96 @@ +import { json } from "@remix-run/server-runtime"; +import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { z } from "zod"; +import { toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencySystem } from "~/v3/services/concurrencySystemInstance.server"; + +const BodySchema = z.object({ + type: RetrieveQueueType.default("id"), + concurrencyLimit: z.number().int().min(0).max(100000), +}); + +const route = createActionApiRoute( + { + body: BodySchema, + params: z.object({ + queueParam: z.string().transform((val) => val.replace(/%2F/g, "/")), + }), + authorization: { + action: "write", + resource: () => ({ type: "queues" }), + }, + }, + async ({ params, body, authentication }) => { + const input: RetrieveQueueParam = + body.type === "id" + ? params.queueParam + : { + type: body.type, + name: decodeURIComponent(params.queueParam).replace(/%2F/g, "/"), + }; + + return concurrencySystem.queues + .overrideTotalConcurrencyLimit(authentication.environment, input, body.concurrencyLimit) + .match( + (queue) => { + return json( + toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + running: queue.running, + queued: queue.queued, + concurrencyLimit: queue.concurrencyLimit, + concurrencyLimitBase: queue.concurrencyLimitBase, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, + concurrencyLimitOverriddenBy: null, + paused: queue.paused, + }), + { status: 200 } + ); + }, + (error) => { + switch (error.type) { + case "queue_not_found": { + return json({ error: "Queue not found" }, { status: 404 }); + } + case "invalid_override": + case "concurrency_limit_exceeds_maximum": { + return json({ error: error.message }, { status: 400 }); + } + case "queue_update_failed": { + return json( + { error: "Failed to update queue total concurrency limit" }, + { status: 500 } + ); + } + case "sync_queue_concurrency_to_engine_failed": { + return json({ error: "Failed to sync the total concurrency limit" }, { status: 500 }); + } + case "get_queue_stats_failed": { + return json({ error: "Failed to read queue stats" }, { status: 500 }); + } + case "other": { + return json( + { error: "Failed to update queue total concurrency limit" }, + { + status: 500, + } + ); + } + default: { + return json( + { error: "Failed to update queue total concurrency limit" }, + { + status: 500, + } + ); + } + } + } + ); + } +); + +export const { action } = route; diff --git a/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.total.reset.ts b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.total.reset.ts new file mode 100644 index 00000000000..2eacbb3ed6a --- /dev/null +++ b/apps/webapp/app/routes/api.v1.queues.$queueParam.concurrency.total.reset.ts @@ -0,0 +1,97 @@ +import { json } from "@remix-run/server-runtime"; +import { type RetrieveQueueParam, RetrieveQueueType } from "@trigger.dev/core/v3"; +import { z } from "zod"; +import { toQueueItem } from "~/presenters/v3/QueueRetrievePresenter.server"; +import { createActionApiRoute } from "~/services/routeBuilders/apiBuilder.server"; +import { concurrencySystem } from "~/v3/services/concurrencySystemInstance.server"; + +const BodySchema = z.object({ + type: RetrieveQueueType.default("id"), +}); + +const route = createActionApiRoute( + { + body: BodySchema, + params: z.object({ + queueParam: z.string().transform((val) => val.replace(/%2F/g, "/")), + }), + authorization: { + action: "write", + resource: () => ({ type: "queues" }), + }, + }, + async ({ params, body, authentication }) => { + const input: RetrieveQueueParam = + body.type === "id" + ? params.queueParam + : { + type: body.type, + name: decodeURIComponent(params.queueParam).replace(/%2F/g, "/"), + }; + + return concurrencySystem.queues + .resetTotalConcurrencyLimit(authentication.environment, input) + .match( + (queue) => { + return json( + toQueueItem({ + friendlyId: queue.friendlyId, + name: queue.name, + type: queue.type, + running: queue.running, + queued: queue.queued, + concurrencyLimit: queue.concurrencyLimit, + concurrencyLimitBase: queue.concurrencyLimitBase, + concurrencyLimitOverriddenAt: queue.concurrencyLimitOverriddenAt, + concurrencyLimitOverriddenBy: null, + paused: queue.paused, + }), + { status: 200 } + ); + }, + (error) => { + switch (error.type) { + case "queue_not_found": { + return json({ error: "Queue not found" }, { status: 404 }); + } + case "queue_not_overridden": { + return json( + { error: "The queue total concurrency limit is not overridden" }, + { status: 400 } + ); + } + case "queue_update_failed": { + return json( + { error: "Failed to reset the queue total concurrency limit" }, + { status: 500 } + ); + } + case "sync_queue_concurrency_to_engine_failed": { + return json({ error: "Failed to sync the total concurrency limit" }, { status: 500 }); + } + case "get_queue_stats_failed": { + return json({ error: "Failed to read queue stats" }, { status: 500 }); + } + case "other": { + return json( + { error: "Failed to reset the queue total concurrency limit" }, + { + status: 500, + } + ); + } + default: { + return json( + { error: "Failed to reset the queue total concurrency limit" }, + { + status: 500, + } + ); + } + } + } + ); + } +); + +export const { action } = route; diff --git a/apps/webapp/app/v3/services/concurrencySystem.server.ts b/apps/webapp/app/v3/services/concurrencySystem.server.ts index 51c51674234..571ed5daf64 100644 --- a/apps/webapp/app/v3/services/concurrencySystem.server.ts +++ b/apps/webapp/app/v3/services/concurrencySystem.server.ts @@ -3,7 +3,13 @@ import { errAsync, fromPromise, okAsync } from "neverthrow"; import type { PrismaClientOrTransaction } from "~/db.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; -import { removeQueueConcurrencyLimits, updateQueueConcurrencyLimits } from "../runQueue.server"; +import { + removeQueueConcurrencyLimits, + removeQueueTotalConcurrencyLimits, + updateQueueConcurrencyLimits, + updateQueueTotalConcurrencyLimits, +} from "../runQueue.server"; +import { RunQueueConcurrencyKeyLimitExceededError } from "@internal/run-engine"; import { engine } from "../runEngine.server"; export type ConcurrencySystemOptions = { @@ -77,6 +83,63 @@ export class ConcurrencySystem { .andThen((queue) => syncQueueConcurrencyToEngine(environment, queue)) .andThen((queue) => getQueueStats(environment, queue)); }, + overrideTotalConcurrencyLimit: ( + environment: AuthenticatedEnvironment, + queue: QueueInput, + totalConcurrencyLimit: number, + overriddenBy?: User + ) => { + return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => + overrideQueueTotalConcurrencyLimit( + this.db, + environment, + queue, + totalConcurrencyLimit, + overriddenBy + ) + ) + .andThen((queue) => syncQueueTotalConcurrencyToEngine(environment, queue)) + .andThen((queue) => getQueueStats(environment, queue)); + }, + resetTotalConcurrencyLimit: (environment: AuthenticatedEnvironment, queue: QueueInput) => { + return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => syncQueueTotalConcurrencyResetToEngine(environment, queue)) + .andThen((queue) => resetQueueTotalConcurrencyLimit(this.db, queue)) + .andThen((queue) => syncQueueTotalConcurrencyToEngine(environment, queue)) + .andThen((queue) => getQueueStats(environment, queue)); + }, + overrideConcurrencyKeyLimit: ( + environment: AuthenticatedEnvironment, + queue: QueueInput, + concurrencyKey: string, + concurrencyLimit: number, + overriddenBy?: User + ) => { + return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => + overrideQueueConcurrencyKeyLimit( + this.db, + environment, + queue, + concurrencyKey, + concurrencyLimit, + overriddenBy + ) + ) + .andThen((queue) => getQueueStats(environment, queue)); + }, + resetConcurrencyKeyLimit: ( + environment: AuthenticatedEnvironment, + queue: QueueInput, + concurrencyKey: string + ) => { + return findQueueFromInput(this.db, environment, queue) + .andThen((queue) => + resetQueueConcurrencyKeyLimit(this.db, environment, queue, concurrencyKey) + ) + .andThen((queue) => getQueueStats(environment, queue)); + }, /** * Recalculates the materialized limit of every percent-based override in the environment * against its CURRENT maximumConcurrencyLimit and syncs changed queues to the run engine. @@ -316,6 +379,272 @@ function syncQueueConcurrencyToEngine(environment: AuthenticatedEnvironment, que } } +function overrideQueueTotalConcurrencyLimit( + db: PrismaClientOrTransaction, + environment: AuthenticatedEnvironment, + queue: TaskQueue, + totalConcurrencyLimit: number, + overriddenBy?: User +) { + const maximum = environment.maximumConcurrencyLimit; + + if (!Number.isFinite(totalConcurrencyLimit) || totalConcurrencyLimit < 0) { + return errAsync({ + type: "invalid_override" as const, + message: "Total concurrency limit must be a non-negative number", + }); + } + + if (totalConcurrencyLimit > maximum) { + return errAsync({ + type: "concurrency_limit_exceeds_maximum" as const, + message: `Total concurrency limit (${totalConcurrencyLimit}) cannot exceed the environment limit (${maximum})`, + }); + } + + const totalConcurrencyLimitBase = queue.totalConcurrencyLimitOverriddenAt + ? queue.totalConcurrencyLimitBase + : queue.totalConcurrencyLimit; + + return fromPromise( + db.taskQueue.update({ + where: { id: queue.id }, + data: { + totalConcurrencyLimit, + totalConcurrencyLimitBase: totalConcurrencyLimitBase ?? null, + totalConcurrencyLimitOverriddenAt: new Date(), + totalConcurrencyLimitOverriddenBy: overriddenBy?.id ?? null, + }, + }), + (error) => ({ + type: "queue_update_failed" as const, + cause: error, + }) + ); +} + +/** + * Enforce first, then persist: syncs the engine to the declared base BEFORE clearing + * the override marker, so an engine failure leaves the marker set and a retry + * converges instead of being rejected while the overridden limit stays enforced. + */ +function syncQueueTotalConcurrencyResetToEngine( + environment: AuthenticatedEnvironment, + queue: TaskQueue +) { + if (queue.totalConcurrencyLimitOverriddenAt === null) { + return errAsync({ type: "queue_not_overridden" as const }); + } + + if (typeof queue.totalConcurrencyLimitBase === "number") { + return fromPromise( + updateQueueTotalConcurrencyLimits(environment, queue.name, queue.totalConcurrencyLimitBase), + (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + }) + ).andThen(() => okAsync(queue)); + } + + return fromPromise(removeQueueTotalConcurrencyLimits(environment, queue.name), (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + })).andThen(() => okAsync(queue)); +} + +function resetQueueTotalConcurrencyLimit(db: PrismaClientOrTransaction, queue: TaskQueue) { + if (queue.totalConcurrencyLimitOverriddenAt === null) { + return errAsync({ type: "queue_not_overridden" as const }); + } + + return fromPromise( + db.taskQueue.update({ + where: { id: queue.id }, + data: { + totalConcurrencyLimit: queue.totalConcurrencyLimitBase, + totalConcurrencyLimitBase: null, + totalConcurrencyLimitOverriddenAt: null, + totalConcurrencyLimitOverriddenBy: null, + }, + }), + (error) => ({ + type: "queue_update_failed" as const, + cause: error, + }) + ); +} + +/** + * The total limit key is separate from the per-queue limit key that pause zeroes, + * so it syncs regardless of the paused state. + */ +function syncQueueTotalConcurrencyToEngine( + environment: AuthenticatedEnvironment, + queue: TaskQueue +) { + if (typeof queue.totalConcurrencyLimit === "number") { + return fromPromise( + updateQueueTotalConcurrencyLimits(environment, queue.name, queue.totalConcurrencyLimit), + (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + }) + ).andThen(() => okAsync(queue)); + } + + return fromPromise(removeQueueTotalConcurrencyLimits(environment, queue.name), (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + })).andThen(() => okAsync(queue)); +} + +/** + * Persist first, then enforce: a database failure leaves nothing enforced and the + * request errors cleanly, while an engine failure after persistence leaves a durable + * record and a retry converges (the upsert is idempotent). A cap rejection deletes + * the row unconditionally: the cap only rejects keys absent from the engine hash + * (updates to present keys always succeed), so a cap-rejected row is never backing + * an enforced limit and must not survive as an authoritative override. + */ +function overrideQueueConcurrencyKeyLimit( + db: PrismaClientOrTransaction, + environment: AuthenticatedEnvironment, + queue: TaskQueue, + concurrencyKey: string, + concurrencyLimit: number, + overriddenBy?: User +) { + const maximum = environment.maximumConcurrencyLimit; + + if (!Number.isFinite(concurrencyLimit) || concurrencyLimit < 0) { + return errAsync({ + type: "invalid_override" as const, + message: "Concurrency limit must be a non-negative number", + }); + } + + if (concurrencyLimit > maximum) { + return errAsync({ + type: "concurrency_limit_exceeds_maximum" as const, + message: `Concurrency limit (${concurrencyLimit}) cannot exceed the environment limit (${maximum})`, + }); + } + + return fromPromise( + db.taskQueueConcurrencyKeyOverride.findFirst({ + where: { taskQueueId: queue.id, concurrencyKey }, + select: { id: true }, + }), + (error) => ({ type: "other" as const, cause: error }) + ) + .andThen(() => + fromPromise( + db.taskQueueConcurrencyKeyOverride.upsert({ + where: { taskQueueId_concurrencyKey: { taskQueueId: queue.id, concurrencyKey } }, + create: { + taskQueueId: queue.id, + concurrencyKey, + concurrencyLimit, + overriddenBy: overriddenBy?.id ?? null, + }, + update: { + concurrencyLimit, + overriddenAt: new Date(), + overriddenBy: overriddenBy?.id ?? null, + }, + }), + (error) => ({ + type: "queue_update_failed" as const, + cause: error, + }) + ) + ) + .andThen((written) => + fromPromise( + engine.runQueue.updateQueueConcurrencyKeyLimit( + environment, + queue.name, + concurrencyKey, + concurrencyLimit + ), + (error) => { + if (error instanceof RunQueueConcurrencyKeyLimitExceededError) { + return { type: "too_many_key_overrides" as const, message: error.message }; + } + return { type: "sync_queue_concurrency_to_engine_failed" as const, cause: error }; + } + ).orElse((error) => { + if (error.type === "too_many_key_overrides") { + return fromPromise( + /** + * Deletes only the exact row generation this request wrote, so a + * concurrent request that succeeded after capacity freed keeps its + * durable record. + */ + db.taskQueueConcurrencyKeyOverride.deleteMany({ + where: { id: written.id, overriddenAt: written.overriddenAt }, + }), + () => error + ).andThen(() => errAsync(error)); + } + return errAsync(error); + }) + ) + .andThen(() => okAsync(queue)); +} + +/** + * Enforce first, then persist: removing the engine limit is idempotent, so an engine + * failure leaves the override row in place and a retry converges instead of being + * rejected as not overridden while the old limit is still enforced. + */ +function resetQueueConcurrencyKeyLimit( + db: PrismaClientOrTransaction, + environment: AuthenticatedEnvironment, + queue: TaskQueue, + concurrencyKey: string +) { + return fromPromise( + db.taskQueueConcurrencyKeyOverride.findFirst({ + where: { taskQueueId: queue.id, concurrencyKey }, + select: { id: true, overriddenAt: true }, + }), + (error) => ({ type: "other" as const, cause: error }) + ).andThen((existing) => { + if (!existing) { + return errAsync({ type: "queue_not_overridden" as const }); + } + + return fromPromise( + engine.runQueue.removeQueueConcurrencyKeyLimit(environment, queue.name, concurrencyKey), + (error) => ({ + type: "sync_queue_concurrency_to_engine_failed" as const, + cause: error, + }) + ) + .andThen(() => + fromPromise( + /** + * Deletes only the exact row generation this reset read, so a concurrent + * override that re-wrote the row after the reset began keeps its record + * (its next write, or the deploy-time restore, re-syncs the engine). + */ + db.taskQueueConcurrencyKeyOverride.deleteMany({ + where: { + id: existing.id, + overriddenAt: existing.overriddenAt, + }, + }), + (error) => ({ + type: "queue_update_failed" as const, + cause: error, + }) + ) + ) + .andThen(() => okAsync(queue)); + }); +} + function getQueueStats(environment: AuthenticatedEnvironment, queue: TaskQueue) { return fromPromise( Promise.all([ diff --git a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts index 87777379744..b0d9ab4ee57 100644 --- a/apps/webapp/app/v3/services/createBackgroundWorker.server.ts +++ b/apps/webapp/app/v3/services/createBackgroundWorker.server.ts @@ -582,6 +582,31 @@ async function createWorkerQueue( await removeQueueTotalConcurrencyLimits(environment, taskQueue.name); } + /** + * Restore per-key limit overrides into the engine so a fresh or flushed Redis + * converges back to the durable records on the next deploy. Row-level failures + * are logged rather than failing the deploy; the next deploy retries them. + */ + const keyOverrides = await prisma.taskQueueConcurrencyKeyOverride.findMany({ + where: { taskQueueId: taskQueue.id }, + }); + for (const keyOverride of keyOverrides) { + try { + await engine.runQueue.updateQueueConcurrencyKeyLimit( + environment, + taskQueue.name, + keyOverride.concurrencyKey, + keyOverride.concurrencyLimit + ); + } catch (error) { + logger.error("createWorkerQueue: failed to restore concurrency key override", { + taskQueueId: taskQueue.id, + concurrencyKey: keyOverride.concurrencyKey, + error, + }); + } + } + if (!taskQueue.paused) { if (typeof newConcurrencyLimit === "number") { logger.debug("createWorkerQueue: updating concurrency limit", { @@ -660,6 +685,7 @@ async function upsertWorkerQueueRecord( }); } else { const hasOverride = taskQueue.concurrencyLimitOverriddenAt !== null; + const hasTotalOverride = taskQueue.totalConcurrencyLimitOverriddenAt !== null; taskQueue = await prisma.taskQueue.update({ where: { @@ -672,7 +698,8 @@ async function upsertWorkerQueueRecord( // If overridden, keep current limit and update base; otherwise update limit normally concurrencyLimit: hasOverride ? undefined : concurrencyLimit, concurrencyLimitBase: hasOverride ? concurrencyLimit : undefined, - totalConcurrencyLimit, + totalConcurrencyLimit: hasTotalOverride ? undefined : totalConcurrencyLimit, + totalConcurrencyLimitBase: hasTotalOverride ? totalConcurrencyLimit : undefined, }, }); } diff --git a/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql b/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql new file mode 100644 index 00000000000..d95778f276d --- /dev/null +++ b/internal-packages/database/prisma/migrations/20260829150000_add_concurrency_overrides/migration.sql @@ -0,0 +1,24 @@ +-- AlterTable +ALTER TABLE "TaskQueue" ADD COLUMN "totalConcurrencyLimitOverriddenAt" TIMESTAMP(3); +ALTER TABLE "TaskQueue" ADD COLUMN "totalConcurrencyLimitOverriddenBy" TEXT; +ALTER TABLE "TaskQueue" ADD COLUMN "totalConcurrencyLimitBase" INTEGER; + +-- CreateTable +CREATE TABLE "TaskQueueConcurrencyKeyOverride" ( + "id" TEXT NOT NULL, + "taskQueueId" TEXT NOT NULL, + "concurrencyKey" TEXT NOT NULL, + "concurrencyLimit" INTEGER NOT NULL, + "overriddenAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "overriddenBy" TEXT, + "createdAt" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updatedAt" TIMESTAMP(3) NOT NULL, + + CONSTRAINT "TaskQueueConcurrencyKeyOverride_pkey" PRIMARY KEY ("id") +); + +-- CreateIndex +CREATE UNIQUE INDEX "TaskQueueConcurrencyKeyOverride_taskQueueId_concurrencyKey_key" ON "TaskQueueConcurrencyKeyOverride"("taskQueueId", "concurrencyKey"); + +-- AddForeignKey +ALTER TABLE "TaskQueueConcurrencyKeyOverride" ADD CONSTRAINT "TaskQueueConcurrencyKeyOverride_taskQueueId_fkey" FOREIGN KEY ("taskQueueId") REFERENCES "TaskQueue"("id") ON DELETE CASCADE ON UPDATE CASCADE; diff --git a/internal-packages/database/prisma/schema.prisma b/internal-packages/database/prisma/schema.prisma index c88ba5b5887..65a84ff1645 100644 --- a/internal-packages/database/prisma/schema.prisma +++ b/internal-packages/database/prisma/schema.prisma @@ -1983,7 +1983,13 @@ model TaskQueue { concurrencyLimitOverridePercent Decimal? @db.Decimal(5, 2) /// Caps total concurrent runs across ALL concurrencyKey values of this queue /// (concurrencyLimit applies per key value). Null = no total cap. - totalConcurrencyLimit Int? + totalConcurrencyLimit Int? + /// When the total concurrency limit was overridden + totalConcurrencyLimitOverriddenAt DateTime? + /// Who overrode the total concurrency limit (null when overridden via the API) + totalConcurrencyLimitOverriddenBy String? + /// If totalConcurrencyLimit is overridden, the declared value it reverts to on reset + totalConcurrencyLimitBase Int? rateLimit Json? paused Boolean @default(false) @@ -1995,9 +2001,31 @@ model TaskQueue { tasks BackgroundWorkerTask[] workers BackgroundWorker[] + concurrencyKeyOverrides TaskQueueConcurrencyKeyOverride[] + @@unique([runtimeEnvironmentId, name]) } +/// A per-concurrency-key limit override for a queue: the named key value gets this +/// limit instead of the queue's concurrencyLimit. Deleting the row resets the key. +model TaskQueueConcurrencyKeyOverride { + id String @id @default(cuid()) + + taskQueue TaskQueue @relation(fields: [taskQueueId], references: [id], onDelete: Cascade, onUpdate: Cascade) + taskQueueId String + + concurrencyKey String + concurrencyLimit Int + + overriddenAt DateTime @default(now()) + overriddenBy String? + + createdAt DateTime @default(now()) + updatedAt DateTime @updatedAt + + @@unique([taskQueueId, concurrencyKey]) +} + enum TaskQueueType { VIRTUAL NAMED diff --git a/internal-packages/run-engine/src/index.ts b/internal-packages/run-engine/src/index.ts index 2c98edf6866..c9b4535b280 100644 --- a/internal-packages/run-engine/src/index.ts +++ b/internal-packages/run-engine/src/index.ts @@ -1,4 +1,5 @@ export { RunEngine } from "./engine/index.js"; +export { RunQueueConcurrencyKeyLimitExceededError } from "./run-queue/index.js"; export { RunDuplicateIdempotencyKeyError, RunOneTimeUseTokenError, diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index b6ef93a9ea3..7f6493dfc2b 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -108,12 +108,18 @@ local function __gateReconcile(setKey, msgKeyPrefix, reconcileKeyPrefix) end end -local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msgKeyPrefix) +local function __gatesHaveCapacity(gatesKeyPrefix, msg, messageId, envLimit, msgKeyPrefix, ckOverridesEnabled) if not msg.gates then return true end for _, gate in ipairs(msg.gates) do local base, variant, gateKey = __gateKeys(gatesKeyPrefix, msg, gate) local occupancy = tonumber(redis.call('SCARD', variant .. ':currentConcurrency') or '0') local perKeyLimit = math.min(tonumber(redis.call('GET', base .. ':concurrency') or '1000000'), envLimit) + if ckOverridesEnabled and gateKey and gateKey ~= '' then + local gateOverride = redis.call('HGET', base .. ':ckLimits', string.sub(variant, #gatesKeyPrefix + 1)) + if gateOverride then + perKeyLimit = math.min(tonumber(gateOverride), envLimit) + end + end if occupancy >= perKeyLimit and redis.call('SISMEMBER', variant .. ':currentConcurrency', messageId) == 0 then __gateReconcile(variant .. ':currentConcurrency', msgKeyPrefix, gatesKeyPrefix) return false @@ -251,6 +257,13 @@ export interface RunQueueMetricsEmitter { emitGauge(shardKey: string, fields: Record): void; } +export class RunQueueConcurrencyKeyLimitExceededError extends Error { + constructor(message: string) { + super(message); + this.name = "RunQueueConcurrencyKeyLimitExceededError"; + } +} + export type RunQueueOptions = { name: string; tracer: Tracer; @@ -297,6 +310,10 @@ export type RunQueueOptions = { * exists are pruned, so such leaks clear within seconds instead of blocking the * queue. Enabling only after every instance runs this build avoids the noise but is * no longer load-bearing for correctness. + * + * Per-concurrency-key limit overrides are part of the same concurrency-limits + * feature and are deliberately enforced behind this flag too: writes are always + * accepted and durable, and enforcement of both arrives together. */ totalConcurrencyEnabled?: boolean; /** @@ -308,6 +325,8 @@ export type RunQueueOptions = { * the total cap covering releases from builds without the mirror. */ gatesEnabled?: boolean; + /** Cap on per-concurrency-key limit overrides stored per queue. Default 1000. */ + maxConcurrencyKeyOverridesPerQueue?: number; workerOptions?: { pollIntervalMs?: number; immediatePollIntervalMs?: number; @@ -421,6 +440,7 @@ export class RunQueue { private queueSelectionStrategy: RunQueueSelectionStrategy; private shardCount: number; private counterTtlSeconds: number; + private maxConcurrencyKeyOverridesPerQueue: number; private abortController: AbortController; private worker: Worker; private workerQueueResolver: WorkerQueueResolver; @@ -431,6 +451,7 @@ export class RunQueue { constructor(public readonly options: RunQueueOptions) { this.shardCount = options.shardCount ?? 2; this.counterTtlSeconds = options.counterTtlSeconds ?? 86400; + this.maxConcurrencyKeyOverridesPerQueue = options.maxConcurrencyKeyOverridesPerQueue ?? 1000; this.retryOptions = options.retryOptions ?? defaultRetrySettings; this.redis = createRedisClient(options.redis, { onError: (error) => { @@ -625,6 +646,62 @@ export class RunQueue { return this.redis.scard(this.keys.queueGroupConcurrencyKey(env, queue)); } + /** + * Sets a per-concurrency-key limit override for a queue. The stored value is the + * raw requested limit; admit paths clamp to the environment limit at read time. + * Throws RunQueueConcurrencyKeyLimitExceededError when a NEW key would push the + * queue past maxConcurrencyKeyOverridesPerQueue (updates to existing keys always + * succeed). + */ + public async updateQueueConcurrencyKeyLimit( + env: MinimalAuthenticatedEnvironment, + queue: string, + concurrencyKey: string, + limit: number + ) { + const result = await this.redis.setQueueConcurrencyKeyLimit( + this.keys.queueCkLimitsKey(env, queue), + this.keys.queueKey(env, queue, concurrencyKey), + String(limit), + String(this.maxConcurrencyKeyOverridesPerQueue) + ); + + if (result === 0) { + throw new RunQueueConcurrencyKeyLimitExceededError( + `Cannot add a concurrency key override to queue ${queue}: the queue already has ${this.maxConcurrencyKeyOverridesPerQueue} overrides` + ); + } + } + + public async removeQueueConcurrencyKeyLimit( + env: MinimalAuthenticatedEnvironment, + queue: string, + concurrencyKey: string + ) { + return this.redis.hdel( + this.keys.queueCkLimitsKey(env, queue), + this.keys.queueKey(env, queue, concurrencyKey) + ); + } + + /** Returns the raw per-concurrency-key limit overrides for a queue, keyed by concurrency key value. */ + public async getQueueConcurrencyKeyLimits( + env: MinimalAuthenticatedEnvironment, + queue: string + ): Promise> { + const raw = await this.redis.hgetall(this.keys.queueCkLimitsKey(env, queue)); + + const limits: Record = {}; + for (const [variantName, value] of Object.entries(raw)) { + const ckIndex = variantName.indexOf(":ck:"); + if (ckIndex === -1) { + continue; + } + limits[variantName.slice(ckIndex + 4)] = Number(value); + } + return limits; + } + public async updateEnvConcurrencyLimits(env: MinimalAuthenticatedEnvironment) { await this.#callUpdateEnvironmentConcurrencyLimits({ envConcurrencyLimitKey: this.keys.envConcurrencyLimitKey(env), @@ -2358,6 +2435,7 @@ export class RunQueue { const totalConcurrencyLimitKey = this.keys.queueTotalConcurrencyLimitKeyFromQueue( message.queue ); + const ckLimitsKey = this.keys.queueCkLimitsKeyFromQueue(message.queue); const totalConcurrencyEnabledArg = this.options.totalConcurrencyEnabled ? "1" : "0"; if (ttlInfo) { @@ -2381,6 +2459,7 @@ export class RunQueue { baseQueueKey, groupConcurrencyKey, totalConcurrencyLimitKey, + ckLimitsKey, // args queueName, messageId, @@ -2420,6 +2499,7 @@ export class RunQueue { baseQueueKey, groupConcurrencyKey, totalConcurrencyLimitKey, + ckLimitsKey, // args queueName, messageId, @@ -2469,6 +2549,7 @@ export class RunQueue { enableFastPathArg, this.options.redis.keyPrefix ?? "", this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); } else { @@ -2498,6 +2579,7 @@ export class RunQueue { enableFastPathArg, this.options.redis.keyPrefix ?? "", this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); } @@ -2579,6 +2661,7 @@ export class RunQueue { this.options.redis.keyPrefix ?? "", String(maxCount), this.options.gatesEnabled ? "1" : "0", + this.options.totalConcurrencyEnabled ? "1" : "0", metricsGaugeArg ); @@ -2707,6 +2790,7 @@ export class RunQueue { runningCounterKey, this.keys.queueGroupConcurrencyKeyFromQueue(ckWildcardQueue), this.keys.queueTotalConcurrencyLimitKeyFromQueue(ckWildcardQueue), + this.keys.queueCkLimitsKeyFromQueue(ckWildcardQueue), //args ckWildcardQueue, String(Date.now()), @@ -3575,6 +3659,7 @@ local currentTime = ARGV[8] local enableFastPath = ARGV[9] local keyPrefix = ARGV[10] local gatesEnabled = ARGV[11] == '1' +local totalConcurrencyEnabled = ARGV[12] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} ${QUEUE_GATES_LUA_HELPERS} @@ -3602,7 +3687,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil, totalConcurrencyEnabled) end end @@ -3689,6 +3774,7 @@ local currentTime = ARGV[10] local enableFastPath = ARGV[11] local keyPrefix = ARGV[12] local gatesEnabled = ARGV[13] == '1' +local totalConcurrencyEnabled = ARGV[14] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} ${QUEUE_GATES_LUA_HELPERS} @@ -3716,7 +3802,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil, totalConcurrencyEnabled) end end @@ -3995,7 +4081,7 @@ return __qmret(0) // *Tracked variants of dequeueMessageFromKey and the ack/nack/dlq/release/clear // scripts. this.redis.defineCommand("enqueueMessageCkTracked", { - numberOfKeys: 17, + numberOfKeys: 18, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -4017,6 +4103,7 @@ local baseQueueKey = KEYS[15] -- Total-cap keys (KEYS 16-17) local groupConcurrencyKey = KEYS[16] local totalConcurrencyLimitKey = KEYS[17] +local ckLimitsKey = KEYS[18] local queueName = ARGV[1] local messageId = ARGV[2] @@ -4054,6 +4141,12 @@ if enableFastPath == '1' then tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000'), envLimit ) + if totalConcurrencyEnabled then + local perKeyOverride = redis.call('HGET', ckLimitsKey, queueName) + if perKeyOverride then + queueLimit = math.min(tonumber(perKeyOverride), envLimit) + end + end if queueCurrent < queueLimit then -- Total-cap gate: a fast-path admit consumes a group slot, so it must @@ -4076,7 +4169,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil, totalConcurrencyEnabled) end end @@ -4165,7 +4258,7 @@ return __qmret(0) }); this.redis.defineCommand("enqueueMessageWithTtlCkTracked", { - numberOfKeys: 18, + numberOfKeys: 19, lua: ` local masterQueueKey = KEYS[1] local queueKey = KEYS[2] @@ -4188,6 +4281,7 @@ local baseQueueKey = KEYS[16] -- Total-cap keys (KEYS 17-18) local groupConcurrencyKey = KEYS[17] local totalConcurrencyLimitKey = KEYS[18] +local ckLimitsKey = KEYS[19] local queueName = ARGV[1] local messageId = ARGV[2] @@ -4227,6 +4321,12 @@ if enableFastPath == '1' then tonumber(redis.call('GET', queueConcurrencyLimitKey) or '1000000'), envLimit ) + if totalConcurrencyEnabled then + local perKeyOverride = redis.call('HGET', ckLimitsKey, queueName) + if perKeyOverride then + queueLimit = math.min(tonumber(perKeyOverride), envLimit) + end + end if queueCurrent < queueLimit then -- Total-cap gate: see enqueueMessageCkTracked. @@ -4247,7 +4347,7 @@ if enableFastPath == '1' then local okDecode, decoded = pcall(cjson.decode, messageData) if okDecode and type(decoded) == 'table' and decoded.gates then gateMsg = decoded - gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil) + gatesAllowFastPath = __gatesHaveCapacity(keyPrefix, decoded, messageId, envLimit, nil, totalConcurrencyEnabled) end end @@ -4577,6 +4677,7 @@ local defaultEnvConcurrencyBurstFactor = ARGV[4] local keyPrefix = ARGV[5] local maxCount = tonumber(ARGV[6] or '1') local gatesEnabled = ARGV[7] == '1' +local totalConcurrencyEnabled = ARGV[8] == '1' ${QUEUE_METRICS_GAUGE_PRELUDE} ${QUEUE_GATES_LUA_HELPERS} ${QUEUE_METRICS_GAUGE_LUA} @@ -4649,7 +4750,7 @@ for i = 1, #messages, 2 do else local gatesAllow = true if gatesEnabled then - gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix, totalConcurrencyEnabled) end if gatesAllow then @@ -4851,7 +4952,7 @@ return results // (normal dequeue, TTL-expired, or stale-orphan path — all of which were // counted at enqueue time). this.redis.defineCommand("dequeueMessagesFromCkQueueTracked", { - numberOfKeys: 13, + numberOfKeys: 14, lua: ` local ckIndexKey = KEYS[1] local queueConcurrencyLimitKey = KEYS[2] @@ -4866,6 +4967,7 @@ local lengthCounterKey = KEYS[10] local runningCounterKey = KEYS[11] local groupConcurrencyKey = KEYS[12] local totalConcurrencyLimitKey = KEYS[13] +local ckLimitsKey = KEYS[14] local ckWildcardName = ARGV[1] local currentTime = tonumber(ARGV[2]) @@ -4952,11 +5054,28 @@ for _, ckQueueName in ipairs(ckQueues) do end local fullQueueKey = keyPrefix .. ckQueueName + local blockedByGates = false local ckConcurrencyKey = fullQueueKey .. ':currentConcurrency' local ckCurrentConcurrency = tonumber(redis.call('SCARD', ckConcurrencyKey) or '0') - if ckCurrentConcurrency < queueConcurrencyLimit then + local perKeyLimit = queueConcurrencyLimit + if totalConcurrencyEnabled then + local perKeyOverride = redis.call('HGET', ckLimitsKey, ckQueueName) + if perKeyOverride then + perKeyLimit = math.min(tonumber(perKeyOverride), envConcurrencyLimit) + end + end + + if ckCurrentConcurrency >= perKeyLimit then + -- Back a blocked variant off so it cannot pin the bounded candidate window + -- and starve later keys (acute with a zero per-key override, which never + -- self-clears). Acks and nacks rebalance the score back to the oldest + -- message, so the key is eligible again the moment capacity frees. + redis.call('ZADD', ckIndexKey, currentTime + 1000, ckQueueName) + end + + if ckCurrentConcurrency < perKeyLimit then local messages = redis.call('ZRANGEBYSCORE', fullQueueKey, '-inf', tostring(currentTime), 'WITHSCORES', 'LIMIT', 0, 1) if #messages >= 2 then @@ -4981,7 +5100,10 @@ for _, ckQueueName in ipairs(ckQueues) do else local gatesAllow = true if gatesEnabled then - gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix) + gatesAllow = __gatesHaveCapacity(keyPrefix, messageData, messageId, envConcurrencyLimit, messageKeyPrefix, totalConcurrencyEnabled) + end + if not gatesAllow then + blockedByGates = true end local alreadyInGroup = false @@ -5025,11 +5147,15 @@ for _, ckQueueName in ipairs(ckQueues) do decrLengthCounter() end - local earliest = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') - if #earliest == 0 then - redis.call('ZREM', ckIndexKey, ckQueueName) + if blockedByGates then + redis.call('ZADD', ckIndexKey, currentTime + 1000, ckQueueName) else - redis.call('ZADD', ckIndexKey, earliest[2], ckQueueName) + local earliest = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') + if #earliest == 0 then + redis.call('ZREM', ckIndexKey, ckQueueName) + else + redis.call('ZADD', ckIndexKey, earliest[2], ckQueueName) + end end else local any = redis.call('ZRANGE', fullQueueKey, 0, 0, 'WITHSCORES') @@ -5872,6 +5998,26 @@ __gatesRelease(keyPrefix, redis.call('GET', messageKey), messageId) `, }); + this.redis.defineCommand("setQueueConcurrencyKeyLimit", { + numberOfKeys: 1, + lua: ` +local ckLimitsKey = KEYS[1] + +local fieldName = ARGV[1] +local limit = ARGV[2] +local maxFields = tonumber(ARGV[3]) + +if redis.call('HEXISTS', ckLimitsKey, fieldName) == 0 then + if redis.call('HLEN', ckLimitsKey) >= maxFields then + return 0 + end +end + +redis.call('HSET', ckLimitsKey, fieldName, limit) +return 1 +`, + }); + this.redis.defineCommand("updateEnvironmentConcurrencyLimits", { numberOfKeys: 2, lua: ` @@ -6062,6 +6208,7 @@ declare module "@internal/redis" { enableFastPath: string, keyPrefix: string, gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6095,6 +6242,7 @@ declare module "@internal/redis" { enableFastPath: string, keyPrefix: string, gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[number, number[] | null]> ): Result<[number, number[] | null], Context>; @@ -6133,6 +6281,7 @@ declare module "@internal/redis" { keyPrefix: string, maxCount: string, gatesEnabled: string, + totalConcurrencyEnabled: string, metricsEnabled: string, callback?: Callback<[string[] | null, number[] | null]> ): Result<[string[] | null, number[] | null], Context>; @@ -6233,6 +6382,14 @@ declare module "@internal/redis" { callback?: Callback ): Result; + setQueueConcurrencyKeyLimit( + ckLimitsKey: string, + fieldName: string, + limit: string, + maxFields: string, + callback?: Callback + ): Result; + updateEnvironmentConcurrencyLimits( // keys envConcurrencyLimitKey: string, @@ -6419,6 +6576,7 @@ declare module "@internal/redis" { baseQueueKey: string, groupConcurrencyKey: string, totalConcurrencyLimitKey: string, + ckLimitsKey: string, queueName: string, messageId: string, messageData: string, @@ -6456,6 +6614,7 @@ declare module "@internal/redis" { baseQueueKey: string, groupConcurrencyKey: string, totalConcurrencyLimitKey: string, + ckLimitsKey: string, queueName: string, messageId: string, messageData: string, @@ -6490,6 +6649,7 @@ declare module "@internal/redis" { runningCounterKey: string, groupConcurrencyKey: string, totalConcurrencyLimitKey: string, + ckLimitsKey: string, ckWildcardName: string, currentTime: string, defaultEnvConcurrencyLimit: string, diff --git a/internal-packages/run-engine/src/run-queue/keyProducer.ts b/internal-packages/run-engine/src/run-queue/keyProducer.ts index 98028f5af7b..7b997043244 100644 --- a/internal-packages/run-engine/src/run-queue/keyProducer.ts +++ b/internal-packages/run-engine/src/run-queue/keyProducer.ts @@ -26,6 +26,7 @@ const constants = { RUNNING_COUNTER_PART: "runningCounter", GROUP_CONCURRENCY_PART: "groupConcurrency", TOTAL_CONCURRENCY_LIMIT_PART: "totalConcurrency", + CK_LIMITS_PART: "ckLimits", } as const; export class RunQueueFullKeyProducer implements RunQueueKeyProducer { @@ -366,6 +367,20 @@ export class RunQueueFullKeyProducer implements RunQueueKeyProducer { return `${this.baseQueueKeyFromQueue(queue)}:${constants.TOTAL_CONCURRENCY_LIMIT_PART}`; } + /** + * HASH of per-concurrency-key limit overrides for a queue. Lives at the base + * queue; each field is the EXACT full ck-variant queue name (the ckIndex ZSET + * member), so reads need no parsing, and values are the raw requested limits + * (readers clamp to the environment limit). + */ + queueCkLimitsKey(env: RunQueueKeyProducerEnvironment, queue: string): string { + return `${this.queueKey(env, queue)}:${constants.CK_LIMITS_PART}`; + } + + queueCkLimitsKeyFromQueue(queue: string): string { + return `${this.baseQueueKeyFromQueue(queue)}:${constants.CK_LIMITS_PART}`; + } + isCkWildcard(queue: string): boolean { return queue.endsWith(":ck:*"); } diff --git a/internal-packages/run-engine/src/run-queue/tests/concurrencyKeyOverrides.test.ts b/internal-packages/run-engine/src/run-queue/tests/concurrencyKeyOverrides.test.ts new file mode 100644 index 00000000000..8c10a8ac2bf --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/concurrencyKeyOverrides.test.ts @@ -0,0 +1,369 @@ +import { redisTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { setTimeout } from "node:timers/promises"; +import { describe } from "vitest"; +import { FairQueueSelectionStrategy } from "../fairQueueSelectionStrategy.js"; +import { RunQueue, RunQueueConcurrencyKeyLimitExceededError } from "../index.js"; +import { RunQueueFullKeyProducer } from "../keyProducer.js"; +import type { InputPayload } from "../types.js"; +import { Decimal } from "@trigger.dev/database"; + +const testOptions = { + name: "rq", + tracer: trace.getTracer("rq"), + workers: 1, + defaultEnvConcurrency: 25, + retryOptions: { + maxAttempts: 5, + factor: 1.1, + minTimeoutInMs: 100, + maxTimeoutInMs: 1_000, + randomize: true, + }, + keys: new RunQueueFullKeyProducer(), +}; + +const authenticatedEnvDev = { + id: "e1234", + type: "DEVELOPMENT" as const, + maximumConcurrencyLimit: 10, + concurrencyLimitBurstFactor: new Decimal(2.0), + project: { id: "p1234" }, + organization: { id: "o1234" }, +}; + +function createQueue( + redisContainer: any, + totalConcurrencyEnabled: boolean, + maxOverrides?: number, + dequeueCount?: number, + gatesEnabled?: boolean +) { + return new RunQueue({ + ...testOptions, + totalConcurrencyEnabled, + gatesEnabled, + maxConcurrencyKeyOverridesPerQueue: maxOverrides, + masterQueueConsumerDequeueCount: dequeueCount, + queueSelectionStrategy: new FairQueueSelectionStrategy({ + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + keys: testOptions.keys, + }), + redis: { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }, + }); +} + +function makeMessage(overrides: Partial = {}): InputPayload { + return { + runId: "r1", + taskIdentifier: "task/my-task", + orgId: "o1234", + projectId: "p1234", + environmentId: "e1234", + environmentType: "DEVELOPMENT", + queue: "task/my-task", + timestamp: Date.now(), + attempt: 0, + ...overrides, + }; +} + +async function waitFor(condition: () => Promise, timeoutMs = 20_000): Promise { + const deadline = Date.now() + timeoutMs; + while (Date.now() < deadline) { + if (await condition()) { + return true; + } + await setTimeout(250); + } + return condition(); +} + +vi.setConfig({ testTimeout: 60_000 }); + +describe("RunQueue per-concurrency-key limit overrides", () => { + redisTest( + "a lowered key is capped while other keys keep the queue limit", + async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 2); + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-a", 1); + + const now = Date.now(); + const messages = [ + ["ck-a", "a0"], + ["ck-a", "a1"], + ["ck-b", "b0"], + ["ck-b", "b1"], + ] as const; + for (const [i, [ck, id]] of messages.entries()) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: id, concurrencyKey: ck, timestamp: now - 1000 + i }), + workerQueue: "main", + }); + } + + const settled = await waitFor(async () => { + const a = await queue.currentConcurrencyOfQueue( + authenticatedEnvDev, + "task/my-task", + "ck-a" + ); + const b = await queue.currentConcurrencyOfQueue( + authenticatedEnvDev, + "task/my-task", + "ck-b" + ); + return a === 1 && b === 2; + }); + expect(settled).toBe(true); + + await setTimeout(2000); + expect( + await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task", "ck-a") + ).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + } + ); + + redisTest("a raised key admits past the queue limit", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 1); + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-a", 3); + + const now = Date.now(); + for (const i of [0, 1, 2]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `a${i}`, + concurrencyKey: "ck-a", + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "b0", concurrencyKey: "ck-b", timestamp: now - 500 }), + workerQueue: "main", + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ runId: "b1", concurrencyKey: "ck-b", timestamp: now - 499 }), + workerQueue: "main", + }); + + const settled = await waitFor(async () => { + const a = await queue.currentConcurrencyOfQueue( + authenticatedEnvDev, + "task/my-task", + "ck-a" + ); + const b = await queue.currentConcurrencyOfQueue( + authenticatedEnvDev, + "task/my-task", + "ck-b" + ); + return a === 3 && b === 1; + }); + expect(settled).toBe(true); + + await setTimeout(2000); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest("removing an override restores the queue limit", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 2); + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-a", 1); + expect(await queue.getQueueConcurrencyKeyLimits(authenticatedEnvDev, "task/my-task")).toEqual( + { "ck-a": 1 } + ); + + await queue.removeQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-a"); + expect(await queue.getQueueConcurrencyKeyLimits(authenticatedEnvDev, "task/my-task")).toEqual( + {} + ); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `a${i}`, + concurrencyKey: "ck-a", + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + + const settled = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task", "ck-a")) === 2 + ); + expect(settled).toBe(true); + } finally { + await queue.quit(); + } + }); + + redisTest("the per-queue override count is capped", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true, 2); + try { + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-a", 1); + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-b", 1); + + await expect( + queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-c", 1) + ).rejects.toThrow(RunQueueConcurrencyKeyLimitExceededError); + + /** Updates to existing keys always succeed at the cap. */ + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-a", 4); + expect(await queue.getQueueConcurrencyKeyLimits(authenticatedEnvDev, "task/my-task")).toEqual( + { "ck-a": 4, "ck-b": 1 } + ); + } finally { + await queue.quit(); + } + }); + + redisTest( + "blocked keys cannot pin the candidate window and starve later keys", + async ({ redisContainer }) => { + /** dequeueCount 2 makes the candidate window 6 variants wide. */ + const queue = createQueue(redisContainer, true, undefined, 2); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + + /** + * Ten zero-limit keys with OLDER messages fill the window many times over; + * without the blocked-key backoff the runnable key behind them would never + * be examined. + */ + const now = Date.now(); + for (let i = 0; i < 10; i++) { + const ck = `blocked-${i}`; + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", ck, 0); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `b${i}`, + concurrencyKey: ck, + timestamp: now - 10_000 + i, + }), + workerQueue: "main", + }); + } + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: "good-0", + concurrencyKey: "ck-good", + timestamp: now - 500, + }), + workerQueue: "main", + }); + + const goodAdmitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue( + authenticatedEnvDev, + "task/my-task", + "ck-good" + )) === 1, + 30_000 + ); + expect(goodAdmitted).toBe(true); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(10); + } finally { + await queue.quit(); + } + } + ); + + redisTest("gate admission honors the gate queue per-key override", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, true, undefined, undefined, true); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 5); + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "tenant", 1); + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "tenant", "acme", 2); + + const now = Date.now(); + for (const [i, ck] of ["ck-a", "ck-b", "ck-c"].entries()) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `r${i}`, + concurrencyKey: ck, + timestamp: now - 1000 + i, + gates: [{ queue: "tenant", concurrencyKey: "acme" }], + }), + workerQueue: "main", + }); + } + + /** The declared gate limit is 1; the override raises acme to 2. */ + const twoAdmitted = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")) === 2 + ); + expect(twoAdmitted).toBe(true); + + await setTimeout(2000); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "tenant", "acme")).toBe(2); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await queue.quit(); + } + }); + + redisTest("overrides are ignored when disabled", async ({ redisContainer }) => { + const queue = createQueue(redisContainer, false); + try { + await queue.updateQueueConcurrencyLimits(authenticatedEnvDev, "task/my-task", 2); + await queue.updateQueueConcurrencyKeyLimit(authenticatedEnvDev, "task/my-task", "ck-a", 1); + + const now = Date.now(); + for (const i of [0, 1]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: makeMessage({ + runId: `a${i}`, + concurrencyKey: "ck-a", + timestamp: now - 1000 + i, + }), + workerQueue: "main", + }); + } + + const settled = await waitFor( + async () => + (await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task", "ck-a")) === 2 + ); + expect(settled).toBe(true); + } finally { + await queue.quit(); + } + }); +}); diff --git a/internal-packages/run-engine/src/run-queue/types.ts b/internal-packages/run-engine/src/run-queue/types.ts index 2cbfe40c775..b21a3f66368 100644 --- a/internal-packages/run-engine/src/run-queue/types.ts +++ b/internal-packages/run-engine/src/run-queue/types.ts @@ -110,6 +110,8 @@ export interface RunQueueKeyProducer { queueGroupConcurrencyKeyFromQueue(queue: string): string; queueTotalConcurrencyLimitKey(env: RunQueueKeyProducerEnvironment, queue: string): string; queueTotalConcurrencyLimitKeyFromQueue(queue: string): string; + queueCkLimitsKey(env: RunQueueKeyProducerEnvironment, queue: string): string; + queueCkLimitsKeyFromQueue(queue: string): string; //env oncurrency envCurrentConcurrencyKey(env: EnvDescriptor): string; diff --git a/packages/core/src/v3/apiClient/index.ts b/packages/core/src/v3/apiClient/index.ts index c270f86ea62..b46f0e6108b 100644 --- a/packages/core/src/v3/apiClient/index.ts +++ b/packages/core/src/v3/apiClient/index.ts @@ -1684,6 +1684,103 @@ export class ApiClient { ); } + overrideQueueTotalConcurrencyLimit( + queue: RetrieveQueueParam, + concurrencyLimit: number, + requestOptions?: ZodFetchOptions + ) { + const type = typeof queue === "string" ? "id" : queue.type; + const value = typeof queue === "string" ? queue : queue.name; + + const encodedValue = encodeURIComponent(value.replace(/\//g, "%2F")); + + return zodfetch( + QueueItem, + `${this.baseUrl}/api/v1/queues/${encodedValue}/concurrency/total/override`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify({ + type, + concurrencyLimit, + }), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + resetQueueTotalConcurrencyLimit(queue: RetrieveQueueParam, requestOptions?: ZodFetchOptions) { + const type = typeof queue === "string" ? "id" : queue.type; + const value = typeof queue === "string" ? queue : queue.name; + + const encodedValue = encodeURIComponent(value.replace(/\//g, "%2F")); + + return zodfetch( + QueueItem, + `${this.baseUrl}/api/v1/queues/${encodedValue}/concurrency/total/reset`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify({ + type, + }), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + overrideQueueConcurrencyKeyLimit( + queue: RetrieveQueueParam, + concurrencyKey: string, + concurrencyLimit: number, + requestOptions?: ZodFetchOptions + ) { + const type = typeof queue === "string" ? "id" : queue.type; + const value = typeof queue === "string" ? queue : queue.name; + + const encodedValue = encodeURIComponent(value.replace(/\//g, "%2F")); + + return zodfetch( + QueueItem, + `${this.baseUrl}/api/v1/queues/${encodedValue}/concurrency/key/override`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify({ + type, + concurrencyKey, + concurrencyLimit, + }), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + + resetQueueConcurrencyKeyLimit( + queue: RetrieveQueueParam, + concurrencyKey: string, + requestOptions?: ZodFetchOptions + ) { + const type = typeof queue === "string" ? "id" : queue.type; + const value = typeof queue === "string" ? queue : queue.name; + + const encodedValue = encodeURIComponent(value.replace(/\//g, "%2F")); + + return zodfetch( + QueueItem, + `${this.baseUrl}/api/v1/queues/${encodedValue}/concurrency/key/reset`, + { + method: "POST", + headers: this.#getHeaders(false), + body: JSON.stringify({ + type, + concurrencyKey, + }), + }, + mergeRequestOptions(this.defaultRequestOptions, requestOptions) + ); + } + subscribeToRun( runId: string, options?: { diff --git a/packages/trigger-sdk/src/v3/queues.ts b/packages/trigger-sdk/src/v3/queues.ts index 7e76c5f940b..7f16d96e58b 100644 --- a/packages/trigger-sdk/src/v3/queues.ts +++ b/packages/trigger-sdk/src/v3/queues.ts @@ -144,9 +144,10 @@ export function pause( export function overrideConcurrencyLimit( queue: RetrieveQueueParam, concurrencyLimit: number, - requestOptions?: ApiRequestOptions + options?: ApiRequestOptions & { concurrencyKey?: string } ): ApiPromise { const apiClient = apiClientManager.clientOrThrow(); + const { concurrencyKey, ...requestOptions } = options ?? {}; const $requestOptions = mergeRequestOptions( { @@ -154,7 +155,7 @@ export function overrideConcurrencyLimit( name: "queues.overrideConcurrencyLimit()", icon: "queue", attributes: { - ...flattenAttributes({ queue }), + ...flattenAttributes({ queue, concurrencyKey }), ...accessoryAttributes({ items: [ { @@ -169,9 +170,93 @@ export function overrideConcurrencyLimit( requestOptions ); + if (concurrencyKey !== undefined) { + return apiClient.overrideQueueConcurrencyKeyLimit( + queue, + concurrencyKey, + concurrencyLimit, + $requestOptions + ); + } + return apiClient.overrideQueueConcurrencyLimit(queue, concurrencyLimit, $requestOptions); } +/** + * Overrides the total concurrency limit of a queue: the cap on concurrent runs across + * all of its `concurrencyKey` values. + * + * @param queue - The ID of the queue, or the type and name + * @param concurrencyLimit - The total concurrency limit to apply + * @returns The updated queue state + */ +export function overrideTotalConcurrencyLimit( + queue: RetrieveQueueParam, + concurrencyLimit: number, + requestOptions?: ApiRequestOptions +): ApiPromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "queues.overrideTotalConcurrencyLimit()", + icon: "queue", + attributes: { + ...flattenAttributes({ queue }), + ...accessoryAttributes({ + items: [ + { + text: typeof queue === "string" ? queue : queue.name, + variant: "normal", + }, + ], + style: "codepath", + }), + }, + }, + requestOptions + ); + + return apiClient.overrideQueueTotalConcurrencyLimit(queue, concurrencyLimit, $requestOptions); +} + +/** + * Resets the total concurrency limit of a queue back to its declared value. + * + * @param queue - The ID of the queue, or the type and name + * @returns The updated queue state + */ +export function resetTotalConcurrencyLimit( + queue: RetrieveQueueParam, + requestOptions?: ApiRequestOptions +): ApiPromise { + const apiClient = apiClientManager.clientOrThrow(); + + const $requestOptions = mergeRequestOptions( + { + tracer, + name: "queues.resetTotalConcurrencyLimit()", + icon: "queue", + attributes: { + ...flattenAttributes({ queue }), + ...accessoryAttributes({ + items: [ + { + text: typeof queue === "string" ? queue : queue.name, + variant: "normal", + }, + ], + style: "codepath", + }), + }, + }, + requestOptions + ); + + return apiClient.resetQueueTotalConcurrencyLimit(queue, $requestOptions); +} + /** * Resets the concurrency limit of a queue to the base value. * @@ -180,9 +265,10 @@ export function overrideConcurrencyLimit( */ export function resetConcurrencyLimit( queue: RetrieveQueueParam, - requestOptions?: ApiRequestOptions + options?: ApiRequestOptions & { concurrencyKey?: string } ): ApiPromise { const apiClient = apiClientManager.clientOrThrow(); + const { concurrencyKey, ...requestOptions } = options ?? {}; const $requestOptions = mergeRequestOptions( { @@ -190,7 +276,7 @@ export function resetConcurrencyLimit( name: "queues.resetConcurrencyLimit()", icon: "queue", attributes: { - ...flattenAttributes({ queue }), + ...flattenAttributes({ queue, concurrencyKey }), ...accessoryAttributes({ items: [ { @@ -205,6 +291,10 @@ export function resetConcurrencyLimit( requestOptions ); + if (concurrencyKey !== undefined) { + return apiClient.resetQueueConcurrencyKeyLimit(queue, concurrencyKey, $requestOptions); + } + return apiClient.resetQueueConcurrencyLimit(queue, $requestOptions); }