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