From 07588f306649b9df235aaf65de4842e99532ed6a 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 20:06:08 +0200
Subject: [PATCH] feat: add support for async kill handler
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>
---
 src/worker/abstract-worker.ts        | 11 +++++++++--
 src/worker/worker-options.ts         |  2 +-
 tests/worker/abstract-worker.test.js | 17 +++++++++--------
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/src/worker/abstract-worker.ts b/src/worker/abstract-worker.ts
index 512a822d..493ca0b2 100644
--- a/src/worker/abstract-worker.ts
+++ b/src/worker/abstract-worker.ts
@@ -326,8 +326,15 @@ export abstract class AbstractWorker<
    */
   protected handleKillMessage (message: MessageValue<Data>): void {
     this.stopCheckActive()
-    this.opts.killHandler?.()
-    this.emitDestroy()
+    if (isAsyncFunction(this.opts.killHandler)) {
+      (this.opts.killHandler?.() as Promise<void>)
+        .then(() => this.emitDestroy())
+        .catch(EMPTY_FUNCTION)
+    } else {
+      // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
+      this.opts.killHandler?.() as void
+      this.emitDestroy()
+    }
   }
 
   /**
diff --git a/src/worker/worker-options.ts b/src/worker/worker-options.ts
index 97d4a30a..f829dcb3 100644
--- a/src/worker/worker-options.ts
+++ b/src/worker/worker-options.ts
@@ -20,7 +20,7 @@ export type KillBehavior = keyof typeof KillBehaviors
 /**
  * Handler called when a worker is killed.
  */
-export type KillHandler = () => void
+export type KillHandler = () => void | Promise<void>
 
 /**
  * Options for workers.
diff --git a/tests/worker/abstract-worker.test.js b/tests/worker/abstract-worker.test.js
index 6582961f..46370777 100644
--- a/tests/worker/abstract-worker.test.js
+++ b/tests/worker/abstract-worker.test.js
@@ -136,14 +136,15 @@ describe('Abstract worker test suite', () => {
     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 async kill handler is called when worker is killed', () => {
+    const killHandlerStub = sinon.stub().returns()
+    const worker = new ClusterWorker(() => {}, {
+      killHandler: async () => Promise.resolve(killHandlerStub())
+    })
+    worker.isMain = false
+    worker.handleKillMessage()
+    expect(killHandlerStub.calledOnce).toBe(true)
+  })
 
   it('Verify that handleError() method works properly', () => {
     const error = new Error('Error as an error')
-- 
2.34.1