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 response of execution. 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   /** @inheritDoc */
 231   public abstract get full (): boolean
 232
 233   /** @inheritDoc */
 234   public abstract get busy (): boolean
 235
 236   protected internalBusy (): boolean {
 237     return this.findFreeWorkerNodeKey() === -1
 238   }
 239
 240   /** @inheritDoc */
 241   public findFreeWorkerNodeKey (): number {
 242     return this.workerNodes.findIndex(workerNode => {
 243       return workerNode.tasksUsage?.running === 0
 244     })
 245   }
 246
 247   /** @inheritDoc */
 248   public async execute (data: Data): Promise<Response> {
 249     const [workerNodeKey, workerNode] = this.chooseWorkerNode()
 250     const submittedTask: Task<Data> = {
 251       // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
 252       data: data ?? ({} as Data),
 253       id: crypto.randomUUID()
 254     }
 255     const res = new Promise<Response>((resolve, reject) => {
 256       this.promiseResponseMap.set(submittedTask.id, {
 257         resolve,
 258         reject,
 259         worker: workerNode.worker
 260       })
 261     })
 262     if (
 263       this.opts.enableTasksQueue === true &&
 264       (this.busy ||
 265         this.workerNodes[workerNodeKey].tasksUsage.running >=
 266           ((this.opts.tasksQueueOptions as TasksQueueOptions)
 267             .concurrency as number))
 268     ) {
 269       this.enqueueTask(workerNodeKey, submittedTask)
 270     } else {
 271       this.executeTask(workerNodeKey, submittedTask)
 272     }
 273     this.checkAndEmitEvents()
 274     // eslint-disable-next-line @typescript-eslint/return-await
 275     return res
 276   }
 277
 278   /** @inheritDoc */
 279   public async destroy (): Promise<void> {
 280     await Promise.all(
 281       this.workerNodes.map(async (workerNode, workerNodeKey) => {
 282         this.flushTasksQueue(workerNodeKey)
 283         await this.destroyWorker(workerNode.worker)
 284       })
 285     )
 286   }
 287
 288   /**
 289    * Shutdowns the given worker.
 290    *
 291    * @param worker - A worker within `workerNodes`.
 292    */
 293   protected abstract destroyWorker (worker: Worker): void | Promise<void>
 294
 295   /**
 296    * Setup hook to execute code before worker node are created in the abstract constructor.
 297    * Can be overridden
 298    *
 299    * @virtual
 300    */
 301   protected setupHook (): void {
 302     // Intentionally empty
 303   }
 304
 305   /**
 306    * Should return whether the worker is the main worker or not.
 307    */
 308   protected abstract isMain (): boolean
 309
 310   /**
 311    * Hook executed before the worker task execution.
 312    * Can be overridden.
 313    *
 314    * @param workerNodeKey - The worker node key.
 315    */
 316   protected beforeTaskExecutionHook (workerNodeKey: number): void {
 317     ++this.workerNodes[workerNodeKey].tasksUsage.running
 318   }
 319
 320   /**
 321    * Hook executed after the worker task execution.
 322    * Can be overridden.
 323    *
 324    * @param worker - The worker.
 325    * @param message - The received message.
 326    */
 327   protected afterTaskExecutionHook (
 328     worker: Worker,
 329     message: MessageValue<Response>
 330   ): void {
 331     const workerTasksUsage = this.getWorkerTasksUsage(worker) as TasksUsage
 332     --workerTasksUsage.running
 333     ++workerTasksUsage.run
 334     if (message.error != null) {
 335       ++workerTasksUsage.error
 336     }
 337     if (this.workerChoiceStrategyContext.getRequiredStatistics().runTime) {
 338       workerTasksUsage.runTime += message.runTime ?? 0
 339       if (
 340         this.workerChoiceStrategyContext.getRequiredStatistics().avgRunTime &&
 341         workerTasksUsage.run !== 0
 342       ) {
 343         workerTasksUsage.avgRunTime =
 344           workerTasksUsage.runTime / workerTasksUsage.run
 345       }
 346       if (this.workerChoiceStrategyContext.getRequiredStatistics().medRunTime) {
 347         workerTasksUsage.runTimeHistory.push(message.runTime ?? 0)
 348         workerTasksUsage.medRunTime = median(workerTasksUsage.runTimeHistory)
 349       }
 350     }
 351   }
 352
 353   /**
 354    * Chooses a worker node for the next task.
 355    *
 356    * The default uses a round robin algorithm to distribute the load.
 357    *
 358    * @returns [worker node key, worker node].
 359    */
 360   protected chooseWorkerNode (): [number, WorkerNode<Worker, Data>] {
 361     let workerNodeKey: number
 362     if (this.type === PoolType.DYNAMIC && !this.full && this.internalBusy()) {
 363       const workerCreated = this.createAndSetupWorker()
 364       this.registerWorkerMessageListener(workerCreated, message => {
 365         if (
 366           isKillBehavior(KillBehaviors.HARD, message.kill) ||
 367           (message.kill != null &&
 368             this.getWorkerTasksUsage(workerCreated)?.running === 0)
 369         ) {
 370           // Kill message received from the worker: no new tasks are submitted to that worker for a while ( > maxInactiveTime)
 371           this.flushTasksQueueByWorker(workerCreated)
 372           void this.destroyWorker(workerCreated)
 373         }
 374       })
 375       workerNodeKey = this.getWorkerNodeKey(workerCreated)
 376     } else {
 377       workerNodeKey = this.workerChoiceStrategyContext.execute()
 378     }
 379     return [workerNodeKey, this.workerNodes[workerNodeKey]]
 380   }
 381
 382   /**
 383    * Sends a message to the given worker.
 384    *
 385    * @param worker - The worker which should receive the message.
 386    * @param message - The message.
 387    */
 388   protected abstract sendToWorker (
 389     worker: Worker,
 390     message: MessageValue<Data>
 391   ): void
 392
 393   /**
 394    * Registers a listener callback on the given worker.
 395    *
 396    * @param worker - The worker which should register a listener.
 397    * @param listener - The message listener callback.
 398    */
 399   protected abstract registerWorkerMessageListener<
 400     Message extends Data | Response
 401   >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
 402
 403   /**
 404    * Returns a newly created worker.
 405    */
 406   protected abstract createWorker (): Worker
 407
 408   /**
 409    * Function that can be hooked up when a worker has been newly created and moved to the pool worker nodes.
 410    *
 411    * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
 412    *
 413    * @param worker - The newly created worker.
 414    */
 415   protected abstract afterWorkerSetup (worker: Worker): void
 416
 417   /**
 418    * Creates a new worker and sets it up completely in the pool worker nodes.
 419    *
 420    * @returns New, completely set up worker.
 421    */
 422   protected createAndSetupWorker (): Worker {
 423     const worker = this.createWorker()
 424
 425     worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
 426     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 427     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 428     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 429     worker.once('exit', () => {
 430       this.removeWorkerNode(worker)
 431     })
 432
 433     this.pushWorkerNode(worker)
 434
 435     this.afterWorkerSetup(worker)
 436
 437     return worker
 438   }
 439
 440   /**
 441    * This function is the listener registered for each worker message.
 442    *
 443    * @returns The listener function to execute when a message is received from a worker.
 444    */
 445   protected workerListener (): (message: MessageValue<Response>) => void {
 446     return message => {
 447       if (message.id != null) {
 448         // Task execution response received
 449         const promiseResponse = this.promiseResponseMap.get(message.id)
 450         if (promiseResponse != null) {
 451           if (message.error != null) {
 452             promiseResponse.reject(message.error)
 453           } else {
 454             promiseResponse.resolve(message.data as Response)
 455           }
 456           this.afterTaskExecutionHook(promiseResponse.worker, message)
 457           this.promiseResponseMap.delete(message.id)
 458           const workerNodeKey = this.getWorkerNodeKey(promiseResponse.worker)
 459           if (
 460             this.opts.enableTasksQueue === true &&
 461             this.tasksQueueSize(workerNodeKey) > 0
 462           ) {
 463             this.executeTask(
 464               workerNodeKey,
 465               this.dequeueTask(workerNodeKey) as Task<Data>
 466             )
 467           }
 468         }
 469       }
 470     }
 471   }
 472
 473   private checkAndEmitEvents (): void {
 474     if (this.opts.enableEvents === true) {
 475       if (this.busy) {
 476         this.emitter?.emit(PoolEvents.busy)
 477       }
 478       if (this.type === PoolType.DYNAMIC && this.full) {
 479         this.emitter?.emit(PoolEvents.full)
 480       }
 481     }
 482   }
 483
 484   /**
 485    * Sets the given worker node its tasks usage in the pool.
 486    *
 487    * @param workerNode - The worker node.
 488    * @param tasksUsage - The worker node tasks usage.
 489    */
 490   private setWorkerNodeTasksUsage (
 491     workerNode: WorkerNode<Worker, Data>,
 492     tasksUsage: TasksUsage
 493   ): void {
 494     workerNode.tasksUsage = tasksUsage
 495   }
 496
 497   /**
 498    * Gets the given worker its tasks usage in the pool.
 499    *
 500    * @param worker - The worker.
 501    * @returns The worker tasks usage.
 502    */
 503   private getWorkerTasksUsage (worker: Worker): TasksUsage | undefined {
 504     const workerNodeKey = this.getWorkerNodeKey(worker)
 505     if (workerNodeKey !== -1) {
 506       return this.workerNodes[workerNodeKey].tasksUsage
 507     }
 508     throw new Error('Worker could not be found in the pool worker nodes')
 509   }
 510
 511   /**
 512    * Pushes the given worker in the pool worker nodes.
 513    *
 514    * @param worker - The worker.
 515    * @returns The worker nodes length.
 516    */
 517   private pushWorkerNode (worker: Worker): number {
 518     return this.workerNodes.push({
 519       worker,
 520       tasksUsage: {
 521         run: 0,
 522         running: 0,
 523         runTime: 0,
 524         runTimeHistory: new CircularArray(),
 525         avgRunTime: 0,
 526         medRunTime: 0,
 527         error: 0
 528       },
 529       tasksQueue: []
 530     })
 531   }
 532
 533   /**
 534    * Sets the given worker in the pool worker nodes.
 535    *
 536    * @param workerNodeKey - The worker node key.
 537    * @param worker - The worker.
 538    * @param tasksUsage - The worker tasks usage.
 539    * @param tasksQueue - The worker task queue.
 540    */
 541   private setWorkerNode (
 542     workerNodeKey: number,
 543     worker: Worker,
 544     tasksUsage: TasksUsage,
 545     tasksQueue: Array<Task<Data>>
 546   ): void {
 547     this.workerNodes[workerNodeKey] = {
 548       worker,
 549       tasksUsage,
 550       tasksQueue
 551     }
 552   }
 553
 554   /**
 555    * Removes the given worker from the pool worker nodes.
 556    *
 557    * @param worker - The worker.
 558    */
 559   private removeWorkerNode (worker: Worker): void {
 560     const workerNodeKey = this.getWorkerNodeKey(worker)
 561     this.workerNodes.splice(workerNodeKey, 1)
 562     this.workerChoiceStrategyContext.remove(workerNodeKey)
 563   }
 564
 565   private executeTask (workerNodeKey: number, task: Task<Data>): void {
 566     this.beforeTaskExecutionHook(workerNodeKey)
 567     this.sendToWorker(this.workerNodes[workerNodeKey].worker, task)
 568   }
 569
 570   private enqueueTask (workerNodeKey: number, task: Task<Data>): number {
 571     return this.workerNodes[workerNodeKey].tasksQueue.push(task)
 572   }
 573
 574   private dequeueTask (workerNodeKey: number): Task<Data> | undefined {
 575     return this.workerNodes[workerNodeKey].tasksQueue.shift()
 576   }
 577
 578   private tasksQueueSize (workerNodeKey: number): number {
 579     return this.workerNodes[workerNodeKey].tasksQueue.length
 580   }
 581
 582   private flushTasksQueue (workerNodeKey: number): void {
 583     if (this.tasksQueueSize(workerNodeKey) > 0) {
 584       for (const task of this.workerNodes[workerNodeKey].tasksQueue) {
 585         this.executeTask(workerNodeKey, task)
 586       }
 587     }
 588   }
 589
 590   private flushTasksQueueByWorker (worker: Worker): void {
 591     const workerNodeKey = this.getWorkerNodeKey(worker)
 592     this.flushTasksQueue(workerNodeKey)
 593   }
 594 }