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

   1 import { DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS } from '../../utils'
   2 import type { IPool } from '../pool'
   3 import type { IWorker } from '../worker'
   4 import { FairShareWorkerChoiceStrategy } from './fair-share-worker-choice-strategy'
   5 import { InterleavedWeightedRoundRobinWorkerChoiceStrategy } from './interleaved-weighted-round-robin-worker-choice-strategy'
   6 import { LeastBusyWorkerChoiceStrategy } from './least-busy-worker-choice-strategy'
   7 import { LeastUsedWorkerChoiceStrategy } from './least-used-worker-choice-strategy'
   8 import { LeastEluWorkerChoiceStrategy } from './least-elu-worker-choice-strategy'
   9 import { RoundRobinWorkerChoiceStrategy } from './round-robin-worker-choice-strategy'
  10 import type {
  11   IWorkerChoiceStrategy,
  12   StrategyPolicy,
  13   TaskStatisticsRequirements,
  14   WorkerChoiceStrategy,
  15   WorkerChoiceStrategyOptions
  16 } from './selection-strategies-types'
  17 import { WorkerChoiceStrategies } from './selection-strategies-types'
  18 import { WeightedRoundRobinWorkerChoiceStrategy } from './weighted-round-robin-worker-choice-strategy'
  19
  20 /**
  21  * The worker choice strategy context.
  22  *
  23  * @typeParam Worker - Type of worker.
  24  * @typeParam Data - Type of data sent to the worker. This can only be structured-cloneable data.
  25  * @typeParam Response - Type of execution response. This can only be structured-cloneable data.
  26  */
  27 export class WorkerChoiceStrategyContext<
  28   Worker extends IWorker,
  29   Data = unknown,
  30   Response = unknown
  31 > {
  32   private readonly workerChoiceStrategies: Map<
  33   WorkerChoiceStrategy,
  34   IWorkerChoiceStrategy
  35   >
  36
  37   /**
  38    * The number of times the worker choice strategy in the context has been retried.
  39    */
  40   private retriesCount = 0
  41
  42   /**
  43    * Worker choice strategy context constructor.
  44    *
  45    * @param pool - The pool instance.
  46    * @param workerChoiceStrategy - The worker choice strategy.
  47    * @param opts - The worker choice strategy options.
  48    */
  49   public constructor (
  50     pool: IPool<Worker, Data, Response>,
  51     private workerChoiceStrategy: WorkerChoiceStrategy = WorkerChoiceStrategies.ROUND_ROBIN,
  52     private opts: WorkerChoiceStrategyOptions = DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
  53   ) {
  54     this.opts = { ...DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS, ...opts }
  55     this.execute = this.execute.bind(this)
  56     this.workerChoiceStrategies = new Map<
  57     WorkerChoiceStrategy,
  58     IWorkerChoiceStrategy
  59     >([
  60       [
  61         WorkerChoiceStrategies.ROUND_ROBIN,
  62         new (RoundRobinWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  63           pool,
  64           opts
  65         )
  66       ],
  67       [
  68         WorkerChoiceStrategies.LEAST_USED,
  69         new (LeastUsedWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  70           pool,
  71           opts
  72         )
  73       ],
  74       [
  75         WorkerChoiceStrategies.LEAST_BUSY,
  76         new (LeastBusyWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  77           pool,
  78           opts
  79         )
  80       ],
  81       [
  82         WorkerChoiceStrategies.LEAST_ELU,
  83         new (LeastEluWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  84           pool,
  85           opts
  86         )
  87       ],
  88       [
  89         WorkerChoiceStrategies.FAIR_SHARE,
  90         new (FairShareWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  91           pool,
  92           opts
  93         )
  94       ],
  95       [
  96         WorkerChoiceStrategies.WEIGHTED_ROUND_ROBIN,
  97         new (WeightedRoundRobinWorkerChoiceStrategy.bind(this))<
  98         Worker,
  99         Data,
 100         Response
 101         >(pool, opts)
 102       ],
 103       [
 104         WorkerChoiceStrategies.INTERLEAVED_WEIGHTED_ROUND_ROBIN,
 105         new (InterleavedWeightedRoundRobinWorkerChoiceStrategy.bind(this))<
 106         Worker,
 107         Data,
 108         Response
 109         >(pool, opts)
 110       ]
 111     ])
 112   }
 113
 114   /**
 115    * Gets the strategy policy in the context.
 116    *
 117    * @returns The strategy policy.
 118    */
 119   public getStrategyPolicy (): StrategyPolicy {
 120     return (
 121       this.workerChoiceStrategies.get(
 122         this.workerChoiceStrategy
 123       ) as IWorkerChoiceStrategy
 124     ).strategyPolicy
 125   }
 126
 127   /**
 128    * Gets the worker choice strategy in the context task statistics requirements.
 129    *
 130    * @returns The task statistics requirements.
 131    */
 132   public getTaskStatisticsRequirements (): TaskStatisticsRequirements {
 133     return (
 134       this.workerChoiceStrategies.get(
 135         this.workerChoiceStrategy
 136       ) as IWorkerChoiceStrategy
 137     ).taskStatisticsRequirements
 138   }
 139
 140   /**
 141    * Sets the worker choice strategy to use in the context.
 142    *
 143    * @param workerChoiceStrategy - The worker choice strategy to set.
 144    */
 145   public setWorkerChoiceStrategy (
 146     workerChoiceStrategy: WorkerChoiceStrategy
 147   ): void {
 148     if (this.workerChoiceStrategy !== workerChoiceStrategy) {
 149       this.workerChoiceStrategy = workerChoiceStrategy
 150     }
 151     this.workerChoiceStrategies.get(this.workerChoiceStrategy)?.reset()
 152   }
 153
 154   /**
 155    * Updates the worker node key in the worker choice strategy in the context internals.
 156    *
 157    * @returns `true` if the update is successful, `false` otherwise.
 158    */
 159   public update (workerNodeKey: number): boolean {
 160     return (
 161       this.workerChoiceStrategies.get(
 162         this.workerChoiceStrategy
 163       ) as IWorkerChoiceStrategy
 164     ).update(workerNodeKey)
 165   }
 166
 167   /**
 168    * Executes the worker choice strategy in the context algorithm.
 169    *
 170    * @returns The key of the worker node.
 171    * @throws {@link https://nodejs.org/api/errors.html#class-error} If after configured retries the worker node key is null or undefined .
 172    */
 173   public execute (): number {
 174     const workerNodeKey = (
 175       this.workerChoiceStrategies.get(
 176         this.workerChoiceStrategy
 177       ) as IWorkerChoiceStrategy
 178     ).choose()
 179     if (
 180       workerNodeKey == null &&
 181       this.retriesCount < (this.opts.retries as number)
 182     ) {
 183       this.retriesCount++
 184       return this.execute()
 185     } else if (workerNodeKey == null) {
 186       throw new Error(
 187         `Worker node key chosen is null or undefined after ${this.retriesCount} retries`
 188       )
 189     }
 190     this.retriesCount = 0
 191     return workerNodeKey
 192   }
 193
 194   /**
 195    * Removes the worker node key from the worker choice strategy in the context.
 196    *
 197    * @param workerNodeKey - The worker node key.
 198    * @returns `true` if the removal is successful, `false` otherwise.
 199    */
 200   public remove (workerNodeKey: number): boolean {
 201     return (
 202       this.workerChoiceStrategies.get(
 203         this.workerChoiceStrategy
 204       ) as IWorkerChoiceStrategy
 205     ).remove(workerNodeKey)
 206   }
 207
 208   /**
 209    * Sets the worker choice strategies in the context options.
 210    *
 211    * @param opts - The worker choice strategy options.
 212    */
 213   public setOptions (opts: WorkerChoiceStrategyOptions): void {
 214     this.opts = { ...DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS, ...opts }
 215     for (const workerChoiceStrategy of this.workerChoiceStrategies.values()) {
 216       workerChoiceStrategy.setOptions(opts)
 217     }
 218   }
 219 }