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