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