src/pools/selection-strategies/weighted-round-robin-worker-choice-strategy.ts

   1 import { cpus } from 'node:os'
   2 import type { IWorker } from '../worker'
   3 import type { IPool } from '../pool'
   4 import { DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS } from '../../utils'
   5 import { AbstractWorkerChoiceStrategy } from './abstract-worker-choice-strategy'
   6 import type {
   7   IWorkerChoiceStrategy,
   8   RequiredStatistics,
   9   WorkerChoiceStrategyOptions
  10 } from './selection-strategies-types'
  11
  12 /**
  13  * Selects the next worker with a weighted round robin scheduling algorithm.
  14  * Loosely modeled after the weighted round robin queueing algorithm: https://en.wikipedia.org/wiki/Weighted_round_robin.
  15  *
  16  * @typeParam Worker - Type of worker which manages the strategy.
  17  * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
  18  * @typeParam Response - Type of execution response. This can only be serializable data.
  19  */
  20 export class WeightedRoundRobinWorkerChoiceStrategy<
  21     Worker extends IWorker,
  22     Data = unknown,
  23     Response = unknown
  24   >
  25   extends AbstractWorkerChoiceStrategy<Worker, Data, Response>
  26   implements IWorkerChoiceStrategy {
  27   /** @inheritDoc */
  28   public readonly requiredStatistics: RequiredStatistics = {
  29     runTime: true,
  30     avgRunTime: true,
  31     medRunTime: false
  32   }
  33
  34   /**
  35    * Worker node id where the current task will be submitted.
  36    */
  37   private currentWorkerNodeId: number = 0
  38   /**
  39    * Default worker weight.
  40    */
  41   private readonly defaultWorkerWeight: number
  42   /**
  43    * Worker virtual task runtime.
  44    */
  45   private workerVirtualTaskRunTime: number = 0
  46
  47   /** @inheritDoc */
  48   public constructor (
  49     pool: IPool<Worker, Data, Response>,
  50     opts: WorkerChoiceStrategyOptions = DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
  51   ) {
  52     super(pool, opts)
  53     this.checkOptions(this.opts)
  54     this.defaultWorkerWeight = this.computeDefaultWorkerWeight()
  55   }
  56
  57   /** @inheritDoc */
  58   public reset (): boolean {
  59     this.currentWorkerNodeId = 0
  60     this.workerVirtualTaskRunTime = 0
  61     return true
  62   }
  63
  64   /** @inheritDoc */
  65   public choose (): number {
  66     const chosenWorkerNodeKey = this.currentWorkerNodeId
  67     const workerTaskRunTime = this.workerVirtualTaskRunTime ?? 0
  68     const workerTaskWeight =
  69       this.opts.weights?.[chosenWorkerNodeKey] ?? this.defaultWorkerWeight
  70     if (workerTaskRunTime < workerTaskWeight) {
  71       this.workerVirtualTaskRunTime =
  72         workerTaskRunTime +
  73         (this.getWorkerVirtualTaskRunTime(chosenWorkerNodeKey) ?? 0)
  74     } else {
  75       this.currentWorkerNodeId =
  76         this.currentWorkerNodeId === this.pool.workerNodes.length - 1
  77           ? 0
  78           : this.currentWorkerNodeId + 1
  79       this.workerVirtualTaskRunTime = 0
  80     }
  81     return chosenWorkerNodeKey
  82   }
  83
  84   /** @inheritDoc */
  85   public remove (workerNodeKey: number): boolean {
  86     if (this.currentWorkerNodeId === workerNodeKey) {
  87       if (this.pool.workerNodes.length === 0) {
  88         this.currentWorkerNodeId = 0
  89       } else {
  90         this.currentWorkerNodeId =
  91           this.currentWorkerNodeId > this.pool.workerNodes.length - 1
  92             ? this.pool.workerNodes.length - 1
  93             : this.currentWorkerNodeId
  94       }
  95       this.workerVirtualTaskRunTime = 0
  96     }
  97     return true
  98   }
  99
 100   private getWorkerVirtualTaskRunTime (workerNodeKey: number): number {
 101     return this.requiredStatistics.medRunTime
 102       ? this.pool.workerNodes[workerNodeKey].tasksUsage.medRunTime
 103       : this.pool.workerNodes[workerNodeKey].tasksUsage.avgRunTime
 104   }
 105
 106   private computeDefaultWorkerWeight (): number {
 107     let cpusCycleTimeWeight = 0
 108     for (const cpu of cpus()) {
 109       // CPU estimated cycle time
 110       const numberOfDigits = cpu.speed.toString().length - 1
 111       const cpuCycleTime = 1 / (cpu.speed / Math.pow(10, numberOfDigits))
 112       cpusCycleTimeWeight += cpuCycleTime * Math.pow(10, numberOfDigits)
 113     }
 114     return Math.round(cpusCycleTimeWeight / cpus().length)
 115   }
 116 }