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     waitTime: false,
  33     avgWaitTime: false,
  34     medWaitTime: false
  35   }
  36
  37   /**
  38    * Worker node id where the current task will be submitted.
  39    */
  40   private currentWorkerNodeId: number = 0
  41   /**
  42    * Default worker weight.
  43    */
  44   private readonly defaultWorkerWeight: number
  45   /**
  46    * Worker virtual task runtime.
  47    */
  48   private workerVirtualTaskRunTime: number = 0
  49
  50   /** @inheritDoc */
  51   public constructor (
  52     pool: IPool<Worker, Data, Response>,
  53     opts: WorkerChoiceStrategyOptions = DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
  54   ) {
  55     super(pool, opts)
  56     this.setRequiredStatistics(this.opts)
  57     this.defaultWorkerWeight = this.computeDefaultWorkerWeight()
  58   }
  59
  60   /** @inheritDoc */
  61   public reset (): boolean {
  62     this.currentWorkerNodeId = 0
  63     this.workerVirtualTaskRunTime = 0
  64     return true
  65   }
  66
  67   /** @inheritDoc */
  68   public update (): boolean {
  69     return true
  70   }
  71
  72   /** @inheritDoc */
  73   public choose (): number {
  74     const chosenWorkerNodeKey = this.currentWorkerNodeId
  75     const workerVirtualTaskRunTime = this.workerVirtualTaskRunTime
  76     const workerWeight =
  77       this.opts.weights?.[chosenWorkerNodeKey] ?? this.defaultWorkerWeight
  78     if (workerVirtualTaskRunTime < workerWeight) {
  79       this.workerVirtualTaskRunTime =
  80         workerVirtualTaskRunTime +
  81         this.getWorkerTaskRunTime(chosenWorkerNodeKey)
  82     } else {
  83       this.currentWorkerNodeId =
  84         this.currentWorkerNodeId === this.pool.workerNodes.length - 1
  85           ? 0
  86           : this.currentWorkerNodeId + 1
  87       this.workerVirtualTaskRunTime = 0
  88     }
  89     return chosenWorkerNodeKey
  90   }
  91
  92   /** @inheritDoc */
  93   public remove (workerNodeKey: number): boolean {
  94     if (this.currentWorkerNodeId === workerNodeKey) {
  95       if (this.pool.workerNodes.length === 0) {
  96         this.currentWorkerNodeId = 0
  97       } else if (this.currentWorkerNodeId > this.pool.workerNodes.length - 1) {
  98         this.currentWorkerNodeId = this.pool.workerNodes.length - 1
  99       }
 100       this.workerVirtualTaskRunTime = 0
 101     }
 102     return true
 103   }
 104
 105   private computeDefaultWorkerWeight (): number {
 106     let cpusCycleTimeWeight = 0
 107     for (const cpu of cpus()) {
 108       // CPU estimated cycle time
 109       const numberOfDigits = cpu.speed.toString().length - 1
 110       const cpuCycleTime = 1 / (cpu.speed / Math.pow(10, numberOfDigits))
 111       cpusCycleTimeWeight += cpuCycleTime * Math.pow(10, numberOfDigits)
 112     }
 113     return Math.round(cpusCycleTimeWeight / cpus().length)
 114   }
 115 }