fix: improve graceful shutdown in Worker

Author: RomainLanz

src/types/main.ts

@@ -89,6 +89,20 @@ export interface WorkerConfig {
    * @default 1
    */
   maxStalledCount?: number
+
+  /**
+   * Whether to automatically stop the worker on SIGINT/SIGTERM signals.
+   * When enabled, the worker will wait for running jobs to complete
+   * before stopping.
+   * @default true
+   */
+  gracefulShutdown?: boolean
+
+  /**
+   * Callback invoked when a shutdown signal is received.
+   * Called before the worker starts stopping.
+   */
+  onShutdownSignal?: () => void | Promise<void>
 }
 
 export type WorkerCycle =

src/worker.ts

@@ -19,13 +19,16 @@ export class Worker {
   readonly #stalledThreshold: number
   readonly #maxStalledCount: number
   readonly #concurrency: number
+  readonly #gracefulShutdown: boolean
+  readonly #onShutdownSignal?: () => void | Promise<void>
 
   #adapter!: Adapter
   #running = false
   #initialized = false
   #generator?: AsyncGenerator<WorkerCycle, void, unknown>
   #pool?: JobPool
   #lastStalledCheck = 0
+  #shutdownHandler?: () => Promise<void>
 
   get id() {
     return this.#id
@@ -41,6 +44,8 @@ export class Worker {
     this.#stalledThreshold = parse(config.worker?.stalledThreshold ?? '30s')
     this.#maxStalledCount = config.worker?.maxStalledCount ?? 1
     this.#concurrency = config.worker?.concurrency ?? 1
+    this.#gracefulShutdown = config.worker?.gracefulShutdown ?? true
+    this.#onShutdownSignal = config.worker?.onShutdownSignal
 
     debug('created worker with id %s and config %O', this.#id, config)
   }
@@ -74,7 +79,7 @@ export class Worker {
 
     debug('starting worker %s on queues: %O', this.#id, queues)
 
-    await this.#setupGracefulShutdown()
+    this.#setupGracefulShutdown()
 
     for await (const cycle of this.process(queues)) {
       if (['started', 'completed'].includes(cycle.type)) {
@@ -109,6 +114,8 @@ export class Worker {
     if (this.#adapter) {
       await this.#adapter.destroy()
     }
+
+    this.#removeShutdownHandlers()
   }
 
   async processCycle(queues: string[]): Promise<WorkerCycle | null> {
@@ -317,14 +324,30 @@ export class Worker {
     }
   }
 
-  async #setupGracefulShutdown() {
-    const shutdown = async () => {
+  #setupGracefulShutdown() {
+    if (!this.#gracefulShutdown) {
+      return
+    }
+
+    this.#shutdownHandler = async () => {
       debug('received shutdown signal, stopping worker...')
+
+      if (this.#onShutdownSignal) {
+        await this.#onShutdownSignal()
+      }
+
       await this.stop()
-      process.exit(0)
     }
 
-    process.on('SIGINT', shutdown)
-    process.on('SIGTERM', shutdown)
+    process.on('SIGINT', this.#shutdownHandler)
+    process.on('SIGTERM', this.#shutdownHandler)
+  }
+
+  #removeShutdownHandlers() {
+    if (this.#shutdownHandler) {
+      process.off('SIGINT', this.#shutdownHandler)
+      process.off('SIGTERM', this.#shutdownHandler)
+      this.#shutdownHandler = undefined
+    }
   }
 }

tests/worker.spec.ts

@@ -1051,4 +1051,58 @@ test.group('Worker', () => {
       assert.equal(count, 1, `${jobId} should be executed exactly once`)
     }
   })
+
+  test('onShutdownSignal callback is invoked on SIGTERM', async ({ assert }) => {
+    let callbackInvoked = false
+
+    const localConfig = {
+      default: 'memory',
+      adapters: { memory: memory() },
+      worker: {
+        gracefulShutdown: true,
+        onShutdownSignal: () => {
+          callbackInvoked = true
+        },
+      },
+    }
+
+    const worker = new Worker(localConfig)
+    const startPromise = worker.start(['default'])
+    await setTimeout(10)
+
+    // Emit SIGTERM to trigger the shutdown handler
+    process.emit('SIGTERM')
+
+    // Wait for the worker to stop
+    await Promise.race([startPromise, setTimeout(500)])
+
+    assert.isTrue(callbackInvoked, 'onShutdownSignal should be called on SIGTERM')
+  })
+
+  test('onShutdownSignal callback is invoked on SIGINT', async ({ assert }) => {
+    let callbackInvoked = false
+
+    const localConfig = {
+      default: 'memory',
+      adapters: { memory: memory() },
+      worker: {
+        gracefulShutdown: true,
+        onShutdownSignal: () => {
+          callbackInvoked = true
+        },
+      },
+    }
+
+    const worker = new Worker(localConfig)
+    const startPromise = worker.start(['default'])
+    await setTimeout(10)
+
+    // Emit SIGINT to trigger the shutdown handler
+    process.emit('SIGINT')
+
+    // Wait for the worker to stop
+    await Promise.race([startPromise, setTimeout(500)])
+
+    assert.isTrue(callbackInvoked, 'onShutdownSignal should be called on SIGINT')
+  })
 })