fix: avoid out of bound at worker node removal

[poolifier.git] / src / pools / selection-strategies / interleaved-weighted-round-robin-worker-choice-strategy.ts

diff --git a/src/pools/selection-strategies/interleaved-weighted-round-robin-worker-choice-strategy.ts b/src/pools/selection-strategies/interleaved-weighted-round-robin-worker-choice-strategy.ts
index ca5b7db9d914d0531fade68eccac9fc7d614bf11..7fba0f1b5dad0972af4ae9630b3c8565301a7863 100644 (file)
--- a/src/pools/selection-strategies/interleaved-weighted-round-robin-worker-choice-strategy.ts
+++ b/src/pools/selection-strategies/interleaved-weighted-round-robin-worker-choice-strategy.ts
@@ -96,9 +96,6 @@ export class InterleavedWeightedRoundRobinWorkerChoiceStrategy<
         workerNodeKey++
       ) {
         this.workerNodeId = workerNodeKey
-        if (!this.isWorkerNodeEligible(workerNodeKey)) {
-          continue
-        }
         if (
           this.workerNodeId !== this.nextWorkerNodeKey &&
           this.workerVirtualTaskRunTime !== 0
@@ -114,8 +111,7 @@ export class InterleavedWeightedRoundRobinWorkerChoiceStrategy<
           this.workerVirtualTaskRunTime =
             this.workerVirtualTaskRunTime +
             this.getWorkerTaskRunTime(workerNodeKey)
-          this.previousWorkerNodeKey =
-            this.nextWorkerNodeKey ?? this.previousWorkerNodeKey
+          this.setPreviousWorkerNodeKey(this.nextWorkerNodeKey)
           this.nextWorkerNodeKey = workerNodeKey
           return this.nextWorkerNodeKey
         }
@@ -141,15 +137,20 @@ export class InterleavedWeightedRoundRobinWorkerChoiceStrategy<
 
   /** @inheritDoc */
   public remove (workerNodeKey: number): boolean {
-    if (this.nextWorkerNodeKey === workerNodeKey) {
-      if (this.pool.workerNodes.length === 0) {
-        this.nextWorkerNodeKey = 0
-      } else if (this.nextWorkerNodeKey > this.pool.workerNodes.length - 1) {
-        this.roundId =
-          this.roundId === this.roundWeights.length - 1 ? 0 : this.roundId + 1
-        this.nextWorkerNodeKey = this.pool.workerNodes.length - 1
-      }
-      this.workerVirtualTaskRunTime = 0
+    if (this.pool.workerNodes.length === 0) {
+      this.reset()
+    }
+    if (
+      this.workerNodeId === workerNodeKey &&
+      this.workerNodeId > this.pool.workerNodes.length - 1
+    ) {
+      this.workerNodeId = this.pool.workerNodes.length - 1
+    }
+    if (
+      this.previousWorkerNodeKey === workerNodeKey &&
+      this.previousWorkerNodeKey > this.pool.workerNodes.length - 1
+    ) {
+      this.previousWorkerNodeKey = this.pool.workerNodes.length - 1
     }
     return true
   }