From: Jérôme Benoit <jerome.benoit@piment-noir.org>
Date: Tue, 19 Dec 2023 22:47:23 +0000 (+0100)
Subject: fix: ensure the number of worker choice retries is enough for WRR
X-Git-Tag: v3.1.7~1^2~20
X-Git-Url: https://git.piment-noir.org/?a=commitdiff_plain;h=449cd15417158c0d1675a8aec8bcc514ff33b5a3;p=poolifier.git

fix: ensure the number of worker choice retries is enough for WRR

Signed-off-by: Jérôme Benoit <jerome.benoit@piment-noir.org>
---

diff --git a/src/pools/selection-strategies/abstract-worker-choice-strategy.ts b/src/pools/selection-strategies/abstract-worker-choice-strategy.ts
index 34473ab4..7324c511 100644
--- a/src/pools/selection-strategies/abstract-worker-choice-strategy.ts
+++ b/src/pools/selection-strategies/abstract-worker-choice-strategy.ts
@@ -113,7 +113,10 @@ export abstract class AbstractWorkerChoiceStrategy<
   /** @inheritDoc */
   public setOptions (opts: InternalWorkerChoiceStrategyOptions): void {
     this.opts = {
-      ...getDefaultInternalWorkerChoiceStrategyOptions(this.pool.info.maxSize),
+      ...getDefaultInternalWorkerChoiceStrategyOptions(
+        this.pool.info.maxSize +
+          Object.keys((opts?.weights as Record<number, number>) ?? {}).length
+      ),
       ...opts
     }
     this.setTaskStatisticsRequirements(this.opts)
diff --git a/src/pools/selection-strategies/worker-choice-strategy-context.ts b/src/pools/selection-strategies/worker-choice-strategy-context.ts
index 1de78aba..cf1b0051 100644
--- a/src/pools/selection-strategies/worker-choice-strategy-context.ts
+++ b/src/pools/selection-strategies/worker-choice-strategy-context.ts
@@ -47,7 +47,11 @@ export class WorkerChoiceStrategyContext<
     private opts?: InternalWorkerChoiceStrategyOptions
   ) {
     this.opts = {
-      ...getDefaultInternalWorkerChoiceStrategyOptions(pool.info.maxSize),
+      ...getDefaultInternalWorkerChoiceStrategyOptions(
+        pool.info.maxSize +
+          Object.keys((this.opts?.weights as Record<number, number>) ?? {})
+            .length
+      ),
       ...this.opts
     }
     this.execute = this.execute.bind(this)
@@ -232,7 +236,10 @@ export class WorkerChoiceStrategyContext<
     opts?: InternalWorkerChoiceStrategyOptions
   ): void {
     this.opts = {
-      ...getDefaultInternalWorkerChoiceStrategyOptions(pool.info.maxSize),
+      ...getDefaultInternalWorkerChoiceStrategyOptions(
+        pool.info.maxSize +
+          Object.keys((opts?.weights as Record<number, number>) ?? {}).length
+      ),
       ...opts
     }
     for (const workerChoiceStrategy of this.workerChoiceStrategies.values()) {
diff --git a/src/utils.ts b/src/utils.ts
index 25233a02..2dd3947c 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -22,16 +22,16 @@ export const EMPTY_FUNCTION: () => void = Object.freeze(() => {
 })
 
 /**
- * Default worker choice strategy options.
+ * Gets default worker choice strategy options.
  *
- * @param poolMaxSize - The pool maximum size.
+ * @param retries - The number of worker choice retries.
  * @returns The default worker choice strategy options.
  */
 export const getDefaultInternalWorkerChoiceStrategyOptions = (
-  poolMaxSize: number
+  retries: number
 ): InternalWorkerChoiceStrategyOptions => {
   return {
-    retries: poolMaxSize,
+    retries,
     runTime: { median: false },
     waitTime: { median: false },
     elu: { median: false }
diff --git a/tests/pools/abstract-pool.test.mjs b/tests/pools/abstract-pool.test.mjs
index 44aca399..29f038c7 100644
--- a/tests/pools/abstract-pool.test.mjs
+++ b/tests/pools/abstract-pool.test.mjs
@@ -288,7 +288,9 @@ describe('Abstract pool test suite', () => {
       exitHandler: testHandler
     })
     expect(pool.workerChoiceStrategyContext.opts).toStrictEqual({
-      retries: pool.info.maxSize,
+      retries:
+        pool.info.maxSize +
+        Object.keys(pool.opts.workerChoiceStrategyOptions.weights).length,
       runTime: { median: true },
       waitTime: { median: false },
       elu: { median: false },
@@ -297,7 +299,9 @@ describe('Abstract pool test suite', () => {
     for (const [, workerChoiceStrategy] of pool.workerChoiceStrategyContext
       .workerChoiceStrategies) {
       expect(workerChoiceStrategy.opts).toStrictEqual({
-        retries: pool.info.maxSize,
+        retries:
+          pool.info.maxSize +
+          Object.keys(pool.opts.workerChoiceStrategyOptions.weights).length,
         runTime: { median: true },
         waitTime: { median: false },
         elu: { median: false },