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 { RoundRobinWorkerChoiceStrategy } from './round-robin-worker-choice-strategy'
   9 import type {
  10   IWorkerChoiceStrategy,
  11   TaskStatistics,
  12   WorkerChoiceStrategy,
  13   WorkerChoiceStrategyOptions
  14 } from './selection-strategies-types'
  15 import { WorkerChoiceStrategies } from './selection-strategies-types'
  16 import { WeightedRoundRobinWorkerChoiceStrategy } from './weighted-round-robin-worker-choice-strategy'
  17
  18 /**
  19  * The worker choice strategy context.
  20  *
  21  * @typeParam Worker - Type of worker.
  22  * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
  23  * @typeParam Response - Type of execution response. This can only be serializable data.
  24  */
  25 export class WorkerChoiceStrategyContext<
  26   Worker extends IWorker,
  27   Data = unknown,
  28   Response = unknown
  29 > {
  30   private readonly workerChoiceStrategies: Map<
  31   WorkerChoiceStrategy,
  32   IWorkerChoiceStrategy
  33   >
  34
  35   /**
  36    * Worker choice strategy context constructor.
  37    *
  38    * @param pool - The pool instance.
  39    * @param workerChoiceStrategy - The worker choice strategy.
  40    * @param opts - The worker choice strategy options.
  41    */
  42   public constructor (
  43     pool: IPool<Worker, Data, Response>,
  44     private workerChoiceStrategy: WorkerChoiceStrategy = WorkerChoiceStrategies.ROUND_ROBIN,
  45     opts: WorkerChoiceStrategyOptions = DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
  46   ) {
  47     this.execute = this.execute.bind(this)
  48     this.workerChoiceStrategies = new Map<
  49     WorkerChoiceStrategy,
  50     IWorkerChoiceStrategy
  51     >([
  52       [
  53         WorkerChoiceStrategies.ROUND_ROBIN,
  54         new (RoundRobinWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  55           pool,
  56           opts
  57         )
  58       ],
  59       [
  60         WorkerChoiceStrategies.LEAST_USED,
  61         new (LeastUsedWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  62           pool,
  63           opts
  64         )
  65       ],
  66       [
  67         WorkerChoiceStrategies.LEAST_BUSY,
  68         new (LeastBusyWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  69           pool,
  70           opts
  71         )
  72       ],
  73       [
  74         WorkerChoiceStrategies.FAIR_SHARE,
  75         new (FairShareWorkerChoiceStrategy.bind(this))<Worker, Data, Response>(
  76           pool,
  77           opts
  78         )
  79       ],
  80       [
  81         WorkerChoiceStrategies.WEIGHTED_ROUND_ROBIN,
  82         new (WeightedRoundRobinWorkerChoiceStrategy.bind(this))<
  83         Worker,
  84         Data,
  85         Response
  86         >(pool, opts)
  87       ],
  88       [
  89         WorkerChoiceStrategies.INTERLEAVED_WEIGHTED_ROUND_ROBIN,
  90         new (InterleavedWeightedRoundRobinWorkerChoiceStrategy.bind(this))<
  91         Worker,
  92         Data,
  93         Response
  94         >(pool, opts)
  95       ]
  96     ])
  97   }
  98
  99   /**
 100    * Gets the worker choice strategy task statistics in the context.
 101    *
 102    * @returns The task statistics.
 103    */
 104   public getTaskStatistics (): TaskStatistics {
 105     return (
 106       this.workerChoiceStrategies.get(
 107         this.workerChoiceStrategy
 108       ) as IWorkerChoiceStrategy
 109     ).taskStatistics
 110   }
 111
 112   /**
 113    * Sets the worker choice strategy to use in the context.
 114    *
 115    * @param workerChoiceStrategy - The worker choice strategy to set.
 116    */
 117   public setWorkerChoiceStrategy (
 118     workerChoiceStrategy: WorkerChoiceStrategy
 119   ): void {
 120     if (this.workerChoiceStrategy !== workerChoiceStrategy) {
 121       this.workerChoiceStrategy = workerChoiceStrategy
 122     }
 123     this.workerChoiceStrategies.get(this.workerChoiceStrategy)?.reset()
 124   }
 125
 126   /**
 127    * Updates the worker node key in the worker choice strategy internals in the context.
 128    *
 129    * @returns `true` if the update is successful, `false` otherwise.
 130    */
 131   public update (workerNodeKey: number): boolean {
 132     return (
 133       this.workerChoiceStrategies.get(
 134         this.workerChoiceStrategy
 135       ) as IWorkerChoiceStrategy
 136     ).update(workerNodeKey)
 137   }
 138
 139   /**
 140    * Executes the worker choice strategy algorithm in the context.
 141    *
 142    * @returns The key of the worker node.
 143    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the worker node key is null or undefined.
 144    */
 145   public execute (): number {
 146     const workerNodeKey = (
 147       this.workerChoiceStrategies.get(
 148         this.workerChoiceStrategy
 149       ) as IWorkerChoiceStrategy
 150     ).choose()
 151     if (workerNodeKey == null) {
 152       throw new Error('Worker node key chosen is null or undefined')
 153     }
 154     return workerNodeKey
 155   }
 156
 157   /**
 158    * Removes the worker node key from the worker choice strategy in the context.
 159    *
 160    * @param workerNodeKey - The key of the worker node.
 161    * @returns `true` if the removal is successful, `false` otherwise.
 162    */
 163   public remove (workerNodeKey: number): boolean {
 164     return (
 165       this.workerChoiceStrategies.get(
 166         this.workerChoiceStrategy
 167       ) as IWorkerChoiceStrategy
 168     ).remove(workerNodeKey)
 169   }
 170
 171   /**
 172    * Sets the worker choice strategies in the context options.
 173    *
 174    * @param opts - The worker choice strategy options.
 175    */
 176   public setOptions (opts: WorkerChoiceStrategyOptions): void {
 177     for (const workerChoiceStrategy of this.workerChoiceStrategies.values()) {
 178       workerChoiceStrategy.setOptions(opts)
 179     }
 180   }
 181 }