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