From febcf8e370ebf9ed2e45ba25296e937277d61e93 Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=B4me=20Benoit?= <jerome.benoit@piment-noir.org>
Date: Sat, 31 Aug 2024 13:51:09 +0200
Subject: [PATCH] refactor: add `isWorkerNodeStealing()` helper
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 | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/pools/abstract-pool.ts b/src/pools/abstract-pool.ts
index 098ed233..1a2c8bf3 100644
--- a/src/pools/abstract-pool.ts
+++ b/src/pools/abstract-pool.ts
@@ -641,15 +641,12 @@ export abstract class AbstractPool<
       const localWorkerNodeKey = this.getWorkerNodeKeyByWorkerId(
         message.workerId
       )
-      const workerInfo = this.getWorkerInfo(localWorkerNodeKey)
       // Kill message received from worker
       if (
         isKillBehavior(KillBehaviors.HARD, message.kill) ||
         (isKillBehavior(KillBehaviors.SOFT, message.kill) &&
           this.isWorkerNodeIdle(localWorkerNodeKey) &&
-          workerInfo != null &&
-          !workerInfo.continuousStealing &&
-          !workerInfo.backPressureStealing)
+          !this.isWorkerNodeStealing(localWorkerNodeKey))
       ) {
         // Flag the worker node as not ready immediately
         this.flagWorkerNodeAsNotReady(localWorkerNodeKey)
@@ -1363,6 +1360,15 @@ export abstract class AbstractPool<
     return workerNode.info.ready && workerNode.usage.tasks.executing === 0
   }
 
+  private isWorkerNodeStealing (workerNodeKey: number): boolean {
+    const workerNode = this.workerNodes[workerNodeKey]
+    return (
+      workerNode.info.ready &&
+      (workerNode.info.continuousStealing ||
+        workerNode.info.backPressureStealing)
+    )
+  }
+
   private redistributeQueuedTasks (sourceWorkerNodeKey: number): void {
     if (sourceWorkerNodeKey === -1 || this.cannotStealTask()) {
       return
@@ -2078,9 +2084,8 @@ export abstract class AbstractPool<
           0
         ),
         stealingWorkerNodes: this.workerNodes.reduce(
-          (accumulator, workerNode) =>
-            workerNode.info.continuousStealing ||
-            workerNode.info.backPressureStealing
+          (accumulator, _, workerNodeKey) =>
+            this.isWorkerNodeStealing(workerNodeKey)
               ? accumulator + 1
               : accumulator,
           0
-- 
2.34.1