|
| 1 | +import { RunEngine } from "@internal/run-engine"; |
| 2 | +import { setupAuthenticatedEnvironment, setupBackgroundWorker } from "@internal/run-engine/tests"; |
| 3 | +import { assertNonNullable, containerTest } from "@internal/testcontainers"; |
| 4 | +import { trace } from "@internal/tracing"; |
| 5 | +import { generateFriendlyId } from "@trigger.dev/core/v3/isomorphic"; |
| 6 | +import type { WorkerQueueSubscription } from "@trigger.dev/core/v3/workers"; |
| 7 | +import { expect } from "vitest"; |
| 8 | +import { dequeueWorkerQueues } from "~/runEngine/concerns/workerQueueDequeue.server"; |
| 9 | +import { createWorkerQueueConsumer } from "~/runEngine/concerns/workerQueueSubscriptions.server"; |
| 10 | + |
| 11 | +containerTest( |
| 12 | + "authorizes the entire subscription set before pop and retains legacy dequeue", |
| 13 | + async ({ prisma, redisOptions }) => { |
| 14 | + const engine = new RunEngine({ |
| 15 | + prisma, |
| 16 | + worker: { redis: redisOptions, workers: 1, tasksPerWorker: 10, pollIntervalMs: 100 }, |
| 17 | + queue: { redis: redisOptions, masterQueueConsumersDisabled: true }, |
| 18 | + runLock: { redis: redisOptions }, |
| 19 | + machines: { |
| 20 | + defaultMachine: "small-1x", |
| 21 | + machines: { |
| 22 | + "small-1x": { name: "small-1x", cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 }, |
| 23 | + }, |
| 24 | + baseCostInCents: 0.0005, |
| 25 | + }, |
| 26 | + tracer: trace.getTracer("multi-queue-test"), |
| 27 | + }); |
| 28 | + |
| 29 | + try { |
| 30 | + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); |
| 31 | + const taskIdentifier = "test-task"; |
| 32 | + await setupBackgroundWorker(engine, environment, taskIdentifier); |
| 33 | + const stable: WorkerQueueSubscription = { |
| 34 | + class: "ondemand", |
| 35 | + phase: "fresh", |
| 36 | + compat: "container", |
| 37 | + channel: "stable", |
| 38 | + }; |
| 39 | + const canary: WorkerQueueSubscription = { ...stable, channel: "canary" }; |
| 40 | + const scheduled: WorkerQueueSubscription = { ...stable, class: "scheduled" }; |
| 41 | + const worker = createWorkerQueueConsumer({ |
| 42 | + workerGroupId: environment.project.defaultWorkerGroupId!, |
| 43 | + workerInstanceId: "test-instance", |
| 44 | + masterQueue: "legacy-region", |
| 45 | + region: "us-east-1", |
| 46 | + workloadType: "CONTAINER" as const, |
| 47 | + allowedSubscriptions: [stable, scheduled], |
| 48 | + }); |
| 49 | + |
| 50 | + const lanes = [ |
| 51 | + "us-east-1:v2:ondemand:fresh:container:stable", |
| 52 | + "us-east-1:v2:scheduled:fresh:container:stable", |
| 53 | + "legacy-region", |
| 54 | + "legacy-region:scheduled", |
| 55 | + ]; |
| 56 | + const runs = []; |
| 57 | + for (const [index, workerQueue] of lanes.entries()) { |
| 58 | + runs.push( |
| 59 | + await engine.trigger( |
| 60 | + { |
| 61 | + number: index + 1, |
| 62 | + friendlyId: generateFriendlyId("run"), |
| 63 | + environment, |
| 64 | + taskIdentifier, |
| 65 | + payload: "{}", |
| 66 | + payloadType: "application/json", |
| 67 | + context: {}, |
| 68 | + traceContext: {}, |
| 69 | + traceId: "test-trace", |
| 70 | + spanId: "test-span", |
| 71 | + workerQueue, |
| 72 | + enableFastPath: true, |
| 73 | + queue: `task/${taskIdentifier}`, |
| 74 | + isTest: false, |
| 75 | + tags: [], |
| 76 | + }, |
| 77 | + prisma |
| 78 | + ) |
| 79 | + ); |
| 80 | + } |
| 81 | + |
| 82 | + expect(() => |
| 83 | + dequeueWorkerQueues({ engine, worker, subscriptions: [stable, canary] }) |
| 84 | + ).toThrow("not authorized"); |
| 85 | + const unchanged = await engine.getRunExecutionData({ runId: runs[0]!.id }); |
| 86 | + expect(unchanged?.snapshot.executionStatus).toBe("QUEUED"); |
| 87 | + |
| 88 | + const deliveredRunIds = new Set<string>(); |
| 89 | + for (const remaining of [1, 0]) { |
| 90 | + const delivery = await dequeueWorkerQueues({ |
| 91 | + engine, |
| 92 | + worker, |
| 93 | + subscriptions: [stable, scheduled], |
| 94 | + }); |
| 95 | + expect(delivery).toHaveLength(1); |
| 96 | + assertNonNullable(delivery[0]); |
| 97 | + deliveredRunIds.add(delivery[0].run.id); |
| 98 | + expect(delivery[0].snapshot.executionStatus).toBe("PENDING_EXECUTING"); |
| 99 | + expect(delivery[0].workerQueueLength).toBe(remaining); |
| 100 | + } |
| 101 | + expect(deliveredRunIds).toEqual(new Set([runs[0]!.id, runs[1]!.id])); |
| 102 | + |
| 103 | + expect((await dequeueWorkerQueues({ engine, worker }))[0]?.run.id).toBe(runs[2]!.id); |
| 104 | + expect( |
| 105 | + (await dequeueWorkerQueues({ engine, worker, queueClass: "scheduled" }))[0]?.run.id |
| 106 | + ).toBe(runs[3]!.id); |
| 107 | + } finally { |
| 108 | + await engine.quit(); |
| 109 | + } |
| 110 | + } |
| 111 | +); |
0 commit comments