add shutdown timeout to redis-worker

Author: nicktrn

apps/webapp/app/env.server.ts

@@ -424,6 +424,7 @@ const EnvironmentSchema = z.object({
   RUN_ENGINE_QUEUE_AGE_RANDOMIZATION_BIAS: z.coerce.number().default(0.25),
   RUN_ENGINE_REUSE_SNAPSHOT_COUNT: z.coerce.number().int().default(0),
   RUN_ENGINE_MAXIMUM_ENV_COUNT: z.coerce.number().int().optional(),
+  RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS: z.coerce.number().int().default(60_000),
 
   RUN_ENGINE_WORKER_REDIS_HOST: z
     .string()
@@ -585,6 +586,7 @@ const EnvironmentSchema = z.object({
   LEGACY_RUN_ENGINE_WORKER_POLL_INTERVAL: z.coerce.number().int().default(1000),
   LEGACY_RUN_ENGINE_WORKER_IMMEDIATE_POLL_INTERVAL: z.coerce.number().int().default(50),
   LEGACY_RUN_ENGINE_WORKER_CONCURRENCY_LIMIT: z.coerce.number().int().default(100),
+  LEGACY_RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS: z.coerce.number().int().default(60_000),
 
   LEGACY_RUN_ENGINE_WORKER_REDIS_HOST: z
     .string()
@@ -627,6 +629,7 @@ const EnvironmentSchema = z.object({
   COMMON_WORKER_POLL_INTERVAL: z.coerce.number().int().default(1000),
   COMMON_WORKER_IMMEDIATE_POLL_INTERVAL: z.coerce.number().int().default(50),
   COMMON_WORKER_CONCURRENCY_LIMIT: z.coerce.number().int().default(100),
+  COMMON_WORKER_SHUTDOWN_TIMEOUT_MS: z.coerce.number().int().default(60_000),
 
   COMMON_WORKER_REDIS_HOST: z
     .string()

apps/webapp/app/v3/commonWorker.server.ts

@@ -80,6 +80,7 @@ function initializeWorker() {
     },
     pollIntervalMs: env.COMMON_WORKER_POLL_INTERVAL,
     immediatePollIntervalMs: env.COMMON_WORKER_IMMEDIATE_POLL_INTERVAL,
+    shutdownTimeoutMs: env.COMMON_WORKER_SHUTDOWN_TIMEOUT_MS,
     logger: new Logger("CommonWorker", "debug"),
     jobs: {
       "v3.deliverAlert": async ({ payload }) => {

apps/webapp/app/v3/legacyRunEngineWorker.server.ts

@@ -67,6 +67,7 @@ function initializeWorker() {
     },
     pollIntervalMs: env.LEGACY_RUN_ENGINE_WORKER_POLL_INTERVAL,
     immediatePollIntervalMs: env.LEGACY_RUN_ENGINE_WORKER_IMMEDIATE_POLL_INTERVAL,
+    shutdownTimeoutMs: env.LEGACY_RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS,
     logger: new Logger("LegacyRunEngineWorker", "debug"),
     jobs: {
       runHeartbeat: async ({ payload }) => {

apps/webapp/app/v3/runEngine.server.ts

@@ -17,6 +17,7 @@ function createRunEngine() {
       workers: env.RUN_ENGINE_WORKER_COUNT,
       tasksPerWorker: env.RUN_ENGINE_TASKS_PER_WORKER,
       pollIntervalMs: env.RUN_ENGINE_WORKER_POLL_INTERVAL,
+      shutdownTimeoutMs: env.RUN_ENGINE_WORKER_SHUTDOWN_TIMEOUT_MS,
       redis: {
         keyPrefix: "engine:",
         port: env.RUN_ENGINE_WORKER_REDIS_PORT ?? undefined,

internal-packages/redis-worker/src/worker.ts

@@ -8,6 +8,7 @@ import { AnyQueueItem, SimpleQueue } from "./queue.js";
 import { nanoid } from "nanoid";
 import pLimit from "p-limit";
 import { createRedisClient } from "@internal/redis";
+import { shutdownManager } from "@trigger.dev/core/v3/serverOnly";
 
 export type WorkerCatalog = {
   [key: string]: {
@@ -44,6 +45,7 @@ type WorkerOptions<TCatalog extends WorkerCatalog> = {
   concurrency?: WorkerConcurrencyOptions;
   pollIntervalMs?: number;
   immediatePollIntervalMs?: number;
+  shutdownTimeoutMs?: number;
   logger?: Logger;
   tracer?: Tracer;
 };
@@ -69,6 +71,7 @@ class Worker<TCatalog extends WorkerCatalog> {
   private workerLoops: Promise<void>[] = [];
   private isShuttingDown = false;
   private concurrency: Required<NonNullable<WorkerOptions<TCatalog>["concurrency"]>>;
+  private shutdownTimeoutMs: number;
 
   // The p-limit limiter to control overall concurrency.
   private limiter: ReturnType<typeof pLimit>;
@@ -77,6 +80,8 @@ class Worker<TCatalog extends WorkerCatalog> {
     this.logger = options.logger ?? new Logger("Worker", "debug");
     this.tracer = options.tracer ?? trace.getTracer(options.name);
 
+    this.shutdownTimeoutMs = options.shutdownTimeoutMs ?? 60_000;
+
     const schema: QueueCatalogFromWorkerCatalog<TCatalog> = Object.fromEntries(
       Object.entries(this.options.catalog).map(([key, value]) => [key, value.schema])
     ) as QueueCatalogFromWorkerCatalog<TCatalog>;
@@ -386,22 +391,33 @@ class Worker<TCatalog extends WorkerCatalog> {
   }
 
   private setupShutdownHandlers() {
-    process.on("SIGTERM", this.shutdown.bind(this));
-    process.on("SIGINT", this.shutdown.bind(this));
+    shutdownManager.register("redis-worker", this.shutdown.bind(this));
   }
 
-  private async shutdown() {
-    if (this.isShuttingDown) return;
+  private async shutdown(signal?: NodeJS.Signals) {
+    if (this.isShuttingDown) {
+      this.logger.log("Worker already shutting down", { signal });
+      return;
+    }
+
     this.isShuttingDown = true;
-    this.logger.log("Shutting down worker loops...");
+    this.logger.log("Shutting down worker loops...", { signal });
 
     // Wait for all worker loops to finish.
-    await Promise.all(this.workerLoops);
+    await Promise.race([
+      Promise.all(this.workerLoops),
+      Worker.delay(this.shutdownTimeoutMs).then(() => {
+        this.logger.error("Worker shutdown timed out", {
+          signal,
+          shutdownTimeoutMs: this.shutdownTimeoutMs,
+        });
+      }),
+    ]);
 
     await this.subscriber?.unsubscribe();
     await this.subscriber?.quit();
     await this.queue.close();
-    this.logger.log("All workers and subscribers shut down.");
+    this.logger.log("All workers and subscribers shut down.", { signal });
   }
 
   public async stop() {

internal-packages/run-engine/src/engine/index.ts

@@ -180,6 +180,7 @@ export class RunEngine {
       concurrency: options.worker,
       pollIntervalMs: options.worker.pollIntervalMs,
       immediatePollIntervalMs: options.worker.immediatePollIntervalMs,
+      shutdownTimeoutMs: options.worker.shutdownTimeoutMs,
       logger: new Logger("RunEngineWorker", "debug"),
       jobs: {
         finishWaitpoint: async ({ payload }) => {

internal-packages/run-engine/src/engine/types.ts

@@ -12,6 +12,7 @@ export type RunEngineOptions = {
     redis: RedisOptions;
     pollIntervalMs?: number;
     immediatePollIntervalMs?: number;
+    shutdownTimeoutMs?: number;
   };
   machines: {
     defaultMachine: MachinePresetName;