From 82888165d01dba6a57f5047d4a4c27c6da6025f3 Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=B4me=20Benoit?= <jerome.benoit@sap.com>
Date: Fri, 5 May 2023 15:54:20 +0200
Subject: [PATCH] fix: fix default task function handling with multiple task
 functions
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              | 21 +++++++++++++--------
 src/worker/thread-worker.ts                | 14 +++++++++++---
 tests/pools/abstract/abstract-pool.test.js |  2 ++
 tests/worker/abstract-worker.test.js       |  1 +
 4 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/src/worker/abstract-worker.ts b/src/worker/abstract-worker.ts
index 3ca52b69..df263b6d 100644
--- a/src/worker/abstract-worker.ts
+++ b/src/worker/abstract-worker.ts
@@ -45,7 +45,7 @@ export abstract class AbstractWorker<
    *
    * @param type - The type of async event.
    * @param isMain - Whether this is the main worker or not.
-   * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked.
+   * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked. The first function is the default function.
    * @param mainWorker - Reference to main worker.
    * @param opts - Options for the worker.
    */
@@ -80,12 +80,7 @@ export abstract class AbstractWorker<
       this.checkAlive.bind(this)()
     }
 
-    this.mainWorker?.on(
-      'message',
-      (message: MessageValue<Data, MainWorker>) => {
-        this.messageListener(message)
-      }
-    )
+    this.mainWorker?.on('message', this.messageListener.bind(this))
   }
 
   private checkWorkerOptions (opts: WorkerOptions): void {
@@ -98,7 +93,7 @@ export abstract class AbstractWorker<
   /**
    * Checks if the `taskFunctions` parameter is passed to the constructor.
    *
-   * @param taskFunctions - The task function(s) that should be defined.
+   * @param taskFunctions - The task function(s) parameter that should be checked.
    */
   private checkTaskFunctions (
     taskFunctions:
@@ -123,6 +118,7 @@ export abstract class AbstractWorker<
     }
     this.taskFunctions = new Map<string, WorkerFunction<Data, Response>>()
     if (typeof taskFunctions !== 'function') {
+      let firstEntry = true
       for (const [name, fn] of Object.entries(taskFunctions)) {
         if (typeof fn !== 'function') {
           throw new Error(
@@ -130,6 +126,10 @@ export abstract class AbstractWorker<
           )
         }
         this.taskFunctions.set(name, fn.bind(this))
+        if (firstEntry) {
+          this.taskFunctions.set(DEFAULT_FUNCTION_NAME, fn.bind(this))
+          firstEntry = false
+        }
       }
     } else {
       this.taskFunctions.set(DEFAULT_FUNCTION_NAME, taskFunctions.bind(this))
@@ -259,6 +259,11 @@ export abstract class AbstractWorker<
       .catch(EMPTY_FUNCTION)
   }
 
+  /**
+   * Gets the task function in the given scope.
+   *
+   * @param name - Name of the function that will be returned.
+   */
   private getTaskFunction (name?: string): WorkerFunction<Data, Response> {
     name = name ?? DEFAULT_FUNCTION_NAME
     const fn = this.taskFunctions.get(name)
diff --git a/src/worker/thread-worker.ts b/src/worker/thread-worker.ts
index ce17c59a..75cd9da5 100644
--- a/src/worker/thread-worker.ts
+++ b/src/worker/thread-worker.ts
@@ -29,14 +29,22 @@ export class ThreadWorker<
   /**
    * Constructs a new poolifier thread worker.
    *
-   * @param fn - Function processed by the worker when the pool's `execution` function is invoked.
+   * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked.
    * @param opts - Options for the worker.
    */
   public constructor (
-    fn: WorkerFunction<Data, Response> | TaskFunctions<Data, Response>,
+    taskFunctions:
+    | WorkerFunction<Data, Response>
+    | TaskFunctions<Data, Response>,
     opts: WorkerOptions = {}
   ) {
-    super('worker-thread-pool:poolifier', isMainThread, fn, parentPort, opts)
+    super(
+      'worker-thread-pool:poolifier',
+      isMainThread,
+      taskFunctions,
+      parentPort,
+      opts
+    )
   }
 
   /** @inheritDoc */
diff --git a/tests/pools/abstract/abstract-pool.test.js b/tests/pools/abstract/abstract-pool.test.js
index 972587b5..81ae3649 100644
--- a/tests/pools/abstract/abstract-pool.test.js
+++ b/tests/pools/abstract/abstract-pool.test.js
@@ -407,6 +407,8 @@ describe('Abstract pool test suite', () => {
       './tests/worker-files/cluster/testMultiTasksWorker.js'
     )
     const data = { n: 10 }
+    const result0 = await pool.execute(data)
+    expect(result0).toBe(false)
     const result1 = await pool.execute(data, 'jsonIntegerSerialization')
     expect(result1).toBe(false)
     const result2 = await pool.execute(data, 'factorial')
diff --git a/tests/worker/abstract-worker.test.js b/tests/worker/abstract-worker.test.js
index 25e09723..5b42eb9d 100644
--- a/tests/worker/abstract-worker.test.js
+++ b/tests/worker/abstract-worker.test.js
@@ -71,6 +71,7 @@ describe('Abstract worker test suite', () => {
       return 2
     }
     const worker = new ClusterWorker({ fn1, fn2 })
+    expect(typeof worker.taskFunctions.get('default') === 'function').toBe(true)
     expect(typeof worker.taskFunctions.get('fn1') === 'function').toBe(true)
     expect(typeof worker.taskFunctions.get('fn2') === 'function').toBe(true)
   })
-- 
2.34.1