diff --git a/.server-changes/dequeue-worker-version-freshness.md b/.server-changes/dequeue-worker-version-freshness.md new file mode 100644 index 0000000000..d584ef2c56 --- /dev/null +++ b/.server-changes/dequeue-worker-version-freshness.md @@ -0,0 +1,6 @@ +--- +area: webapp +type: fix +--- + +Fixed a brief window after promoting or rolling back a deployment where newly triggered runs could still execute on the previous version. New runs now pick up the current version immediately. diff --git a/apps/webapp/app/env.server.ts b/apps/webapp/app/env.server.ts index 3ce98dd521..5467aaca7d 100644 --- a/apps/webapp/app/env.server.ts +++ b/apps/webapp/app/env.server.ts @@ -179,6 +179,7 @@ const EnvironmentSchema = z // Explicit positive opt-in. Split behavior is unreachable unless this is true // AND the distinct-DB sentinel confirms the two URLs are physically distinct DBs. RUN_OPS_SPLIT_ENABLED: BoolEnv.default(false), + RUN_OPS_WORKER_VERSION_FRESH_READ_ENABLED: BoolEnv.default(true), // Canonical connection URL for the dedicated NEW run-ops DB — drives the runtime pool, the split // decision, replication, and migrations. Optional so single-DB installs never set it. RUN_OPS_DATABASE_URL: z diff --git a/apps/webapp/app/v3/runOpsMigration/controlPlaneCache.server.ts b/apps/webapp/app/v3/runOpsMigration/controlPlaneCache.server.ts index c62d6d9593..fdfd6d92ce 100644 --- a/apps/webapp/app/v3/runOpsMigration/controlPlaneCache.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/controlPlaneCache.server.ts @@ -228,7 +228,6 @@ export class ControlPlaneCache { this.#bump(`env:${id}`); } - // worker version: key = `${environmentId}:${backgroundWorkerId ?? "current"}` getWorkerVersion(key: string): (ResolvedWorkerVersion | null) | undefined { return this.#read(this.#version, `version:${key}`); } diff --git a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.dispatchFreshness.test.ts b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.dispatchFreshness.test.ts new file mode 100644 index 0000000000..af1d30288e --- /dev/null +++ b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.dispatchFreshness.test.ts @@ -0,0 +1,320 @@ +/** + * TRI-13291: the dequeue worker-version resolve reads the currently-promoted worker fresh on every + * call. The previous env-keyed TTL cache served the superseded worker for up to the TTL after a + * deployment promotion / dev re-register (nothing invalidated it); it has been removed, so a + * promotion or re-register is reflected on the very next resolve. The DB is never mocked: every + * query runs against the real Postgres container. + */ +import { postgresTest } from "@internal/testcontainers"; +import { describe, expect } from "vitest"; +import type { PrismaClient, PrismaReplicaClient } from "@trigger.dev/database"; +import { CURRENT_DEPLOYMENT_LABEL } from "@trigger.dev/core/v3/isomorphic"; +import { ControlPlaneCache } from "./controlPlaneCache.server"; +import { ControlPlaneResolver } from "./controlPlaneResolver.server"; + +let n = 0; + +async function seedEnv(prisma: PrismaClient, type: "PRODUCTION" | "DEVELOPMENT" = "PRODUCTION") { + const s = n++; + const organization = await prisma.organization.create({ + data: { title: `Org ${s}`, slug: `org-${s}` }, + }); + const project = await prisma.project.create({ + data: { + name: `P ${s}`, + slug: `p-${s}`, + externalRef: `proj_${s}`, + organizationId: organization.id, + }, + }); + const environment = await prisma.runtimeEnvironment.create({ + data: { + type, + slug: `env-${s}`, + projectId: project.id, + organizationId: organization.id, + apiKey: `tr_${s}`, + pkApiKey: `pk_${s}`, + shortcode: `sc_${s}`, + }, + }); + return { organization, project, environment }; +} + +async function seedWorkerWithTask( + prisma: PrismaClient, + ctx: { projectId: string; runtimeEnvironmentId: string }, + version: string, + taskSlug: string +) { + const s = n++; + const worker = await prisma.backgroundWorker.create({ + data: { + friendlyId: `worker_${s}`, + version, + contentHash: `hash_${s}`, + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.runtimeEnvironmentId, + metadata: {}, + }, + }); + await prisma.backgroundWorkerTask.create({ + data: { + friendlyId: `task_${s}`, + slug: taskSlug, + filePath: `src/${taskSlug}.ts`, + workerId: worker.id, + projectId: ctx.projectId, + runtimeEnvironmentId: ctx.runtimeEnvironmentId, + }, + }); + return worker; +} + +async function seedManagedDeployment( + prisma: PrismaClient, + ctx: { projectId: string; environmentId: string; workerId: string }, + version: string +) { + const s = n++; + return prisma.workerDeployment.create({ + data: { + friendlyId: `deploy_${s}`, + shortCode: `dep${s}`, + contentHash: `dhash_${s}`, + version, + type: "MANAGED", + status: "DEPLOYED", + projectId: ctx.projectId, + environmentId: ctx.environmentId, + workerId: ctx.workerId, + }, + }); +} + +async function promote(prisma: PrismaClient, environmentId: string, deploymentId: string) { + await prisma.workerDeploymentPromotion.upsert({ + where: { environmentId_label: { environmentId, label: CURRENT_DEPLOYMENT_LABEL } }, + create: { deploymentId, environmentId, label: CURRENT_DEPLOYMENT_LABEL }, + update: { deploymentId }, + }); +} + +function makeResolver(prisma: PrismaClient, freshRead = true) { + return new ControlPlaneResolver({ + controlPlanePrimary: prisma, + controlPlaneReplica: prisma as unknown as PrismaReplicaClient, + cache: new ControlPlaneCache(), + splitEnabled: () => true, + workerVersionFreshReadEnabled: () => freshRead, + }); +} + +describe("ControlPlaneResolver worker-version dispatch freshness (TRI-13291)", () => { + postgresTest( + "deployed :current: resolves the newly-promoted worker on the next call (no stale cache)", + async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma); + const ctx = { projectId: project.id, runtimeEnvironmentId: environment.id }; + + const workerV1 = await seedWorkerWithTask(prisma, ctx, "20240101.1", "task-a"); + const workerV2 = await seedWorkerWithTask(prisma, ctx, "20240101.2", "task-a"); + const depV1 = await seedManagedDeployment( + prisma, + { projectId: project.id, environmentId: environment.id, workerId: workerV1.id }, + "20240101.1" + ); + const depV2 = await seedManagedDeployment( + prisma, + { projectId: project.id, environmentId: environment.id, workerId: workerV2.id }, + "20240101.2" + ); + + await promote(prisma, environment.id, depV1.id); + + const resolver = makeResolver(prisma); + + const first = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + type: "PRODUCTION", + taskIdentifier: "task-a", + }); + expect(first?.worker.id).toBe(workerV1.id); + expect(first?.worker.version).toBe("20240101.1"); + + await promote(prisma, environment.id, depV2.id); + + const afterPromotion = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + type: "PRODUCTION", + taskIdentifier: "task-a", + }); + expect(afterPromotion?.worker.id).toBe(workerV2.id); + expect(afterPromotion?.worker.version).toBe("20240101.2"); + }, + 30_000 + ); + + postgresTest( + "deployed :current: resolves the rolled-back worker on the next call", + async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma); + const ctx = { projectId: project.id, runtimeEnvironmentId: environment.id }; + + const workerV1 = await seedWorkerWithTask(prisma, ctx, "20240101.1", "task-a"); + const workerV2 = await seedWorkerWithTask(prisma, ctx, "20240101.2", "task-a"); + const depV1 = await seedManagedDeployment( + prisma, + { projectId: project.id, environmentId: environment.id, workerId: workerV1.id }, + "20240101.1" + ); + const depV2 = await seedManagedDeployment( + prisma, + { projectId: project.id, environmentId: environment.id, workerId: workerV2.id }, + "20240101.2" + ); + + await promote(prisma, environment.id, depV2.id); + const resolver = makeResolver(prisma); + const onV2 = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + type: "PRODUCTION", + taskIdentifier: "task-a", + }); + expect(onV2?.worker.id).toBe(workerV2.id); + + await promote(prisma, environment.id, depV1.id); + const rolledBack = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + type: "PRODUCTION", + taskIdentifier: "task-a", + }); + expect(rolledBack?.worker.id).toBe(workerV1.id); + }, + 30_000 + ); + + postgresTest( + "dev :current: resolves the re-registered worker on the next call", + async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma, "DEVELOPMENT"); + const ctx = { projectId: project.id, runtimeEnvironmentId: environment.id }; + + const workerV1 = await seedWorkerWithTask(prisma, ctx, "20240101.1", "task-a"); + + const resolver = makeResolver(prisma); + const first = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + type: "DEVELOPMENT", + taskIdentifier: "task-a", + }); + expect(first?.worker.id).toBe(workerV1.id); + + const workerV2 = await seedWorkerWithTask(prisma, ctx, "20240101.2", "task-a"); + + const afterReregister = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + type: "DEVELOPMENT", + taskIdentifier: "task-a", + }); + expect(afterReregister?.worker.id).toBe(workerV2.id); + }, + 30_000 + ); + + postgresTest( + "resolves only the matched task and queue (per-slug/per-queue dispatch)", + async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma); + const ctx = { projectId: project.id, runtimeEnvironmentId: environment.id }; + + const worker = await seedWorkerWithTask(prisma, ctx, "20240101.1", "task-a"); + await prisma.backgroundWorkerTask.create({ + data: { + friendlyId: `task_extra_${n++}`, + slug: "task-b", + filePath: "src/task-b.ts", + workerId: worker.id, + projectId: project.id, + runtimeEnvironmentId: environment.id, + }, + }); + const qA = await prisma.taskQueue.create({ + data: { + friendlyId: `q_${n++}`, + name: "queue-a", + runtimeEnvironmentId: environment.id, + projectId: project.id, + workers: { connect: { id: worker.id } }, + }, + }); + await prisma.taskQueue.create({ + data: { + friendlyId: `q_${n++}`, + name: "queue-b", + runtimeEnvironmentId: environment.id, + projectId: project.id, + workers: { connect: { id: worker.id } }, + }, + }); + const dep = await seedManagedDeployment( + prisma, + { projectId: project.id, environmentId: environment.id, workerId: worker.id }, + "20240101.1" + ); + await promote(prisma, environment.id, dep.id); + + const resolved = await makeResolver(prisma).resolveWorkerVersion({ + environmentId: environment.id, + type: "PRODUCTION", + taskIdentifier: "task-a", + queue: { name: "queue-a" }, + }); + + expect(resolved?.tasks.map((t) => t.slug)).toEqual(["task-a"]); + expect(resolved?.queues.map((q) => q.id)).toEqual([qA.id]); + }, + 30_000 + ); + + postgresTest( + "kill-switch off falls back to the legacy env-keyed cache (serves the pre-promotion worker)", + async ({ prisma }) => { + const { project, environment } = await seedEnv(prisma); + const ctx = { projectId: project.id, runtimeEnvironmentId: environment.id }; + + const workerV1 = await seedWorkerWithTask(prisma, ctx, "20240101.1", "task-a"); + const workerV2 = await seedWorkerWithTask(prisma, ctx, "20240101.2", "task-a"); + const depV1 = await seedManagedDeployment( + prisma, + { projectId: project.id, environmentId: environment.id, workerId: workerV1.id }, + "20240101.1" + ); + const depV2 = await seedManagedDeployment( + prisma, + { projectId: project.id, environmentId: environment.id, workerId: workerV2.id }, + "20240101.2" + ); + + await promote(prisma, environment.id, depV1.id); + + const resolver = makeResolver(prisma, false); + const first = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + type: "PRODUCTION", + taskIdentifier: "task-a", + }); + expect(first?.worker.id).toBe(workerV1.id); + + await promote(prisma, environment.id, depV2.id); + + const afterPromotion = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + type: "PRODUCTION", + taskIdentifier: "task-a", + }); + expect(afterPromotion?.worker.id).toBe(workerV1.id); + }, + 30_000 + ); +}); diff --git a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts index 573160b9da..4a92b227a8 100644 --- a/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/controlPlaneResolver.server.ts @@ -58,10 +58,27 @@ export type ControlPlaneResolverOptions = { controlPlaneReplica: PrismaReplicaClient; cache: ControlPlaneCache; splitEnabled: () => boolean; + /** + * When true (the default when omitted), the dequeue worker-version resolve reads the matched + * task/queue fresh on every call (no cache). When false, it falls back to the legacy env-keyed + * cache over the whole task/queue set — a kill-switch, retained only so the read shape can be + * reverted via config. + */ + workerVersionFreshReadEnabled?: () => boolean; }; type CpClient = PrismaClient | PrismaReplicaClient; +type WorkerVersionWheres = { + taskWhere: { slug: string } | undefined; + queueWhere: { id: string } | { name: string } | undefined; +}; + +const ALL_WORKER_VERSION_WHERES: WorkerVersionWheres = { + taskWhere: undefined, + queueWhere: undefined, +}; + function workerVersionKey( environmentId: string, backgroundWorkerId: string | undefined, @@ -88,12 +105,14 @@ export class ControlPlaneResolver { private readonly controlPlaneReplica: PrismaReplicaClient; private readonly cache: ControlPlaneCache; private readonly splitEnabled: () => boolean; + private readonly workerVersionFreshReadEnabled: () => boolean; constructor(opts: ControlPlaneResolverOptions) { this.controlPlanePrimary = opts.controlPlanePrimary; this.controlPlaneReplica = opts.controlPlaneReplica; this.cache = opts.cache; this.splitEnabled = opts.splitEnabled; + this.workerVersionFreshReadEnabled = opts.workerVersionFreshReadEnabled ?? (() => true); } async resolveEnv(environmentId: string): Promise { @@ -346,47 +365,66 @@ export class ControlPlaneResolver { * When omitted, the original app behavior applies (worker-by-id, else current promotion). */ type?: RuntimeEnvironmentType; + taskIdentifier?: string; + queue?: { lockedQueueId?: string | null; name: string }; }): Promise { const { environmentId, backgroundWorkerId, type } = args; - if (!this.splitEnabled()) { - return this.#queryWorkerVersion( - this.controlPlanePrimary, + if (!this.workerVersionFreshReadEnabled()) { + if (!this.splitEnabled()) { + return this.#queryWorkerVersion( + this.controlPlanePrimary, + environmentId, + backgroundWorkerId, + type, + ALL_WORKER_VERSION_WHERES + ); + } + + const key = workerVersionKey(environmentId, backgroundWorkerId, type); + const cached = this.cache.getWorkerVersion(key); + if (cached !== undefined) { + return cached; + } + + const resolved = await this.#queryWorkerVersion( + this.controlPlaneReplica, environmentId, backgroundWorkerId, - type + type, + ALL_WORKER_VERSION_WHERES ); + this.cache.setWorkerVersion(key, resolved); + return resolved; } - const key = workerVersionKey(environmentId, backgroundWorkerId, type); - const cached = this.cache.getWorkerVersion(key); - if (cached !== undefined) { - return cached; - } + const wheres: WorkerVersionWheres = { + taskWhere: args.taskIdentifier ? { slug: args.taskIdentifier } : undefined, + queueWhere: args.queue + ? args.queue.lockedQueueId + ? { id: args.queue.lockedQueueId } + : { name: args.queue.name } + : undefined, + }; - const resolved = await this.#queryWorkerVersion( - this.controlPlaneReplica, - environmentId, - backgroundWorkerId, - type - ); - this.cache.setWorkerVersion(key, resolved); - return resolved; + const client = this.splitEnabled() ? this.controlPlaneReplica : this.controlPlanePrimary; + return this.#queryWorkerVersion(client, environmentId, backgroundWorkerId, type, wheres); } async #queryWorkerVersion( client: CpClient, environmentId: string, - backgroundWorkerId?: string, - type?: RuntimeEnvironmentType + backgroundWorkerId: string | undefined, + type: RuntimeEnvironmentType | undefined, + wheres: WorkerVersionWheres ): Promise { // Full run-engine dequeue dispatch (mirrors dequeueSystem's four helpers) when the env type is // known. DEVELOPMENT envs resolve by most-recent worker; deployed envs resolve the promoted // MANAGED deployment. if (type === "DEVELOPMENT") { return backgroundWorkerId - ? this.#queryWorkerById(client, backgroundWorkerId) - : this.#queryMostRecentWorker(client, environmentId); + ? this.#queryWorkerById(client, backgroundWorkerId, wheres) + : this.#queryMostRecentWorker(client, environmentId, wheres); } if (backgroundWorkerId) { @@ -394,8 +432,8 @@ export class ControlPlaneResolver { where: { id: backgroundWorkerId }, include: { deployment: { select: resolvedWorkerDeploymentSelect }, - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, }); @@ -423,8 +461,8 @@ export class ControlPlaneResolver { type: true, worker: { include: { - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, }, }, @@ -460,8 +498,8 @@ export class ControlPlaneResolver { ...resolvedWorkerDeploymentSelect, worker: { include: { - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, }, }, @@ -486,14 +524,15 @@ export class ControlPlaneResolver { async #queryWorkerById( client: CpClient, - workerId: string + workerId: string, + wheres: WorkerVersionWheres ): Promise { const worker = await client.backgroundWorker.findFirst({ where: { id: workerId }, include: { deployment: { select: resolvedWorkerDeploymentSelect }, - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, }); @@ -506,13 +545,14 @@ export class ControlPlaneResolver { async #queryMostRecentWorker( client: CpClient, - environmentId: string + environmentId: string, + wheres: WorkerVersionWheres ): Promise { const worker = await client.backgroundWorker.findFirst({ where: { runtimeEnvironmentId: environmentId }, include: { - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, orderBy: [{ createdAt: "desc" }, { id: "desc" }], }); @@ -592,4 +632,5 @@ export const controlPlaneResolver = new ControlPlaneResolver({ maxEntries: env.CONTROL_PLANE_CACHE_MAX_ENTRIES ?? DEFAULT_CP_CACHE_MAX_ENTRIES, }), splitEnabled: () => SPLIT_ENABLED, + workerVersionFreshReadEnabled: () => env.RUN_OPS_WORKER_VERSION_FRESH_READ_ENABLED, }); diff --git a/apps/webapp/app/v3/runOpsMigration/runEngineControlPlaneResolver.server.ts b/apps/webapp/app/v3/runOpsMigration/runEngineControlPlaneResolver.server.ts index 014e446464..71e7eafa88 100644 --- a/apps/webapp/app/v3/runOpsMigration/runEngineControlPlaneResolver.server.ts +++ b/apps/webapp/app/v3/runOpsMigration/runEngineControlPlaneResolver.server.ts @@ -49,11 +49,15 @@ export class RunEngineControlPlaneResolver implements EngineControlPlaneResolver environmentId: string; type: RuntimeEnvironmentType; workerId?: string; + taskIdentifier?: string; + queue?: { lockedQueueId?: string | null; name: string }; }): Promise { return this.#resolver.resolveWorkerVersion({ environmentId: args.environmentId, backgroundWorkerId: args.workerId, type: args.type, + taskIdentifier: args.taskIdentifier, + queue: args.queue, }); } diff --git a/apps/webapp/test/engine/dequeueWorkerVersionFreshness.test.ts b/apps/webapp/test/engine/dequeueWorkerVersionFreshness.test.ts new file mode 100644 index 0000000000..da906e8940 --- /dev/null +++ b/apps/webapp/test/engine/dequeueWorkerVersionFreshness.test.ts @@ -0,0 +1,149 @@ +import { describe, expect, onTestFinished, vi } from "vitest"; + +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; + 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 { setTimeout } from "node:timers/promises"; +import type { PrismaClient, PrismaReplicaClient } from "@trigger.dev/database"; +import { IdempotencyKeyConcern } from "~/runEngine/concerns/idempotencyKeys.server"; +import { DefaultQueueManager } from "~/runEngine/concerns/queues.server"; +import { RunEngineTriggerTaskService } from "../../app/runEngine/services/triggerTask.server"; +import { ControlPlaneCache } from "~/v3/runOpsMigration/controlPlaneCache.server"; +import { ControlPlaneResolver } from "~/v3/runOpsMigration/controlPlaneResolver.server"; +import { RunEngineControlPlaneResolver } from "~/v3/runOpsMigration/runEngineControlPlaneResolver.server"; +import { + MockPayloadProcessor, + MockTraceEventConcern, + MockTriggerTaskValidator, +} from "./triggerTaskTestHelpers"; + +vi.setConfig({ testTimeout: 60_000, hookTimeout: 60_000 }); + +function buildEngine(prisma: PrismaClient, redisOptions: any) { + const appResolver = new ControlPlaneResolver({ + controlPlanePrimary: prisma, + controlPlaneReplica: prisma as unknown as PrismaReplicaClient, + cache: new ControlPlaneCache(), + splitEnabled: () => true, + workerVersionFreshReadEnabled: () => true, + }); + + return 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"), + controlPlaneResolver: new RunEngineControlPlaneResolver(appResolver), + }); +} + +function buildTriggerService(engine: RunEngine, prisma: PrismaClient) { + return new RunEngineTriggerTaskService({ + engine, + prisma, + payloadProcessor: new MockPayloadProcessor(), + queueConcern: new DefaultQueueManager(prisma, engine), + idempotencyKeyConcern: new IdempotencyKeyConcern(prisma, engine, new MockTraceEventConcern()), + validator: new MockTriggerTaskValidator(), + traceEventConcern: new MockTraceEventConcern(), + tracer: trace.getTracer("test", "0.0.0"), + metadataMaximumSize: 1024 * 1024 * 1, + }); +} + +describe("Dequeue worker-version dispatch freshness end-to-end (TRI-13291)", () => { + containerTest( + "a run triggered after a mid-stream promotion dequeues onto the NEWLY-promoted version", + async ({ prisma, redisOptions }) => { + const engine = buildEngine(prisma as unknown as PrismaClient, redisOptions); + onTestFinished(() => engine.quit()); + + const environment = await setupAuthenticatedEnvironment(prisma, "PRODUCTION"); + const taskIdentifier = "fresh-task"; + + const v1 = await setupBackgroundWorker(engine, environment, taskIdentifier); + const triggerService = buildTriggerService(engine, prisma as unknown as PrismaClient); + + const run1 = await triggerService.call({ + taskId: taskIdentifier, + environment, + body: { payload: { n: 1 } }, + }); + assertNonNullable(run1); + + await engine.runQueue.processMasterQueueForEnvironment(environment.id, 10); + await setTimeout(500); + + const run1Row = await prisma.taskRun.findUniqueOrThrow({ where: { id: run1.run.id } }); + const dequeued1 = await engine.dequeueFromWorkerQueue({ + consumerId: "c1", + workerQueue: run1Row.workerQueue!, + }); + expect(dequeued1.length).toBe(1); + assertNonNullable(dequeued1[0]); + expect(dequeued1[0].run.id).toBe(run1.run.id); + expect(dequeued1[0].backgroundWorker.id).toBe(v1.worker.id); + expect(dequeued1[0].backgroundWorker.version).toBe(v1.worker.version); + + const v2 = await setupBackgroundWorker(engine, environment, taskIdentifier); + expect(v2.worker.id).not.toBe(v1.worker.id); + expect(v2.worker.version).not.toBe(v1.worker.version); + + const currentPromotion = await prisma.workerDeploymentPromotion.findFirstOrThrow({ + where: { environmentId: environment.id, label: "current" }, + include: { deployment: true }, + }); + expect(currentPromotion.deployment.workerId).toBe(v2.worker.id); + + const run2 = await triggerService.call({ + taskId: taskIdentifier, + environment, + body: { payload: { n: 2 } }, + }); + assertNonNullable(run2); + + await engine.runQueue.processMasterQueueForEnvironment(environment.id, 10); + await setTimeout(500); + + const run2Row = await prisma.taskRun.findUniqueOrThrow({ where: { id: run2.run.id } }); + const dequeued2 = await engine.dequeueFromWorkerQueue({ + consumerId: "c2", + workerQueue: run2Row.workerQueue!, + }); + expect(dequeued2.length).toBe(1); + assertNonNullable(dequeued2[0]); + expect(dequeued2[0].run.id).toBe(run2.run.id); + + expect(dequeued2[0].backgroundWorker.id).toBe(v2.worker.id); + expect(dequeued2[0].backgroundWorker.version).toBe(v2.worker.version); + expect(dequeued2[0].deployment.id).toBe(v2.deployment.id); + expect(dequeued2[0].image).toContain(v2.worker.version); + } + ); +}); diff --git a/apps/webapp/test/v3/runOpsMigration/controlPlaneResolver.server.test.ts b/apps/webapp/test/v3/runOpsMigration/controlPlaneResolver.server.test.ts index af734ca1f7..08b383afbc 100644 --- a/apps/webapp/test/v3/runOpsMigration/controlPlaneResolver.server.test.ts +++ b/apps/webapp/test/v3/runOpsMigration/controlPlaneResolver.server.test.ts @@ -224,7 +224,7 @@ heteroPostgresTest( // --- resolveWorkerVersion --------------------------------------------------- heteroPostgresTest( - "resolveWorkerVersion (pinned) returns worker/tasks/queues and caches it", + "resolveWorkerVersion (pinned) returns worker/tasks/queues and reads fresh each call", async ({ prisma14 }) => { const { environment, project } = await seedControlPlane(prisma14); const { worker, task, queue } = await seedWorker(prisma14, { @@ -250,6 +250,37 @@ heteroPostgresTest( const readsAfterFirst = reads(); expect(readsAfterFirst).toBeGreaterThanOrEqual(1); + const second = await resolver.resolveWorkerVersion({ + environmentId: environment.id, + backgroundWorkerId: worker.id, + }); + expect(second?.worker.id).toBe(worker.id); + expect(reads()).toBeGreaterThan(readsAfterFirst); + } +); + +heteroPostgresTest( + "resolveWorkerVersion (pinned) serves from cache when the kill-switch is off", + async ({ prisma14 }) => { + const { environment, project } = await seedControlPlane(prisma14); + const { worker } = await seedWorker(prisma14, { + projectId: project.id, + environmentId: environment.id, + }); + const { client: counting, reads } = countQueries(prisma14); + const resolver = new ControlPlaneResolver({ + controlPlaneReplica: counting, + controlPlanePrimary: counting, + cache: new ControlPlaneCache(), + splitEnabled: () => true, + workerVersionFreshReadEnabled: () => false, + }); + + await resolver.resolveWorkerVersion({ + environmentId: environment.id, + backgroundWorkerId: worker.id, + }); + const readsAfterFirst = reads(); const second = await resolver.resolveWorkerVersion({ environmentId: environment.id, backgroundWorkerId: worker.id, @@ -283,7 +314,7 @@ heteroPostgresTest( const second = await resolver.resolveWorkerVersion({ environmentId: environment.id }); expect(second?.worker.id).toBe(worker.id); - expect(reads()).toBe(readsAfterFirst); + expect(reads()).toBeGreaterThan(readsAfterFirst); } ); diff --git a/internal-packages/run-engine/src/engine/controlPlaneResolver.ts b/internal-packages/run-engine/src/engine/controlPlaneResolver.ts index 256314b1d9..114095cc67 100644 --- a/internal-packages/run-engine/src/engine/controlPlaneResolver.ts +++ b/internal-packages/run-engine/src/engine/controlPlaneResolver.ts @@ -114,17 +114,47 @@ export type ResolvedWorkerVersion = { deployment: ResolvedWorkerDeployment | null; }; +/** + * Optional dispatch filter threaded from the dequeue consumer. When present, the resolve fetches + * ONLY the task matching `taskIdentifier` and the queue matching `queue` (index point-lookups via + * `@@unique([workerId, slug])` / `@@unique([runtimeEnvironmentId, name])`) instead of the worker's + * whole task/queue set. Absent (any non-dequeue caller) keeps the full-set behaviour. + */ +export type WorkerVersionDispatchFilter = { + taskIdentifier?: string; + queue?: { lockedQueueId?: string | null; name: string }; +}; + export interface ControlPlaneResolver { resolveEnv(environmentId: string): Promise; resolveAuthenticatedEnv(environmentId: string): Promise; - resolveWorkerVersion(args: { - environmentId: string; - type: RuntimeEnvironmentType; - workerId?: string; - }): Promise; + resolveWorkerVersion( + args: { + environmentId: string; + type: RuntimeEnvironmentType; + workerId?: string; + } & WorkerVersionDispatchFilter + ): Promise; assertEnvExists(environmentId: string): Promise; } +type WorkerVersionWheres = { + taskWhere: { slug: string } | undefined; + queueWhere: { id: string } | { name: string } | undefined; +}; + +/** Build the nested-include `where`s for a dispatch filter (undefined = fetch the whole set). */ +export function workerVersionWheres(filter: WorkerVersionDispatchFilter): WorkerVersionWheres { + return { + taskWhere: filter.taskIdentifier ? { slug: filter.taskIdentifier } : undefined, + queueWhere: filter.queue + ? filter.queue.lockedQueueId + ? { id: filter.queue.lockedQueueId } + : { name: filter.queue.name } + : undefined, + }; +} + export class PassthroughControlPlaneResolver implements ControlPlaneResolver { readonly #prisma: PrismaClient; @@ -240,31 +270,39 @@ export class PassthroughControlPlaneResolver implements ControlPlaneResolver { // no cross-seam FK to replace (matches main, which dropped the TaskRun env FK). } - async resolveWorkerVersion(args: { - environmentId: string; - type: RuntimeEnvironmentType; - workerId?: string; - }): Promise { + async resolveWorkerVersion( + args: { + environmentId: string; + type: RuntimeEnvironmentType; + workerId?: string; + } & WorkerVersionDispatchFilter + ): Promise { const { environmentId, type, workerId } = args; + const wheres = workerVersionWheres(args); if (type === "DEVELOPMENT") { - return workerId ? this.#getWorkerById(workerId) : this.#getMostRecentWorker(environmentId); + return workerId + ? this.#getWorkerById(workerId, wheres) + : this.#getMostRecentWorker(environmentId, wheres); } return workerId - ? this.#getWorkerDeploymentFromWorker(workerId) - : this.#getManagedWorkerFromCurrentlyPromotedDeployment(environmentId); + ? this.#getWorkerDeploymentFromWorker(workerId, wheres) + : this.#getManagedWorkerFromCurrentlyPromotedDeployment(environmentId, wheres); } - async #getWorkerDeploymentFromWorker(workerId: string): Promise { + async #getWorkerDeploymentFromWorker( + workerId: string, + wheres: WorkerVersionWheres + ): Promise { const worker = await this.#prisma.backgroundWorker.findFirst({ where: { id: workerId, }, include: { deployment: { select: resolvedWorkerDeploymentSelect }, - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, }); @@ -280,14 +318,17 @@ export class PassthroughControlPlaneResolver implements ControlPlaneResolver { }; } - async #getMostRecentWorker(environmentId: string): Promise { + async #getMostRecentWorker( + environmentId: string, + wheres: WorkerVersionWheres + ): Promise { const worker = await this.#prisma.backgroundWorker.findFirst({ where: { runtimeEnvironmentId: environmentId, }, include: { - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, orderBy: [{ createdAt: "desc" }, { id: "desc" }], }); @@ -299,15 +340,18 @@ export class PassthroughControlPlaneResolver implements ControlPlaneResolver { return { worker, tasks: worker.tasks, queues: worker.queues, deployment: null }; } - async #getWorkerById(workerId: string): Promise { + async #getWorkerById( + workerId: string, + wheres: WorkerVersionWheres + ): Promise { const worker = await this.#prisma.backgroundWorker.findFirst({ where: { id: workerId, }, include: { deployment: { select: resolvedWorkerDeploymentSelect }, - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, }); @@ -324,7 +368,8 @@ export class PassthroughControlPlaneResolver implements ControlPlaneResolver { } async #getManagedWorkerFromCurrentlyPromotedDeployment( - environmentId: string + environmentId: string, + wheres: WorkerVersionWheres ): Promise { const promotion = await this.#prisma.workerDeploymentPromotion.findFirst({ where: { @@ -338,8 +383,8 @@ export class PassthroughControlPlaneResolver implements ControlPlaneResolver { type: true, worker: { include: { - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, }, }, @@ -378,8 +423,8 @@ export class PassthroughControlPlaneResolver implements ControlPlaneResolver { ...resolvedWorkerDeploymentSelect, worker: { include: { - tasks: { select: resolvedWorkerTaskSelect }, - queues: { select: resolvedTaskQueueSelect }, + tasks: { where: wheres.taskWhere, select: resolvedWorkerTaskSelect }, + queues: { where: wheres.queueWhere, select: resolvedTaskQueueSelect }, }, }, }, diff --git a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts index f8f31d848a..b1584cfb69 100644 --- a/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts +++ b/internal-packages/run-engine/src/engine/systems/dequeueSystem.ts @@ -895,6 +895,8 @@ export class DequeueSystem { environmentId: run.runtimeEnvironmentId, type: env.type, workerId: workerId ?? undefined, + taskIdentifier: run.taskIdentifier, + queue: { lockedQueueId: run.lockedQueueId, name: run.queue }, }); if (!workerWithTasks) {