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

   1 import { cpus } from 'node:os'
   2 import type { IPoolInternal } from '../pool-internal'
   3 import type { IPoolWorker } from '../pool-worker'
   4 import { AbstractWorkerChoiceStrategy } from './abstract-worker-choice-strategy'
   5 import type { RequiredStatistics } from './selection-strategies-types'
   6
   7 /**
   8  * Virtual task runtime.
   9  */
  10 interface TaskRunTime {
  11   weight: number
  12   runTime: number
  13 }
  14
  15 /**
  16  * Selects the next worker with a weighted round robin scheduling algorithm.
  17  * Loosely modeled after the weighted round robin queueing algorithm: https://en.wikipedia.org/wiki/Weighted_round_robin.
  18  *
  19  * @typeParam Worker - Type of worker which manages the strategy.
  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 WeightedRoundRobinWorkerChoiceStrategy<
  24   Worker extends IPoolWorker,
  25   Data,
  26   Response
  27 > extends AbstractWorkerChoiceStrategy<Worker, Data, Response> {
  28   /** {@inheritDoc} */
  29   public readonly requiredStatistics: RequiredStatistics = {
  30     runTime: true
  31   }
  32
  33   /**
  34    * Worker id where the current task will be submitted.
  35    */
  36   private currentWorkerId: number = 0
  37   /**
  38    * Default worker weight.
  39    */
  40   private readonly defaultWorkerWeight: number
  41   /**
  42    * Per worker virtual task runtime map.
  43    */
  44   private readonly workersTaskRunTime: Map<number, TaskRunTime> = new Map<
  45   number,
  46   TaskRunTime
  47   >()
  48
  49   /**
  50    * Constructs a worker choice strategy that selects with a weighted round robin scheduling algorithm.
  51    *
  52    * @param pool - The pool instance.
  53    */
  54   public constructor (pool: IPoolInternal<Worker, Data, Response>) {
  55     super(pool)
  56     this.defaultWorkerWeight = this.computeWorkerWeight()
  57     this.initWorkersTaskRunTime()
  58   }
  59
  60   /** {@inheritDoc} */
  61   public reset (): boolean {
  62     this.currentWorkerId = 0
  63     this.workersTaskRunTime.clear()
  64     this.initWorkersTaskRunTime()
  65     return true
  66   }
  67
  68   /** {@inheritDoc} */
  69   public choose (): number {
  70     const chosenWorkerKey = this.currentWorkerId
  71     if (this.isDynamicPool && !this.workersTaskRunTime.has(chosenWorkerKey)) {
  72       this.initWorkerTaskRunTime(chosenWorkerKey)
  73     }
  74     const workerTaskRunTime =
  75       this.workersTaskRunTime.get(chosenWorkerKey)?.runTime ?? 0
  76     const workerTaskWeight =
  77       this.workersTaskRunTime.get(chosenWorkerKey)?.weight ??
  78       this.defaultWorkerWeight
  79     if (workerTaskRunTime < workerTaskWeight) {
  80       this.setWorkerTaskRunTime(
  81         chosenWorkerKey,
  82         workerTaskWeight,
  83         workerTaskRunTime +
  84           (this.getWorkerVirtualTaskRunTime(chosenWorkerKey) ?? 0)
  85       )
  86     } else {
  87       this.currentWorkerId =
  88         this.currentWorkerId === this.pool.workers.length - 1
  89           ? 0
  90           : this.currentWorkerId + 1
  91       this.setWorkerTaskRunTime(this.currentWorkerId, workerTaskWeight, 0)
  92     }
  93     return chosenWorkerKey
  94   }
  95
  96   /** {@inheritDoc} */
  97   public remove (workerKey: number): boolean {
  98     if (this.currentWorkerId === workerKey) {
  99       this.currentWorkerId =
 100         this.currentWorkerId > this.pool.workers.length - 1
 101           ? this.pool.workers.length - 1
 102           : this.currentWorkerId
 103     }
 104     const workerDeleted = this.workersTaskRunTime.delete(workerKey)
 105     for (const [key, value] of this.workersTaskRunTime) {
 106       if (key > workerKey) {
 107         this.workersTaskRunTime.set(key - 1, value)
 108       }
 109     }
 110     return workerDeleted
 111   }
 112
 113   private initWorkersTaskRunTime (): void {
 114     for (const [index] of this.pool.workers.entries()) {
 115       this.initWorkerTaskRunTime(index)
 116     }
 117   }
 118
 119   private initWorkerTaskRunTime (workerKey: number): void {
 120     this.setWorkerTaskRunTime(workerKey, this.defaultWorkerWeight, 0)
 121   }
 122
 123   private setWorkerTaskRunTime (
 124     workerKey: number,
 125     weight: number,
 126     runTime: number
 127   ): void {
 128     this.workersTaskRunTime.set(workerKey, {
 129       weight,
 130       runTime
 131     })
 132   }
 133
 134   private getWorkerVirtualTaskRunTime (workerKey: number): number {
 135     return this.pool.workers[workerKey].tasksUsage.avgRunTime
 136   }
 137
 138   private computeWorkerWeight (): number {
 139     let cpusCycleTimeWeight = 0
 140     for (const cpu of cpus()) {
 141       // CPU estimated cycle time
 142       const numberOfDigits = cpu.speed.toString().length - 1
 143       const cpuCycleTime = 1 / (cpu.speed / Math.pow(10, numberOfDigits))
 144       cpusCycleTimeWeight += cpuCycleTime * Math.pow(10, numberOfDigits)
 145     }
 146     return Math.round(cpusCycleTimeWeight / cpus().length)
 147   }
 148 }