Update delays to use a redis worker and work with the new reserve concurrency system

This commit is contained in:
Eric Allam
2025-03-12 15:44:55 +00:00
parent 1a233efd42
commit e9fa4ce3a2
4 changed files with 390 additions and 57 deletions
+1 -1
View File
@@ -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
},
+103 -53
View File
@@ -118,6 +118,12 @@ const workerCatalog = {
}),
visibilityTimeoutMs: 10_000,
},
enqueueDelayedRun: {
schema: z.object({
runId: z.string(),
}),
visibilityTimeoutMs: 10_000,
},
};
type EngineWorker = Worker<typeof workerCatalog>;
@@ -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
@@ -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();
}
}
);
});
@@ -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,