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

   1 import type { IPoolInternal } from '../pool-internal'
   2 import type { IPoolWorker } from '../pool-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 } from './selection-strategies-types'
  12 import { WorkerChoiceStrategies } from './selection-strategies-types'
  13 import { WeightedRoundRobinWorkerChoiceStrategy } from './weighted-round-robin-worker-choice-strategy'
  14
  15 /**
  16  * The worker choice strategy context.
  17  *
  18  * @typeParam Worker - Type of worker.
  19  * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
  20  * @typeParam Response - Type of response of execution. This can only be serializable data.
  21  */
  22 export class WorkerChoiceStrategyContext<
  23   Worker extends IPoolWorker,
  24   Data = unknown,
  25   Response = unknown
  26 > {
  27   private readonly workerChoiceStrategies = new Map<
  28   WorkerChoiceStrategy,
  29   IWorkerChoiceStrategy<Worker, Data, Response>
  30   >()
  31
  32   /**
  33    * Worker choice strategy context constructor.
  34    *
  35    * @param pool - The pool instance.
  36    * @param createWorkerCallback - The worker creation callback for dynamic pool.
  37    * @param workerChoiceStrategy - The worker choice strategy.
  38    */
  39   public constructor (
  40     pool: IPoolInternal<Worker, Data, Response>,
  41     private readonly createWorkerCallback: () => number,
  42     private workerChoiceStrategyType: WorkerChoiceStrategy = WorkerChoiceStrategies.ROUND_ROBIN
  43   ) {
  44     this.execute.bind(this)
  45     this.registerWorkerChoiceStrategy(pool, workerChoiceStrategyType)
  46   }
  47
  48   /**
  49    * Gets the worker choice strategy in the context required statistics.
  50    *
  51    * @returns The required statistics.
  52    */
  53   public getRequiredStatistics (): RequiredStatistics {
  54     return (
  55       this.workerChoiceStrategies.get(
  56         this.workerChoiceStrategyType
  57       ) as IWorkerChoiceStrategy<Worker, Data, Response>
  58     ).requiredStatistics
  59   }
  60
  61   /**
  62    * Sets the worker choice strategy to use in the context.
  63    *
  64    * @param workerChoiceStrategy - The worker choice strategy to set.
  65    */
  66   public setWorkerChoiceStrategy (
  67     pool: IPoolInternal<Worker, Data, Response>,
  68     workerChoiceStrategy: WorkerChoiceStrategy
  69   ): void {
  70     if (this.workerChoiceStrategyType === workerChoiceStrategy) {
  71       this.workerChoiceStrategies.get(workerChoiceStrategy)?.reset()
  72     } else {
  73       this.workerChoiceStrategyType = workerChoiceStrategy
  74       this.registerWorkerChoiceStrategy(pool, workerChoiceStrategy)
  75     }
  76   }
  77
  78   /**
  79    * Executes the worker choice strategy algorithm in the context.
  80    *
  81    * @returns The key of the chosen one.
  82    */
  83   public execute (): number {
  84     const workerChoiceStrategy = this.workerChoiceStrategies.get(
  85       this.workerChoiceStrategyType
  86     ) as IWorkerChoiceStrategy<Worker, Data, Response>
  87     if (
  88       workerChoiceStrategy.isDynamicPool &&
  89       !workerChoiceStrategy.pool.full &&
  90       workerChoiceStrategy.pool.findFreeWorkerKey() === -1
  91     ) {
  92       return this.createWorkerCallback()
  93     }
  94     return workerChoiceStrategy.choose()
  95   }
  96
  97   /**
  98    * Removes a worker from the worker choice strategy in the context.
  99    *
 100    * @param workerKey - The key of the worker to remove.
 101    * @returns `true` if the removal is successful, `false` otherwise.
 102    */
 103   public remove (workerKey: number): boolean {
 104     return (
 105       this.workerChoiceStrategies.get(
 106         this.workerChoiceStrategyType
 107       ) as IWorkerChoiceStrategy<Worker, Data, Response>
 108     ).remove(workerKey)
 109   }
 110
 111   private registerWorkerChoiceStrategy (
 112     pool: IPoolInternal<Worker, Data, Response>,
 113     workerChoiceStrategy: WorkerChoiceStrategy
 114   ): void {
 115     if (!this.workerChoiceStrategies.has(workerChoiceStrategy)) {
 116       this.workerChoiceStrategies.set(
 117         workerChoiceStrategy,
 118         this.getWorkerChoiceStrategy(pool, workerChoiceStrategy)
 119       )
 120     }
 121   }
 122
 123   /**
 124    * Gets the worker choice strategy instance.
 125    *
 126    * @param pool - The pool instance.
 127    * @param workerChoiceStrategy - The worker choice strategy.
 128    * @returns The worker choice strategy instance.
 129    */
 130   private getWorkerChoiceStrategy (
 131     pool: IPoolInternal<Worker, Data, Response>,
 132     workerChoiceStrategy: WorkerChoiceStrategy = WorkerChoiceStrategies.ROUND_ROBIN
 133   ): IWorkerChoiceStrategy<Worker, Data, Response> {
 134     switch (workerChoiceStrategy) {
 135       case WorkerChoiceStrategies.ROUND_ROBIN:
 136         return new RoundRobinWorkerChoiceStrategy<Worker, Data, Response>(pool)
 137       case WorkerChoiceStrategies.LESS_USED:
 138         return new LessUsedWorkerChoiceStrategy<Worker, Data, Response>(pool)
 139       case WorkerChoiceStrategies.LESS_BUSY:
 140         return new LessBusyWorkerChoiceStrategy<Worker, Data, Response>(pool)
 141       case WorkerChoiceStrategies.FAIR_SHARE:
 142         return new FairShareWorkerChoiceStrategy<Worker, Data, Response>(pool)
 143       case WorkerChoiceStrategies.WEIGHTED_ROUND_ROBIN:
 144         return new WeightedRoundRobinWorkerChoiceStrategy<
 145         Worker,
 146         Data,
 147         Response
 148         >(pool)
 149       default:
 150         throw new Error(
 151           // eslint-disable-next-line @typescript-eslint/restrict-template-expressions
 152           `Worker choice strategy '${workerChoiceStrategy}' not found`
 153         )
 154     }
 155   }
 156 }