From b2559a74817131f03cd8184013084cc8314a6167 Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=B4me=20Benoit?= <jerome.benoit@piment-noir.org>
Date: Fri, 25 Aug 2023 17:42:56 +0200
Subject: [PATCH] refactor: ensure stolen tasks account is done safely
MIME-Version: 1.0
Content-Type: text/plain; charset=utf8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jérôme Benoit <jerome.benoit@piment-noir.org>
---
 src/pools/abstract-pool.ts | 31 +++++++++++++++++++++++++++----
 1 file changed, 27 insertions(+), 4 deletions(-)

diff --git a/src/pools/abstract-pool.ts b/src/pools/abstract-pool.ts
index 4abbd615..3cc8b220 100644
--- a/src/pools/abstract-pool.ts
+++ b/src/pools/abstract-pool.ts
@@ -1238,6 +1238,17 @@ export abstract class AbstractPool<
     }
   }
 
+  // private incrementTasksStolen (workerNodeKey: number, workerNode: WorkerNode): void {
+  //   ++workerNode.usage.tasks.stolen
+  //   if (this.shallUpdateTaskFunctionWorkerUsage(workerNodeKey)) {
+  //     const taskFunctionWorkerUsage =
+  //       workerNode.getTaskFunctionWorkerUsage(
+  //         task.name as string
+  //       ) as WorkerUsage
+  //     ++taskFunctionWorkerUsage.tasks.stolen
+  //   }
+  // }
+
   private taskStealingOnEmptyQueue (workerId: number): void {
     const destinationWorkerNodeKey = this.getWorkerNodeKeyByWorkerId(workerId)
     const destinationWorkerNode = this.workerNodes[destinationWorkerNodeKey]
@@ -1269,8 +1280,15 @@ export abstract class AbstractPool<
         } else {
           this.enqueueTask(destinationWorkerNodeKey, task)
         }
-        ++destinationWorkerNode.usage.tasks.stolen
-        if (this.shallUpdateTaskFunctionWorkerUsage(destinationWorkerNodeKey)) {
+        if (destinationWorkerNode?.usage != null) {
+          ++destinationWorkerNode.usage.tasks.stolen
+        }
+        if (
+          this.shallUpdateTaskFunctionWorkerUsage(destinationWorkerNodeKey) &&
+          destinationWorkerNode.getTaskFunctionWorkerUsage(
+            task.name as string
+          ) != null
+        ) {
           const taskFunctionWorkerUsage =
             destinationWorkerNode.getTaskFunctionWorkerUsage(
               task.name as string
@@ -1312,8 +1330,13 @@ export abstract class AbstractPool<
         } else {
           this.enqueueTask(workerNodeKey, task)
         }
-        ++workerNode.usage.tasks.stolen
-        if (this.shallUpdateTaskFunctionWorkerUsage(workerNodeKey)) {
+        if (workerNode?.usage != null) {
+          ++workerNode.usage.tasks.stolen
+        }
+        if (
+          this.shallUpdateTaskFunctionWorkerUsage(workerNodeKey) &&
+          workerNode.getTaskFunctionWorkerUsage(task.name as string) != null
+        ) {
           const taskFunctionWorkerUsage = workerNode.getTaskFunctionWorkerUsage(
             task.name as string
           ) as WorkerUsage
-- 
2.34.1