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 { LessBusyWorkerChoiceStrategy } from './less-busy-worker-choice-strategy'
   6 import { LessUsedWorkerChoiceStrategy } from './less-used-worker-choice-strategy'
   7 import { RoundRobinWorkerChoiceStrategy } from './round-robin-worker-choice-strategy'
   8 import type {
   9   IWorkerChoiceStrategy,
  10   RequiredStatistics,
  11   WorkerChoiceStrategy,
  12   WorkerChoiceStrategyOptions
  13 } from './selection-strategies-types'
  14 import { WorkerChoiceStrategies } from './selection-strategies-types'
  15 import { WeightedRoundRobinWorkerChoiceStrategy } from './weighted-round-robin-worker-choice-strategy'
  16
  17 /**
  18  * The worker choice strategy context.
  19  *
  20  * @typeParam Worker - Type of worker.
  21  * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
  22  * @typeParam Response - Type of execution response. This can only be serializable data.
  23  */
  24 export class WorkerChoiceStrategyContext<
  25   Worker extends IWorker,
  26   Data = unknown,
  27   Response = unknown
  28 > {
  29   private readonly workerChoiceStrategies: Map<
  30   WorkerChoiceStrategy,
  31   IWorkerChoiceStrategy
  32   >
  33
  34   /**
  35    * Worker choice strategy context constructor.
  36    *
  37    * @param pool - The pool instance.
  38    * @param workerChoiceStrategy - The worker choice strategy.
  39    * @param opts - The worker choice strategy options.
  40    */
  41   public constructor (
  42     pool: IPool<Worker, Data, Response>,
  43     /**
  44      * @defaultValue WorkerChoiceStrategies.ROUND_ROBIN
  45      */
  46     private workerChoiceStrategy: WorkerChoiceStrategy = WorkerChoiceStrategies.ROUND_ROBIN,
  47     opts: WorkerChoiceStrategyOptions = DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
  48   ) {
  49     this.execute.bind(this)
  50     this.workerChoiceStrategies = new Map<
  51     WorkerChoiceStrategy,
  52     IWorkerChoiceStrategy
  53     >([
  54       [
  55         WorkerChoiceStrategies.ROUND_ROBIN,
  56         new RoundRobinWorkerChoiceStrategy<Worker, Data, Response>(pool, opts)
  57       ],
  58       [
  59         WorkerChoiceStrategies.LESS_USED,
  60         new LessUsedWorkerChoiceStrategy<Worker, Data, Response>(pool, opts)
  61       ],
  62       [
  63         WorkerChoiceStrategies.LESS_BUSY,
  64         new LessBusyWorkerChoiceStrategy<Worker, Data, Response>(pool, opts)
  65       ],
  66       [
  67         WorkerChoiceStrategies.FAIR_SHARE,
  68         new FairShareWorkerChoiceStrategy<Worker, Data, Response>(pool, opts)
  69       ],
  70       [
  71         WorkerChoiceStrategies.WEIGHTED_ROUND_ROBIN,
  72         new WeightedRoundRobinWorkerChoiceStrategy<Worker, Data, Response>(
  73           pool,
  74           opts
  75         )
  76       ]
  77     ])
  78   }
  79
  80   /**
  81    * Gets the worker choice strategy in the context required statistics.
  82    *
  83    * @returns The required statistics.
  84    */
  85   public getRequiredStatistics (): RequiredStatistics {
  86     return (
  87       this.workerChoiceStrategies.get(
  88         this.workerChoiceStrategy
  89       ) as IWorkerChoiceStrategy
  90     ).requiredStatistics
  91   }
  92
  93   /**
  94    * Sets the worker choice strategy to use in the context.
  95    *
  96    * @param workerChoiceStrategy - The worker choice strategy to set.
  97    */
  98   public setWorkerChoiceStrategy (
  99     workerChoiceStrategy: WorkerChoiceStrategy
 100   ): void {
 101     if (this.workerChoiceStrategy !== workerChoiceStrategy) {
 102       this.workerChoiceStrategy = workerChoiceStrategy
 103     }
 104     this.workerChoiceStrategies.get(this.workerChoiceStrategy)?.reset()
 105   }
 106
 107   /**
 108    * Executes the worker choice strategy algorithm in the context.
 109    *
 110    * @returns The key of the worker node.
 111    */
 112   public execute (): number {
 113     return (
 114       this.workerChoiceStrategies.get(
 115         this.workerChoiceStrategy
 116       ) as IWorkerChoiceStrategy
 117     ).choose()
 118   }
 119
 120   /**
 121    * Removes a worker node key from the worker choice strategy in the context.
 122    *
 123    * @param workerNodeKey - The key of the worker node.
 124    * @returns `true` if the removal is successful, `false` otherwise.
 125    */
 126   public remove (workerNodeKey: number): boolean {
 127     return (
 128       this.workerChoiceStrategies.get(
 129         this.workerChoiceStrategy
 130       ) as IWorkerChoiceStrategy
 131     ).remove(workerNodeKey)
 132   }
 133
 134   /**
 135    * Sets the worker choice strategies in the context options.
 136    *
 137    * @param opts - The worker choice strategy options.
 138    */
 139   public setOptions (opts: WorkerChoiceStrategyOptions): void {
 140     this.workerChoiceStrategies.forEach(workerChoiceStrategy => {
 141       workerChoiceStrategy.setOptions(opts)
 142     })
 143   }
 144 }