From 352c612c77b9bf138732563fd626ded4be6015ee Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Mon, 27 Jul 2026 17:47:51 +0100 Subject: [PATCH 1/3] fix(run-engine,webapp): return unclaimed runs to the queue when pausing Pausing a queue or environment set its concurrency limit to 0, which stops the queue admitting new runs but does nothing about runs it had already admitted. Those runs sit in a per-region worker queue that supervisors drain with a plain BLPOP, with no pause or concurrency check on the pop path, so they went on to execute during the pause. An idle environment with headroom can have its whole concurrency limit's worth of runs in that state, since the trigger-time fast path pushes straight to the worker queue. Adds RunQueue.returnUnclaimedMessagesToQueue, which moves runs the queue has admitted but no worker has claimed yet back into the pending queue. Candidates come from the difference between currentConcurrency and currentDequeued, which identifies worker queue residents without scanning the region-wide list. Each move runs in a Lua script whose LREM doubles as the claim: if it removes nothing then a worker already owns that run and every other key is left alone, so a concurrent BLPOP can never produce a duplicate dispatch. Ordering is preserved by re-adding at the message's original timestamp rather than a fresh score the way nack does, so a returned run keeps its place in line. Callers must set the concurrency limit to 0 first: returning a run costs it its position in the worker queue FIFO and it can only re-enter at the back, so a queue that can still admit work could let a newer run overtake one on its way back. All three pause callers already do this ordering. Covers the manual queue pause, the manual environment pause, and the billing-limit environment pause. --- .server-changes/pause-returns-waiting-runs.md | 6 + apps/webapp/app/v3/runQueue.server.ts | 16 + ...billingLimitConvergeEnvironments.server.ts | 27 +- .../v3/services/pauseEnvironment.server.ts | 3 +- .../app/v3/services/pauseQueue.server.ts | 7 +- .../billingLimitConvergeEnvironments.test.ts | 72 +++ .../run-engine/src/engine/index.ts | 18 + .../run-engine/src/run-queue/index.ts | 369 ++++++++++++++ .../dequeueMessageFromWorkerQueue.test.ts | 53 ++ .../returnUnclaimedMessagesToQueue.test.ts | 454 ++++++++++++++++++ 10 files changed, 1019 insertions(+), 6 deletions(-) create mode 100644 .server-changes/pause-returns-waiting-runs.md create mode 100644 internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts diff --git a/.server-changes/pause-returns-waiting-runs.md b/.server-changes/pause-returns-waiting-runs.md new file mode 100644 index 00000000000..e5f7c198754 --- /dev/null +++ b/.server-changes/pause-returns-waiting-runs.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: fix +--- + +Pausing a queue or an environment now also holds back runs that were already waiting to start. Previously those runs would still go ahead, so a pause could take effect a little later than expected. diff --git a/apps/webapp/app/v3/runQueue.server.ts b/apps/webapp/app/v3/runQueue.server.ts index 0ff28fa0889..fc716f0f1a8 100644 --- a/apps/webapp/app/v3/runQueue.server.ts +++ b/apps/webapp/app/v3/runQueue.server.ts @@ -30,3 +30,19 @@ export async function removeQueueConcurrencyLimits( ) { await engine.runQueue.removeQueueConcurrencyLimits(environment, queueName); } + +/** + * Returns runs waiting to be picked up by a worker back into their queue. + * + * Only safe once the concurrency limit has been set to 0, otherwise a newer run can be + * admitted and overtake one on its way back. + */ +export async function returnUnclaimedMessagesToQueue({ + environment, + queue, +}: { + environment: AuthenticatedEnvironment; + queue?: string; +}) { + return engine.returnUnclaimedMessagesToQueue({ environment, queue }); +} diff --git a/apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts b/apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts index 0b59d4c7fae..abde4a0f1c1 100644 --- a/apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts +++ b/apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts @@ -29,6 +29,8 @@ type UpdateEnvConcurrency = ( maximumConcurrencyLimit?: number ) => Promise; +type ReturnUnclaimedMessages = (environment: EnvironmentWithRelations) => Promise; + export async function convergeBillingLimitEnvironmentsForOrg( organizationId: string, targetState: BillingLimitConvergeTargetState, @@ -36,6 +38,7 @@ export async function convergeBillingLimitEnvironmentsForOrg( batchSize?: number; prismaClient?: PrismaClient; updateConcurrency?: UpdateEnvConcurrency; + returnUnclaimed?: ReturnUnclaimedMessages; } ): Promise { const db = options?.prismaClient ?? prisma; @@ -51,18 +54,32 @@ export async function convergeBillingLimitEnvironmentsForOrg( return updateEnvConcurrencyLimits(environment, maximumConcurrencyLimit); }); + const returnUnclaimed = + options?.returnUnclaimed ?? + (async (environment) => { + const { returnUnclaimedMessagesToQueue } = await import("~/v3/runQueue.server"); + await returnUnclaimedMessagesToQueue({ environment }); + }); + if (targetState === "ok") { return unpauseBillingLimitEnvironments(organizationId, db, batchSize, updateConcurrency); } - return pauseBillingLimitEnvironments(organizationId, db, batchSize, updateConcurrency); + return pauseBillingLimitEnvironments( + organizationId, + db, + batchSize, + updateConcurrency, + returnUnclaimed + ); } async function pauseBillingLimitEnvironments( organizationId: string, db: PrismaClient, batchSize: number, - updateConcurrency: UpdateEnvConcurrency + updateConcurrency: UpdateEnvConcurrency, + returnUnclaimed: ReturnUnclaimedMessages ): Promise { let paused = 0; let cursor: string | undefined; @@ -88,7 +105,7 @@ async function pauseBillingLimitEnvironments( } for (const environment of environments) { - await pauseEnvironmentForBillingLimit(environment, db, updateConcurrency); + await pauseEnvironmentForBillingLimit(environment, db, updateConcurrency, returnUnclaimed); paused++; } @@ -156,7 +173,8 @@ async function unpauseBillingLimitEnvironments( async function pauseEnvironmentForBillingLimit( environment: EnvironmentWithRelations, db: PrismaClient, - updateConcurrency: UpdateEnvConcurrency + updateConcurrency: UpdateEnvConcurrency, + returnUnclaimed: ReturnUnclaimedMessages ) { const updated = await db.runtimeEnvironment.update({ where: { id: environment.id }, @@ -172,6 +190,7 @@ async function pauseEnvironmentForBillingLimit( try { await updateConcurrency(updated, 0); + await returnUnclaimed(updated); } catch (error) { await db.runtimeEnvironment.update({ where: { id: environment.id }, diff --git a/apps/webapp/app/v3/services/pauseEnvironment.server.ts b/apps/webapp/app/v3/services/pauseEnvironment.server.ts index af9edff856c..ab10e08ca0a 100644 --- a/apps/webapp/app/v3/services/pauseEnvironment.server.ts +++ b/apps/webapp/app/v3/services/pauseEnvironment.server.ts @@ -2,7 +2,7 @@ import { EnvironmentPauseSource, type PrismaClientOrTransaction } from "@trigger import { prisma } from "~/db.server"; import { logger } from "~/services/logger.server"; import { getManualPauseEnvironmentResult } from "~/v3/services/billingLimit/manualPauseEnvironmentGuard.server"; -import { updateEnvConcurrencyLimits } from "../runQueue.server"; +import { returnUnclaimedMessagesToQueue, updateEnvConcurrencyLimits } from "../runQueue.server"; import { WithRunEngine } from "./baseService.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; @@ -114,6 +114,7 @@ export class PauseEnvironmentService extends WithRunEngine { environmentId: environment.id, }); await updateEnvConcurrencyLimits(environment, 0); + await returnUnclaimedMessagesToQueue({ environment }); } else { logger.debug("PauseEnvironmentService: resuming environment", { environmentId: environment.id, diff --git a/apps/webapp/app/v3/services/pauseQueue.server.ts b/apps/webapp/app/v3/services/pauseQueue.server.ts index aa3e21f9727..19c7abfb3da 100644 --- a/apps/webapp/app/v3/services/pauseQueue.server.ts +++ b/apps/webapp/app/v3/services/pauseQueue.server.ts @@ -4,7 +4,11 @@ import { type AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { logger } from "~/services/logger.server"; import { BaseService } from "./baseService.server"; import { determineEngineVersion } from "../engineVersion.server"; -import { removeQueueConcurrencyLimits, updateQueueConcurrencyLimits } from "../runQueue.server"; +import { + removeQueueConcurrencyLimits, + returnUnclaimedMessagesToQueue, + updateQueueConcurrencyLimits, +} from "../runQueue.server"; import { engine } from "../runEngine.server"; export type PauseStatus = "paused" | "resumed"; @@ -59,6 +63,7 @@ export class PauseQueueService extends BaseService { if (action === "paused") { await updateQueueConcurrencyLimits(environment, queue.name, 0); + await returnUnclaimedMessagesToQueue({ environment, queue: queue.name }); } else { if (queue.concurrencyLimit) { await updateQueueConcurrencyLimits(environment, queue.name, queue.concurrencyLimit); diff --git a/apps/webapp/test/billingLimitConvergeEnvironments.test.ts b/apps/webapp/test/billingLimitConvergeEnvironments.test.ts index 0854ea42ef4..f1ec8dff2be 100644 --- a/apps/webapp/test/billingLimitConvergeEnvironments.test.ts +++ b/apps/webapp/test/billingLimitConvergeEnvironments.test.ts @@ -49,6 +49,78 @@ describe("convergeBillingLimitEnvironmentsForOrg", () => { expect(envAfter.pauseSource).toBeNull(); }); + postgresTest("returns unclaimed runs after pausing for a billing limit", async ({ prisma }) => { + const { organization, project } = await createTestOrgProjectWithMember(prisma); + const environment = await createRuntimeEnvironment(prisma, { + projectId: project.id, + organizationId: organization.id, + type: "PRODUCTION", + slug: uniqueId("prod"), + }); + + const calls: Array<{ concurrency?: number; returnedFor?: string }> = []; + + const result = await convergeBillingLimitEnvironmentsForOrg(organization.id, "grace", { + prismaClient: prisma, + updateConcurrency: async (_env, maximumConcurrencyLimit) => { + calls.push({ concurrency: maximumConcurrencyLimit }); + }, + returnUnclaimed: async (env) => { + calls.push({ returnedFor: env.id }); + }, + }); + + expect(result).toEqual({ paused: 1, unpaused: 0 }); + + expect(calls).toEqual([{ concurrency: 0 }, { returnedFor: environment.id }]); + + const envAfter = await prisma.runtimeEnvironment.findUniqueOrThrow({ + where: { id: environment.id }, + }); + expect(envAfter.paused).toBe(true); + expect(envAfter.pauseSource).toBe(EnvironmentPauseSource.BILLING_LIMIT); + }); + + postgresTest("does not return unclaimed runs when unpausing", async ({ prisma }) => { + const { organization } = await createBillingPausedProductionEnv(prisma); + + const returnUnclaimed = vi.fn(async () => undefined); + + await convergeBillingLimitEnvironmentsForOrg(organization.id, "ok", { + prismaClient: prisma, + updateConcurrency: async () => undefined, + returnUnclaimed, + }); + + expect(returnUnclaimed).not.toHaveBeenCalled(); + }); + + postgresTest("rolls back pause when returning unclaimed runs fails", async ({ prisma }) => { + const { organization, project } = await createTestOrgProjectWithMember(prisma); + const environment = await createRuntimeEnvironment(prisma, { + projectId: project.id, + organizationId: organization.id, + type: "PRODUCTION", + slug: uniqueId("prod"), + }); + + await expect( + convergeBillingLimitEnvironmentsForOrg(organization.id, "grace", { + prismaClient: prisma, + updateConcurrency: async () => undefined, + returnUnclaimed: async () => { + throw new Error("run queue unavailable"); + }, + }) + ).rejects.toThrow("run queue unavailable"); + + const envAfter = await prisma.runtimeEnvironment.findUniqueOrThrow({ + where: { id: environment.id }, + }); + expect(envAfter.paused).toBe(false); + expect(envAfter.pauseSource).toBeNull(); + }); + postgresTest("rolls back pause when concurrency update fails", async ({ prisma }) => { const { organization, project } = await createTestOrgProjectWithMember(prisma); const environment = await createRuntimeEnvironment(prisma, { diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 3c1f4330a0f..3fe2d325033 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -1634,6 +1634,24 @@ export class RunEngine { return this.runQueue.currentConcurrencyOfQueues(environment, queues); } + /** + * Returns runs that have been admitted for execution but that no worker has claimed yet + * back into their queue, at the position they held before. + * + * Pausing only stops runs being admitted; anything already handed to a worker queue would + * otherwise still execute. Call this after the queue has been made ineligible for + * dequeuing, never before — see `RunQueue.returnUnclaimedMessagesToQueue`. + */ + async returnUnclaimedMessagesToQueue({ + environment, + queue, + }: { + environment: MinimalAuthenticatedEnvironment; + queue?: string; + }): Promise<{ returned: number; skipped: number }> { + return this.runQueue.returnUnclaimedMessagesToQueue({ env: environment, queue }); + } + async removeEnvironmentQueuesFromMasterQueue({ runtimeEnvironmentId, organizationId, diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index a0571206538..ad951ed770c 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -1033,6 +1033,118 @@ export class RunQueue { ); } + /** + * Returns runs that the queue has admitted but that no worker has claimed yet back into + * the pending queue, restoring the position they held before they were admitted. + * + * A run sitting in the worker queue is in `currentConcurrency` but not yet in + * `currentDequeued`, so the difference of those two sets identifies the candidates without + * having to scan the (region-wide) worker queue list. + * + * Callers must make the queue ineligible for dequeuing (concurrency limit 0) *before* + * calling this. Returning a run costs it its place in the worker queue FIFO and it can only + * re-enter at the back, so if the queue can still admit work then a newer run can overtake + * one that is on its way back. + * + * @param queue - Restrict to a single queue (including its concurrency key variants). + * Omit to cover every queue in the environment. + */ + public async returnUnclaimedMessagesToQueue({ + env, + queue, + }: { + env: MinimalAuthenticatedEnvironment; + queue?: string; + }): Promise<{ returned: number; skipped: number }> { + return this.#trace( + "returnUnclaimedMessagesToQueue", + async (span) => { + const unclaimedRunIds = await this.redis.sdiff( + this.keys.envCurrentConcurrencyKey(env), + this.keys.envCurrentDequeuedKey(env) + ); + + span.setAttribute("unclaimed_count", unclaimedRunIds.length); + + if (unclaimedRunIds.length === 0) { + return { returned: 0, skipped: 0 }; + } + + const targetBaseQueueKey = queue ? this.keys.queueKey(env, queue) : undefined; + const messages = await this.#readMessages(env.organization.id, unclaimedRunIds); + + let returned = 0; + let skipped = 0; + + for (const message of messages) { + if ( + targetBaseQueueKey && + this.keys.baseQueueKeyFromQueue(message.queue) !== targetBaseQueueKey + ) { + continue; + } + + if (await this.#callReturnMessageToQueue(message)) { + returned++; + } else { + skipped++; + } + } + + span.setAttribute("returned_count", returned); + span.setAttribute("skipped_count", skipped); + + this.logger.info("returnUnclaimedMessagesToQueue", { + service: this.name, + environmentId: env.id, + queue, + candidates: unclaimedRunIds.length, + returned, + skipped, + }); + + return { returned, skipped }; + }, + { + kind: SpanKind.INTERNAL, + attributes: { + [SEMATTRS_MESSAGING_SYSTEM]: "runqueue", + ...attributesFromAuthenticatedEnv(env), + }, + } + ); + } + + async #readMessages(orgId: string, runIds: string[]): Promise { + const rawMessages = await this.redis.mget( + runIds.map((runId) => this.keys.messageKey(orgId, runId)) + ); + + const messages: OutputPayload[] = []; + + for (const rawMessage of rawMessages) { + if (!rawMessage) { + continue; + } + + const [error, message] = parseRawMessage(rawMessage); + + if (error) { + this.logger.error(`[${this.name}] Failed to parse message`, { + error, + service: this.name, + message: message ?? rawMessage, + }); + } + + if (message) { + messages.push(message); + } + } + + return messages; + } + public async removeEnvironmentQueuesFromMasterQueue( runtimeEnvironmentId: string, organizationId: string, @@ -2623,6 +2735,80 @@ export class RunQueue { } } + /** + * @returns true if the message was returned to its queue, false if a worker had already + * claimed it (or its payload was gone) and it was left alone. + */ + async #callReturnMessageToQueue(message: OutputPayload): Promise { + const messageId = message.runId; + const messageKey = this.keys.messageKey(message.orgId, message.runId); + const messageQueue = message.queue; + const queueCurrentConcurrencyKey = this.keys.queueCurrentConcurrencyKeyFromQueue(message.queue); + const envCurrentConcurrencyKey = this.keys.envCurrentConcurrencyKeyFromQueue(message.queue); + const queueCurrentDequeuedKey = this.keys.queueCurrentDequeuedKeyFromQueue(message.queue); + const envCurrentDequeuedKey = this.keys.envCurrentDequeuedKeyFromQueue(message.queue); + const envQueueKey = this.keys.envQueueKeyFromQueue(message.queue); + const masterQueueKey = this.keys.masterQueueKeyForEnvironment( + message.environmentId, + this.shardCount + ); + const workerQueueKey = this.keys.workerQueueKey(this.#getWorkerQueueFromMessage(message)); + + this.logger.debug("Calling returnMessageToQueue", { + messageKey, + messageQueue, + masterQueueKey, + workerQueueKey, + messageId, + messageScore: message.timestamp, + service: this.name, + }); + + let result: number; + + if (message.concurrencyKey) { + result = await this.redis.returnMessageToQueueCkTracked( + masterQueueKey, + messageKey, + messageQueue, + queueCurrentConcurrencyKey, + envCurrentConcurrencyKey, + queueCurrentDequeuedKey, + envCurrentDequeuedKey, + envQueueKey, + workerQueueKey, + this.keys.ckIndexKeyFromQueue(message.queue), + this.keys.queueLengthCounterKeyFromQueue(message.queue), + this.keys.queueRunningCounterKeyFromQueue(message.queue), + messageId, + messageQueue, + String(message.timestamp), + messageKey, + this.keys.toCkWildcard(message.queue), + this.options.redis.keyPrefix ?? "", + String(this.counterTtlSeconds) + ); + } else { + result = await this.redis.returnMessageToQueue( + masterQueueKey, + messageKey, + messageQueue, + queueCurrentConcurrencyKey, + envCurrentConcurrencyKey, + queueCurrentDequeuedKey, + envCurrentDequeuedKey, + envQueueKey, + workerQueueKey, + messageId, + messageQueue, + String(message.timestamp), + messageKey + ); + } + + return result === 1; + } + async #callMoveToDeadLetterQueue({ message }: { message: OutputPayload }) { const messageId = message.runId; const messageKey = this.keys.messageKey(message.orgId, message.runId); @@ -4496,6 +4682,149 @@ end `, }); + /** + * Return an admitted-but-unclaimed message to its queue at its original score. + * + * The LREM is the claim: if it removes nothing then a worker has already popped the + * entry and owns the run, so every other key must be left untouched. + * + * Returns 1 when the message was returned, 0 when it was left alone. + */ + this.redis.defineCommand("returnMessageToQueue", { + numberOfKeys: 9, + lua: ` +-- Keys: +local masterQueueKey = KEYS[1] +local messageKey = KEYS[2] +local messageQueueKey = KEYS[3] +local queueCurrentConcurrencyKey = KEYS[4] +local envCurrentConcurrencyKey = KEYS[5] +local queueCurrentDequeuedKey = KEYS[6] +local envCurrentDequeuedKey = KEYS[7] +local envQueueKey = KEYS[8] +local workerQueueKey = KEYS[9] + +-- Args: +local messageId = ARGV[1] +local messageQueueName = ARGV[2] +local messageScore = tonumber(ARGV[3]) +local messageKeyValue = ARGV[4] + +if redis.call('LREM', workerQueueKey, 0, messageKeyValue) == 0 then + return 0 +end + +-- An orphaned list entry (payload already acknowledged): the LREM above cleaned it up. +if redis.call('EXISTS', messageKey) == 0 then + return 0 +end + +redis.call('SREM', queueCurrentConcurrencyKey, messageId) +redis.call('SREM', envCurrentConcurrencyKey, messageId) +redis.call('SREM', queueCurrentDequeuedKey, messageId) +redis.call('SREM', envCurrentDequeuedKey, messageId) + +redis.call('ZADD', messageQueueKey, messageScore, messageId) +redis.call('ZADD', envQueueKey, messageScore, messageId) + +local earliestMessage = redis.call('ZRANGE', messageQueueKey, 0, 0, 'WITHSCORES') +if #earliestMessage == 0 then + redis.call('ZREM', masterQueueKey, messageQueueName) +else + redis.call('ZADD', masterQueueKey, earliestMessage[2], messageQueueName) +end + +return 1 +`, + }); + + /** + * Tracked CK variant of returnMessageToQueue. Mirrors nackMessageCkTracked's counter + * bookkeeping: SREM currentDequeued may DECR runningCounter, and the ZADD back into the + * variant zset INCRs lengthCounter only when it reported a new entry. + */ + this.redis.defineCommand("returnMessageToQueueCkTracked", { + numberOfKeys: 12, + lua: ` +-- Keys: +local masterQueueKey = KEYS[1] +local messageKey = KEYS[2] +local messageQueueKey = KEYS[3] +local queueCurrentConcurrencyKey = KEYS[4] +local envCurrentConcurrencyKey = KEYS[5] +local queueCurrentDequeuedKey = KEYS[6] +local envCurrentDequeuedKey = KEYS[7] +local envQueueKey = KEYS[8] +local workerQueueKey = KEYS[9] +local ckIndexKey = KEYS[10] +local lengthCounterKey = KEYS[11] +local runningCounterKey = KEYS[12] + +-- Args: +local messageId = ARGV[1] +local messageQueueName = ARGV[2] +local messageScore = tonumber(ARGV[3]) +local messageKeyValue = ARGV[4] +local ckWildcardName = ARGV[5] +local keyPrefix = ARGV[6] +local counterTtl = ARGV[7] + +local function decrFloored(key) + if tonumber(redis.call('GET', key) or '0') > 0 then + redis.call('DECR', key) + end +end + +if redis.call('LREM', workerQueueKey, 0, messageKeyValue) == 0 then + return 0 +end + +-- An orphaned list entry (payload already acknowledged): the LREM above cleaned it up. +if redis.call('EXISTS', messageKey) == 0 then + return 0 +end + +redis.call('SREM', queueCurrentConcurrencyKey, messageId) +redis.call('SREM', envCurrentConcurrencyKey, messageId) +local removedFromDequeued = redis.call('SREM', queueCurrentDequeuedKey, messageId) +redis.call('SREM', envCurrentDequeuedKey, messageId) +if removedFromDequeued == 1 then + decrFloored(runningCounterKey) +end + +if redis.call('EXISTS', lengthCounterKey) == 0 then + local total = 0 + local variants = redis.call('ZRANGE', ckIndexKey, 0, -1) + for _, v in ipairs(variants) do + total = total + tonumber(redis.call('ZCARD', keyPrefix .. v) or '0') + end + redis.call('SET', lengthCounterKey, total, 'EX', counterTtl) +end + +local added = redis.call('ZADD', messageQueueKey, messageScore, messageId) +redis.call('ZADD', envQueueKey, messageScore, messageId) +if added == 1 then + redis.call('INCR', lengthCounterKey) +end + +local earliest = redis.call('ZRANGE', messageQueueKey, 0, 0, 'WITHSCORES') +if #earliest > 0 then + redis.call('ZADD', ckIndexKey, earliest[2], messageQueueName) +end + +local earliestIdx = redis.call('ZRANGE', ckIndexKey, 0, 0, 'WITHSCORES') +if #earliestIdx == 0 then + redis.call('ZREM', masterQueueKey, ckWildcardName) +else + redis.call('ZADD', masterQueueKey, earliestIdx[2], ckWildcardName) +end + +redis.call('ZREM', masterQueueKey, messageQueueName) + +return 1 +`, + }); + this.redis.defineCommand("moveToDeadLetterQueue", { numberOfKeys: 9, lua: ` @@ -5636,6 +5965,46 @@ declare module "@internal/redis" { callback?: Callback ): Result; + returnMessageToQueue( + masterQueueKey: string, + messageKey: string, + messageQueue: string, + queueCurrentConcurrencyKey: string, + envCurrentConcurrencyKey: string, + queueCurrentDequeuedKey: string, + envCurrentDequeuedKey: string, + envQueueKey: string, + workerQueueKey: string, + messageId: string, + messageQueueName: string, + messageScore: string, + messageKeyValue: string, + callback?: Callback + ): Result; + + returnMessageToQueueCkTracked( + masterQueueKey: string, + messageKey: string, + messageQueue: string, + queueCurrentConcurrencyKey: string, + envCurrentConcurrencyKey: string, + queueCurrentDequeuedKey: string, + envCurrentDequeuedKey: string, + envQueueKey: string, + workerQueueKey: string, + ckIndexKey: string, + lengthCounterKey: string, + runningCounterKey: string, + messageId: string, + messageQueueName: string, + messageScore: string, + messageKeyValue: string, + ckWildcardName: string, + keyPrefix: string, + counterTtl: string, + callback?: Callback + ): Result; + nackMessageCkTracked( masterQueueKey: string, messageKey: string, diff --git a/internal-packages/run-engine/src/run-queue/tests/dequeueMessageFromWorkerQueue.test.ts b/internal-packages/run-engine/src/run-queue/tests/dequeueMessageFromWorkerQueue.test.ts index 83d8c7a723f..59fb91719ca 100644 --- a/internal-packages/run-engine/src/run-queue/tests/dequeueMessageFromWorkerQueue.test.ts +++ b/internal-packages/run-engine/src/run-queue/tests/dequeueMessageFromWorkerQueue.test.ts @@ -137,6 +137,59 @@ describe("RunQueue.dequeueMessageFromWorkerQueue", () => { } }); + redisTest( + "tracks a message waiting in the worker queue as concurrent but not yet dequeued", + async ({ redisContainer }) => { + const queue = new RunQueue({ + ...testOptions, + 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(), + }, + }); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageDev, + workerQueue: "main", + }); + + await setTimeout(1000); + + expect(await queue.peekAllOnWorkerQueue("main")).toHaveLength(1); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, messageDev.queue)).toBe( + 1 + ); + expect(await queue.currentDequeuedOfQueue(authenticatedEnvDev, messageDev.queue)).toBe(0); + expect(await queue.operationalCurrentConcurrencyOfEnvironment(authenticatedEnvDev)).toBe(1); + expect(await queue.currentConcurrencyOfEnvironment(authenticatedEnvDev)).toBe(0); + + const dequeued = await queue.dequeueMessageFromWorkerQueue("test_12345", "main"); + assertNonNullable(dequeued); + + expect(await queue.peekAllOnWorkerQueue("main")).toHaveLength(0); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, messageDev.queue)).toBe( + 1 + ); + expect(await queue.currentDequeuedOfQueue(authenticatedEnvDev, messageDev.queue)).toBe(1); + expect(await queue.operationalCurrentConcurrencyOfEnvironment(authenticatedEnvDev)).toBe(1); + expect(await queue.currentConcurrencyOfEnvironment(authenticatedEnvDev)).toBe(1); + } finally { + await queue.quit(); + } + } + ); + redisTest( "should not dequeue when env current concurrency equals env concurrency limit", async ({ redisContainer }) => { diff --git a/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts b/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts new file mode 100644 index 00000000000..01638e10bd1 --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts @@ -0,0 +1,454 @@ +import { assertNonNullable, redisTest } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { describe } from "node:test"; +import { FairQueueSelectionStrategy } from "../fairQueueSelectionStrategy.js"; +import { RunQueue } 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" }, +}; + +const otherEnv = { + id: "e5678", + type: "DEVELOPMENT" as const, + maximumConcurrencyLimit: 10, + concurrencyLimitBurstFactor: new Decimal(2.0), + project: { id: "p1234" }, + organization: { id: "o1234" }, +}; + +const baseTimestamp = 1_745_000_000_000; + +function messageFor( + overrides: Partial & Pick, + env = authenticatedEnvDev +): InputPayload { + return { + taskIdentifier: "task/my-task", + orgId: env.organization.id, + projectId: env.project.id, + environmentId: env.id, + environmentType: "DEVELOPMENT", + queue: "task/my-task", + timestamp: baseTimestamp, + attempt: 0, + ...overrides, + }; +} + +function createRunQueue(redisContainer: { getHost(): string; getPort(): number }) { + const redisOptions = { + keyPrefix: "runqueue:test:", + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }; + + return new RunQueue({ + ...testOptions, + queueSelectionStrategy: new FairQueueSelectionStrategy({ + redis: redisOptions, + keys: testOptions.keys, + }), + redis: redisOptions, + }); +} + +vi.setConfig({ testTimeout: 60_000 }); + +describe("RunQueue.returnUnclaimedMessagesToQueue", () => { + redisTest( + "returns worker queue messages to the pending queue and releases their concurrency", + async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + for (const runId of ["r1", "r2", "r3"]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + } + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(3); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(3); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + }); + + expect(result).toEqual({ returned: 3, skipped: 0 }); + + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(0); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(3); + expect(await queue.lengthOfEnvQueue(authenticatedEnvDev)).toBe(3); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + expect(await queue.operationalCurrentConcurrencyOfEnvironment(authenticatedEnvDev)).toBe(0); + expect(await queue.currentDequeuedOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + } finally { + await queue.quit(); + } + } + ); + + redisTest( + "restores the original queue score so returned runs keep their place in line", + async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1", timestamp: baseTimestamp }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r2", timestamp: baseTimestamp + 1_000 }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r3", timestamp: baseTimestamp + 2_000 }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + }); + + expect(await queue.oldestMessageInQueue(authenticatedEnvDev, "task/my-task")).toBe( + baseTimestamp + ); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + const first = await queue.dequeueMessageFromWorkerQueue("consumer", authenticatedEnvDev.id); + const second = await queue.dequeueMessageFromWorkerQueue( + "consumer", + authenticatedEnvDev.id + ); + const third = await queue.dequeueMessageFromWorkerQueue("consumer", authenticatedEnvDev.id); + + assertNonNullable(first); + assertNonNullable(second); + assertNonNullable(third); + + expect([first.messageId, second.messageId, third.messageId]).toEqual(["r1", "r2", "r3"]); + } finally { + await queue.quit(); + } + } + ); + + redisTest("does not return a run a worker has already claimed", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + for (const runId of ["r1", "r2"]) { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + } + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + const claimed = await queue.dequeueMessageFromWorkerQueue("consumer", authenticatedEnvDev.id); + assertNonNullable(claimed); + expect(claimed.messageId).toBe("r1"); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + }); + + expect(result.returned).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.oldestMessageInQueue(authenticatedEnvDev, "task/my-task")).toBe( + baseTimestamp + ); + + expect(await queue.currentDequeuedOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.getCurrentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toEqual( + ["r1"] + ); + } finally { + await queue.quit(); + } + }); + + redisTest("returns runs across every queue in the environment", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1", queue: "task/queue-a" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r2", queue: "task/queue-b" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(2); + + const result = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); + + expect(result.returned).toBe(2); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(0); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/queue-a")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/queue-b")).toBe(1); + expect(await queue.operationalCurrentConcurrencyOfEnvironment(authenticatedEnvDev)).toBe(0); + } finally { + await queue.quit(); + } + }); + + redisTest("leaves other queues in the environment untouched", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1", queue: "task/paused" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r2", queue: "task/running" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/paused", + }); + + expect(result.returned).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/paused")).toBe(1); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/paused")).toBe(0); + + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/running")).toBe(0); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/running")).toBe(1); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(1); + } finally { + await queue.quit(); + } + }); + + redisTest("leaves other environments untouched", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + await queue.enqueueMessage({ + env: otherEnv, + message: messageFor({ runId: "r2" }, otherEnv), + workerQueue: otherEnv.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + await queue.processMasterQueueForEnvironment(otherEnv.id, 10); + + await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); + + expect(await queue.operationalCurrentConcurrencyOfEnvironment(authenticatedEnvDev)).toBe(0); + expect(await queue.operationalCurrentConcurrencyOfEnvironment(otherEnv)).toBe(1); + expect(await queue.peekAllOnWorkerQueue(otherEnv.id)).toHaveLength(1); + } finally { + await queue.quit(); + } + }); + + redisTest( + "returns runs that skipped the pending queue via the fast path", + async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + enableFastPath: true, + }); + + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + }); + + expect(result.returned).toBe(1); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(0); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + expect(await queue.oldestMessageInQueue(authenticatedEnvDev, "task/my-task")).toBe( + baseTimestamp + ); + } finally { + await queue.quit(); + } + } + ); + + redisTest("returns runs on concurrency key queues", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1", concurrencyKey: "ck-a" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r2", concurrencyKey: "ck-b" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(2); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + }); + + expect(result.returned).toBe(2); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(0); + + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task", "ck-a")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task", "ck-b")).toBe(1); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(2); + + expect( + await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task", "ck-a") + ).toBe(0); + expect( + await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task", "ck-b") + ).toBe(0); + expect(await queue.operationalCurrentConcurrencyOfEnvironment(authenticatedEnvDev)).toBe(0); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(2); + } finally { + await queue.quit(); + } + }); + + redisTest( + "leaves runs that have not reached the worker queue alone", + async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + const empty = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); + expect(empty).toEqual({ returned: 0, skipped: 0 }); + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + const pendingOnly = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + }); + expect(pendingOnly).toEqual({ returned: 0, skipped: 0 }); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.oldestMessageInQueue(authenticatedEnvDev, "task/my-task")).toBe( + baseTimestamp + ); + } finally { + await queue.quit(); + } + } + ); + + redisTest("is idempotent when run twice", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + const first = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); + expect(first.returned).toBe(1); + + const second = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); + expect(second).toEqual({ returned: 0, skipped: 0 }); + + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfEnvQueue(authenticatedEnvDev)).toBe(1); + } finally { + await queue.quit(); + } + }); +}); From 0a4aa5860d530185b4b49de5c61eb5c30d120338 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Mon, 27 Jul 2026 18:09:45 +0100 Subject: [PATCH 2/3] fix(run-engine,webapp): close gaps found in adversarial review of the pause sweep MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Restore the TTL registration when returning a run. Admission removes the run's member from the TTL sorted set, so returning it to the pending queue without re-adding left a run that could never expire: on the next admission the dequeue script takes its already-expired branch, which deliberately leaves cleanup to the TTL consumer, and that consumer no longer has an entry to find. The per-run expireRun job is not a backstop because it is only armed for development environments. Take the pause sweep out of the rollback window. It ran inside the try that compensates a failed pause by clearing the DB flag, but by then the Redis limit was already 0, so a sweep failure left the environment with no concurrency and no paused flag — dark, and unrecoverable without a manual pause/resume cycle. For the billing-limit path it also aborted the remaining environments in the org. The pause is durable and the sweep is idempotent, so it now runs after the pause has been committed and a failure is logged rather than propagated. That also stops the queue pause reporting failure for a pause that took effect. Retry passes while runs remain unclaimed. Admission is not atomic with the push onto the worker queue: the dequeue script claims the concurrency slot and a separate round trip does the RPUSH. A run caught in that gap fails the claim and would previously have been abandoned, then executed during the pause. Claim from both the resolver-mapped and the raw worker queue key. The trigger- time fast path pushes to the raw workerQueue on the message while the master queue consumer pushes to the resolver-mapped one, so with a region override active the two producers disagree and a single-key LREM missed fast-pathed runs. Batch the sweep. It read every payload in one MGET and then made one sequential round trip per run, on a request path, with no error isolation — one failure abandoned the remaining runs. Tests: cover the LREM claim directly by emptying the worker queue without marking the run dequeued (the previous test was excluded by the candidate set before ever reaching the script), the TTL round trip, and the concurrency key running counter across a return. --- apps/webapp/app/v3/runQueue.server.ts | 29 +++ ...billingLimitConvergeEnvironments.server.ts | 10 +- .../v3/services/pauseEnvironment.server.ts | 7 +- .../app/v3/services/pauseQueue.server.ts | 7 +- .../billingLimitConvergeEnvironments.test.ts | 38 +-- .../run-engine/src/run-queue/index.ts | 231 ++++++++++++++---- .../returnUnclaimedMessagesToQueue.test.ts | 162 +++++++++++- 7 files changed, 407 insertions(+), 77 deletions(-) diff --git a/apps/webapp/app/v3/runQueue.server.ts b/apps/webapp/app/v3/runQueue.server.ts index fc716f0f1a8..2759252163f 100644 --- a/apps/webapp/app/v3/runQueue.server.ts +++ b/apps/webapp/app/v3/runQueue.server.ts @@ -1,4 +1,5 @@ import { type AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { logger } from "~/services/logger.server"; import { engine } from "./runEngine.server"; /** Updates the RunQueue env concurrency limits */ @@ -46,3 +47,31 @@ export async function returnUnclaimedMessagesToQueue({ }) { return engine.returnUnclaimedMessagesToQueue({ environment, queue }); } + +/** + * Best-effort sweep for the pause paths. + * + * By the time this runs the pause is already durable and in force, so a failure here must + * not roll it back or surface as a failed pause — the worst case is that some already-admitted + * runs still execute, which is the behaviour that existed before the sweep. + */ +export async function sweepUnclaimedRuns( + environment: AuthenticatedEnvironment, + queue?: string +): Promise { + try { + const result = await returnUnclaimedMessagesToQueue({ environment, queue }); + + logger.debug("sweepUnclaimedRuns", { + environmentId: environment.id, + queue, + ...result, + }); + } catch (error) { + logger.error("sweepUnclaimedRuns failed", { + environmentId: environment.id, + queue, + error, + }); + } +} diff --git a/apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts b/apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts index abde4a0f1c1..1066943e428 100644 --- a/apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts +++ b/apps/webapp/app/v3/services/billingLimit/billingLimitConvergeEnvironments.server.ts @@ -190,7 +190,6 @@ async function pauseEnvironmentForBillingLimit( try { await updateConcurrency(updated, 0); - await returnUnclaimed(updated); } catch (error) { await db.runtimeEnvironment.update({ where: { id: environment.id }, @@ -201,6 +200,15 @@ async function pauseEnvironmentForBillingLimit( // The env's paused state changed (or was rolled back); drop any cached copy either way. controlPlaneResolver.invalidateEnvironment(environment.id); } + + try { + await returnUnclaimed(updated); + } catch (error) { + logger.error("Billing limit converge failed to return unclaimed runs", { + environmentId: environment.id, + error, + }); + } } async function resumeEnvironmentFromBillingLimit( diff --git a/apps/webapp/app/v3/services/pauseEnvironment.server.ts b/apps/webapp/app/v3/services/pauseEnvironment.server.ts index ab10e08ca0a..96440d7f144 100644 --- a/apps/webapp/app/v3/services/pauseEnvironment.server.ts +++ b/apps/webapp/app/v3/services/pauseEnvironment.server.ts @@ -2,7 +2,7 @@ import { EnvironmentPauseSource, type PrismaClientOrTransaction } from "@trigger import { prisma } from "~/db.server"; import { logger } from "~/services/logger.server"; import { getManualPauseEnvironmentResult } from "~/v3/services/billingLimit/manualPauseEnvironmentGuard.server"; -import { returnUnclaimedMessagesToQueue, updateEnvConcurrencyLimits } from "../runQueue.server"; +import { sweepUnclaimedRuns, updateEnvConcurrencyLimits } from "../runQueue.server"; import { WithRunEngine } from "./baseService.server"; import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; import { controlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; @@ -114,7 +114,6 @@ export class PauseEnvironmentService extends WithRunEngine { environmentId: environment.id, }); await updateEnvConcurrencyLimits(environment, 0); - await returnUnclaimedMessagesToQueue({ environment }); } else { logger.debug("PauseEnvironmentService: resuming environment", { environmentId: environment.id, @@ -137,6 +136,10 @@ export class PauseEnvironmentService extends WithRunEngine { // The env's `paused` state changed in the control-plane; drop any cached copy. controlPlaneResolver.invalidateEnvironment(environment.id); + if (action === "paused") { + await sweepUnclaimedRuns(environment); + } + return { success: true, state: action, diff --git a/apps/webapp/app/v3/services/pauseQueue.server.ts b/apps/webapp/app/v3/services/pauseQueue.server.ts index 19c7abfb3da..ab54d3f2b96 100644 --- a/apps/webapp/app/v3/services/pauseQueue.server.ts +++ b/apps/webapp/app/v3/services/pauseQueue.server.ts @@ -6,7 +6,7 @@ import { BaseService } from "./baseService.server"; import { determineEngineVersion } from "../engineVersion.server"; import { removeQueueConcurrencyLimits, - returnUnclaimedMessagesToQueue, + sweepUnclaimedRuns, updateQueueConcurrencyLimits, } from "../runQueue.server"; import { engine } from "../runEngine.server"; @@ -63,7 +63,6 @@ export class PauseQueueService extends BaseService { if (action === "paused") { await updateQueueConcurrencyLimits(environment, queue.name, 0); - await returnUnclaimedMessagesToQueue({ environment, queue: queue.name }); } else { if (queue.concurrencyLimit) { await updateQueueConcurrencyLimits(environment, queue.name, queue.concurrencyLimit); @@ -78,6 +77,10 @@ export class PauseQueueService extends BaseService { environmentId: environment.id, }); + if (action === "paused") { + await sweepUnclaimedRuns(environment, queue.name); + } + const results = await Promise.all([ engine.lengthOfQueues(environment, [queue.name]), engine.currentConcurrencyOfQueues(environment, [queue.name]), diff --git a/apps/webapp/test/billingLimitConvergeEnvironments.test.ts b/apps/webapp/test/billingLimitConvergeEnvironments.test.ts index f1ec8dff2be..62f2640b349 100644 --- a/apps/webapp/test/billingLimitConvergeEnvironments.test.ts +++ b/apps/webapp/test/billingLimitConvergeEnvironments.test.ts @@ -95,7 +95,7 @@ describe("convergeBillingLimitEnvironmentsForOrg", () => { expect(returnUnclaimed).not.toHaveBeenCalled(); }); - postgresTest("rolls back pause when returning unclaimed runs fails", async ({ prisma }) => { + postgresTest("keeps the pause when returning unclaimed runs fails", async ({ prisma }) => { const { organization, project } = await createTestOrgProjectWithMember(prisma); const environment = await createRuntimeEnvironment(prisma, { projectId: project.id, @@ -103,22 +103,30 @@ describe("convergeBillingLimitEnvironmentsForOrg", () => { type: "PRODUCTION", slug: uniqueId("prod"), }); + const second = await createRuntimeEnvironment(prisma, { + projectId: project.id, + organizationId: organization.id, + type: "PRODUCTION", + slug: uniqueId("prod"), + }); - await expect( - convergeBillingLimitEnvironmentsForOrg(organization.id, "grace", { - prismaClient: prisma, - updateConcurrency: async () => undefined, - returnUnclaimed: async () => { - throw new Error("run queue unavailable"); - }, - }) - ).rejects.toThrow("run queue unavailable"); - - const envAfter = await prisma.runtimeEnvironment.findUniqueOrThrow({ - where: { id: environment.id }, + const result = await convergeBillingLimitEnvironmentsForOrg(organization.id, "grace", { + prismaClient: prisma, + updateConcurrency: async () => undefined, + returnUnclaimed: async () => { + throw new Error("run queue unavailable"); + }, }); - expect(envAfter.paused).toBe(false); - expect(envAfter.pauseSource).toBeNull(); + + expect(result).toEqual({ paused: 2, unpaused: 0 }); + + for (const env of [environment, second]) { + const envAfter = await prisma.runtimeEnvironment.findUniqueOrThrow({ + where: { id: env.id }, + }); + expect(envAfter.paused).toBe(true); + expect(envAfter.pauseSource).toBe(EnvironmentPauseSource.BILLING_LIMIT); + } }); postgresTest("rolls back pause when concurrency update fails", async ({ prisma }) => { diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index ad951ed770c..3bd757b2a03 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -137,6 +137,16 @@ type MarkedRun = { score: number; }; +const RETURN_UNCLAIMED_BATCH_SIZE = 50; +const READ_MESSAGES_BATCH_SIZE = 200; + +export type ReturnUnclaimedMessagesResult = { + returned: number; + skipped: number; + errors: number; + passes: number; +}; + const defaultRetrySettings = { maxAttempts: 12, factor: 2, @@ -1046,64 +1056,68 @@ export class RunQueue { * re-enter at the back, so if the queue can still admit work then a newer run can overtake * one that is on its way back. * + * Admission is not atomic with the push onto the worker queue: the dequeue script claims + * the concurrency slot and a separate round trip does the RPUSH. A run caught in that gap + * is a candidate whose claim fails, so a pass reporting skipped runs is retried — by then + * the run has either landed on the worker queue (and is claimable) or been picked up by a + * worker (and has left the candidate set). + * * @param queue - Restrict to a single queue (including its concurrency key variants). * Omit to cover every queue in the environment. */ public async returnUnclaimedMessagesToQueue({ env, queue, + maxPasses = 3, + passDelayMs = 250, }: { env: MinimalAuthenticatedEnvironment; queue?: string; - }): Promise<{ returned: number; skipped: number }> { + maxPasses?: number; + passDelayMs?: number; + }): Promise { return this.#trace( "returnUnclaimedMessagesToQueue", async (span) => { - const unclaimedRunIds = await this.redis.sdiff( - this.keys.envCurrentConcurrencyKey(env), - this.keys.envCurrentDequeuedKey(env) - ); - - span.setAttribute("unclaimed_count", unclaimedRunIds.length); - - if (unclaimedRunIds.length === 0) { - return { returned: 0, skipped: 0 }; - } - const targetBaseQueueKey = queue ? this.keys.queueKey(env, queue) : undefined; - const messages = await this.#readMessages(env.organization.id, unclaimedRunIds); let returned = 0; let skipped = 0; + let errors = 0; + let passes = 0; - for (const message of messages) { - if ( - targetBaseQueueKey && - this.keys.baseQueueKeyFromQueue(message.queue) !== targetBaseQueueKey - ) { - continue; - } + for (let pass = 0; pass < Math.max(1, maxPasses); pass++) { + passes++; - if (await this.#callReturnMessageToQueue(message)) { - returned++; - } else { - skipped++; + const passResult = await this.#returnUnclaimedMessagesPass(env, targetBaseQueueKey); + + returned += passResult.returned; + errors += passResult.errors; + skipped = passResult.skipped; + + if (passResult.skipped === 0 || pass === Math.max(1, maxPasses) - 1) { + break; } + + await setTimeout(passDelayMs); } span.setAttribute("returned_count", returned); span.setAttribute("skipped_count", skipped); + span.setAttribute("error_count", errors); + span.setAttribute("passes", passes); this.logger.info("returnUnclaimedMessagesToQueue", { service: this.name, environmentId: env.id, queue, - candidates: unclaimedRunIds.length, returned, skipped, + errors, + passes, }); - return { returned, skipped }; + return { returned, skipped, errors, passes }; }, { kind: SpanKind.INTERNAL, @@ -1115,30 +1129,86 @@ export class RunQueue { ); } - async #readMessages(orgId: string, runIds: string[]): Promise { - const rawMessages = await this.redis.mget( - runIds.map((runId) => this.keys.messageKey(orgId, runId)) + async #returnUnclaimedMessagesPass( + env: MinimalAuthenticatedEnvironment, + targetBaseQueueKey: string | undefined + ): Promise<{ returned: number; skipped: number; errors: number }> { + const unclaimedRunIds = await this.redis.sdiff( + this.keys.envCurrentConcurrencyKey(env), + this.keys.envCurrentDequeuedKey(env) ); - const messages: OutputPayload[] = []; + if (unclaimedRunIds.length === 0) { + return { returned: 0, skipped: 0, errors: 0 }; + } - for (const rawMessage of rawMessages) { - if (!rawMessage) { - continue; - } + const messages = (await this.#readMessages(env.organization.id, unclaimedRunIds)).filter( + (message) => + !targetBaseQueueKey || this.keys.baseQueueKeyFromQueue(message.queue) === targetBaseQueueKey + ); - const [error, message] = parseRawMessage(rawMessage); + let returned = 0; + let skipped = 0; + let errors = 0; - if (error) { - this.logger.error(`[${this.name}] Failed to parse message`, { - error, - service: this.name, - message: message ?? rawMessage, - }); + for (let i = 0; i < messages.length; i += RETURN_UNCLAIMED_BATCH_SIZE) { + const batch = messages.slice(i, i + RETURN_UNCLAIMED_BATCH_SIZE); + + const results = await Promise.allSettled( + batch.map((message) => this.#callReturnMessageToQueue(message)) + ); + + for (const [index, result] of results.entries()) { + if (result.status === "rejected") { + errors++; + this.logger.error("returnUnclaimedMessagesToQueue failed for a run", { + service: this.name, + environmentId: env.id, + runId: batch[index]?.runId, + error: result.reason, + }); + continue; + } + + if (result.value) { + returned++; + } else { + skipped++; + } } + } - if (message) { - messages.push(message); + return { returned, skipped, errors }; + } + + async #readMessages(orgId: string, runIds: string[]): Promise { + const messages: OutputPayload[] = []; + + for (let i = 0; i < runIds.length; i += READ_MESSAGES_BATCH_SIZE) { + const rawMessages = await this.redis.mget( + runIds + .slice(i, i + READ_MESSAGES_BATCH_SIZE) + .map((runId) => this.keys.messageKey(orgId, runId)) + ); + + for (const rawMessage of rawMessages) { + if (!rawMessage) { + continue; + } + + const [error, message] = parseRawMessage(rawMessage); + + if (error) { + this.logger.error(`[${this.name}] Failed to parse message`, { + error, + service: this.name, + message: message ?? rawMessage, + }); + } + + if (message) { + messages.push(message); + } } } @@ -2753,12 +2823,21 @@ export class RunQueue { this.shardCount ); const workerQueueKey = this.keys.workerQueueKey(this.#getWorkerQueueFromMessage(message)); + const rawWorkerQueueKey = + message.version === "2" ? this.keys.workerQueueKey(message.workerQueue) : workerQueueKey; + + const returnTtl = Boolean(message.ttlExpiresAt) && Boolean(this.options.ttlSystem); + const ttlQueueKey = this.keys.ttlQueueKeyForShard(this.#getTtlShardForQueue(message.queue)); + const ttlMember = `${message.queue}|${message.runId}|${message.orgId}`; + const ttlScore = returnTtl ? String(message.ttlExpiresAt) : "0"; this.logger.debug("Calling returnMessageToQueue", { messageKey, messageQueue, masterQueueKey, workerQueueKey, + rawWorkerQueueKey, + ttlQueueKey, messageId, messageScore: message.timestamp, service: this.name, @@ -2777,6 +2856,8 @@ export class RunQueue { envCurrentDequeuedKey, envQueueKey, workerQueueKey, + rawWorkerQueueKey, + ttlQueueKey, this.keys.ckIndexKeyFromQueue(message.queue), this.keys.queueLengthCounterKeyFromQueue(message.queue), this.keys.queueRunningCounterKeyFromQueue(message.queue), @@ -2786,7 +2867,9 @@ export class RunQueue { messageKey, this.keys.toCkWildcard(message.queue), this.options.redis.keyPrefix ?? "", - String(this.counterTtlSeconds) + String(this.counterTtlSeconds), + ttlMember, + ttlScore ); } else { result = await this.redis.returnMessageToQueue( @@ -2799,10 +2882,14 @@ export class RunQueue { envCurrentDequeuedKey, envQueueKey, workerQueueKey, + rawWorkerQueueKey, + ttlQueueKey, messageId, messageQueue, String(message.timestamp), - messageKey + messageKey, + ttlMember, + ttlScore ); } @@ -4691,7 +4778,7 @@ end * Returns 1 when the message was returned, 0 when it was left alone. */ this.redis.defineCommand("returnMessageToQueue", { - numberOfKeys: 9, + numberOfKeys: 11, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -4703,14 +4790,25 @@ local queueCurrentDequeuedKey = KEYS[6] local envCurrentDequeuedKey = KEYS[7] local envQueueKey = KEYS[8] local workerQueueKey = KEYS[9] +local rawWorkerQueueKey = KEYS[10] +local ttlQueueKey = KEYS[11] -- Args: local messageId = ARGV[1] local messageQueueName = ARGV[2] local messageScore = tonumber(ARGV[3]) local messageKeyValue = ARGV[4] +local ttlMember = ARGV[5] +local ttlScore = tonumber(ARGV[6]) -if redis.call('LREM', workerQueueKey, 0, messageKeyValue) == 0 then +-- The resolver-mapped key and the raw one can differ while a region override is active, +-- and the two producers disagree on which they push to. +local claimed = redis.call('LREM', workerQueueKey, 0, messageKeyValue) +if rawWorkerQueueKey ~= workerQueueKey then + claimed = claimed + redis.call('LREM', rawWorkerQueueKey, 0, messageKeyValue) +end + +if claimed == 0 then return 0 end @@ -4727,6 +4825,11 @@ redis.call('SREM', envCurrentDequeuedKey, messageId) redis.call('ZADD', messageQueueKey, messageScore, messageId) redis.call('ZADD', envQueueKey, messageScore, messageId) +-- Admission removed the TTL member; restore it so the run can still expire while queued. +if ttlScore > 0 then + redis.call('ZADD', ttlQueueKey, ttlScore, ttlMember) +end + local earliestMessage = redis.call('ZRANGE', messageQueueKey, 0, 0, 'WITHSCORES') if #earliestMessage == 0 then redis.call('ZREM', masterQueueKey, messageQueueName) @@ -4744,7 +4847,7 @@ return 1 * variant zset INCRs lengthCounter only when it reported a new entry. */ this.redis.defineCommand("returnMessageToQueueCkTracked", { - numberOfKeys: 12, + numberOfKeys: 14, lua: ` -- Keys: local masterQueueKey = KEYS[1] @@ -4756,9 +4859,11 @@ local queueCurrentDequeuedKey = KEYS[6] local envCurrentDequeuedKey = KEYS[7] local envQueueKey = KEYS[8] local workerQueueKey = KEYS[9] -local ckIndexKey = KEYS[10] -local lengthCounterKey = KEYS[11] -local runningCounterKey = KEYS[12] +local rawWorkerQueueKey = KEYS[10] +local ttlQueueKey = KEYS[11] +local ckIndexKey = KEYS[12] +local lengthCounterKey = KEYS[13] +local runningCounterKey = KEYS[14] -- Args: local messageId = ARGV[1] @@ -4768,6 +4873,8 @@ local messageKeyValue = ARGV[4] local ckWildcardName = ARGV[5] local keyPrefix = ARGV[6] local counterTtl = ARGV[7] +local ttlMember = ARGV[8] +local ttlScore = tonumber(ARGV[9]) local function decrFloored(key) if tonumber(redis.call('GET', key) or '0') > 0 then @@ -4775,7 +4882,14 @@ local function decrFloored(key) end end -if redis.call('LREM', workerQueueKey, 0, messageKeyValue) == 0 then +-- The resolver-mapped key and the raw one can differ while a region override is active, +-- and the two producers disagree on which they push to. +local claimed = redis.call('LREM', workerQueueKey, 0, messageKeyValue) +if rawWorkerQueueKey ~= workerQueueKey then + claimed = claimed + redis.call('LREM', rawWorkerQueueKey, 0, messageKeyValue) +end + +if claimed == 0 then return 0 end @@ -4807,6 +4921,11 @@ if added == 1 then redis.call('INCR', lengthCounterKey) end +-- Admission removed the TTL member; restore it so the run can still expire while queued. +if ttlScore > 0 then + redis.call('ZADD', ttlQueueKey, ttlScore, ttlMember) +end + local earliest = redis.call('ZRANGE', messageQueueKey, 0, 0, 'WITHSCORES') if #earliest > 0 then redis.call('ZADD', ckIndexKey, earliest[2], messageQueueName) @@ -5975,10 +6094,14 @@ declare module "@internal/redis" { envCurrentDequeuedKey: string, envQueueKey: string, workerQueueKey: string, + rawWorkerQueueKey: string, + ttlQueueKey: string, messageId: string, messageQueueName: string, messageScore: string, messageKeyValue: string, + ttlMember: string, + ttlScore: string, callback?: Callback ): Result; @@ -5992,6 +6115,8 @@ declare module "@internal/redis" { envCurrentDequeuedKey: string, envQueueKey: string, workerQueueKey: string, + rawWorkerQueueKey: string, + ttlQueueKey: string, ckIndexKey: string, lengthCounterKey: string, runningCounterKey: string, @@ -6002,6 +6127,8 @@ declare module "@internal/redis" { ckWildcardName: string, keyPrefix: string, counterTtl: string, + ttlMember: string, + ttlScore: string, callback?: Callback ): Result; diff --git a/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts b/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts index 01638e10bd1..9809725f290 100644 --- a/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts +++ b/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts @@ -6,6 +6,7 @@ import { RunQueue } from "../index.js"; import { RunQueueFullKeyProducer } from "../keyProducer.js"; import type { InputPayload } from "../types.js"; import { Decimal } from "@trigger.dev/database"; +import { createRedisClient as createRawRedisClient } from "@internal/redis"; const testOptions = { name: "rq", @@ -59,7 +60,10 @@ function messageFor( }; } -function createRunQueue(redisContainer: { getHost(): string; getPort(): number }) { +function createRunQueue( + redisContainer: { getHost(): string; getPort(): number }, + overrides: Partial[0]> = {} +) { const redisOptions = { keyPrefix: "runqueue:test:", host: redisContainer.getHost(), @@ -73,9 +77,21 @@ function createRunQueue(redisContainer: { getHost(): string; getPort(): number } keys: testOptions.keys, }), redis: redisOptions, + ...overrides, }); } +function createRedisClient(redisContainer: { getHost(): string; getPort(): number }) { + return createRawRedisClient({ + host: redisContainer.getHost(), + port: redisContainer.getPort(), + }); +} + +function queueKeyFor(queue: string) { + return testOptions.keys.queueKey(authenticatedEnvDev, queue); +} + vi.setConfig({ testTimeout: 60_000 }); describe("RunQueue.returnUnclaimedMessagesToQueue", () => { @@ -105,7 +121,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { queue: "task/my-task", }); - expect(result).toEqual({ returned: 3, skipped: 0 }); + expect(result).toEqual({ returned: 3, skipped: 0, errors: 0, passes: 1 }); expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(0); expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(3); @@ -215,6 +231,88 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { } }); + redisTest( + "leaves a run whose worker queue entry is gone but is not yet marked dequeued", + async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + const redis = createRedisClient(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + const popped = await redis.lpop(`runqueue:test:workerQueue:${authenticatedEnvDev.id}`); + expect(popped).not.toBeNull(); + + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.currentDequeuedOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + maxPasses: 1, + }); + + expect(result.returned).toBe(0); + expect(result.skipped).toBe(1); + + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + } finally { + await redis.quit(); + await queue.quit(); + } + } + ); + + redisTest("restores the TTL registration of a returned run", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer, { + ttlSystem: { + shardCount: 1, + consumersDisabled: true, + workerQueueSuffix: "ttlWorker", + workerItemsSuffix: "ttlWorkerItems", + }, + }); + const redis = createRedisClient(redisContainer); + + try { + const ttlExpiresAt = Date.now() + 600_000; + + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1", ttlExpiresAt }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + const ttlKey = "runqueue:test:ttl:shard:0"; + const ttlMember = `${queueKeyFor("task/my-task")}|r1|o1234`; + + expect(await redis.zscore(ttlKey, ttlMember)).toBe(String(ttlExpiresAt)); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + expect(await redis.zscore(ttlKey, ttlMember)).toBeNull(); + + await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + }); + + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await redis.zscore(ttlKey, ttlMember)).toBe(String(ttlExpiresAt)); + } finally { + await redis.quit(); + await queue.quit(); + } + }); + redisTest("returns runs across every queue in the environment", async ({ redisContainer }) => { const queue = createRunQueue(redisContainer); @@ -396,6 +494,60 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { } }); + redisTest( + "keeps the concurrency key running counter correct across a return", + async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1", concurrencyKey: "ck-a" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r2", concurrencyKey: "ck-a" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + const claimed = await queue.dequeueMessageFromWorkerQueue( + "consumer", + authenticatedEnvDev.id + ); + assertNonNullable(claimed); + + expect( + await queue.currentConcurrencyOfQueues(authenticatedEnvDev, ["task/my-task"]) + ).toEqual({ "task/my-task": 1 }); + + await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + }); + + expect( + await queue.currentConcurrencyOfQueues(authenticatedEnvDev, ["task/my-task"]) + ).toEqual({ "task/my-task": 1 }); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, claimed.messageId, { + skipDequeueProcessing: true, + }); + + expect( + await queue.currentConcurrencyOfQueues(authenticatedEnvDev, ["task/my-task"]) + ).toEqual({ "task/my-task": 0 }); + } finally { + await queue.quit(); + } + } + ); + redisTest( "leaves runs that have not reached the worker queue alone", async ({ redisContainer }) => { @@ -403,7 +555,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { try { const empty = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); - expect(empty).toEqual({ returned: 0, skipped: 0 }); + expect(empty).toEqual({ returned: 0, skipped: 0, errors: 0, passes: 1 }); await queue.enqueueMessage({ env: authenticatedEnvDev, @@ -415,7 +567,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { const pendingOnly = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev, }); - expect(pendingOnly).toEqual({ returned: 0, skipped: 0 }); + expect(pendingOnly).toEqual({ returned: 0, skipped: 0, errors: 0, passes: 1 }); expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); expect(await queue.oldestMessageInQueue(authenticatedEnvDev, "task/my-task")).toBe( baseTimestamp @@ -443,7 +595,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { expect(first.returned).toBe(1); const second = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); - expect(second).toEqual({ returned: 0, skipped: 0 }); + expect(second).toEqual({ returned: 0, skipped: 0, errors: 0, passes: 1 }); expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); expect(await queue.lengthOfEnvQueue(authenticatedEnvDev)).toBe(1); From b717937e842373adbd0a73f7e5c608522e7d7437 Mon Sep 17 00:00:00 2001 From: Matt Aitken Date: Mon, 27 Jul 2026 21:00:55 +0100 Subject: [PATCH 3/3] fix(run-engine,webapp): harden the pause sweep retry and report it honestly MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Guard the pass count against non-finite input. Math.max(1, NaN) is NaN, so a NaN maxPasses skipped the loop entirely and reported a successful no-op sweep, and Infinity spun at one pass per 50ms inside an awaited request. The count is now clamped to a ceiling and falls back to the default when it is not finite. Retry a pass that errored, not just one that skipped. Rejected script calls never increment skipped, so a transient Redis failure part-way through a batch ended the sweep with those runs still admitted and still on the worker queue — the exact outcome the sweep exists to prevent — while the result read as a success with a non-zero error count nobody was gated on. Rename the residual count to skippedLastPass. It is the final pass's value rather than a total, and it conflates a run a worker claimed first (which escaped the pause and is executing) with a leaked concurrency slot that has no worker queue entry behind it. The name and its doc comment now say so instead of implying it partitions the candidate set alongside returned. Tests: pin the non-finite pass-count guard, cover that an acknowledged run is never resurrected by a later sweep, and stop the worker queue override test leaking its environment variable if the queue constructor throws. --- apps/webapp/test/pauseSweepWiring.test.ts | 94 +++++++++++ apps/webapp/test/sweepUnclaimedRuns.test.ts | 39 +++++ .../run-engine/src/engine/index.ts | 4 +- .../run-engine/src/run-queue/index.ts | 72 ++++++--- .../returnUnclaimedMessagesToQueue.test.ts | 147 +++++++++++++++++- 5 files changed, 328 insertions(+), 28 deletions(-) create mode 100644 apps/webapp/test/pauseSweepWiring.test.ts create mode 100644 apps/webapp/test/sweepUnclaimedRuns.test.ts diff --git a/apps/webapp/test/pauseSweepWiring.test.ts b/apps/webapp/test/pauseSweepWiring.test.ts new file mode 100644 index 00000000000..72924953cc6 --- /dev/null +++ b/apps/webapp/test/pauseSweepWiring.test.ts @@ -0,0 +1,94 @@ +import { type PrismaClient } from "@trigger.dev/database"; +import { describe, expect, vi } from "vitest"; +import { postgresTest } from "@internal/testcontainers"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; +import { + createRuntimeEnvironment, + createTestOrgProjectWithMember, + uniqueId, +} from "./fixtures/environmentVariablesFixtures"; + +vi.setConfig({ testTimeout: 60_000 }); + +const calls: string[] = []; + +vi.mock("~/v3/runQueue.server", () => ({ + updateEnvConcurrencyLimits: vi.fn(async (_env: unknown, limit?: number) => { + calls.push(`updateEnvConcurrencyLimits:${limit}`); + }), + updateQueueConcurrencyLimits: vi.fn(async (_env: unknown, name: string, limit: number) => { + calls.push(`updateQueueConcurrencyLimits:${name}:${limit}`); + }), + removeQueueConcurrencyLimits: vi.fn(async () => { + calls.push("removeQueueConcurrencyLimits"); + }), + returnUnclaimedMessagesToQueue: vi.fn(async () => ({ + returned: 0, + skipped: 0, + errors: 0, + passes: 1, + })), + sweepUnclaimedRuns: vi.fn(async (_env: unknown, queue?: string) => { + calls.push(`sweepUnclaimedRuns:${queue ?? "*"}`); + }), +})); + +async function loadServices() { + const [{ PauseEnvironmentService }, { authIncludeBase, toAuthenticated }] = await Promise.all([ + import("~/v3/services/pauseEnvironment.server"), + import("~/models/runtimeEnvironment.server"), + ]); + return { PauseEnvironmentService, authIncludeBase, toAuthenticated }; +} + +type Loaded = Awaited>; + +async function seedEnv( + loaded: Loaded, + prisma: PrismaClient +): Promise<{ environment: AuthenticatedEnvironment; environmentId: string }> { + const { organization, project } = await createTestOrgProjectWithMember(prisma); + const created = await createRuntimeEnvironment(prisma, { + projectId: project.id, + organizationId: organization.id, + type: "PRODUCTION", + slug: uniqueId("prod"), + }); + + const row = await prisma.runtimeEnvironment.findFirstOrThrow({ + where: { id: created.id }, + include: loaded.authIncludeBase, + }); + + return { environment: loaded.toAuthenticated(row), environmentId: created.id }; +} + +describe("pause sweep wiring", () => { + postgresTest("pausing an environment sweeps after the limit is zeroed", async ({ prisma }) => { + calls.length = 0; + const loaded = await loadServices(); + const { environment } = await seedEnv(loaded, prisma); + + const result = await new loaded.PauseEnvironmentService(prisma).call(environment, "paused"); + + expect(result).toEqual({ success: true, state: "paused" }); + expect(calls).toEqual(["updateEnvConcurrencyLimits:0", "sweepUnclaimedRuns:*"]); + }); + + postgresTest("resuming an environment does not sweep", async ({ prisma }) => { + const loaded = await loadServices(); + const { environment, environmentId } = await seedEnv(loaded, prisma); + + await prisma.runtimeEnvironment.update({ + where: { id: environmentId }, + data: { paused: true }, + }); + + calls.length = 0; + + const result = await new loaded.PauseEnvironmentService(prisma).call(environment, "resumed"); + + expect(result).toEqual({ success: true, state: "resumed" }); + expect(calls).toEqual(["updateEnvConcurrencyLimits:undefined"]); + }); +}); diff --git a/apps/webapp/test/sweepUnclaimedRuns.test.ts b/apps/webapp/test/sweepUnclaimedRuns.test.ts new file mode 100644 index 00000000000..de8043d53f9 --- /dev/null +++ b/apps/webapp/test/sweepUnclaimedRuns.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it, vi } from "vitest"; +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; + +const returnUnclaimedMessagesToQueue = vi.fn(); + +vi.mock("~/v3/runEngine.server", () => ({ + engine: { returnUnclaimedMessagesToQueue }, +})); + +const environment = { id: "env_1234" } as AuthenticatedEnvironment; + +describe("sweepUnclaimedRuns", () => { + it("swallows a failing sweep so a pause that is already in force is not reported as failed", async () => { + const { sweepUnclaimedRuns } = await import("~/v3/runQueue.server"); + + returnUnclaimedMessagesToQueue.mockRejectedValueOnce(new Error("run queue unavailable")); + + await expect(sweepUnclaimedRuns(environment)).resolves.toBeUndefined(); + expect(returnUnclaimedMessagesToQueue).toHaveBeenCalledWith({ environment, queue: undefined }); + }); + + it("passes the queue through when one is targeted", async () => { + const { sweepUnclaimedRuns } = await import("~/v3/runQueue.server"); + + returnUnclaimedMessagesToQueue.mockResolvedValueOnce({ + returned: 2, + skipped: 0, + errors: 0, + passes: 1, + }); + + await sweepUnclaimedRuns(environment, "task/my-task"); + + expect(returnUnclaimedMessagesToQueue).toHaveBeenCalledWith({ + environment, + queue: "task/my-task", + }); + }); +}); diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 3fe2d325033..342f8f90d00 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -49,7 +49,7 @@ import type { BatchCompletionCallback, } from "../batch-queue/types.js"; import { FairQueueSelectionStrategy } from "../run-queue/fairQueueSelectionStrategy.js"; -import { RunQueue } from "../run-queue/index.js"; +import { RunQueue, type ReturnUnclaimedMessagesResult } from "../run-queue/index.js"; import { RunQueueFullKeyProducer } from "../run-queue/keyProducer.js"; import type { AuthenticatedEnvironment, MinimalAuthenticatedEnvironment } from "../shared/index.js"; import { BillingCache } from "./billingCache.js"; @@ -1648,7 +1648,7 @@ export class RunEngine { }: { environment: MinimalAuthenticatedEnvironment; queue?: string; - }): Promise<{ returned: number; skipped: number }> { + }): Promise { return this.runQueue.returnUnclaimedMessagesToQueue({ env: environment, queue }); } diff --git a/internal-packages/run-engine/src/run-queue/index.ts b/internal-packages/run-engine/src/run-queue/index.ts index 3bd757b2a03..448d74b800f 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -140,9 +140,25 @@ type MarkedRun = { const RETURN_UNCLAIMED_BATCH_SIZE = 50; const READ_MESSAGES_BATCH_SIZE = 200; +/** + * One retry covers the admission-to-RPUSH gap, which is a single round trip. `skipped` cannot + * distinguish that gap from a leaked concurrency slot, which never becomes claimable, so the + * pass count is kept low enough that a leak costs one short delay rather than a stall — these + * sweeps run per environment inside bulk pause loops. + */ +const RETURN_UNCLAIMED_MAX_PASSES = 2; +const RETURN_UNCLAIMED_PASS_DELAY_MS = 50; +const RETURN_UNCLAIMED_MAX_PASS_CEILING = 10; + export type ReturnUnclaimedMessagesResult = { + /** Distinct runs put back on their queue. */ returned: number; - skipped: number; + /** + * Candidates the final pass could not claim. Each is either a run a worker took first — + * which means it escaped the pause and is executing — or a leaked concurrency slot with no + * worker queue entry behind it. The two are indistinguishable from here. + */ + skippedLastPass: number; errors: number; passes: number; }; @@ -1058,8 +1074,8 @@ export class RunQueue { * * Admission is not atomic with the push onto the worker queue: the dequeue script claims * the concurrency slot and a separate round trip does the RPUSH. A run caught in that gap - * is a candidate whose claim fails, so a pass reporting skipped runs is retried — by then - * the run has either landed on the worker queue (and is claimable) or been picked up by a + * is a candidate whose claim fails, so a pass that skips or errors is retried — by then the + * run has either landed on the worker queue (and is claimable) or been picked up by a * worker (and has left the candidate set). * * @param queue - Restrict to a single queue (including its concurrency key variants). @@ -1068,8 +1084,8 @@ export class RunQueue { public async returnUnclaimedMessagesToQueue({ env, queue, - maxPasses = 3, - passDelayMs = 250, + maxPasses = RETURN_UNCLAIMED_MAX_PASSES, + passDelayMs = RETURN_UNCLAIMED_PASS_DELAY_MS, }: { env: MinimalAuthenticatedEnvironment; queue?: string; @@ -1080,30 +1096,39 @@ export class RunQueue { "returnUnclaimedMessagesToQueue", async (span) => { const targetBaseQueueKey = queue ? this.keys.queueKey(env, queue) : undefined; + const totalPasses = Number.isFinite(maxPasses) + ? Math.min(RETURN_UNCLAIMED_MAX_PASS_CEILING, Math.max(1, Math.floor(maxPasses))) + : RETURN_UNCLAIMED_MAX_PASSES; - let returned = 0; - let skipped = 0; + const returnedRunIds = new Set(); + let skippedLastPass = 0; let errors = 0; let passes = 0; - for (let pass = 0; pass < Math.max(1, maxPasses); pass++) { + for (let pass = 0; pass < totalPasses; pass++) { passes++; const passResult = await this.#returnUnclaimedMessagesPass(env, targetBaseQueueKey); - returned += passResult.returned; + for (const runId of passResult.returnedRunIds) { + returnedRunIds.add(runId); + } errors += passResult.errors; - skipped = passResult.skipped; + skippedLastPass = passResult.skipped; - if (passResult.skipped === 0 || pass === Math.max(1, maxPasses) - 1) { + const settled = passResult.skipped === 0 && passResult.errors === 0; + + if (settled || pass === totalPasses - 1) { break; } await setTimeout(passDelayMs); } + const returned = returnedRunIds.size; + span.setAttribute("returned_count", returned); - span.setAttribute("skipped_count", skipped); + span.setAttribute("skipped_last_pass", skippedLastPass); span.setAttribute("error_count", errors); span.setAttribute("passes", passes); @@ -1112,12 +1137,12 @@ export class RunQueue { environmentId: env.id, queue, returned, - skipped, + skippedLastPass, errors, passes, }); - return { returned, skipped, errors, passes }; + return { returned, skippedLastPass, errors, passes }; }, { kind: SpanKind.INTERNAL, @@ -1129,17 +1154,23 @@ export class RunQueue { ); } + /** + * Candidates come from the environment-level sets even when a single queue is targeted. + * The per-queue sets would be cheaper but `ckIndex` only tracks concurrency key variants + * that still have pending messages, so a variant whose only runs are already in flight is + * not discoverable from it. + */ async #returnUnclaimedMessagesPass( env: MinimalAuthenticatedEnvironment, targetBaseQueueKey: string | undefined - ): Promise<{ returned: number; skipped: number; errors: number }> { + ): Promise<{ returnedRunIds: string[]; skipped: number; errors: number }> { const unclaimedRunIds = await this.redis.sdiff( this.keys.envCurrentConcurrencyKey(env), this.keys.envCurrentDequeuedKey(env) ); if (unclaimedRunIds.length === 0) { - return { returned: 0, skipped: 0, errors: 0 }; + return { returnedRunIds: [], skipped: 0, errors: 0 }; } const messages = (await this.#readMessages(env.organization.id, unclaimedRunIds)).filter( @@ -1147,7 +1178,7 @@ export class RunQueue { !targetBaseQueueKey || this.keys.baseQueueKeyFromQueue(message.queue) === targetBaseQueueKey ); - let returned = 0; + const returnedRunIds: string[] = []; let skipped = 0; let errors = 0; @@ -1171,14 +1202,17 @@ export class RunQueue { } if (result.value) { - returned++; + const runId = batch[index]?.runId; + if (runId) { + returnedRunIds.push(runId); + } } else { skipped++; } } } - return { returned, skipped, errors }; + return { returnedRunIds, skipped, errors }; } async #readMessages(orgId: string, runIds: string[]): Promise { diff --git a/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts b/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts index 9809725f290..4132a897f65 100644 --- a/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts +++ b/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts @@ -121,7 +121,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { queue: "task/my-task", }); - expect(result).toEqual({ returned: 3, skipped: 0, errors: 0, passes: 1 }); + expect(result).toEqual({ returned: 3, skippedLastPass: 0, errors: 0, passes: 1 }); expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(0); expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(3); @@ -260,7 +260,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { }); expect(result.returned).toBe(0); - expect(result.skipped).toBe(1); + expect(result.skippedLastPass).toBe(1); expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); @@ -271,6 +271,41 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { } ); + redisTest("never resurrects an acknowledged run", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(1); + + await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, "r1", { + skipDequeueProcessing: true, + }); + + expect(await queue.messageExists(authenticatedEnvDev.organization.id, "r1")).toBe(0); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + maxPasses: 1, + }); + + expect(result.returned).toBe(0); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + expect(await queue.lengthOfEnvQueue(authenticatedEnvDev)).toBe(0); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + } finally { + await queue.quit(); + } + }); + redisTest("restores the TTL registration of a returned run", async ({ redisContainer }) => { const queue = createRunQueue(redisContainer, { ttlSystem: { @@ -520,20 +555,27 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { authenticatedEnvDev.id ); assertNonNullable(claimed); + expect(claimed.messageId).toBe("r1"); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(1); expect( await queue.currentConcurrencyOfQueues(authenticatedEnvDev, ["task/my-task"]) ).toEqual({ "task/my-task": 1 }); - await queue.returnUnclaimedMessagesToQueue({ + const result = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev, queue: "task/my-task", }); + expect(result.returned).toBe(1); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(0); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect( await queue.currentConcurrencyOfQueues(authenticatedEnvDev, ["task/my-task"]) ).toEqual({ "task/my-task": 1 }); - expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); await queue.acknowledgeMessage(authenticatedEnvDev.organization.id, claimed.messageId, { skipDequeueProcessing: true, @@ -548,6 +590,97 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { } ); + redisTest( + "claims a fast-pathed run pushed to the unresolved worker queue name", + async ({ redisContainer }) => { + vi.stubEnv( + "RUN_ENGINE_WORKER_QUEUE_OVERRIDES", + JSON.stringify({ environmentId: { [authenticatedEnvDev.id]: "rerouted" } }) + ); + + let queue: RunQueue | undefined; + + try { + queue = createRunQueue(redisContainer); + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + enableFastPath: true, + }); + + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(1); + expect(await queue.peekAllOnWorkerQueue("rerouted")).toHaveLength(0); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + }); + + expect(result.returned).toBe(1); + expect(await queue.peekAllOnWorkerQueue(authenticatedEnvDev.id)).toHaveLength(0); + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.currentConcurrencyOfQueue(authenticatedEnvDev, "task/my-task")).toBe(0); + } finally { + vi.unstubAllEnvs(); + await queue?.quit(); + } + } + ); + + redisTest("retries a pass while runs remain unclaimed", async ({ redisContainer }) => { + const queue = createRunQueue(redisContainer); + const redis = createRedisClient(redisContainer); + + try { + await queue.enqueueMessage({ + env: authenticatedEnvDev, + message: messageFor({ runId: "r1" }), + workerQueue: authenticatedEnvDev.id, + skipDequeueProcessing: true, + }); + + await queue.processMasterQueueForEnvironment(authenticatedEnvDev.id, 10); + + const entry = await redis.lpop(`runqueue:test:workerQueue:${authenticatedEnvDev.id}`); + assertNonNullable(entry); + + const result = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + passDelayMs: 10, + }); + + expect(result.passes).toBe(2); + expect(result.skippedLastPass).toBe(1); + expect(result.returned).toBe(0); + + const single = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + maxPasses: 1, + }); + + expect(single.passes).toBe(1); + + for (const maxPasses of [Number.NaN, Number.POSITIVE_INFINITY, 0, -1]) { + const guarded = await queue.returnUnclaimedMessagesToQueue({ + env: authenticatedEnvDev, + queue: "task/my-task", + maxPasses, + passDelayMs: 1, + }); + + expect(guarded.passes).toBeGreaterThanOrEqual(1); + expect(guarded.passes).toBeLessThanOrEqual(10); + } + } finally { + await redis.quit(); + await queue.quit(); + } + }); + redisTest( "leaves runs that have not reached the worker queue alone", async ({ redisContainer }) => { @@ -555,7 +688,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { try { const empty = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); - expect(empty).toEqual({ returned: 0, skipped: 0, errors: 0, passes: 1 }); + expect(empty).toEqual({ returned: 0, skippedLastPass: 0, errors: 0, passes: 1 }); await queue.enqueueMessage({ env: authenticatedEnvDev, @@ -567,7 +700,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { const pendingOnly = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev, }); - expect(pendingOnly).toEqual({ returned: 0, skipped: 0, errors: 0, passes: 1 }); + expect(pendingOnly).toEqual({ returned: 0, skippedLastPass: 0, errors: 0, passes: 1 }); expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); expect(await queue.oldestMessageInQueue(authenticatedEnvDev, "task/my-task")).toBe( baseTimestamp @@ -595,7 +728,7 @@ describe("RunQueue.returnUnclaimedMessagesToQueue", () => { expect(first.returned).toBe(1); const second = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); - expect(second).toEqual({ returned: 0, skipped: 0, errors: 0, passes: 1 }); + expect(second).toEqual({ returned: 0, skippedLastPass: 0, errors: 0, passes: 1 }); expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); expect(await queue.lengthOfEnvQueue(authenticatedEnvDev)).toBe(1);