src/pools/abstract-pool.ts

   1 import crypto from 'node:crypto'
   2 import type { MessageValue, PromiseResponseWrapper } from '../utility-types'
   3 import {
   4   DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS,
   5   EMPTY_FUNCTION,
   6   median
   7 } from '../utils'
   8 import { KillBehaviors, isKillBehavior } from '../worker/worker-options'
   9 import { PoolEvents, type PoolOptions } from './pool'
  10 import { PoolEmitter } from './pool'
  11 import type { IPoolInternal } from './pool-internal'
  12 import { PoolType } from './pool-internal'
  13 import type { IWorker, Task, TasksUsage, WorkerNode } from './worker'
  14 import {
  15   WorkerChoiceStrategies,
  16   type WorkerChoiceStrategy
  17 } from './selection-strategies/selection-strategies-types'
  18 import { WorkerChoiceStrategyContext } from './selection-strategies/worker-choice-strategy-context'
  19 import { CircularArray } from '../circular-array'
  20
  21 /**
  22  * Base class that implements some shared logic for all poolifier pools.
  23  *
  24  * @typeParam Worker - Type of worker which manages this pool.
  25  * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
  26  * @typeParam Response - Type of response of execution. This can only be serializable data.
  27  */
  28 export abstract class AbstractPool<
  29   Worker extends IWorker,
  30   Data = unknown,
  31   Response = unknown
  32 > implements IPoolInternal<Worker, Data, Response> {
  33   /** @inheritDoc */
  34   public readonly workerNodes: Array<WorkerNode<Worker, Data>> = []
  35
  36   /** @inheritDoc */
  37   public readonly emitter?: PoolEmitter
  38
  39   /**
  40    * The execution response promise map.
  41    *
  42    * - `key`: The message id of each submitted task.
  43    * - `value`: An object that contains the worker, the execution response promise resolve and reject callbacks.
  44    *
  45    * When we receive a message from the worker, we get a map entry with the promise resolve/reject bound to the message id.
  46    */
  47   protected promiseResponseMap: Map<
  48   string,
  49   PromiseResponseWrapper<Worker, Response>
  50   > = new Map<string, PromiseResponseWrapper<Worker, Response>>()
  51
  52   /**
  53    * Worker choice strategy context referencing a worker choice algorithm implementation.
  54    *
  55    * Default to a round robin algorithm.
  56    */
  57   protected workerChoiceStrategyContext: WorkerChoiceStrategyContext<
  58   Worker,
  59   Data,
  60   Response
  61   >
  62
  63   /**
  64    * Constructs a new poolifier pool.
  65    *
  66    * @param numberOfWorkers - Number of workers that this pool should manage.
  67    * @param filePath - Path to the worker-file.
  68    * @param opts - Options for the pool.
  69    */
  70   public constructor (
  71     public readonly numberOfWorkers: number,
  72     public readonly filePath: string,
  73     public readonly opts: PoolOptions<Worker>
  74   ) {
  75     if (!this.isMain()) {
  76       throw new Error('Cannot start a pool from a worker!')
  77     }
  78     this.checkNumberOfWorkers(this.numberOfWorkers)
  79     this.checkFilePath(this.filePath)
  80     this.checkPoolOptions(this.opts)
  81
  82     this.chooseWorkerNode.bind(this)
  83     this.executeTask.bind(this)
  84     this.enqueueTask.bind(this)
  85     this.checkAndEmitEvents.bind(this)
  86
  87     this.setupHook()
  88
  89     for (let i = 1; i <= this.numberOfWorkers; i++) {
  90       this.createAndSetupWorker()
  91     }
  92
  93     if (this.opts.enableEvents === true) {
  94       this.emitter = new PoolEmitter()
  95     }
  96     this.workerChoiceStrategyContext = new WorkerChoiceStrategyContext<
  97     Worker,
  98     Data,
  99     Response
 100     >(
 101       this,
 102       this.opts.workerChoiceStrategy,
 103       this.opts.workerChoiceStrategyOptions
 104     )
 105   }
 106
 107   private checkFilePath (filePath: string): void {
 108     if (
 109       filePath == null ||
 110       (typeof filePath === 'string' && filePath.trim().length === 0)
 111     ) {
 112       throw new Error('Please specify a file with a worker implementation')
 113     }
 114   }
 115
 116   private checkNumberOfWorkers (numberOfWorkers: number): void {
 117     if (numberOfWorkers == null) {
 118       throw new Error(
 119         'Cannot instantiate a pool without specifying the number of workers'
 120       )
 121     } else if (!Number.isSafeInteger(numberOfWorkers)) {
 122       throw new TypeError(
 123         'Cannot instantiate a pool with a non integer number of workers'
 124       )
 125     } else if (numberOfWorkers < 0) {
 126       throw new RangeError(
 127         'Cannot instantiate a pool with a negative number of workers'
 128       )
 129     } else if (this.type === PoolType.FIXED && numberOfWorkers === 0) {
 130       throw new Error('Cannot instantiate a fixed pool with no worker')
 131     }
 132   }
 133
 134   private checkPoolOptions (opts: PoolOptions<Worker>): void {
 135     this.opts.workerChoiceStrategy =
 136       opts.workerChoiceStrategy ?? WorkerChoiceStrategies.ROUND_ROBIN
 137     this.checkValidWorkerChoiceStrategy(this.opts.workerChoiceStrategy)
 138     this.opts.workerChoiceStrategyOptions =
 139       opts.workerChoiceStrategyOptions ?? DEFAULT_WORKER_CHOICE_STRATEGY_OPTIONS
 140     this.opts.enableEvents = opts.enableEvents ?? true
 141     this.opts.enableTasksQueue = opts.enableTasksQueue ?? false
 142   }
 143
 144   private checkValidWorkerChoiceStrategy (
 145     workerChoiceStrategy: WorkerChoiceStrategy
 146   ): void {
 147     if (!Object.values(WorkerChoiceStrategies).includes(workerChoiceStrategy)) {
 148       throw new Error(
 149         `Invalid worker choice strategy '${workerChoiceStrategy}'`
 150       )
 151     }
 152   }
 153
 154   /** @inheritDoc */
 155   public abstract get type (): PoolType
 156
 157   /**
 158    * Number of tasks running in the pool.
 159    */
 160   private get numberOfRunningTasks (): number {
 161     return this.workerNodes.reduce(
 162       (accumulator, workerNode) => accumulator + workerNode.tasksUsage.running,
 163       0
 164     )
 165   }
 166
 167   /**
 168    * Number of tasks queued in the pool.
 169    */
 170   private get numberOfQueuedTasks (): number {
 171     if (this.opts.enableTasksQueue === false) {
 172       return 0
 173     }
 174     return this.workerNodes.reduce(
 175       (accumulator, workerNode) => accumulator + workerNode.tasksQueue.length,
 176       0
 177     )
 178   }
 179
 180   /**
 181    * Gets the given worker its worker node key.
 182    *
 183    * @param worker - The worker.
 184    * @returns The worker node key if the worker is found in the pool worker nodes, `-1` otherwise.
 185    */
 186   private getWorkerNodeKey (worker: Worker): number {
 187     return this.workerNodes.findIndex(
 188       workerNode => workerNode.worker === worker
 189     )
 190   }
 191
 192   /** @inheritDoc */
 193   public setWorkerChoiceStrategy (
 194     workerChoiceStrategy: WorkerChoiceStrategy
 195   ): void {
 196     this.checkValidWorkerChoiceStrategy(workerChoiceStrategy)
 197     this.opts.workerChoiceStrategy = workerChoiceStrategy
 198     for (const workerNode of this.workerNodes) {
 199       this.setWorkerNodeTasksUsage(workerNode, {
 200         run: 0,
 201         running: 0,
 202         runTime: 0,
 203         runTimeHistory: new CircularArray(),
 204         avgRunTime: 0,
 205         medRunTime: 0,
 206         error: 0
 207       })
 208     }
 209     this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
 210       workerChoiceStrategy
 211     )
 212   }
 213
 214   /** @inheritDoc */
 215   public abstract get full (): boolean
 216
 217   /** @inheritDoc */
 218   public abstract get busy (): boolean
 219
 220   protected internalBusy (): boolean {
 221     return this.findFreeWorkerNodeKey() === -1
 222   }
 223
 224   /** @inheritDoc */
 225   public findFreeWorkerNodeKey (): number {
 226     return this.workerNodes.findIndex(workerNode => {
 227       return workerNode.tasksUsage?.running === 0
 228     })
 229   }
 230
 231   /** @inheritDoc */
 232   public async execute (data: Data): Promise<Response> {
 233     const [workerNodeKey, workerNode] = this.chooseWorkerNode()
 234     const submittedTask: Task<Data> = {
 235       // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
 236       data: data ?? ({} as Data),
 237       id: crypto.randomUUID()
 238     }
 239     const res = new Promise<Response>((resolve, reject) => {
 240       this.promiseResponseMap.set(submittedTask.id, {
 241         resolve,
 242         reject,
 243         worker: workerNode.worker
 244       })
 245     })
 246     if (
 247       this.opts.enableTasksQueue === true &&
 248       (this.busy || this.workerNodes[workerNodeKey].tasksUsage.running > 0)
 249     ) {
 250       this.enqueueTask(workerNodeKey, submittedTask)
 251     } else {
 252       this.executeTask(workerNodeKey, submittedTask)
 253     }
 254     this.checkAndEmitEvents()
 255     // eslint-disable-next-line @typescript-eslint/return-await
 256     return res
 257   }
 258
 259   /** @inheritDoc */
 260   public async destroy (): Promise<void> {
 261     await Promise.all(
 262       this.workerNodes.map(async (workerNode, workerNodeKey) => {
 263         this.flushTasksQueue(workerNodeKey)
 264         await this.destroyWorker(workerNode.worker)
 265       })
 266     )
 267   }
 268
 269   /**
 270    * Shutdowns the given worker.
 271    *
 272    * @param worker - A worker within `workerNodes`.
 273    */
 274   protected abstract destroyWorker (worker: Worker): void | Promise<void>
 275
 276   /**
 277    * Setup hook to execute code before worker node are created in the abstract constructor.
 278    * Can be overridden
 279    *
 280    * @virtual
 281    */
 282   protected setupHook (): void {
 283     // Intentionally empty
 284   }
 285
 286   /**
 287    * Should return whether the worker is the main worker or not.
 288    */
 289   protected abstract isMain (): boolean
 290
 291   /**
 292    * Hook executed before the worker task execution.
 293    * Can be overridden.
 294    *
 295    * @param workerNodeKey - The worker node key.
 296    */
 297   protected beforeTaskExecutionHook (workerNodeKey: number): void {
 298     ++this.workerNodes[workerNodeKey].tasksUsage.running
 299   }
 300
 301   /**
 302    * Hook executed after the worker task execution.
 303    * Can be overridden.
 304    *
 305    * @param worker - The worker.
 306    * @param message - The received message.
 307    */
 308   protected afterTaskExecutionHook (
 309     worker: Worker,
 310     message: MessageValue<Response>
 311   ): void {
 312     const workerTasksUsage = this.getWorkerTasksUsage(worker) as TasksUsage
 313     --workerTasksUsage.running
 314     ++workerTasksUsage.run
 315     if (message.error != null) {
 316       ++workerTasksUsage.error
 317     }
 318     if (this.workerChoiceStrategyContext.getRequiredStatistics().runTime) {
 319       workerTasksUsage.runTime += message.runTime ?? 0
 320       if (
 321         this.workerChoiceStrategyContext.getRequiredStatistics().avgRunTime &&
 322         workerTasksUsage.run !== 0
 323       ) {
 324         workerTasksUsage.avgRunTime =
 325           workerTasksUsage.runTime / workerTasksUsage.run
 326       }
 327       if (this.workerChoiceStrategyContext.getRequiredStatistics().medRunTime) {
 328         workerTasksUsage.runTimeHistory.push(message.runTime ?? 0)
 329         workerTasksUsage.medRunTime = median(workerTasksUsage.runTimeHistory)
 330       }
 331     }
 332   }
 333
 334   /**
 335    * Chooses a worker node for the next task.
 336    *
 337    * The default uses a round robin algorithm to distribute the load.
 338    *
 339    * @returns [worker node key, worker node].
 340    */
 341   protected chooseWorkerNode (): [number, WorkerNode<Worker, Data>] {
 342     let workerNodeKey: number
 343     if (this.type === PoolType.DYNAMIC && !this.full && this.internalBusy()) {
 344       const workerCreated = this.createAndSetupWorker()
 345       this.registerWorkerMessageListener(workerCreated, message => {
 346         if (
 347           isKillBehavior(KillBehaviors.HARD, message.kill) ||
 348           (message.kill != null &&
 349             this.getWorkerTasksUsage(workerCreated)?.running === 0)
 350         ) {
 351           // Kill message received from the worker: no new tasks are submitted to that worker for a while ( > maxInactiveTime)
 352           this.flushTasksQueueByWorker(workerCreated)
 353           void this.destroyWorker(workerCreated)
 354         }
 355       })
 356       workerNodeKey = this.getWorkerNodeKey(workerCreated)
 357     } else {
 358       workerNodeKey = this.workerChoiceStrategyContext.execute()
 359     }
 360     return [workerNodeKey, this.workerNodes[workerNodeKey]]
 361   }
 362
 363   /**
 364    * Sends a message to the given worker.
 365    *
 366    * @param worker - The worker which should receive the message.
 367    * @param message - The message.
 368    */
 369   protected abstract sendToWorker (
 370     worker: Worker,
 371     message: MessageValue<Data>
 372   ): void
 373
 374   /**
 375    * Registers a listener callback on the given worker.
 376    *
 377    * @param worker - The worker which should register a listener.
 378    * @param listener - The message listener callback.
 379    */
 380   protected abstract registerWorkerMessageListener<
 381     Message extends Data | Response
 382   >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
 383
 384   /**
 385    * Returns a newly created worker.
 386    */
 387   protected abstract createWorker (): Worker
 388
 389   /**
 390    * Function that can be hooked up when a worker has been newly created and moved to the pool worker nodes.
 391    *
 392    * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
 393    *
 394    * @param worker - The newly created worker.
 395    */
 396   protected abstract afterWorkerSetup (worker: Worker): void
 397
 398   /**
 399    * Creates a new worker and sets it up completely in the pool worker nodes.
 400    *
 401    * @returns New, completely set up worker.
 402    */
 403   protected createAndSetupWorker (): Worker {
 404     const worker = this.createWorker()
 405
 406     worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
 407     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 408     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 409     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 410     worker.once('exit', () => {
 411       this.removeWorkerNode(worker)
 412     })
 413
 414     this.pushWorkerNode(worker)
 415
 416     this.afterWorkerSetup(worker)
 417
 418     return worker
 419   }
 420
 421   /**
 422    * This function is the listener registered for each worker message.
 423    *
 424    * @returns The listener function to execute when a message is received from a worker.
 425    */
 426   protected workerListener (): (message: MessageValue<Response>) => void {
 427     return message => {
 428       if (message.id != null) {
 429         // Task execution response received
 430         const promiseResponse = this.promiseResponseMap.get(message.id)
 431         if (promiseResponse != null) {
 432           if (message.error != null) {
 433             promiseResponse.reject(message.error)
 434           } else {
 435             promiseResponse.resolve(message.data as Response)
 436           }
 437           this.afterTaskExecutionHook(promiseResponse.worker, message)
 438           this.promiseResponseMap.delete(message.id)
 439           const workerNodeKey = this.getWorkerNodeKey(promiseResponse.worker)
 440           if (
 441             this.opts.enableTasksQueue === true &&
 442             this.tasksQueueSize(workerNodeKey) > 0
 443           ) {
 444             this.executeTask(
 445               workerNodeKey,
 446               this.dequeueTask(workerNodeKey) as Task<Data>
 447             )
 448           }
 449         }
 450       }
 451     }
 452   }
 453
 454   private checkAndEmitEvents (): void {
 455     if (this.opts.enableEvents === true) {
 456       if (this.busy) {
 457         this.emitter?.emit(PoolEvents.busy)
 458       }
 459       if (this.type === PoolType.DYNAMIC && this.full) {
 460         this.emitter?.emit(PoolEvents.full)
 461       }
 462     }
 463   }
 464
 465   /**
 466    * Sets the given worker node its tasks usage in the pool.
 467    *
 468    * @param workerNode - The worker node.
 469    * @param tasksUsage - The worker node tasks usage.
 470    */
 471   private setWorkerNodeTasksUsage (
 472     workerNode: WorkerNode<Worker, Data>,
 473     tasksUsage: TasksUsage
 474   ): void {
 475     workerNode.tasksUsage = tasksUsage
 476   }
 477
 478   /**
 479    * Gets the given worker its tasks usage in the pool.
 480    *
 481    * @param worker - The worker.
 482    * @returns The worker tasks usage.
 483    */
 484   private getWorkerTasksUsage (worker: Worker): TasksUsage | undefined {
 485     const workerNodeKey = this.getWorkerNodeKey(worker)
 486     if (workerNodeKey !== -1) {
 487       return this.workerNodes[workerNodeKey].tasksUsage
 488     }
 489     throw new Error('Worker could not be found in the pool worker nodes')
 490   }
 491
 492   /**
 493    * Pushes the given worker in the pool worker nodes.
 494    *
 495    * @param worker - The worker.
 496    * @returns The worker nodes length.
 497    */
 498   private pushWorkerNode (worker: Worker): number {
 499     return this.workerNodes.push({
 500       worker,
 501       tasksUsage: {
 502         run: 0,
 503         running: 0,
 504         runTime: 0,
 505         runTimeHistory: new CircularArray(),
 506         avgRunTime: 0,
 507         medRunTime: 0,
 508         error: 0
 509       },
 510       tasksQueue: []
 511     })
 512   }
 513
 514   /**
 515    * Sets the given worker in the pool worker nodes.
 516    *
 517    * @param workerNodeKey - The worker node key.
 518    * @param worker - The worker.
 519    * @param tasksUsage - The worker tasks usage.
 520    * @param tasksQueue - The worker task queue.
 521    */
 522   private setWorkerNode (
 523     workerNodeKey: number,
 524     worker: Worker,
 525     tasksUsage: TasksUsage,
 526     tasksQueue: Array<Task<Data>>
 527   ): void {
 528     this.workerNodes[workerNodeKey] = {
 529       worker,
 530       tasksUsage,
 531       tasksQueue
 532     }
 533   }
 534
 535   /**
 536    * Removes the given worker from the pool worker nodes.
 537    *
 538    * @param worker - The worker.
 539    */
 540   private removeWorkerNode (worker: Worker): void {
 541     const workerNodeKey = this.getWorkerNodeKey(worker)
 542     this.workerNodes.splice(workerNodeKey, 1)
 543     this.workerChoiceStrategyContext.remove(workerNodeKey)
 544   }
 545
 546   private executeTask (workerNodeKey: number, task: Task<Data>): void {
 547     this.beforeTaskExecutionHook(workerNodeKey)
 548     this.sendToWorker(this.workerNodes[workerNodeKey].worker, task)
 549   }
 550
 551   private enqueueTask (workerNodeKey: number, task: Task<Data>): number {
 552     return this.workerNodes[workerNodeKey].tasksQueue.push(task)
 553   }
 554
 555   private dequeueTask (workerNodeKey: number): Task<Data> | undefined {
 556     return this.workerNodes[workerNodeKey].tasksQueue.shift()
 557   }
 558
 559   private tasksQueueSize (workerNodeKey: number): number {
 560     return this.workerNodes[workerNodeKey].tasksQueue.length
 561   }
 562
 563   private flushTasksQueue (workerNodeKey: number): void {
 564     if (this.tasksQueueSize(workerNodeKey) > 0) {
 565       for (const task of this.workerNodes[workerNodeKey].tasksQueue) {
 566         this.executeTask(workerNodeKey, task)
 567       }
 568     }
 569   }
 570
 571   private flushTasksQueueByWorker (worker: Worker): void {
 572     const workerNodeKey = this.getWorkerNodeKey(worker)
 573     this.flushTasksQueue(workerNodeKey)
 574   }
 575 }