diff --git a/internal-packages/run-engine/src/engine/index.ts b/internal-packages/run-engine/src/engine/index.ts index 2f7447af7..588c3a147 100644 --- a/internal-packages/run-engine/src/engine/index.ts +++ b/internal-packages/run-engine/src/engine/index.ts @@ -114,6 +114,7 @@ export class RunEngine { private repairSnapshotTimeoutMs: number; private batchQueue: BatchQueue; private workerQueueObserverAbortController?: AbortController; + private quitPromise?: Promise; prisma: PrismaClient; readOnlyPrisma: PrismaReplicaClient; @@ -2312,26 +2313,57 @@ export class RunEngine { } } - async quit() { - try { - this.workerQueueObserverAbortController?.abort(); + quit(): Promise { + this.quitPromise ??= this.#quit(); + return this.quitPromise; + } - await this.runQueue.quit(); - await this.worker.stop(); - await this.ttlWorker.stop(); - await this.runLock.quit(); + async #quit(): Promise { + this.workerQueueObserverAbortController?.abort(); - // This is just a failsafe - await this.runLockRedis.quit(); + // Stop resources that actively process work before closing support resources they may use. + const processingResults = await Promise.allSettled([ + this.runQueue.quit(), + this.worker.stop(), + this.ttlWorker.stop(), + this.batchQueue.close(), + ]); + this.#logShutdownFailures( + ["runQueue.quit", "worker.stop", "ttlWorker.stop", "batchQueue.close"], + processingResults + ); - await this.batchQueue.close(); + const supportResults = await Promise.allSettled([ + this.runLock.quit(), + this.debounceSystem.quit(), + ]); + this.#logShutdownFailures(["runLock.quit", "debounceSystem.quit"], supportResults); - await this.debounceSystem.quit(); - } catch (_error) { - // Best-effort shutdown; ignore quit/close errors. + // RunLocker/Redlock owns this client and normally closes it. Do not send a second QUIT, + // but force-disconnect if Redlock failed to leave the connection in its terminal state. + if (this.runLockRedis.status !== "end") { + try { + this.runLockRedis.disconnect(); + } catch (error) { + this.logger.error("RunEngine shutdown operation failed", { + operation: "runLockRedis.disconnect", + error, + }); + } } } + #logShutdownFailures(operations: string[], results: PromiseSettledResult[]): void { + results.forEach((result, index) => { + if (result.status === "rejected") { + this.logger.error("RunEngine shutdown operation failed", { + operation: operations[index], + error: result.reason, + }); + } + }); + } + async repairEnvironment(environment: AuthenticatedEnvironment, dryRun: boolean) { const runIds = await this.runQueue.getCurrentConcurrencyOfEnvironment(environment); diff --git a/internal-packages/run-engine/src/engine/tests/shutdown.test.ts b/internal-packages/run-engine/src/engine/tests/shutdown.test.ts new file mode 100644 index 000000000..d12e73aee --- /dev/null +++ b/internal-packages/run-engine/src/engine/tests/shutdown.test.ts @@ -0,0 +1,98 @@ +import { createRedisClient, type Redis, type RedisOptions } from "@internal/redis"; +import { containerTestWithIsolatedRedisNoClickhouse } from "@internal/testcontainers"; +import { trace } from "@internal/tracing"; +import { Logger } from "@trigger.dev/core/logger"; +import { expect } from "vitest"; +import { RunEngine } from "../index.js"; + +async function connectedClientCount(redis: Redis): Promise { + const clientsInfo = await redis.info("clients"); + const match = clientsInfo.match(/^connected_clients:(\d+)$/m); + + if (!match) { + throw new Error("Redis INFO clients response did not include connected_clients"); + } + + return Number(match[1]); +} + +function engineOptions(redisOptions: RedisOptions) { + // Keep caches and disabled consumers lazy so every connection opened by this test belongs to a + // shutdown resource. The run-lock client remains eager to exercise Redlock's ownership of it. + const lazyRedisOptions = { ...redisOptions, lazyConnect: true }; + + return { + worker: { + disabled: true, + redis: lazyRedisOptions, + workers: 1, + tasksPerWorker: 1, + pollIntervalMs: 10, + immediatePollIntervalMs: 10, + shutdownTimeoutMs: 30_000, + }, + queue: { + redis: lazyRedisOptions, + masterQueueConsumersDisabled: true, + ttlSystem: { disabled: true }, + logLevel: "error" as const, + }, + runLock: { redis: redisOptions }, + cache: { redis: lazyRedisOptions }, + debounce: { redis: lazyRedisOptions }, + batchQueue: { + redis: lazyRedisOptions, + consumerEnabled: false, + }, + machines: { + defaultMachine: "small-1x" as const, + machines: { + "small-1x": { + name: "small-1x" as const, + cpu: 0.5, + memory: 0.5, + centsPerMs: 0.0001, + }, + }, + baseCostInCents: 0.0001, + }, + tracer: trace.getTracer("run-engine-shutdown-test", "0.0.0"), + logger: new Logger("run-engine-shutdown-test", "error"), + }; +} + +describe("RunEngine.quit", () => { + containerTestWithIsolatedRedisNoClickhouse( + "is concurrency-safe, repeatable, and returns Redis connections to baseline", + { timeout: 60_000 }, + async ({ prisma, redisOptions }) => { + await prisma.$queryRaw`SELECT 1`; + + const observer = createRedisClient(redisOptions); + await observer.ping(); + const baselineConnections = await connectedClientCount(observer); + const engine = new RunEngine({ prisma, ...engineOptions(redisOptions) }); + + try { + await expect + .poll(() => connectedClientCount(observer)) + .toBeGreaterThan(baselineConnections); + + const firstQuit = engine.quit(); + const concurrentQuit = engine.quit(); + expect(concurrentQuit).toBe(firstQuit); + + await Promise.all([firstQuit, concurrentQuit, engine.quit()]); + + const repeatedQuit = engine.quit(); + expect(repeatedQuit).toBe(firstQuit); + await repeatedQuit; + + await expect.poll(() => connectedClientCount(observer)).toBe(baselineConnections); + } finally { + await engine.quit(); + await observer.quit(); + } + } + ); +}); diff --git a/packages/redis-worker/src/worker.test.ts b/packages/redis-worker/src/worker.test.ts index f5659f379..16d4ffe54 100644 --- a/packages/redis-worker/src/worker.test.ts +++ b/packages/redis-worker/src/worker.test.ts @@ -4,7 +4,22 @@ import { describe } from "node:test"; import { expect } from "vitest"; import { z } from "zod"; import { Worker } from "./worker.js"; -import { createRedisClient } from "@internal/redis"; +import { createRedisClient, type Redis } from "@internal/redis"; + +async function connectedClientCount(redis: Redis): Promise { + const clientsInfo = await redis.info("clients"); + const match = clientsInfo.match(/^connected_clients:(\d+)$/m); + + if (!match) { + throw new Error("Redis INFO clients response did not include connected_clients"); + } + + return Number(match[1]); +} + +function activeTimeoutCount(): number { + return process.getActiveResourcesInfo().filter((resource) => resource === "Timeout").length; +} describe("Worker", () => { redisTest("Process items that don't throw", { timeout: 30_000 }, async ({ redisContainer }) => { @@ -549,6 +564,58 @@ describe("Worker", () => { } ); + redisTest( + "clears its shutdown deadline and closes Redis connections after a prompt stop", + { timeout: 30_000 }, + async ({ redisContainer }) => { + const redisOptions = { + host: redisContainer.getHost(), + port: redisContainer.getPort(), + password: redisContainer.getPassword(), + }; + const observer = createRedisClient(redisOptions); + await observer.ping(); + + const baselineConnections = await connectedClientCount(observer); + const baselineTimeouts = activeTimeoutCount(); + const worker = new Worker({ + name: "shutdown-lifecycle-worker", + redisOptions, + catalog: { + testJob: { + schema: z.object({ value: z.number() }), + visibilityTimeoutMs: 5000, + }, + }, + jobs: { + testJob: async () => {}, + }, + concurrency: { workers: 1, tasksPerWorker: 1 }, + pollIntervalMs: 10, + immediatePollIntervalMs: 10, + shutdownTimeoutMs: 30_000, + logger: new Logger("shutdown-lifecycle-test", "error"), + }).start(); + + try { + await expect + .poll(() => connectedClientCount(observer)) + .toBeGreaterThan(baselineConnections); + + // Let the worker enter its polling loop so the loop, rather than the deadline, wins shutdown. + await new Promise((resolve) => setTimeout(resolve, 50)); + await worker.stop(); + + await expect.poll(() => connectedClientCount(observer)).toBe(baselineConnections); + await new Promise((resolve) => setImmediate(resolve)); + expect(activeTimeoutCount()).toBeLessThanOrEqual(baselineTimeouts); + } finally { + await worker.stop(); + await observer.quit(); + } + } + ); + redisTest( "Should allow cancelling a job before it's enqueued, but only if the enqueue.cancellationKey is provided", { timeout: 30_000 }, diff --git a/packages/redis-worker/src/worker.ts b/packages/redis-worker/src/worker.ts index 64268a1c9..0db651a54 100644 --- a/packages/redis-worker/src/worker.ts +++ b/packages/redis-worker/src/worker.ts @@ -1198,16 +1198,26 @@ class Worker { this.isShuttingDown = true; this.logger.log("Shutting down worker loops...", { signal }); - // Wait for all worker loops to finish. - await Promise.race([ - Promise.all(this.workerLoops), - Worker.delay(this.shutdownTimeoutMs).then(() => { + // Wait for all worker loops to finish, retaining ownership of the deadline timer so the + // losing timeout cannot keep the process alive after a prompt shutdown. + let shutdownDeadline: ReturnType | undefined; + const deadlinePromise = new Promise((resolve) => { + shutdownDeadline = setTimeout(() => { this.logger.error("Worker shutdown timed out", { signal, shutdownTimeoutMs: this.shutdownTimeoutMs, }); - }), - ]); + resolve(); + }, this.shutdownTimeoutMs); + }); + + try { + await Promise.race([Promise.all(this.workerLoops), deadlinePromise]); + } finally { + if (shutdownDeadline) { + clearTimeout(shutdownDeadline); + } + } await this.subscriber?.unsubscribe(); await this.subscriber?.quit();