652 lines
20 KiB
TypeScript
652 lines
20 KiB
TypeScript
import { describe, expect, onTestFinished, vi } from "vitest";
|
|
|
|
// db.server + splitMode are mocked so the idempotency dedup client resolves to
|
|
// the container prisma passed into the concern (split stays off).
|
|
vi.mock("~/db.server", () => ({
|
|
prisma: {},
|
|
$replica: {},
|
|
runOpsNewPrisma: {},
|
|
runOpsLegacyPrisma: {},
|
|
}));
|
|
|
|
vi.mock("~/v3/runOpsMigration/splitMode.server", () => ({ isSplitEnabled: async () => false }));
|
|
|
|
vi.mock("~/services/platform.v3.server", async (importOriginal) => {
|
|
const actual = (await importOriginal()) as Record<string, unknown>;
|
|
return {
|
|
...actual,
|
|
getEntitlement: vi.fn(),
|
|
};
|
|
});
|
|
|
|
import { RunEngine } from "@internal/run-engine";
|
|
import { setupAuthenticatedEnvironment, setupBackgroundWorker } from "@internal/run-engine/tests";
|
|
import { assertNonNullable, containerTest } from "@internal/testcontainers";
|
|
import { trace } from "@opentelemetry/api";
|
|
import { IdempotencyKeyConcern } from "~/runEngine/concerns/idempotencyKeys.server";
|
|
import { DefaultQueueManager } from "~/runEngine/concerns/queues.server";
|
|
import { NoopTaskMetadataCache } from "~/services/taskMetadataCache.server";
|
|
import { RunEngineTriggerTaskService } from "../../app/runEngine/services/triggerTask.server";
|
|
import { setTimeout } from "node:timers/promises";
|
|
import {
|
|
MockPayloadProcessor,
|
|
MockTraceEventConcern,
|
|
MockTriggerRacepointSystem,
|
|
MockTriggerTaskValidator,
|
|
} from "./triggerTaskTestHelpers";
|
|
|
|
vi.setConfig({ testTimeout: 60_000, hookTimeout: 60_000 });
|
|
|
|
describe("RunEngineTriggerTaskService", () => {
|
|
containerTest("should handle idempotency keys correctly", async ({ prisma, redisOptions }) => {
|
|
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.0005,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
onTestFinished(() => engine.quit());
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const taskIdentifier = "test-task";
|
|
|
|
await setupBackgroundWorker(engine, authenticatedEnvironment, taskIdentifier);
|
|
|
|
const queuesManager = new DefaultQueueManager(prisma, engine);
|
|
|
|
const idempotencyKeyConcern = new IdempotencyKeyConcern(
|
|
prisma,
|
|
engine,
|
|
new MockTraceEventConcern()
|
|
);
|
|
|
|
const triggerTaskService = new RunEngineTriggerTaskService({
|
|
engine,
|
|
prisma,
|
|
payloadProcessor: new MockPayloadProcessor(),
|
|
queueConcern: queuesManager,
|
|
idempotencyKeyConcern,
|
|
validator: new MockTriggerTaskValidator(),
|
|
traceEventConcern: new MockTraceEventConcern(),
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
metadataMaximumSize: 1024 * 1024 * 1, // 1MB
|
|
});
|
|
|
|
const result = await triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
idempotencyKey: "test-idempotency-key",
|
|
},
|
|
},
|
|
});
|
|
|
|
expect(result).toBeDefined();
|
|
expect(result?.run.friendlyId).toBeDefined();
|
|
expect(result?.run.status).toBe("PENDING");
|
|
expect(result?.isCached).toBe(false);
|
|
|
|
const run = await prisma.taskRun.findFirst({
|
|
where: {
|
|
id: result?.run.id,
|
|
},
|
|
});
|
|
|
|
expect(run).toBeDefined();
|
|
expect(run?.friendlyId).toBe(result?.run.friendlyId);
|
|
expect(run?.engine).toBe("V2");
|
|
expect(run?.queuedAt).toBeDefined();
|
|
expect(run?.queue).toBe(`task/${taskIdentifier}`);
|
|
|
|
// Lets make sure the task is in the queue
|
|
const queueLength = await engine.runQueue.lengthOfQueue(
|
|
authenticatedEnvironment,
|
|
`task/${taskIdentifier}`
|
|
);
|
|
expect(queueLength).toBe(1);
|
|
|
|
// Now lets try to trigger the same task with the same idempotency key
|
|
const cachedResult = await triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
idempotencyKey: "test-idempotency-key",
|
|
},
|
|
},
|
|
});
|
|
|
|
expect(cachedResult).toBeDefined();
|
|
expect(cachedResult?.run.friendlyId).toBe(result?.run.friendlyId);
|
|
expect(cachedResult?.isCached).toBe(true);
|
|
});
|
|
|
|
containerTest(
|
|
"should handle idempotency keys when the engine throws an RunDuplicateIdempotencyKeyError",
|
|
async ({ prisma, redisOptions }) => {
|
|
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.0005,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
logLevel: "debug",
|
|
});
|
|
onTestFinished(() => engine.quit());
|
|
|
|
const parentTask = "parent-task";
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
|
|
const taskIdentifier = "test-task";
|
|
|
|
await setupBackgroundWorker(engine, authenticatedEnvironment, [parentTask, taskIdentifier]);
|
|
|
|
const parentRun1 = await engine.trigger(
|
|
{
|
|
number: 1,
|
|
friendlyId: "run_cmqxvncxq0000kaulzpafkicv",
|
|
environment: authenticatedEnvironment,
|
|
taskIdentifier: parentTask,
|
|
payload: "{}",
|
|
payloadType: "application/json",
|
|
context: {},
|
|
traceContext: {},
|
|
traceId: "t12345",
|
|
spanId: "s12345",
|
|
queue: `task/${parentTask}`,
|
|
isTest: false,
|
|
tags: [],
|
|
workerQueue: "main",
|
|
},
|
|
prisma
|
|
);
|
|
|
|
//dequeue parent and create the attempt
|
|
await setTimeout(500);
|
|
const dequeued = await engine.dequeueFromWorkerQueue({
|
|
consumerId: "test_12345",
|
|
workerQueue: "main",
|
|
});
|
|
await engine.startRunAttempt({
|
|
runId: parentRun1.id,
|
|
snapshotId: dequeued[0].snapshot.id,
|
|
});
|
|
|
|
const parentRun2 = await engine.trigger(
|
|
{
|
|
number: 2,
|
|
friendlyId: "run_cmqxvncxr0001kauldv9mqa9z",
|
|
environment: authenticatedEnvironment,
|
|
taskIdentifier: parentTask,
|
|
payload: "{}",
|
|
payloadType: "application/json",
|
|
context: {},
|
|
traceContext: {},
|
|
traceId: "t12346",
|
|
spanId: "s12346",
|
|
queue: `task/${parentTask}`,
|
|
isTest: false,
|
|
tags: [],
|
|
workerQueue: "main",
|
|
},
|
|
prisma
|
|
);
|
|
|
|
await setTimeout(500);
|
|
const dequeued2 = await engine.dequeueFromWorkerQueue({
|
|
consumerId: "test_12345",
|
|
workerQueue: "main",
|
|
});
|
|
await engine.startRunAttempt({
|
|
runId: parentRun2.id,
|
|
snapshotId: dequeued2[0].snapshot.id,
|
|
});
|
|
|
|
const queuesManager = new DefaultQueueManager(prisma, engine);
|
|
|
|
const idempotencyKeyConcern = new IdempotencyKeyConcern(
|
|
prisma,
|
|
engine,
|
|
new MockTraceEventConcern()
|
|
);
|
|
|
|
const triggerRacepointSystem = new MockTriggerRacepointSystem();
|
|
|
|
const triggerTaskService = new RunEngineTriggerTaskService({
|
|
engine,
|
|
prisma,
|
|
payloadProcessor: new MockPayloadProcessor(),
|
|
queueConcern: queuesManager,
|
|
idempotencyKeyConcern,
|
|
validator: new MockTriggerTaskValidator(),
|
|
traceEventConcern: new MockTraceEventConcern(),
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
metadataMaximumSize: 1024 * 1024 * 1, // 1MB
|
|
triggerRacepointSystem,
|
|
});
|
|
|
|
const idempotencyKey = "test-idempotency-key";
|
|
|
|
const racepoint = triggerRacepointSystem.registerRacepoint("idempotencyKey", idempotencyKey);
|
|
|
|
const childTriggerPromise1 = triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
idempotencyKey,
|
|
parentRunId: parentRun1.friendlyId,
|
|
resumeParentOnCompletion: true,
|
|
},
|
|
},
|
|
});
|
|
|
|
const childTriggerPromise2 = triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
idempotencyKey,
|
|
parentRunId: parentRun2.friendlyId,
|
|
resumeParentOnCompletion: true,
|
|
},
|
|
},
|
|
});
|
|
|
|
await setTimeout(500);
|
|
|
|
// Now we can resolve the racepoint
|
|
racepoint.resolve();
|
|
|
|
const result = await childTriggerPromise1;
|
|
const result2 = await childTriggerPromise2;
|
|
|
|
expect(result).toBeDefined();
|
|
expect(result?.run.friendlyId).toBeDefined();
|
|
expect(result?.run.status).toBe("PENDING");
|
|
|
|
const run = await prisma.taskRun.findFirst({
|
|
where: {
|
|
id: result?.run.id,
|
|
},
|
|
});
|
|
|
|
expect(run).toBeDefined();
|
|
expect(run?.friendlyId).toBe(result?.run.friendlyId);
|
|
expect(run?.engine).toBe("V2");
|
|
expect(run?.queuedAt).toBeDefined();
|
|
expect(run?.queue).toBe(`task/${taskIdentifier}`);
|
|
|
|
expect(result2).toBeDefined();
|
|
expect(result2?.run.friendlyId).toBe(result?.run.friendlyId);
|
|
|
|
const parent1ExecutionData = await engine.getRunExecutionData({ runId: parentRun1.id });
|
|
assertNonNullable(parent1ExecutionData);
|
|
expect(parent1ExecutionData.snapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS");
|
|
|
|
const parent2ExecutionData = await engine.getRunExecutionData({ runId: parentRun2.id });
|
|
assertNonNullable(parent2ExecutionData);
|
|
expect(parent2ExecutionData.snapshot.executionStatus).toBe("EXECUTING_WITH_WAITPOINTS");
|
|
|
|
const parent1RunWaitpoint = await prisma.taskRunWaitpoint.findFirst({
|
|
where: {
|
|
taskRunId: parentRun1.id,
|
|
},
|
|
include: {
|
|
waitpoint: true,
|
|
},
|
|
});
|
|
|
|
assertNonNullable(parent1RunWaitpoint);
|
|
expect(parent1RunWaitpoint.waitpoint.type).toBe("RUN");
|
|
expect(parent1RunWaitpoint.waitpoint.completedByTaskRunId).toBe(result?.run.id);
|
|
|
|
const parent2RunWaitpoint = await prisma.taskRunWaitpoint.findFirst({
|
|
where: {
|
|
taskRunId: parentRun2.id,
|
|
},
|
|
include: {
|
|
waitpoint: true,
|
|
},
|
|
});
|
|
|
|
assertNonNullable(parent2RunWaitpoint);
|
|
expect(parent2RunWaitpoint.waitpoint.type).toBe("RUN");
|
|
expect(parent2RunWaitpoint.waitpoint.completedByTaskRunId).toBe(result2?.run.id);
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should resolve queue names correctly when locked to version",
|
|
async ({ prisma, redisOptions }) => {
|
|
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.0005,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
onTestFinished(() => engine.quit());
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
const taskIdentifier = "test-task";
|
|
|
|
// Create a background worker with a specific version
|
|
const worker = await setupBackgroundWorker(engine, authenticatedEnvironment, taskIdentifier, {
|
|
preset: "small-1x",
|
|
});
|
|
|
|
// Create a specific queue for this worker
|
|
const specificQueue = await prisma.taskQueue.create({
|
|
data: {
|
|
name: "specific-queue",
|
|
friendlyId: "specific-queue",
|
|
projectId: authenticatedEnvironment.projectId,
|
|
runtimeEnvironmentId: authenticatedEnvironment.id,
|
|
workers: {
|
|
connect: {
|
|
id: worker.worker.id,
|
|
},
|
|
},
|
|
},
|
|
});
|
|
|
|
// Associate the task with the queue
|
|
await prisma.backgroundWorkerTask.update({
|
|
where: {
|
|
workerId_slug: {
|
|
workerId: worker.worker.id,
|
|
slug: taskIdentifier,
|
|
},
|
|
},
|
|
data: {
|
|
queueId: specificQueue.id,
|
|
},
|
|
});
|
|
|
|
const queuesManager = new DefaultQueueManager(prisma, engine);
|
|
const idempotencyKeyConcern = new IdempotencyKeyConcern(
|
|
prisma,
|
|
engine,
|
|
new MockTraceEventConcern()
|
|
);
|
|
|
|
const triggerTaskService = new RunEngineTriggerTaskService({
|
|
engine,
|
|
prisma,
|
|
payloadProcessor: new MockPayloadProcessor(),
|
|
queueConcern: queuesManager,
|
|
idempotencyKeyConcern,
|
|
validator: new MockTriggerTaskValidator(),
|
|
traceEventConcern: new MockTraceEventConcern(),
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
metadataMaximumSize: 1024 * 1024 * 1, // 1MB
|
|
});
|
|
|
|
// Test case 1: Trigger with lockToVersion but no specific queue
|
|
const result1 = await triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
lockToVersion: worker.worker.version,
|
|
},
|
|
},
|
|
});
|
|
|
|
expect(result1).toBeDefined();
|
|
expect(result1?.run.queue).toBe("specific-queue");
|
|
|
|
// Test case 2: Trigger with lockToVersion and specific queue
|
|
const result2 = await triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
lockToVersion: worker.worker.version,
|
|
queue: {
|
|
name: "specific-queue",
|
|
},
|
|
},
|
|
},
|
|
});
|
|
|
|
expect(result2).toBeDefined();
|
|
expect(result2?.run.queue).toBe("specific-queue");
|
|
expect(result2?.run.lockedQueueId).toBe(specificQueue.id);
|
|
|
|
// Test case 3: Try to use non-existent queue with locked version (should throw)
|
|
await expect(
|
|
triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
lockToVersion: worker.worker.version,
|
|
queue: {
|
|
name: "non-existent-queue",
|
|
},
|
|
},
|
|
},
|
|
})
|
|
).rejects.toThrow(
|
|
`Specified queue 'non-existent-queue' not found or not associated with locked version '${worker.worker.version}'`
|
|
);
|
|
|
|
// Test case 4: Trigger with a non-existent queue without a locked version
|
|
const result4 = await triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
queue: {
|
|
name: "non-existent-queue",
|
|
},
|
|
},
|
|
},
|
|
});
|
|
|
|
expect(result4).toBeDefined();
|
|
expect(result4?.run.queue).toBe("non-existent-queue");
|
|
expect(result4?.run.status).toBe("PENDING");
|
|
}
|
|
);
|
|
|
|
containerTest(
|
|
"should fall back to the writer when a stale replica returns no row for a locked task",
|
|
async ({ prisma, redisOptions }) => {
|
|
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.0005,
|
|
},
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
});
|
|
onTestFinished(() => engine.quit());
|
|
|
|
const authenticatedEnvironment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION");
|
|
const taskIdentifier = "test-task";
|
|
|
|
const worker = await setupBackgroundWorker(engine, authenticatedEnvironment, taskIdentifier);
|
|
|
|
// A read replica that has not yet caught up to the BackgroundWorkerTask
|
|
// row: it is the real database for every query except the locked-task
|
|
// lookup, which comes back empty (the TRI-10868 false-negative window).
|
|
const staleReplica = new Proxy(prisma, {
|
|
get(target, prop, receiver) {
|
|
if (prop === "backgroundWorkerTask") {
|
|
const delegate = Reflect.get(target, prop, receiver);
|
|
return new Proxy(delegate, {
|
|
get(taskTarget, taskProp, taskReceiver) {
|
|
if (taskProp === "findFirst") {
|
|
return async () => null;
|
|
}
|
|
const value = Reflect.get(taskTarget, taskProp, taskReceiver);
|
|
return typeof value === "function" ? value.bind(taskTarget) : value;
|
|
},
|
|
});
|
|
}
|
|
const value = Reflect.get(target, prop, receiver);
|
|
return typeof value === "function" ? value.bind(target) : value;
|
|
},
|
|
}) as typeof prisma;
|
|
|
|
// Noop cache so every resolve misses the cache and exercises the
|
|
// replica -> writer fallback. The writer is the real `prisma`.
|
|
const queuesManager = new DefaultQueueManager(
|
|
prisma,
|
|
engine,
|
|
staleReplica,
|
|
new NoopTaskMetadataCache()
|
|
);
|
|
|
|
const triggerTaskService = new RunEngineTriggerTaskService({
|
|
engine,
|
|
prisma,
|
|
payloadProcessor: new MockPayloadProcessor(),
|
|
queueConcern: queuesManager,
|
|
idempotencyKeyConcern: new IdempotencyKeyConcern(
|
|
prisma,
|
|
engine,
|
|
new MockTraceEventConcern()
|
|
),
|
|
validator: new MockTriggerTaskValidator(),
|
|
traceEventConcern: new MockTraceEventConcern(),
|
|
tracer: trace.getTracer("test", "0.0.0"),
|
|
metadataMaximumSize: 1024 * 1024 * 1,
|
|
});
|
|
|
|
// The task IS registered on the locked worker, but the replica returns
|
|
// nothing. Before the fix this threw "not found on locked version"; now
|
|
// the writer fallback resolves the registered row.
|
|
const result = await triggerTaskService.call({
|
|
taskId: taskIdentifier,
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
lockToVersion: worker.worker.version,
|
|
},
|
|
},
|
|
});
|
|
|
|
expect(result).toBeDefined();
|
|
expect(result?.run.status).toBe("PENDING");
|
|
expect(result?.run.queue).toBe(`task/${taskIdentifier}`);
|
|
|
|
// A genuinely unregistered task must still throw, even with the writer
|
|
// fallback — the writer has no row either, so the 422 is correct.
|
|
await expect(
|
|
triggerTaskService.call({
|
|
taskId: "not-a-registered-task",
|
|
environment: authenticatedEnvironment,
|
|
body: {
|
|
payload: { test: "test" },
|
|
options: {
|
|
lockToVersion: worker.worker.version,
|
|
},
|
|
},
|
|
})
|
|
).rejects.toThrow(
|
|
`Task 'not-a-registered-task' not found on locked version '${worker.worker.version}'`
|
|
);
|
|
}
|
|
);
|
|
});
|