From 8a1260a3eca1f3fc72c9db255627daccfa7616bc Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=B4me=20Benoit?= <jerome.benoit@sap.com>
Date: Tue, 4 Jul 2023 22:00:44 +0200
Subject: [PATCH] fix: recreate the right worker type on uncaught error
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/pools/abstract-pool.ts | 52 +++++++++++++++++++++++++++-----------
 src/pools/worker.ts        |  4 +++
 2 files changed, 41 insertions(+), 15 deletions(-)

diff --git a/src/pools/abstract-pool.ts b/src/pools/abstract-pool.ts
index ad657b17..e12023f1 100644
--- a/src/pools/abstract-pool.ts
+++ b/src/pools/abstract-pool.ts
@@ -28,6 +28,7 @@ import type {
   IWorker,
   MessageHandler,
   Task,
+  WorkerInfo,
   WorkerNode,
   WorkerUsage
 } from './worker'
@@ -379,10 +380,10 @@ export abstract class AbstractPool<
     if (workerChoiceStrategyOptions != null) {
       this.setWorkerChoiceStrategyOptions(workerChoiceStrategyOptions)
     }
-    for (const [workerNodeKey, workerNode] of this.workerNodes.entries()) {
+    for (const workerNode of this.workerNodes) {
       this.setWorkerNodeTasksUsage(
         workerNode,
-        this.getWorkerUsage(workerNodeKey)
+        this.getInitialWorkerUsage(workerNode.worker)
       )
       this.setWorkerStatistics(workerNode.worker)
     }
@@ -827,7 +828,11 @@ export abstract class AbstractPool<
         }
       }
       if (this.opts.restartWorkerOnError === true) {
-        this.createAndSetupWorker()
+        if (this.getWorkerInfo(this.getWorkerNodeKey(worker)).dynamic) {
+          this.createAndSetupDynamicWorker()
+        } else {
+          this.createAndSetupWorker()
+        }
       }
     })
     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
@@ -852,6 +857,7 @@ export abstract class AbstractPool<
    */
   protected createAndSetupDynamicWorker (): Worker {
     const worker = this.createAndSetupWorker()
+    this.getWorkerInfo(this.getWorkerNodeKey(worker)).dynamic = true
     this.registerWorkerMessageListener(worker, message => {
       const workerNodeKey = this.getWorkerNodeKey(worker)
       if (
@@ -953,6 +959,15 @@ export abstract class AbstractPool<
     workerNode.usage = workerUsage
   }
 
+  /**
+   * Gets the worker information.
+   *
+   * @param workerNodeKey - The worker node key.
+   */
+  private getWorkerInfo (workerNodeKey: number): WorkerInfo {
+    return this.workerNodes[workerNodeKey].info
+  }
+
   /**
    * Pushes the given worker in the pool worker nodes.
    *
@@ -962,14 +977,13 @@ export abstract class AbstractPool<
   private pushWorkerNode (worker: Worker): number {
     this.workerNodes.push({
       worker,
-      info: { id: this.getWorkerId(worker), started: true },
-      usage: this.getWorkerUsage(),
+      info: this.getInitialWorkerInfo(worker),
+      usage: this.getInitialWorkerUsage(),
       tasksQueue: new Queue<Task<Data>>()
     })
-    const workerNodeKey = this.getWorkerNodeKey(worker)
     this.setWorkerNodeTasksUsage(
-      this.workerNodes[workerNodeKey],
-      this.getWorkerUsage(workerNodeKey)
+      this.workerNodes[this.getWorkerNodeKey(worker)],
+      this.getInitialWorkerUsage(worker)
     )
     return this.workerNodes.length
   }
@@ -1074,22 +1088,26 @@ export abstract class AbstractPool<
     })
   }
 
-  private getWorkerUsage (workerNodeKey?: number): WorkerUsage {
-    const getTasksQueueSize = (workerNodeKey?: number): number => {
-      return workerNodeKey != null ? this.tasksQueueSize(workerNodeKey) : 0
+  private getInitialWorkerUsage (worker?: Worker): WorkerUsage {
+    const getTasksQueueSize = (worker?: Worker): number => {
+      return worker != null
+        ? this.tasksQueueSize(this.getWorkerNodeKey(worker))
+        : 0
     }
-    const getTasksMaxQueueSize = (workerNodeKey?: number): number => {
-      return workerNodeKey != null ? this.tasksMaxQueueSize(workerNodeKey) : 0
+    const getTasksMaxQueueSize = (worker?: Worker): number => {
+      return worker != null
+        ? this.tasksMaxQueueSize(this.getWorkerNodeKey(worker))
+        : 0
     }
     return {
       tasks: {
         executed: 0,
         executing: 0,
         get queued (): number {
-          return getTasksQueueSize(workerNodeKey)
+          return getTasksQueueSize(worker)
         },
         get maxQueued (): number {
-          return getTasksMaxQueueSize(workerNodeKey)
+          return getTasksMaxQueueSize(worker)
         },
         failed: 0
       },
@@ -1129,4 +1147,8 @@ export abstract class AbstractPool<
       }
     }
   }
+
+  private getInitialWorkerInfo (worker: Worker): WorkerInfo {
+    return { id: this.getWorkerId(worker), dynamic: false, started: true }
+  }
 }
diff --git a/src/pools/worker.ts b/src/pools/worker.ts
index 7f9cef0a..20fac7cc 100644
--- a/src/pools/worker.ts
+++ b/src/pools/worker.ts
@@ -136,6 +136,10 @@ export interface WorkerInfo {
    * Worker id.
    */
   readonly id: number | undefined
+  /**
+   * Dynamic flag.
+   */
+  dynamic: boolean
   /**
    * Started flag.
    */
-- 
2.34.1