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..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 */ @@ -30,3 +31,47 @@ 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 }); +} + +/** + * 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 0b59d4c7fae..1066943e428 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 }, @@ -182,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 af9edff856c..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 { 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"; @@ -136,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 aa3e21f9727..ab54d3f2b96 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, + sweepUnclaimedRuns, + updateQueueConcurrencyLimits, +} from "../runQueue.server"; import { engine } from "../runEngine.server"; export type PauseStatus = "paused" | "resumed"; @@ -73,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 0854ea42ef4..62f2640b349 100644 --- a/apps/webapp/test/billingLimitConvergeEnvironments.test.ts +++ b/apps/webapp/test/billingLimitConvergeEnvironments.test.ts @@ -49,6 +49,86 @@ 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("keeps the 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"), + }); + const second = await createRuntimeEnvironment(prisma, { + projectId: project.id, + organizationId: organization.id, + type: "PRODUCTION", + slug: uniqueId("prod"), + }); + + const result = await convergeBillingLimitEnvironmentsForOrg(organization.id, "grace", { + prismaClient: prisma, + updateConcurrency: async () => undefined, + returnUnclaimed: async () => { + throw new Error("run queue unavailable"); + }, + }); + + 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 }) => { const { organization, project } = await createTestOrgProjectWithMember(prisma); const environment = await createRuntimeEnvironment(prisma, { 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 3c1f4330a0f..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"; @@ -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 { + 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..448d74b800f 100644 --- a/internal-packages/run-engine/src/run-queue/index.ts +++ b/internal-packages/run-engine/src/run-queue/index.ts @@ -137,6 +137,32 @@ type MarkedRun = { score: number; }; +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; + /** + * 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; +}; + const defaultRetrySettings = { maxAttempts: 12, factor: 2, @@ -1033,6 +1059,196 @@ 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. + * + * 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 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). + * Omit to cover every queue in the environment. + */ + public async returnUnclaimedMessagesToQueue({ + env, + queue, + maxPasses = RETURN_UNCLAIMED_MAX_PASSES, + passDelayMs = RETURN_UNCLAIMED_PASS_DELAY_MS, + }: { + env: MinimalAuthenticatedEnvironment; + queue?: string; + maxPasses?: number; + passDelayMs?: number; + }): Promise { + return this.#trace( + "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; + + const returnedRunIds = new Set(); + let skippedLastPass = 0; + let errors = 0; + let passes = 0; + + for (let pass = 0; pass < totalPasses; pass++) { + passes++; + + const passResult = await this.#returnUnclaimedMessagesPass(env, targetBaseQueueKey); + + for (const runId of passResult.returnedRunIds) { + returnedRunIds.add(runId); + } + errors += passResult.errors; + skippedLastPass = passResult.skipped; + + 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_last_pass", skippedLastPass); + span.setAttribute("error_count", errors); + span.setAttribute("passes", passes); + + this.logger.info("returnUnclaimedMessagesToQueue", { + service: this.name, + environmentId: env.id, + queue, + returned, + skippedLastPass, + errors, + passes, + }); + + return { returned, skippedLastPass, errors, passes }; + }, + { + kind: SpanKind.INTERNAL, + attributes: { + [SEMATTRS_MESSAGING_SYSTEM]: "runqueue", + ...attributesFromAuthenticatedEnv(env), + }, + } + ); + } + + /** + * 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<{ 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 { returnedRunIds: [], skipped: 0, errors: 0 }; + } + + const messages = (await this.#readMessages(env.organization.id, unclaimedRunIds)).filter( + (message) => + !targetBaseQueueKey || this.keys.baseQueueKeyFromQueue(message.queue) === targetBaseQueueKey + ); + + const returnedRunIds: string[] = []; + let skipped = 0; + let errors = 0; + + 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) { + const runId = batch[index]?.runId; + if (runId) { + returnedRunIds.push(runId); + } + } else { + skipped++; + } + } + } + + return { returnedRunIds, 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); + } + } + } + + return messages; + } + public async removeEnvironmentQueuesFromMasterQueue( runtimeEnvironmentId: string, organizationId: string, @@ -2623,6 +2839,97 @@ 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)); + 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, + }); + + let result: number; + + if (message.concurrencyKey) { + result = await this.redis.returnMessageToQueueCkTracked( + masterQueueKey, + messageKey, + messageQueue, + queueCurrentConcurrencyKey, + envCurrentConcurrencyKey, + queueCurrentDequeuedKey, + envCurrentDequeuedKey, + envQueueKey, + workerQueueKey, + rawWorkerQueueKey, + ttlQueueKey, + 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), + ttlMember, + ttlScore + ); + } else { + result = await this.redis.returnMessageToQueue( + masterQueueKey, + messageKey, + messageQueue, + queueCurrentConcurrencyKey, + envCurrentConcurrencyKey, + queueCurrentDequeuedKey, + envCurrentDequeuedKey, + envQueueKey, + workerQueueKey, + rawWorkerQueueKey, + ttlQueueKey, + messageId, + messageQueue, + String(message.timestamp), + messageKey, + ttlMember, + ttlScore + ); + } + + return result === 1; + } + async #callMoveToDeadLetterQueue({ message }: { message: OutputPayload }) { const messageId = message.runId; const messageKey = this.keys.messageKey(message.orgId, message.runId); @@ -4496,6 +4803,181 @@ 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: 11, + 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 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]) + +-- 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 + +-- 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) + +-- 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) +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: 14, + 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 rawWorkerQueueKey = KEYS[10] +local ttlQueueKey = KEYS[11] +local ckIndexKey = KEYS[12] +local lengthCounterKey = KEYS[13] +local runningCounterKey = KEYS[14] + +-- 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 ttlMember = ARGV[8] +local ttlScore = tonumber(ARGV[9]) + +local function decrFloored(key) + if tonumber(redis.call('GET', key) or '0') > 0 then + redis.call('DECR', key) + end +end + +-- 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 + +-- 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 + +-- 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) +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 +6118,54 @@ 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, + rawWorkerQueueKey: string, + ttlQueueKey: string, + messageId: string, + messageQueueName: string, + messageScore: string, + messageKeyValue: string, + ttlMember: string, + ttlScore: string, + callback?: Callback + ): Result; + + returnMessageToQueueCkTracked( + masterQueueKey: string, + messageKey: string, + messageQueue: string, + queueCurrentConcurrencyKey: string, + envCurrentConcurrencyKey: string, + queueCurrentDequeuedKey: string, + envCurrentDequeuedKey: string, + envQueueKey: string, + workerQueueKey: string, + rawWorkerQueueKey: string, + ttlQueueKey: string, + ckIndexKey: string, + lengthCounterKey: string, + runningCounterKey: string, + messageId: string, + messageQueueName: string, + messageScore: string, + messageKeyValue: string, + ckWildcardName: string, + keyPrefix: string, + counterTtl: string, + ttlMember: string, + ttlScore: 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..4132a897f65 --- /dev/null +++ b/internal-packages/run-engine/src/run-queue/tests/returnUnclaimedMessagesToQueue.test.ts @@ -0,0 +1,739 @@ +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"; +import { createRedisClient as createRawRedisClient } from "@internal/redis"; + +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 }, + overrides: Partial[0]> = {} +) { + 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, + ...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", () => { + 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, skippedLastPass: 0, errors: 0, passes: 1 }); + + 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( + "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.skippedLastPass).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("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: { + 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); + + 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( + "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(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 }); + + 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 }); + + 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( + "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 }) => { + const queue = createRunQueue(redisContainer); + + try { + const empty = await queue.returnUnclaimedMessagesToQueue({ env: authenticatedEnvDev }); + expect(empty).toEqual({ returned: 0, skippedLastPass: 0, errors: 0, passes: 1 }); + + 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, 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 + ); + } 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, skippedLastPass: 0, errors: 0, passes: 1 }); + + expect(await queue.lengthOfQueue(authenticatedEnvDev, "task/my-task")).toBe(1); + expect(await queue.lengthOfEnvQueue(authenticatedEnvDev)).toBe(1); + } finally { + await queue.quit(); + } + }); +});