|
| 1 | +import { RunEngine } from "@internal/run-engine"; |
| 2 | +import { containerTest } from "@internal/testcontainers"; |
| 3 | +import { trace } from "@opentelemetry/api"; |
| 4 | +import type { PrismaClient } from "@trigger.dev/database"; |
| 5 | +import type { RedisOptions } from "ioredis"; |
| 6 | +import { describe, expect, onTestFinished, vi } from "vitest"; |
| 7 | +import type { AuthenticatedEnvironment } from "~/services/apiAuth.server"; |
| 8 | +import { |
| 9 | + createRuntimeEnvironment, |
| 10 | + createTestOrgProjectWithMember, |
| 11 | + uniqueId, |
| 12 | +} from "./fixtures/environmentVariablesFixtures"; |
| 13 | + |
| 14 | +vi.setConfig({ testTimeout: 60_000, hookTimeout: 60_000 }); |
| 15 | + |
| 16 | +// test/setup.ts replaces the app's engine singleton with a no-op for every webapp suite, which |
| 17 | +// would make any assertion about the RunQueue limits vacuous. Every test in this file asserts on |
| 18 | +// real RunQueue state, so put a real RunEngine - built on the test's own Redis container - back |
| 19 | +// behind the singleton. No test here uses the no-op default. |
| 20 | +const { engineHolder } = vi.hoisted(() => ({ |
| 21 | + engineHolder: { current: undefined as any }, |
| 22 | +})); |
| 23 | + |
| 24 | +vi.mock("~/v3/runEngine.server", () => ({ |
| 25 | + engine: new Proxy({} as Record<string, any>, { |
| 26 | + get: (_target, prop) => engineHolder.current?.[prop as string], |
| 27 | + }), |
| 28 | +})); |
| 29 | + |
| 30 | +function useEngine(prisma: PrismaClient, redisOptions: RedisOptions) { |
| 31 | + const engine = new RunEngine({ |
| 32 | + prisma, |
| 33 | + worker: { redis: redisOptions, disabled: true }, |
| 34 | + queue: { redis: redisOptions, masterQueueConsumersDisabled: true }, |
| 35 | + runLock: { redis: redisOptions }, |
| 36 | + machines: { |
| 37 | + defaultMachine: "small-1x", |
| 38 | + machines: { |
| 39 | + "small-1x": { name: "small-1x" as const, cpu: 0.5, memory: 0.5, centsPerMs: 0.0001 }, |
| 40 | + }, |
| 41 | + baseCostInCents: 0.0001, |
| 42 | + }, |
| 43 | + tracer: trace.getTracer("test", "0.0.0"), |
| 44 | + }); |
| 45 | + |
| 46 | + engineHolder.current = engine; |
| 47 | + onTestFinished(async () => { |
| 48 | + engineHolder.current = undefined; |
| 49 | + await engine.quit(); |
| 50 | + }); |
| 51 | + |
| 52 | + return engine; |
| 53 | +} |
| 54 | + |
| 55 | +// The import chain reaches module-level singletons that throw at load time when |
| 56 | +// REDIS_HOST/REDIS_PORT are unset (autoIncrementCounter via triggerTaskV1), so the env must point |
| 57 | +// at the redis container BEFORE the modules are imported. Hence dynamic imports; vitest runs each |
| 58 | +// file in its own fork, so the env mutation cannot leak into other suites. |
| 59 | +async function loadServices(redisOptions: RedisOptions) { |
| 60 | + process.env.REDIS_HOST = redisOptions.host; |
| 61 | + process.env.REDIS_PORT = String(redisOptions.port); |
| 62 | + process.env.REDIS_TLS_DISABLED = "true"; |
| 63 | + const [{ updateEnvConcurrencyLimits }, { PauseEnvironmentService }, runtimeEnvironment] = |
| 64 | + await Promise.all([ |
| 65 | + import("~/v3/runQueue.server"), |
| 66 | + import("~/v3/services/pauseEnvironment.server"), |
| 67 | + import("~/models/runtimeEnvironment.server"), |
| 68 | + ]); |
| 69 | + return { |
| 70 | + updateEnvConcurrencyLimits, |
| 71 | + PauseEnvironmentService, |
| 72 | + authIncludeBase: runtimeEnvironment.authIncludeBase, |
| 73 | + toAuthenticated: runtimeEnvironment.toAuthenticated, |
| 74 | + }; |
| 75 | +} |
| 76 | + |
| 77 | +type Loaded = Awaited<ReturnType<typeof loadServices>>; |
| 78 | + |
| 79 | +async function authEnv( |
| 80 | + loaded: Loaded, |
| 81 | + prisma: PrismaClient, |
| 82 | + environmentId: string |
| 83 | +): Promise<AuthenticatedEnvironment> { |
| 84 | + const row = await prisma.runtimeEnvironment.findFirstOrThrow({ |
| 85 | + where: { id: environmentId }, |
| 86 | + include: loaded.authIncludeBase, |
| 87 | + }); |
| 88 | + return loaded.toAuthenticated(row); |
| 89 | +} |
| 90 | + |
| 91 | +async function seedProductionEnv(prisma: PrismaClient, maximumConcurrencyLimit: number) { |
| 92 | + const { organization, project } = await createTestOrgProjectWithMember(prisma); |
| 93 | + const environment = await createRuntimeEnvironment(prisma, { |
| 94 | + projectId: project.id, |
| 95 | + organizationId: organization.id, |
| 96 | + type: "PRODUCTION", |
| 97 | + slug: uniqueId("prod"), |
| 98 | + }); |
| 99 | + |
| 100 | + await prisma.runtimeEnvironment.update({ |
| 101 | + where: { id: environment.id }, |
| 102 | + data: { maximumConcurrencyLimit }, |
| 103 | + }); |
| 104 | + |
| 105 | + return { organization, project, environment }; |
| 106 | +} |
| 107 | + |
| 108 | +// An unset RunQueue limit reads back as the engine default (10), so neither the 0 nor the 17 |
| 109 | +// assertions below can pass just because a push never happened. |
| 110 | +describe("updateEnvConcurrencyLimits", () => { |
| 111 | + containerTest( |
| 112 | + "clamps to 0 when the environment is paused, even though the caller's copy says otherwise", |
| 113 | + async ({ prisma, redisOptions }) => { |
| 114 | + const loaded = await loadServices(redisOptions); |
| 115 | + const engine = useEngine(prisma, redisOptions); |
| 116 | + |
| 117 | + const { environment } = await seedProductionEnv(prisma, 17); |
| 118 | + // What an argument-less caller holds: an environment read when the request authenticated, |
| 119 | + // before the pause landed (finalizing a deployment, registering a background worker). |
| 120 | + const atAuthTime = await authEnv(loaded, prisma, environment.id); |
| 121 | + expect(atAuthTime.paused).toBe(false); |
| 122 | + |
| 123 | + await prisma.runtimeEnvironment.update({ |
| 124 | + where: { id: environment.id }, |
| 125 | + data: { paused: true }, |
| 126 | + }); |
| 127 | + |
| 128 | + await loaded.updateEnvConcurrencyLimits(atAuthTime, undefined, prisma); |
| 129 | + |
| 130 | + // The 0 limit is the only thing stopping dequeues, so the real limit must not go back in. |
| 131 | + expect(await engine.runQueue.getEnvConcurrencyLimit(atAuthTime)).toBe(0); |
| 132 | + } |
| 133 | + ); |
| 134 | + |
| 135 | + containerTest( |
| 136 | + "pushes the real limit for a running environment", |
| 137 | + async ({ prisma, redisOptions }) => { |
| 138 | + const loaded = await loadServices(redisOptions); |
| 139 | + const engine = useEngine(prisma, redisOptions); |
| 140 | + |
| 141 | + const { environment } = await seedProductionEnv(prisma, 17); |
| 142 | + const env = await authEnv(loaded, prisma, environment.id); |
| 143 | + |
| 144 | + await loaded.updateEnvConcurrencyLimits(env, undefined, prisma); |
| 145 | + |
| 146 | + expect(await engine.runQueue.getEnvConcurrencyLimit(env)).toBe(17); |
| 147 | + } |
| 148 | + ); |
| 149 | + |
| 150 | + containerTest( |
| 151 | + "restores the real limit when the environment was resumed while the request was in flight", |
| 152 | + async ({ prisma, redisOptions }) => { |
| 153 | + const loaded = await loadServices(redisOptions); |
| 154 | + const engine = useEngine(prisma, redisOptions); |
| 155 | + |
| 156 | + const { environment } = await seedProductionEnv(prisma, 17); |
| 157 | + await prisma.runtimeEnvironment.update({ |
| 158 | + where: { id: environment.id }, |
| 159 | + data: { paused: true }, |
| 160 | + }); |
| 161 | + |
| 162 | + // Captured while paused, then resumed before the push. Trusting this copy would write 0 over |
| 163 | + // the restored limit and leave the env stalled with `paused: false` and nothing to fix it. |
| 164 | + const whilePaused = await authEnv(loaded, prisma, environment.id); |
| 165 | + expect(whilePaused.paused).toBe(true); |
| 166 | + |
| 167 | + await prisma.runtimeEnvironment.update({ |
| 168 | + where: { id: environment.id }, |
| 169 | + data: { paused: false }, |
| 170 | + }); |
| 171 | + |
| 172 | + await loaded.updateEnvConcurrencyLimits(whilePaused, undefined, prisma); |
| 173 | + |
| 174 | + expect(await engine.runQueue.getEnvConcurrencyLimit(whilePaused)).toBe(17); |
| 175 | + } |
| 176 | + ); |
| 177 | + |
| 178 | + containerTest( |
| 179 | + "an explicit limit wins over the stored pause state", |
| 180 | + async ({ prisma, redisOptions }) => { |
| 181 | + const loaded = await loadServices(redisOptions); |
| 182 | + const engine = useEngine(prisma, redisOptions); |
| 183 | + |
| 184 | + const { environment } = await seedProductionEnv(prisma, 17); |
| 185 | + await prisma.runtimeEnvironment.update({ |
| 186 | + where: { id: environment.id }, |
| 187 | + data: { paused: true }, |
| 188 | + }); |
| 189 | + const env = await authEnv(loaded, prisma, environment.id); |
| 190 | + |
| 191 | + // How billing-limit converge restores a limit as it unpauses: the caller decides, no read. |
| 192 | + await loaded.updateEnvConcurrencyLimits(env, 9, prisma); |
| 193 | + |
| 194 | + expect(await engine.runQueue.getEnvConcurrencyLimit(env)).toBe(9); |
| 195 | + } |
| 196 | + ); |
| 197 | + |
| 198 | + containerTest( |
| 199 | + "a pause writes 0 and a resume restores the limit", |
| 200 | + async ({ prisma, redisOptions }) => { |
| 201 | + const loaded = await loadServices(redisOptions); |
| 202 | + const engine = useEngine(prisma, redisOptions); |
| 203 | + |
| 204 | + const { environment } = await seedProductionEnv(prisma, 17); |
| 205 | + const service = new loaded.PauseEnvironmentService(prisma); |
| 206 | + const env = await authEnv(loaded, prisma, environment.id); |
| 207 | + |
| 208 | + expect(await service.call(env, "paused")).toEqual({ success: true, state: "paused" }); |
| 209 | + expect(await engine.runQueue.getEnvConcurrencyLimit(env)).toBe(0); |
| 210 | + |
| 211 | + // The service holds an environment read before its own resume update, so `env.paused` is |
| 212 | + // stale here too. |
| 213 | + expect(await service.call(env, "resumed")).toEqual({ success: true, state: "resumed" }); |
| 214 | + expect(await engine.runQueue.getEnvConcurrencyLimit(env)).toBe(17); |
| 215 | + } |
| 216 | + ); |
| 217 | +}); |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?v=4&size=40){kind=avatar}
0 commit comments