From df9aaf207bf2e96ed759f9b89fa6573e8c48fb65 Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=B4me=20Benoit?= <jerome.benoit@sap.com>
Date: Mon, 14 Aug 2023 16:22:57 +0200
Subject: [PATCH] feat: add a kill handler to worker to allow to run custom
 code at worker termination
MIME-Version: 1.0
Content-Type: text/plain; charset=utf8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jérôme Benoit <jerome.benoit@sap.com>
---
 rome.json                            |  2 +-
 src/worker/abstract-worker.ts        | 14 ++++++++++---
 src/worker/worker-options.ts         |  9 +++++++++
 tests/worker/abstract-worker.test.js | 30 ++++++++++++++++++++++++++--
 4 files changed, 49 insertions(+), 6 deletions(-)

diff --git a/rome.json b/rome.json
index 9e65641a..1c91a4d7 100644
--- a/rome.json
+++ b/rome.json
@@ -23,8 +23,8 @@
   },
   "files": {
     "ignore": [
-      ".vscode/",
       ".nyc_output/",
+      ".vscode/",
       "benchmarks/internal/results/",
       "coverage/",
       "docs/**/*.css",
diff --git a/src/worker/abstract-worker.ts b/src/worker/abstract-worker.ts
index ae7778c4..508d8689 100644
--- a/src/worker/abstract-worker.ts
+++ b/src/worker/abstract-worker.ts
@@ -84,7 +84,11 @@ export abstract class AbstractWorker<
        * The maximum time to keep this worker active while idle.
        * The pool automatically checks and terminates this worker when the time expires.
        */
-      maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME
+      maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME,
+      /**
+       * The function to call when the worker is killed.
+       */
+      killHandler: EMPTY_FUNCTION
     }
   ) {
     super(type)
@@ -100,6 +104,7 @@ export abstract class AbstractWorker<
     this.opts.maxInactiveTime =
       opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME
     delete this.opts.async
+    this.opts.killHandler = opts.killHandler ?? EMPTY_FUNCTION
   }
 
   /**
@@ -320,8 +325,11 @@ export abstract class AbstractWorker<
    * @param message - The kill message.
    */
   protected handleKillMessage (message: MessageValue<Data>): void {
-    !this.isMain && this.stopCheckActive()
-    this.emitDestroy()
+    if (!this.isMain) {
+      this.stopCheckActive()
+      this.opts.killHandler?.()
+      this.emitDestroy()
+    }
   }
 
   /**
diff --git a/src/worker/worker-options.ts b/src/worker/worker-options.ts
index e6137c7f..97d4a30a 100644
--- a/src/worker/worker-options.ts
+++ b/src/worker/worker-options.ts
@@ -17,6 +17,11 @@ export const KillBehaviors = Object.freeze({
  */
 export type KillBehavior = keyof typeof KillBehaviors
 
+/**
+ * Handler called when a worker is killed.
+ */
+export type KillHandler = () => void
+
 /**
  * Options for workers.
  */
@@ -52,4 +57,8 @@ export interface WorkerOptions {
    * @defaultValue KillBehaviors.SOFT
    */
   killBehavior?: KillBehavior
+  /**
+   * The function to call when a worker is killed.
+   */
+  killHandler?: KillHandler
 }
diff --git a/tests/worker/abstract-worker.test.js b/tests/worker/abstract-worker.test.js
index e9da7a87..6582961f 100644
--- a/tests/worker/abstract-worker.test.js
+++ b/tests/worker/abstract-worker.test.js
@@ -1,5 +1,7 @@
 const { expect } = require('expect')
+const sinon = require('sinon')
 const { ClusterWorker, KillBehaviors, ThreadWorker } = require('../../lib')
+const { EMPTY_FUNCTION } = require('../../lib/utils')
 
 describe('Abstract worker test suite', () => {
   class StubWorkerWithMainWorker extends ThreadWorker {
@@ -13,17 +15,23 @@ describe('Abstract worker test suite', () => {
     const worker = new ThreadWorker(() => {})
     expect(worker.opts.maxInactiveTime).toStrictEqual(60000)
     expect(worker.opts.killBehavior).toBe(KillBehaviors.SOFT)
+    expect(worker.opts.killHandler).toStrictEqual(EMPTY_FUNCTION)
     expect(worker.opts.async).toBe(undefined)
   })
 
   it('Verify that worker options are set at worker creation', () => {
+    const killHandler = () => {
+      console.info('Worker received kill message')
+    }
     const worker = new ClusterWorker(() => {}, {
       maxInactiveTime: 6000,
-      async: true,
-      killBehavior: KillBehaviors.HARD
+      killBehavior: KillBehaviors.HARD,
+      killHandler,
+      async: true
     })
     expect(worker.opts.maxInactiveTime).toStrictEqual(6000)
     expect(worker.opts.killBehavior).toBe(KillBehaviors.HARD)
+    expect(worker.opts.killHandler).toStrictEqual(killHandler)
     expect(worker.opts.async).toBe(undefined)
   })
 
@@ -119,6 +127,24 @@ describe('Abstract worker test suite', () => {
     )
   })
 
+  it('Verify that sync kill handler is called when worker is killed', () => {
+    const worker = new ClusterWorker(() => {}, {
+      killHandler: sinon.stub().returns()
+    })
+    worker.isMain = false
+    worker.handleKillMessage()
+    expect(worker.opts.killHandler.calledOnce).toBe(true)
+  })
+
+  // it('Verify that async kill handler is called when worker is killed', () => {
+  //   const worker = new ClusterWorker(() => {}, {
+  //     killHandler: sinon.stub().resolves()
+  //   })
+  //   worker.isMain = false
+  //   worker.handleKillMessage()
+  //   expect(worker.opts.killHandler.calledOnce).toBe(true)
+  // })
+
   it('Verify that handleError() method works properly', () => {
     const error = new Error('Error as an error')
     const worker = new ClusterWorker(() => {})
-- 
2.34.1