fix: properly account strategy retries on a per strategy basis

Signed-off-by: Jérôme Benoit <[email protected]>

Author: jerome-benoit

src/pools/abstract-pool.ts

@@ -309,7 +309,8 @@ export abstract class AbstractPool<
       started: this.started,
       ready: this.ready,
       defaultStrategy: this.opts.workerChoiceStrategy!,
-      strategyRetries: this.workerChoiceStrategiesContext?.retriesCount ?? 0,
+      strategyRetries:
+        this.workerChoiceStrategiesContext?.getStrategyRetries() ?? 0,
       minSize: this.minimumNumberOfWorkers,
       maxSize: this.maximumNumberOfWorkers ?? this.minimumNumberOfWorkers,
       ...(taskStatisticsRequirements?.runTime.aggregate === true &&

src/pools/selection-strategies/abstract-worker-choice-strategy.ts

@@ -28,15 +28,18 @@ export abstract class AbstractWorkerChoiceStrategy<
   /** @inheritDoc */
   public abstract readonly name: WorkerChoiceStrategy
 
+  /** @inheritDoc */
+  public retriesCount: number
+
   /**
    * The next worker node key.
    */
-  protected nextWorkerNodeKey: number | undefined = 0
+  protected nextWorkerNodeKey: number | undefined
 
   /**
    * The previous worker node key.
    */
-  protected previousWorkerNodeKey = 0
+  protected previousWorkerNodeKey: number
 
   /** @inheritDoc */
   public readonly strategyPolicy: StrategyPolicy = Object.freeze({
@@ -62,6 +65,9 @@ export abstract class AbstractWorkerChoiceStrategy<
     protected readonly pool: IPool<Worker, Data, Response>,
     protected opts?: WorkerChoiceStrategyOptions,
   ) {
+    this.retriesCount = 0
+    this.nextWorkerNodeKey = 0
+    this.previousWorkerNodeKey = 0
     this.choose = this.choose.bind(this)
     this.setOptions(this.opts)
   }

src/pools/selection-strategies/selection-strategies-types.ts

@@ -194,6 +194,10 @@ export interface IWorkerChoiceStrategy {
    * @returns `true` if the reset is successful, `false` otherwise.
    */
   readonly reset: () => boolean
+  /**
+   * The worker choice strategy execution retries count.
+   */
+  retriesCount: number
   /**
    * Updates the worker node key strategy internals.
    * This is called after a task has been executed on a worker node.

src/pools/selection-strategies/worker-choice-strategies-context.ts

@@ -28,11 +28,6 @@ export class WorkerChoiceStrategiesContext<
   Data = unknown,
   Response = unknown,
 > {
-  /**
-   * The number of worker choice strategies execution retries.
-   */
-  public retriesCount: number
-
   /**
    * The default worker choice strategy in the context.
    */
@@ -92,7 +87,6 @@ export class WorkerChoiceStrategiesContext<
       buildWorkerChoiceStrategiesTaskStatisticsRequirements(
         this.workerChoiceStrategies,
       )
-    this.retriesCount = 0
     this.retries = getWorkerChoiceStrategiesRetries<Worker, Data, Response>(
       this.pool,
       opts,
@@ -108,6 +102,17 @@ export class WorkerChoiceStrategiesContext<
     return this.workerChoiceStrategiesPolicy
   }
 
+  /**
+   * Gets the number of worker choice strategies execution retries.
+   * @returns The number of retries.
+   */
+  public getStrategyRetries(): number {
+    return Array.from(
+      this.workerChoiceStrategies,
+      ([_, workerChoiceStrategy]) => workerChoiceStrategy.retriesCount,
+    ).reduce((accumulator, retries) => accumulator + retries, 0)
+  }
+
   /**
    * Gets the active worker choice strategies in the context task statistics requirements.
    *
@@ -173,13 +178,13 @@ export class WorkerChoiceStrategiesContext<
     let workerNodeKey: number | undefined = workerChoiceStrategy.choose()
     let retriesCount = 0
     while (workerNodeKey == null && retriesCount < this.retries) {
-      workerNodeKey = workerChoiceStrategy.choose()
       retriesCount++
-      this.retriesCount++
+      workerNodeKey = workerChoiceStrategy.choose()
     }
+    workerChoiceStrategy.retriesCount = retriesCount
     if (workerNodeKey == null) {
       throw new Error(
-        `Worker node key chosen by ${workerChoiceStrategy.name} is null or undefined after ${retriesCount.toString()} retries (max: ${this.retries.toString()})`,
+        `Worker node key chosen by ${workerChoiceStrategy.name} is null or undefined after ${workerChoiceStrategy.retriesCount.toString()} retries (max: ${this.retries.toString()})`,
       )
     }
     return workerNodeKey

src/worker/abstract-worker.ts

@@ -287,14 +287,10 @@ export abstract class AbstractWorker<
         DEFAULT_TASK_NAME,
         this.taskFunctions.get(DEFAULT_TASK_NAME),
       ),
-      ...(defaultTaskFunctionName !== DEFAULT_TASK_NAME
-        ? [
-          buildTaskFunctionProperties(
-            defaultTaskFunctionName,
-            this.taskFunctions.get(defaultTaskFunctionName),
-          ),
-        ]
-        : []),
+      buildTaskFunctionProperties(
+        defaultTaskFunctionName,
+        this.taskFunctions.get(defaultTaskFunctionName),
+      ),
       ...taskFunctionsProperties,
     ]
   }

tests/pools/selection-strategies/selection-strategies.test.mjs

@@ -97,6 +97,11 @@ describe('Selection strategies test suite', () => {
           workerChoiceStrategy,
         ).name,
       ).toBe(workerChoiceStrategy)
+      expect(
+        pool.workerChoiceStrategiesContext.workerChoiceStrategies.get(
+          workerChoiceStrategy,
+        ).retriesCount,
+      ).toBe(0)
       expect(
         pool.workerChoiceStrategiesContext.workerChoiceStrategies.get(
           workerChoiceStrategy,

tests/pools/selection-strategies/worker-choice-strategies-context.test.mjs

@@ -80,39 +80,58 @@ describe('Worker choice strategies context test suite', () => {
     expect(workerChoiceStrategiesContext.defaultWorkerChoiceStrategy).toBe(
       WorkerChoiceStrategies.ROUND_ROBIN,
     )
-    const workerChoiceStrategyUndefinedStub =
-      new RoundRobinWorkerChoiceStrategy(fixedPool)
+    const workerChoiceStrategyUndefinedStub = workerChoiceStrategiesContext
+      .workerChoiceStrategies.get(
+        workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
+      )
     stub(
       workerChoiceStrategyUndefinedStub,
       'choose',
       returnsNext(Array(7).fill(undefined)),
     )
-    workerChoiceStrategiesContext.workerChoiceStrategies.set(
-      workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
-      workerChoiceStrategyUndefinedStub,
+    let err
+    try {
+      workerChoiceStrategiesContext.execute()
+    } catch (e) {
+      err = e
+    }
+    expect(err).toBeInstanceOf(Error)
+    expect(err.message).toBe(
+      `Worker node key chosen by ${workerChoiceStrategyUndefinedStub.name} is null or undefined after ${workerChoiceStrategyUndefinedStub.retriesCount.toString()} retries (max: ${workerChoiceStrategiesContext.retries.toString()})`,
     )
-    expect(() => workerChoiceStrategiesContext.execute()).toThrow(
-      new Error(
-        `Worker node key chosen by ${workerChoiceStrategyUndefinedStub.name} is null or undefined after ${workerChoiceStrategiesContext.retries.toString()} retries (max: ${workerChoiceStrategiesContext.retries.toString()})`,
-      ),
+    assertSpyCalls(
+      workerChoiceStrategyUndefinedStub.choose,
+      workerChoiceStrategyUndefinedStub.retriesCount + 1,
     )
-    workerChoiceStrategyUndefinedStub.choose.restore()
-    const workerChoiceStrategyNullStub = new RoundRobinWorkerChoiceStrategy(
-      fixedPool,
+    expect(workerChoiceStrategiesContext.getStrategyRetries()).toBe(
+      workerChoiceStrategyUndefinedStub.retriesCount,
     )
+    workerChoiceStrategyUndefinedStub.choose.restore()
+    const workerChoiceStrategyNullStub = workerChoiceStrategiesContext
+      .workerChoiceStrategies.get(
+        workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
+      )
     stub(
       workerChoiceStrategyNullStub,
       'choose',
       returnsNext(Array(7).fill(null)),
     )
-    workerChoiceStrategiesContext.workerChoiceStrategies.set(
-      workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
-      workerChoiceStrategyNullStub,
+    err = undefined
+    try {
+      workerChoiceStrategiesContext.execute()
+    } catch (e) {
+      err = e
+    }
+    expect(err).toBeInstanceOf(Error)
+    expect(err.message).toBe(
+      `Worker node key chosen by ${workerChoiceStrategyNullStub.name} is null or undefined after ${workerChoiceStrategyNullStub.retriesCount.toString()} retries (max: ${workerChoiceStrategiesContext.retries.toString()})`,
     )
-    expect(() => workerChoiceStrategiesContext.execute()).toThrow(
-      new Error(
-        `Worker node key chosen by ${workerChoiceStrategyNullStub.name} is null or undefined after ${workerChoiceStrategiesContext.retries.toString()} retries (max: ${workerChoiceStrategiesContext.retries.toString()})`,
-      ),
+    assertSpyCalls(
+      workerChoiceStrategyNullStub.choose,
+      workerChoiceStrategyUndefinedStub.retriesCount + 1,
+    )
+    expect(workerChoiceStrategiesContext.getStrategyRetries()).toBe(
+      workerChoiceStrategyNullStub.retriesCount,
     )
     workerChoiceStrategyNullStub.choose.restore()
   })
@@ -121,28 +140,21 @@ describe('Worker choice strategies context test suite', () => {
     const workerChoiceStrategiesContext = new WorkerChoiceStrategiesContext(
       fixedPool,
     )
-    const workerChoiceStrategyStub = new RoundRobinWorkerChoiceStrategy(
-      fixedPool,
+    expect(workerChoiceStrategiesContext.defaultWorkerChoiceStrategy).toBe(
+      WorkerChoiceStrategies.ROUND_ROBIN,
     )
+    const workerChoiceStrategyStub = workerChoiceStrategiesContext
+      .workerChoiceStrategies.get(
+        workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
+      )
     stub(
       workerChoiceStrategyStub,
       'choose',
       returnsNext(Array(5).fill(undefined).concat(Array(1).fill(1))),
     )
-    expect(workerChoiceStrategiesContext.defaultWorkerChoiceStrategy).toBe(
-      WorkerChoiceStrategies.ROUND_ROBIN,
-    )
-    workerChoiceStrategiesContext.workerChoiceStrategies.set(
-      workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
-      workerChoiceStrategyStub,
-    )
     const chosenWorkerKey = workerChoiceStrategiesContext.execute()
-    assertSpyCalls(
-      workerChoiceStrategiesContext.workerChoiceStrategies.get(
-        workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
-      ).choose,
-      6,
-    )
+    assertSpyCalls(workerChoiceStrategyStub.choose, 6)
+    expect(workerChoiceStrategiesContext.getStrategyRetries()).toBe(5)
     expect(chosenWorkerKey).toBe(1)
     workerChoiceStrategyStub.choose.restore()
   })
@@ -151,24 +163,17 @@ describe('Worker choice strategies context test suite', () => {
     const workerChoiceStrategiesContext = new WorkerChoiceStrategiesContext(
       fixedPool,
     )
-    const workerChoiceStrategyStub = new RoundRobinWorkerChoiceStrategy(
-      fixedPool,
-    )
-    stub(workerChoiceStrategyStub, 'choose', returnsNext([0]))
     expect(workerChoiceStrategiesContext.defaultWorkerChoiceStrategy).toBe(
       WorkerChoiceStrategies.ROUND_ROBIN,
     )
-    workerChoiceStrategiesContext.workerChoiceStrategies.set(
-      workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
-      workerChoiceStrategyStub,
-    )
-    const chosenWorkerKey = workerChoiceStrategiesContext.execute()
-    assertSpyCalls(
-      workerChoiceStrategiesContext.workerChoiceStrategies.get(
+    const workerChoiceStrategyStub = workerChoiceStrategiesContext
+      .workerChoiceStrategies.get(
         workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
-      ).choose,
-      1,
-    )
+      )
+    stub(workerChoiceStrategyStub, 'choose', returnsNext([0]))
+    const chosenWorkerKey = workerChoiceStrategiesContext.execute()
+    assertSpyCalls(workerChoiceStrategyStub.choose, 1)
+    expect(workerChoiceStrategiesContext.getStrategyRetries()).toBe(0)
     expect(chosenWorkerKey).toBe(0)
     workerChoiceStrategyStub.choose.restore()
   })
@@ -177,24 +182,17 @@ describe('Worker choice strategies context test suite', () => {
     const workerChoiceStrategiesContext = new WorkerChoiceStrategiesContext(
       dynamicPool,
     )
-    const workerChoiceStrategyStub = new RoundRobinWorkerChoiceStrategy(
-      dynamicPool,
-    )
-    stub(workerChoiceStrategyStub, 'choose', returnsNext([0]))
     expect(workerChoiceStrategiesContext.defaultWorkerChoiceStrategy).toBe(
       WorkerChoiceStrategies.ROUND_ROBIN,
     )
-    workerChoiceStrategiesContext.workerChoiceStrategies.set(
-      workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
-      workerChoiceStrategyStub,
-    )
-    const chosenWorkerKey = workerChoiceStrategiesContext.execute()
-    assertSpyCalls(
-      workerChoiceStrategiesContext.workerChoiceStrategies.get(
+    const workerChoiceStrategyStub = workerChoiceStrategiesContext
+      .workerChoiceStrategies.get(
         workerChoiceStrategiesContext.defaultWorkerChoiceStrategy,
-      ).choose,
-      1,
-    )
+      )
+    stub(workerChoiceStrategyStub, 'choose', returnsNext([0]))
+    const chosenWorkerKey = workerChoiceStrategiesContext.execute()
+    assertSpyCalls(workerChoiceStrategyStub.choose, 1)
+    expect(workerChoiceStrategiesContext.getStrategyRetries()).toBe(0)
     expect(chosenWorkerKey).toBe(0)
     workerChoiceStrategyStub.choose.restore()
   })