fix: fix race condition in worker choice strategies
[poolifier.git] / src / pools / selection-strategies / round-robin-worker-choice-strategy.ts
1 import { DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS } from '../../utils'
2 import type { IPool } from '../pool'
3 import type { IWorker } from '../worker'
4 import { AbstractWorkerChoiceStrategy } from './abstract-worker-choice-strategy'
5 import type {
6 IWorkerChoiceStrategy,
7 WorkerChoiceStrategyOptions
8 } from './selection-strategies-types'
9
10 /**
11 * Selects the next worker in a round robin fashion.
12 *
13 * @typeParam Worker - Type of worker which manages the strategy.
14 * @typeParam Data - Type of data sent to the worker. This can only be structured-cloneable data.
15 * @typeParam Response - Type of execution response. This can only be structured-cloneable data.
16 */
17 export class RoundRobinWorkerChoiceStrategy<
18 Worker extends IWorker,
19 Data = unknown,
20 Response = unknown
21 >
22 extends AbstractWorkerChoiceStrategy<Worker, Data, Response>
23 implements IWorkerChoiceStrategy {
24 /** @inheritDoc */
25 public constructor (
26 pool: IPool<Worker, Data, Response>,
27 opts: WorkerChoiceStrategyOptions = DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
28 ) {
29 super(pool, opts)
30 this.setTaskStatisticsRequirements(this.opts)
31 }
32
33 /** @inheritDoc */
34 public reset (): boolean {
35 this.resetWorkerNodeKeyProperties()
36 return true
37 }
38
39 /** @inheritDoc */
40 public update (): boolean {
41 return true
42 }
43
44 /** @inheritDoc */
45 public choose (): number | undefined {
46 const chosenWorkerNodeKey = this.nextWorkerNodeKey
47 this.roundRobinNextWorkerNodeKey()
48 return chosenWorkerNodeKey
49 }
50
51 /** @inheritDoc */
52 public remove (workerNodeKey: number): boolean {
53 if (this.nextWorkerNodeKey === workerNodeKey) {
54 if (this.pool.workerNodes.length === 0) {
55 this.nextWorkerNodeKey = 0
56 } else if (this.nextWorkerNodeKey > this.pool.workerNodes.length - 1) {
57 this.nextWorkerNodeKey = this.pool.workerNodes.length - 1
58 }
59 }
60 return true
61 }
62
63 private roundRobinNextWorkerNodeKey (): number | undefined {
64 this.nextWorkerNodeKey =
65 this.nextWorkerNodeKey === this.pool.workerNodes.length - 1
66 ? 0
67 : (this.nextWorkerNodeKey ?? this.previousWorkerNodeKey) + 1
68 return this.nextWorkerNodeKey
69 }
70 }