From f201a0cd2575be684ab0f291c8590eec538038e0 Mon Sep 17 00:00:00 2001
From: =?utf8?q?J=C3=A9r=C3=B4me=20Benoit?= <jerome.benoit@piment-noir.org>
Date: Sun, 3 Sep 2023 21:31:58 +0200
Subject: [PATCH] perf: optimize task(s) stealing implementation
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                    | 64 +++++++++----------
 .../fair-share-worker-choice-strategy.ts      | 10 +--
 2 files changed, 34 insertions(+), 40 deletions(-)

diff --git a/src/pools/abstract-pool.ts b/src/pools/abstract-pool.ts
index 6668361f..aae08395 100644
--- a/src/pools/abstract-pool.ts
+++ b/src/pools/abstract-pool.ts
@@ -1201,20 +1201,18 @@ export abstract class AbstractPool<
 
   private redistributeQueuedTasks (workerNodeKey: number): void {
     while (this.tasksQueueSize(workerNodeKey) > 0) {
-      let destinationWorkerNodeKey!: number
-      let minQueuedTasks = Infinity
-      for (const [workerNodeId, workerNode] of this.workerNodes.entries()) {
-        if (workerNode.info.ready && workerNodeId !== workerNodeKey) {
-          if (workerNode.usage.tasks.queued === 0) {
-            destinationWorkerNodeKey = workerNodeId
-            break
+      const destinationWorkerNodeKey = this.workerNodes.reduce(
+        (minWorkerNodeKey, workerNode, workerNodeKey, workerNodes) => {
+          if (!workerNode.info.ready) {
+            return minWorkerNodeKey
           }
-          if (workerNode.usage.tasks.queued < minQueuedTasks) {
-            minQueuedTasks = workerNode.usage.tasks.queued
-            destinationWorkerNodeKey = workerNodeId
-          }
-        }
-      }
+          return workerNode.usage.tasks.queued <
+            workerNodes[minWorkerNodeKey].usage.tasks.queued
+            ? workerNodeKey
+            : minWorkerNodeKey
+        },
+        0
+      )
       if (destinationWorkerNodeKey != null) {
         const destinationWorkerNode = this.workerNodes[destinationWorkerNodeKey]
         const task = {
@@ -1258,30 +1256,26 @@ export abstract class AbstractPool<
         (workerNodeA, workerNodeB) =>
           workerNodeB.usage.tasks.queued - workerNodeA.usage.tasks.queued
       )
-    for (const sourceWorkerNode of workerNodes) {
-      if (sourceWorkerNode.usage.tasks.queued === 0) {
-        break
+    const sourceWorkerNode = workerNodes.find(
+      (workerNode) =>
+        workerNode.info.ready &&
+        workerNode.info.id !== workerId &&
+        workerNode.usage.tasks.queued > 0
+    )
+    if (sourceWorkerNode != null) {
+      const task = {
+        ...(sourceWorkerNode.popTask() as Task<Data>),
+        workerId: destinationWorkerNode.info.id as number
       }
-      if (
-        sourceWorkerNode.info.ready &&
-        sourceWorkerNode.info.id !== workerId &&
-        sourceWorkerNode.usage.tasks.queued > 0
-      ) {
-        const task = {
-          ...(sourceWorkerNode.popTask() as Task<Data>),
-          workerId: destinationWorkerNode.info.id as number
-        }
-        if (this.shallExecuteTask(destinationWorkerNodeKey)) {
-          this.executeTask(destinationWorkerNodeKey, task)
-        } else {
-          this.enqueueTask(destinationWorkerNodeKey, task)
-        }
-        this.updateTaskStolenStatisticsWorkerUsage(
-          destinationWorkerNodeKey,
-          task.name as string
-        )
-        break
+      if (this.shallExecuteTask(destinationWorkerNodeKey)) {
+        this.executeTask(destinationWorkerNodeKey, task)
+      } else {
+        this.enqueueTask(destinationWorkerNodeKey, task)
       }
+      this.updateTaskStolenStatisticsWorkerUsage(
+        destinationWorkerNodeKey,
+        task.name as string
+      )
     }
   }
 
diff --git a/src/pools/selection-strategies/fair-share-worker-choice-strategy.ts b/src/pools/selection-strategies/fair-share-worker-choice-strategy.ts
index af3fcfa0..a337278b 100644
--- a/src/pools/selection-strategies/fair-share-worker-choice-strategy.ts
+++ b/src/pools/selection-strategies/fair-share-worker-choice-strategy.ts
@@ -128,12 +128,12 @@ export class FairShareWorkerChoiceStrategy<
   private getWorkerNodeVirtualTaskStartTimestamp (
     workerNodeKey: number
   ): number {
+    const virtualTaskEndTimestamp =
+      this.pool.workerNodes[workerNodeKey]?.strategyData
+        ?.virtualTaskEndTimestamp
     const now = performance.now()
-    return now <
-      (this.pool.workerNodes[workerNodeKey]?.strategyData
-        ?.virtualTaskEndTimestamp ?? -Infinity)
-      ? (this.pool.workerNodes[workerNodeKey]?.strategyData
-          ?.virtualTaskEndTimestamp as number)
+    return now < (virtualTaskEndTimestamp ?? -Infinity)
+      ? (virtualTaskEndTimestamp as number)
       : now
   }
 }
-- 
2.34.1