From e9fa4ce3a2ba443f0d1562fdf787d73a58c6af39 Mon Sep 17 00:00:00 2001 From: Eric Allam Date: Wed, 12 Mar 2025 15:44:55 +0000 Subject: [PATCH] Update delays to use a redis worker and work with the new reserve concurrency system --- .vscode/launch.json | 2 +- .../run-engine/src/engine/index.ts | 156 ++++++---- .../src/engine/tests/delays.test.ts | 287 +++++++++++++++++- .../engine/tests/reserveConcurrency.test.ts | 2 +- 4 files changed, 390 insertions(+), 57 deletions(-) diff --git a/.vscode/launch.json b/.vscode/launch.json index da6e7674a..bb8c931ec 100644 --- a/.vscode/launch.json +++ b/.vscode/launch.json @@ -138,7 +138,7 @@ "type": "node-terminal", "request": "launch", "name": "Debug RunEngine tests", - "command": "pnpm run test ./src/engine/tests/attemptFailures.test.ts -t 'OOM fails after retrying on larger machine'", + "command": "pnpm run test ./src/engine/tests/delays.test.ts -t 'Delayed run with a ttl'", "cwd": "${workspaceFolder}/internal-packages/run-engine", "sourceMaps": true }, diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 95d78af40..1c8a296c6 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -118,6 +118,12 @@ const workerCatalog = { }), visibilityTimeoutMs: 10_000, }, + enqueueDelayedRun: { + schema: z.object({ + runId: z.string(), + }), + visibilityTimeoutMs: 10_000, + }, }; type EngineWorker = Worker; @@ -215,6 +221,9 @@ export class RunEngine { runId: payload.runId, }); }, + enqueueDelayedRun: async ({ payload }) => { + await this.#enqueueDelayedRun({ runId: payload.runId }); + }, }, }).start(); @@ -515,41 +524,18 @@ export class RunEngine { } } - if (taskRun.delayUntil) { - const delayWaitpointResult = await this.createDateTimeWaitpoint({ - projectId: environment.project.id, - environmentId: environment.id, - completedAfter: taskRun.delayUntil, - tx: prisma, - }); - - await prisma.taskRunWaitpoint.create({ - data: { - taskRunId: taskRun.id, - waitpointId: delayWaitpointResult.waitpoint.id, - projectId: delayWaitpointResult.waitpoint.projectId, - }, - }); - } - - if (!taskRun.delayUntil && taskRun.ttl) { - const expireAt = parseNaturalLanguageDuration(taskRun.ttl); - - if (expireAt) { - await this.worker.enqueue({ - id: `expireRun:${taskRun.id}`, - job: "expireRun", - payload: { runId: taskRun.id }, - availableAt: expireAt, - }); - } - } - //Make sure lock extension succeeded signal.throwIfAborted(); - //enqueue the run if it's not delayed - if (!taskRun.delayUntil) { + if (taskRun.delayUntil) { + // Schedule the run to be enqueued at the delayUntil time + await this.worker.enqueue({ + id: `enqueueDelayedRun:${taskRun.id}`, + job: "enqueueDelayedRun", + payload: { runId: taskRun.id }, + availableAt: taskRun.delayUntil, + }); + } else { const { wasEnqueued, error } = await this.#enqueueRun({ run: taskRun, env: environment, @@ -565,12 +551,25 @@ export class RunEngine { taskRun = await prisma.taskRun.update({ where: { id: taskRun.id }, data: { - status: "SYSTEM_FAILURE", + status: runStatusFromError(error), completedAt: new Date(), error, }, }); } + + if (wasEnqueued && taskRun.ttl) { + const expireAt = parseNaturalLanguageDuration(taskRun.ttl); + + if (expireAt) { + await this.worker.enqueue({ + id: `expireRun:${taskRun.id}`, + job: "expireRun", + payload: { runId: taskRun.id }, + availableAt: expireAt, + }); + } + } } }); @@ -1598,7 +1597,7 @@ export class RunEngine { /** * Reschedules a delayed run where the run hasn't been queued yet */ - async rescheduleRun({ + async rescheduleDelayedRun({ runId, delayUntil, tx, @@ -1634,26 +1633,9 @@ export class RunEngine { }, }, }, - include: { - blockedByWaitpoints: true, - }, }); - if (updatedRun.blockedByWaitpoints.length === 0) { - throw new ServiceValidationError( - "Cannot reschedule a run that is not blocked by a waitpoint" - ); - } - - const result = await this.#rescheduleDateTimeWaitpoint( - prisma, - updatedRun.blockedByWaitpoints[0].waitpointId, - delayUntil - ); - - if (!result.success) { - throw new ServiceValidationError("Failed to reschedule waitpoint, too late.", 400); - } + await this.worker.reschedule(`enqueueDelayedRun:${updatedRun.id}`, delayUntil); return updatedRun; }); @@ -3118,7 +3100,7 @@ export class RunEngine { runnerId, }: { runId: string; - snapshotId: string; + snapshotId?: string; failedAt: Date; error: TaskRunError; workerId?: string; @@ -3472,6 +3454,74 @@ export class RunEngine { }); } + async #enqueueDelayedRun({ runId }: { runId: string }) { + const run = await this.prisma.taskRun.findFirst({ + where: { id: runId }, + include: { + runtimeEnvironment: { + include: { + project: true, + organization: true, + }, + }, + }, + }); + + if (!run) { + throw new Error(`#enqueueDelayedRun: run not found: ${runId}`); + } + + let reserveConcurrency: RunQueueReserveConcurrencyOptions | undefined; + + if (run.parentTaskRunId) { + const parentRun = await this.prisma.taskRun.findFirst({ + where: { id: run.parentTaskRunId }, + }); + + if (parentRun) { + reserveConcurrency = { + messageId: parentRun.id, + recursiveQueue: parentRun.queue === run.queue, + }; + } + } + + // Now we need to enqueue the run into the RunQueue + const { wasEnqueued, error } = await this.#enqueueRun({ + run, + env: run.runtimeEnvironment, + timestamp: run.createdAt.getTime() - run.priorityMs, + reserveConcurrency, + }); + + if (error) { + await this.#permanentlyFailRun({ runId, error, failedAt: new Date() }); + } + + if (wasEnqueued) { + await this.prisma.taskRun.update({ + where: { id: runId }, + data: { + status: "PENDING", + queuedAt: new Date(), + }, + }); + + if (run.ttl) { + const expireAt = parseNaturalLanguageDuration(run.ttl); + + if (expireAt) { + await this.worker.enqueue({ + id: `expireRun:${runId}`, + job: "expireRun", + payload: { runId }, + availableAt: expireAt, + }); + } + } + } + } + async #queueRunsWaitingForWorker({ backgroundWorkerId }: { backgroundWorkerId: string }) { //It could be a lot of runs, so we will process them in a batch //if there are still more to process we will enqueue this function again diff --git a/internal-packages/run-engine/src/engine/tests/delays.test.ts b/internal-packages/run-engine/src/engine/tests/delays.test.ts index 0d87d27e4..8dd24bf7e 100644 --- a/internal-packages/run-engine/src/engine/tests/delays.test.ts +++ b/internal-packages/run-engine/src/engine/tests/delays.test.ts @@ -8,6 +8,7 @@ import { trace } from "@internal/tracing"; import { expect } from "vitest"; import { RunEngine } from "../index.js"; import { setTimeout } from "timers/promises"; +import { TaskRunErrorCodes } from "@trigger.dev/core/v3"; vi.setConfig({ testTimeout: 60_000 }); @@ -154,7 +155,7 @@ describe("RunEngine delays", () => { queueName: "task/test-task", isTest: false, tags: [], - delayUntil: new Date(Date.now() + 200), + delayUntil: new Date(Date.now() + 400), }, prisma ); @@ -165,7 +166,10 @@ describe("RunEngine delays", () => { expect(executionData.snapshot.executionStatus).toBe("RUN_CREATED"); const rescheduleTo = new Date(Date.now() + 1_500); - const updatedRun = await engine.rescheduleRun({ runId: run.id, delayUntil: rescheduleTo }); + const updatedRun = await engine.rescheduleDelayedRun({ + runId: run.id, + delayUntil: rescheduleTo, + }); expect(updatedRun.delayUntil?.toISOString()).toBe(rescheduleTo.toISOString()); //wait so the initial delay passes @@ -187,4 +191,283 @@ describe("RunEngine delays", () => { engine.quit(); } }); + + containerTest("Delayed run with a ttl", async ({ prisma, redisOptions }) => { + //create environment + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + + try { + const taskIdentifier = "test-task"; + + //create background worker + const backgroundWorker = await setupBackgroundWorker( + prisma, + authenticatedEnvironment, + taskIdentifier + ); + + //trigger the run + const run = await engine.trigger( + { + number: 1, + friendlyId: "run_1234", + environment: authenticatedEnvironment, + taskIdentifier, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: "task/test-task", + isTest: false, + tags: [], + delayUntil: new Date(Date.now() + 1000), + ttl: "2s", + }, + prisma + ); + + //should be created but not queued yet + const executionData = await engine.getRunExecutionData({ runId: run.id }); + assertNonNullable(executionData); + expect(executionData.snapshot.executionStatus).toBe("RUN_CREATED"); + expect(run.status).toBe("DELAYED"); + + //wait for 1 seconds + await setTimeout(2_500); + + //should now be queued + const executionData2 = await engine.getRunExecutionData({ runId: run.id }); + assertNonNullable(executionData2); + expect(executionData2.snapshot.executionStatus).toBe("QUEUED"); + + const run2 = await prisma.taskRun.findFirstOrThrow({ + where: { id: run.id }, + }); + + expect(run2.status).toBe("PENDING"); + + //wait for 3 seconds + await setTimeout(3_000); + + //should now be expired + const executionData3 = await engine.getRunExecutionData({ runId: run.id }); + assertNonNullable(executionData3); + expect(executionData3.snapshot.executionStatus).toBe("FINISHED"); + + const run3 = await prisma.taskRun.findFirstOrThrow({ + where: { id: run.id }, + }); + + expect(run3.status).toBe("EXPIRED"); + } finally { + engine.quit(); + } + }); + + containerTest( + "Delayed run that fails to enqueue because of a recursive deadlock issue", + async ({ prisma, redisOptions }) => { + //create environment + const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + + const engine = new RunEngine({ + prisma, + worker: { + redis: redisOptions, + workers: 1, + tasksPerWorker: 10, + pollIntervalMs: 100, + }, + queue: { + redis: redisOptions, + }, + runLock: { + redis: redisOptions, + }, + machines: { + defaultMachine: "small-1x", + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("test", "0.0.0"), + }); + + try { + const parentTask = "parent-task"; + const childTask = "child-task"; + + //create background worker + await setupBackgroundWorker(prisma, authenticatedEnvironment, [parentTask, childTask]); + + //trigger the run + const parentRun = await engine.trigger( + { + number: 1, + friendlyId: "run_p1234", + environment: authenticatedEnvironment, + taskIdentifier: parentTask, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: "shared-queue", + queue: { + concurrencyLimit: 1, + }, + isTest: false, + tags: [], + }, + prisma + ); + + const dequeued = await engine.dequeueFromMasterQueue({ + consumerId: "test_12345", + masterQueue: parentRun.masterQueue, + maxRunCount: 10, + }); + + expect(dequeued.length).toBe(1); + + const initialExecutionData = await engine.getRunExecutionData({ runId: parentRun.id }); + assertNonNullable(initialExecutionData); + const attemptResult = await engine.startRunAttempt({ + runId: parentRun.id, + snapshotId: initialExecutionData.snapshot.id, + }); + + expect(attemptResult).toBeDefined(); + + const childRun = await engine.trigger( + { + number: 1, + friendlyId: "run_c1234", + environment: authenticatedEnvironment, + taskIdentifier: childTask, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345", + spanId: "s12345", + masterQueue: "main", + queueName: "shared-queue", + queue: { + concurrencyLimit: 1, + }, + isTest: false, + tags: [], + resumeParentOnCompletion: true, + parentTaskRunId: parentRun.id, + }, + prisma + ); + + const childExecutionData = await engine.getRunExecutionData({ runId: childRun.id }); + assertNonNullable(childExecutionData); + expect(childExecutionData.snapshot.executionStatus).toBe("QUEUED"); + + const parentExecutionData = await engine.getRunExecutionData({ runId: parentRun.id }); + assertNonNullable(parentExecutionData); + expect(parentExecutionData.snapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS"); + + //dequeue the child run + const dequeuedChild = await engine.dequeueFromMasterQueue({ + consumerId: "test_12345", + masterQueue: childRun.masterQueue, + maxRunCount: 10, + }); + + expect(dequeuedChild.length).toBe(1); + + // Now try and trigger another child run on the same queue + const childRun2 = await engine.trigger( + { + number: 1, + friendlyId: "run_c12345", + environment: authenticatedEnvironment, + taskIdentifier: childTask, + payload: "{}", + payloadType: "application/json", + context: {}, + traceContext: {}, + traceId: "t12345_2", + spanId: "s12345_2", + masterQueue: "main", + queueName: "shared-queue", + queue: { + concurrencyLimit: 1, + }, + isTest: false, + tags: [], + resumeParentOnCompletion: true, + parentTaskRunId: parentRun.id, + delayUntil: new Date(Date.now() + 1000), + }, + prisma + ); + + const executionData = await engine.getRunExecutionData({ runId: childRun2.id }); + assertNonNullable(executionData); + expect(executionData.snapshot.executionStatus).toBe("RUN_CREATED"); + + await setTimeout(1_500); + + // Now the run should be failed + const run2 = await prisma.taskRun.findFirstOrThrow({ + where: { id: childRun2.id }, + }); + + expect(run2.status).toBe("COMPLETED_WITH_ERRORS"); + expect(run2.error).toEqual({ + type: "INTERNAL_ERROR", + code: TaskRunErrorCodes.RECURSIVE_WAIT_DEADLOCK, + message: expect.any(String), + }); + } finally { + engine.quit(); + } + } + ); }); diff --git a/internal-packages/run-engine/src/engine/tests/reserveConcurrency.test.ts b/internal-packages/run-engine/src/engine/tests/reserveConcurrency.test.ts index d2ba5f1b3..97b631935 100644 --- a/internal-packages/run-engine/src/engine/tests/reserveConcurrency.test.ts +++ b/internal-packages/run-engine/src/engine/tests/reserveConcurrency.test.ts @@ -571,7 +571,7 @@ describe("Reserve concurrency", () => { prisma ); - expect(childRun2.status).toBe("SYSTEM_FAILURE"); + expect(childRun2.status).toBe("COMPLETED_WITH_ERRORS"); expect(childRun2.error).toEqual({ type: "INTERNAL_ERROR", code: TaskRunErrorCodes.RECURSIVE_WAIT_DEADLOCK,