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

[poolifier.git] / src / pools / selection-strategies / worker-choice-strategy-context.ts

diff --git a/src/pools/selection-strategies/worker-choice-strategy-context.ts b/src/pools/selection-strategies/worker-choice-strategy-context.ts
index 1de78aba9f15f1ca2cddb1cdac474c8470e7991f..cf1b005169d8fbaa04fb074adb4f3f2966cbe335 100644 (file)
--- 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()) {