refactor: cleanup tasks redistribution on error code

[poolifier.git] / src / pools / abstract-pool.ts

diff --git a/src/pools/abstract-pool.ts b/src/pools/abstract-pool.ts
index 3a1a87807cc065df871bb1ca6d2195619672a64d..7b0d541d8bcafaaaab1150fa01067c66acea30c2 100644 (file)
--- a/src/pools/abstract-pool.ts
+++ b/src/pools/abstract-pool.ts
@@ -1156,6 +1156,8 @@ export abstract class AbstractPool<
     // Send the statistics message to worker.
     this.sendStatisticsMessageToWorker(workerNodeKey)
     if (this.opts.enableTasksQueue === true) {
+      this.workerNodes[workerNodeKey].onEmptyQueue =
+        this.taskStealingOnEmptyQueue.bind(this)
       this.workerNodes[workerNodeKey].onBackPressure =
         this.tasksStealingOnBackPressure.bind(this)
     }
@@ -1187,42 +1189,69 @@ export abstract class AbstractPool<
   }
 
   private redistributeQueuedTasks (workerNodeKey: number): void {
-    const workerNodes = this.workerNodes.filter(
-      (_, workerNodeId) => workerNodeId !== workerNodeKey
-    )
     while (this.tasksQueueSize(workerNodeKey) > 0) {
-      let targetWorkerNodeKey: number = workerNodeKey
+      let destinationWorkerNodeKey: number = workerNodeKey
       let minQueuedTasks = Infinity
       let executeTask = false
-      for (const [workerNodeId, workerNode] of workerNodes.entries()) {
-        if (
-          workerNode.usage.tasks.executing <
-          (this.opts.tasksQueueOptions?.concurrency as number)
-        ) {
-          executeTask = true
+      for (const [workerNodeId, workerNode] of this.workerNodes.entries()) {
+        if (workerNode.info.ready && workerNodeId !== workerNodeKey) {
+          if (
+            workerNode.usage.tasks.executing <
+            (this.opts.tasksQueueOptions?.concurrency as number)
+          ) {
+            executeTask = true
+          }
+          if (workerNode.usage.tasks.queued === 0) {
+            destinationWorkerNodeKey = workerNodeId
+            break
+          }
+          if (workerNode.usage.tasks.queued < minQueuedTasks) {
+            minQueuedTasks = workerNode.usage.tasks.queued
+            destinationWorkerNodeKey = workerNodeId
+          }
         }
-        if (workerNode.info.ready && workerNode.usage.tasks.queued === 0) {
-          targetWorkerNodeKey = workerNodeId
-          break
+      }
+      const task = {
+        ...(this.dequeueTask(workerNodeKey) as Task<Data>),
+        workerId: (this.getWorkerInfo(destinationWorkerNodeKey) as WorkerInfo)
+          .id as number
+      }
+      if (executeTask) {
+        this.executeTask(destinationWorkerNodeKey, task)
+      } else {
+        this.enqueueTask(destinationWorkerNodeKey, task)
+      }
+    }
+  }
+
+  private taskStealingOnEmptyQueue (workerId: number): void {
+    const destinationWorkerNodeKey = this.getWorkerNodeKeyByWorkerId(workerId)
+    const destinationWorkerNode = this.workerNodes[destinationWorkerNodeKey]
+    const workerNodes = this.workerNodes
+      .slice()
+      .sort(
+        (workerNodeA, workerNodeB) =>
+          workerNodeB.usage.tasks.queued - workerNodeA.usage.tasks.queued
+      )
+    for (const sourceWorkerNode of workerNodes) {
+      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 (
-          workerNode.info.ready &&
-          workerNode.usage.tasks.queued < minQueuedTasks
+          destinationWorkerNode.usage.tasks.executing <
+          (this.opts.tasksQueueOptions?.concurrency as number)
         ) {
-          minQueuedTasks = workerNode.usage.tasks.queued
-          targetWorkerNodeKey = workerNodeId
+          this.executeTask(destinationWorkerNodeKey, task)
+        } else {
+          this.enqueueTask(destinationWorkerNodeKey, task)
         }
-      }
-      if (executeTask) {
-        this.executeTask(
-          targetWorkerNodeKey,
-          this.dequeueTask(workerNodeKey) as Task<Data>
-        )
-      } else {
-        this.enqueueTask(
-          targetWorkerNodeKey,
-          this.dequeueTask(workerNodeKey) as Task<Data>
-        )
+        break
       }
     }
   }
@@ -1231,7 +1260,7 @@ export abstract class AbstractPool<
     const sourceWorkerNode =
       this.workerNodes[this.getWorkerNodeKeyByWorkerId(workerId)]
     const workerNodes = this.workerNodes
-      .filter((workerNode) => workerNode.info.id !== workerId)
+      .slice()
       .sort(
         (workerNodeA, workerNodeB) =>
           workerNodeA.usage.tasks.queued - workerNodeB.usage.tasks.queued
@@ -1239,22 +1268,21 @@ export abstract class AbstractPool<
     for (const [workerNodeKey, workerNode] of workerNodes.entries()) {
       if (
         workerNode.info.ready &&
+        workerNode.info.id !== workerId &&
         sourceWorkerNode.usage.tasks.queued > 0 &&
         !workerNode.hasBackPressure()
       ) {
+        const task = {
+          ...(sourceWorkerNode.popTask() as Task<Data>),
+          workerId: workerNode.info.id as number
+        }
         if (
           workerNode.usage.tasks.executing <
           (this.opts.tasksQueueOptions?.concurrency as number)
         ) {
-          this.executeTask(
-            workerNodeKey,
-            sourceWorkerNode.popTask() as Task<Data>
-          )
+          this.executeTask(workerNodeKey, task)
         } else {
-          this.enqueueTask(
-            workerNodeKey,
-            sourceWorkerNode.popTask() as Task<Data>
-          )
+          this.enqueueTask(workerNodeKey, task)
         }
       }
     }