fix: fix worker choice strategy retries mechanism on some edge cases
[poolifier.git] / src / pools / selection-strategies / weighted-round-robin-worker-choice-strategy.ts
1 import type { IWorker } from '../worker'
2 import type { IPool } from '../pool'
3 import {
4 DEFAULT_MEASUREMENT_STATISTICS_REQUIREMENTS,
5 DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
6 } from '../../utils'
7 import { AbstractWorkerChoiceStrategy } from './abstract-worker-choice-strategy'
8 import type {
9 IWorkerChoiceStrategy,
10 StrategyPolicy,
11 TaskStatisticsRequirements,
12 WorkerChoiceStrategyOptions
13 } from './selection-strategies-types'
14
15 /**
16 * Selects the next worker with a weighted round robin scheduling algorithm.
17 * Loosely modeled after the weighted round robin queueing algorithm: https://en.wikipedia.org/wiki/Weighted_round_robin.
18 *
19 * @typeParam Worker - Type of worker which manages the strategy.
20 * @typeParam Data - Type of data sent to the worker. This can only be structured-cloneable data.
21 * @typeParam Response - Type of execution response. This can only be structured-cloneable data.
22 */
23 export class WeightedRoundRobinWorkerChoiceStrategy<
24 Worker extends IWorker,
25 Data = unknown,
26 Response = unknown
27 >
28 extends AbstractWorkerChoiceStrategy<Worker, Data, Response>
29 implements IWorkerChoiceStrategy {
30 /** @inheritDoc */
31 public readonly strategyPolicy: StrategyPolicy = {
32 dynamicWorkerUsage: false,
33 dynamicWorkerReady: true
34 }
35
36 /** @inheritDoc */
37 public readonly taskStatisticsRequirements: TaskStatisticsRequirements = {
38 runTime: {
39 aggregate: true,
40 average: true,
41 median: false
42 },
43 waitTime: DEFAULT_MEASUREMENT_STATISTICS_REQUIREMENTS,
44 elu: DEFAULT_MEASUREMENT_STATISTICS_REQUIREMENTS
45 }
46
47 /**
48 * Default worker weight.
49 */
50 private readonly defaultWorkerWeight: number
51 /**
52 * Worker virtual task runtime.
53 */
54 private workerVirtualTaskRunTime: number = 0
55
56 /** @inheritDoc */
57 public constructor (
58 pool: IPool<Worker, Data, Response>,
59 opts: WorkerChoiceStrategyOptions = DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
60 ) {
61 super(pool, opts)
62 this.setTaskStatisticsRequirements(this.opts)
63 this.defaultWorkerWeight = this.computeDefaultWorkerWeight()
64 }
65
66 /** @inheritDoc */
67 public reset (): boolean {
68 this.nextWorkerNodeKey = 0
69 this.workerVirtualTaskRunTime = 0
70 return true
71 }
72
73 /** @inheritDoc */
74 public update (): boolean {
75 return true
76 }
77
78 /** @inheritDoc */
79 public choose (): number | undefined {
80 const chosenWorkerNodeKey = this.nextWorkerNodeKey
81 this.weightedRoundRobinNextWorkerNodeKey()
82 if (!this.isWorkerNodeEligible(this.nextWorkerNodeKey as number)) {
83 this.nextWorkerNodeKey = undefined
84 }
85 return chosenWorkerNodeKey
86 }
87
88 /** @inheritDoc */
89 public remove (workerNodeKey: number): boolean {
90 if (this.nextWorkerNodeKey === workerNodeKey) {
91 if (this.pool.workerNodes.length === 0) {
92 this.nextWorkerNodeKey = 0
93 } else if (this.nextWorkerNodeKey > this.pool.workerNodes.length - 1) {
94 this.nextWorkerNodeKey = this.pool.workerNodes.length - 1
95 }
96 this.workerVirtualTaskRunTime = 0
97 }
98 return true
99 }
100
101 private weightedRoundRobinNextWorkerNodeKey (): number | undefined {
102 const workerVirtualTaskRunTime = this.workerVirtualTaskRunTime
103 const workerWeight =
104 this.opts.weights?.[this.nextWorkerNodeKey ?? 0] ??
105 this.defaultWorkerWeight
106 if (workerVirtualTaskRunTime < workerWeight) {
107 this.workerVirtualTaskRunTime =
108 workerVirtualTaskRunTime +
109 this.getWorkerTaskRunTime(this.nextWorkerNodeKey ?? 0)
110 } else {
111 this.nextWorkerNodeKey =
112 this.nextWorkerNodeKey === this.pool.workerNodes.length - 1
113 ? 0
114 : (this.nextWorkerNodeKey ?? 0) + 1
115 this.workerVirtualTaskRunTime = 0
116 }
117 return this.nextWorkerNodeKey
118 }
119 }