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 workerChoiceStrategyType - The worker choice strategy.
  39    * @param opts - The worker choice strategy options.
  40    */
  41   public constructor (
  42     pool: IPool<Worker, Data, Response>,
  43     private workerChoiceStrategyType: WorkerChoiceStrategy = WorkerChoiceStrategies.ROUND_ROBIN,
  44     opts: WorkerChoiceStrategyOptions = DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
  45   ) {
  46     this.execute.bind(this)
  47     this.workerChoiceStrategies = new Map<
  48     WorkerChoiceStrategy,
  49     IWorkerChoiceStrategy
  50     >([
  51       [
  52         WorkerChoiceStrategies.ROUND_ROBIN,
  53         new RoundRobinWorkerChoiceStrategy<Worker, Data, Response>(pool, opts)
  54       ],
  55       [
  56         WorkerChoiceStrategies.LESS_USED,
  57         new LessUsedWorkerChoiceStrategy<Worker, Data, Response>(pool, opts)
  58       ],
  59       [
  60         WorkerChoiceStrategies.LESS_BUSY,
  61         new LessBusyWorkerChoiceStrategy<Worker, Data, Response>(pool, opts)
  62       ],
  63       [
  64         WorkerChoiceStrategies.FAIR_SHARE,
  65         new FairShareWorkerChoiceStrategy<Worker, Data, Response>(pool, opts)
  66       ],
  67       [
  68         WorkerChoiceStrategies.WEIGHTED_ROUND_ROBIN,
  69         new WeightedRoundRobinWorkerChoiceStrategy<Worker, Data, Response>(
  70           pool,
  71           opts
  72         )
  73       ]
  74     ])
  75   }
  76
  77   /**
  78    * Gets the worker choice strategy in the context required statistics.
  79    *
  80    * @returns The required statistics.
  81    */
  82   public getRequiredStatistics (): RequiredStatistics {
  83     return (
  84       this.workerChoiceStrategies.get(
  85         this.workerChoiceStrategyType
  86       ) as IWorkerChoiceStrategy
  87     ).requiredStatistics
  88   }
  89
  90   /**
  91    * Sets the worker choice strategy to use in the context.
  92    *
  93    * @param workerChoiceStrategy - The worker choice strategy to set.
  94    */
  95   public setWorkerChoiceStrategy (
  96     workerChoiceStrategy: WorkerChoiceStrategy
  97   ): void {
  98     if (this.workerChoiceStrategyType !== workerChoiceStrategy) {
  99       this.workerChoiceStrategyType = workerChoiceStrategy
 100     }
 101     this.workerChoiceStrategies.get(this.workerChoiceStrategyType)?.reset()
 102   }
 103
 104   /**
 105    * Executes the worker choice strategy algorithm in the context.
 106    *
 107    * @returns The key of the worker node.
 108    */
 109   public execute (): number {
 110     return (
 111       this.workerChoiceStrategies.get(
 112         this.workerChoiceStrategyType
 113       ) as IWorkerChoiceStrategy
 114     ).choose()
 115   }
 116
 117   /**
 118    * Removes a worker node key from the worker choice strategy in the context.
 119    *
 120    * @param workerNodeKey - The key of the worker node.
 121    * @returns `true` if the removal is successful, `false` otherwise.
 122    */
 123   public remove (workerNodeKey: number): boolean {
 124     return (
 125       this.workerChoiceStrategies.get(
 126         this.workerChoiceStrategyType
 127       ) as IWorkerChoiceStrategy
 128     ).remove(workerNodeKey)
 129   }
 130 }