From d52d4773cc0ed2936aea894224a7edc3dab8359a Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=B4me=20Benoit?= <jerome.benoit@piment-noir.org>
Date: Tue, 13 Aug 2024 18:42:09 +0200
Subject: [PATCH] refactor: cleanup worker condition checks at task stealing
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 | 73 +++++++++++++++++++-------------------
 1 file changed, 36 insertions(+), 37 deletions(-)

diff --git a/src/pools/abstract-pool.ts b/src/pools/abstract-pool.ts
index 4c387477..29dd42db 100644
--- a/src/pools/abstract-pool.ts
+++ b/src/pools/abstract-pool.ts
@@ -1894,6 +1894,42 @@ export abstract class AbstractPool<
     }
   }
 
+  private readonly stealTask = (
+    sourceWorkerNode: IWorkerNode<Worker, Data>,
+    destinationWorkerNodeKey: number
+  ): Task<Data> | undefined => {
+    const destinationWorkerInfo = this.getWorkerInfo(destinationWorkerNodeKey)
+    if (destinationWorkerInfo == null) {
+      throw new Error(
+        `Worker node with key '${destinationWorkerNodeKey.toString()}' not found in pool`
+      )
+    }
+    // Avoid cross task stealing. Could be smarter by checking stealing/stolen worker ids pair.
+    if (
+      !sourceWorkerNode.info.ready ||
+      sourceWorkerNode.info.stolen ||
+      sourceWorkerNode.info.stealing ||
+      !destinationWorkerInfo.ready ||
+      destinationWorkerInfo.stolen ||
+      destinationWorkerInfo.stealing
+    ) {
+      return
+    }
+    destinationWorkerInfo.stealing = true
+    sourceWorkerNode.info.stolen = true
+    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+    const task = sourceWorkerNode.dequeueLastPrioritizedTask()!
+    sourceWorkerNode.info.stolen = false
+    destinationWorkerInfo.stealing = false
+    this.handleTask(destinationWorkerNodeKey, task)
+    this.updateTaskStolenStatisticsWorkerUsage(
+      destinationWorkerNodeKey,
+      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
+      task.name!
+    )
+    return task
+  }
+
   private readonly handleWorkerNodeIdleEvent = (
     eventDetail: WorkerNodeEventDetail,
     previousStolenTask?: Task<Data>
@@ -1904,12 +1940,6 @@ export abstract class AbstractPool<
         "WorkerNode event detail 'workerNodeKey' property must be defined"
       )
     }
-    const workerInfo = this.getWorkerInfo(workerNodeKey)
-    if (workerInfo == null) {
-      throw new Error(
-        `Worker node with key '${workerNodeKey.toString()}' not found in pool`
-      )
-    }
     if (
       this.cannotStealTask() ||
       (this.info.stealingWorkerNodes ?? 0) >
@@ -1960,31 +1990,6 @@ export abstract class AbstractPool<
       })
   }
 
-  private readonly stealTask = (
-    sourceWorkerNode: IWorkerNode<Worker, Data>,
-    destinationWorkerNodeKey: number
-  ): Task<Data> | undefined => {
-    const destinationWorkerInfo = this.getWorkerInfo(destinationWorkerNodeKey)
-    if (destinationWorkerInfo == null) {
-      throw new Error(
-        `Worker node with key '${destinationWorkerNodeKey.toString()}' not found in pool`
-      )
-    }
-    destinationWorkerInfo.stealing = true
-    sourceWorkerNode.info.stolen = true
-    // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-    const task = sourceWorkerNode.dequeueLastPrioritizedTask()!
-    sourceWorkerNode.info.stolen = false
-    destinationWorkerInfo.stealing = false
-    this.handleTask(destinationWorkerNodeKey, task)
-    this.updateTaskStolenStatisticsWorkerUsage(
-      destinationWorkerNodeKey,
-      // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-      task.name!
-    )
-    return task
-  }
-
   private readonly workerNodeStealTask = (
     workerNodeKey: number
   ): Task<Data> | undefined => {
@@ -1996,9 +2001,6 @@ export abstract class AbstractPool<
       )
     const sourceWorkerNode = workerNodes.find(
       (sourceWorkerNode, sourceWorkerNodeKey) =>
-        sourceWorkerNode.info.ready &&
-        !sourceWorkerNode.info.stolen &&
-        !sourceWorkerNode.info.stealing &&
         sourceWorkerNodeKey !== workerNodeKey &&
         sourceWorkerNode.usage.tasks.queued > 0
     )
@@ -2039,9 +2041,6 @@ export abstract class AbstractPool<
     for (const [workerNodeKey, workerNode] of workerNodes.entries()) {
       if (
         sourceWorkerNode.usage.tasks.queued > 0 &&
-        workerNode.info.ready &&
-        !workerNode.info.stolen &&
-        !workerNode.info.stealing &&
         workerNode.info.id !== workerId &&
         workerNode.usage.tasks.queued <
           // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
-- 
2.34.1