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

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