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.checkValidWorkerChoiceStrategyOptions(
 150         this.opts.workerChoiceStrategyOptions
 151       )
 152       this.opts.enableEvents = opts.enableEvents ?? true
 153       this.opts.enableTasksQueue = opts.enableTasksQueue ?? false
 154       if (this.opts.enableTasksQueue) {
 155         this.checkValidTasksQueueOptions(
 156           opts.tasksQueueOptions as TasksQueueOptions
 157         )
 158         this.opts.tasksQueueOptions = this.buildTasksQueueOptions(
 159           opts.tasksQueueOptions as TasksQueueOptions
 160         )
 161       }
 162     } else {
 163       throw new TypeError('Invalid pool options: must be a plain object')
 164     }
 165   }
 166
 167   private checkValidWorkerChoiceStrategy (
 168     workerChoiceStrategy: WorkerChoiceStrategy
 169   ): void {
 170     if (!Object.values(WorkerChoiceStrategies).includes(workerChoiceStrategy)) {
 171       throw new Error(
 172         `Invalid worker choice strategy '${workerChoiceStrategy}'`
 173       )
 174     }
 175   }
 176
 177   private checkValidWorkerChoiceStrategyOptions (
 178     workerChoiceStrategyOptions: WorkerChoiceStrategyOptions
 179   ): void {
 180     if (!isPlainObject(workerChoiceStrategyOptions)) {
 181       throw new TypeError(
 182         'Invalid worker choice strategy options: must be a plain object'
 183       )
 184     }
 185     if (
 186       workerChoiceStrategyOptions.weights != null &&
 187       Object.keys(workerChoiceStrategyOptions.weights).length !== this.size
 188     ) {
 189       throw new Error(
 190         'Invalid worker choice strategy options: must have a weight for each worker node'
 191       )
 192     }
 193   }
 194
 195   private checkValidTasksQueueOptions (
 196     tasksQueueOptions: TasksQueueOptions
 197   ): void {
 198     if (tasksQueueOptions != null && !isPlainObject(tasksQueueOptions)) {
 199       throw new TypeError('Invalid tasks queue options: must be a plain object')
 200     }
 201     if ((tasksQueueOptions?.concurrency as number) <= 0) {
 202       throw new Error(
 203         `Invalid worker tasks concurrency '${
 204           tasksQueueOptions.concurrency as number
 205         }'`
 206       )
 207     }
 208   }
 209
 210   /** @inheritDoc */
 211   public abstract get type (): PoolType
 212
 213   /** @inheritDoc */
 214   public abstract get size (): number
 215
 216   /**
 217    * Number of tasks running in the pool.
 218    */
 219   private get numberOfRunningTasks (): number {
 220     return this.workerNodes.reduce(
 221       (accumulator, workerNode) => accumulator + workerNode.tasksUsage.running,
 222       0
 223     )
 224   }
 225
 226   /**
 227    * Number of tasks queued in the pool.
 228    */
 229   private get numberOfQueuedTasks (): number {
 230     if (this.opts.enableTasksQueue === false) {
 231       return 0
 232     }
 233     return this.workerNodes.reduce(
 234       (accumulator, workerNode) => accumulator + workerNode.tasksQueue.size,
 235       0
 236     )
 237   }
 238
 239   /**
 240    * Gets the given worker its worker node key.
 241    *
 242    * @param worker - The worker.
 243    * @returns The worker node key if the worker is found in the pool worker nodes, `-1` otherwise.
 244    */
 245   private getWorkerNodeKey (worker: Worker): number {
 246     return this.workerNodes.findIndex(
 247       workerNode => workerNode.worker === worker
 248     )
 249   }
 250
 251   /** @inheritDoc */
 252   public setWorkerChoiceStrategy (
 253     workerChoiceStrategy: WorkerChoiceStrategy,
 254     workerChoiceStrategyOptions?: WorkerChoiceStrategyOptions
 255   ): void {
 256     this.checkValidWorkerChoiceStrategy(workerChoiceStrategy)
 257     this.opts.workerChoiceStrategy = workerChoiceStrategy
 258     for (const workerNode of this.workerNodes) {
 259       this.setWorkerNodeTasksUsage(workerNode, {
 260         run: 0,
 261         running: 0,
 262         runTime: 0,
 263         runTimeHistory: new CircularArray(),
 264         avgRunTime: 0,
 265         medRunTime: 0,
 266         error: 0
 267       })
 268     }
 269     this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
 270       this.opts.workerChoiceStrategy
 271     )
 272     if (workerChoiceStrategyOptions != null) {
 273       this.setWorkerChoiceStrategyOptions(workerChoiceStrategyOptions)
 274     }
 275   }
 276
 277   /** @inheritDoc */
 278   public setWorkerChoiceStrategyOptions (
 279     workerChoiceStrategyOptions: WorkerChoiceStrategyOptions
 280   ): void {
 281     this.checkValidWorkerChoiceStrategyOptions(workerChoiceStrategyOptions)
 282     this.opts.workerChoiceStrategyOptions = workerChoiceStrategyOptions
 283     this.workerChoiceStrategyContext.setOptions(
 284       this.opts.workerChoiceStrategyOptions
 285     )
 286   }
 287
 288   /** @inheritDoc */
 289   public enableTasksQueue (
 290     enable: boolean,
 291     tasksQueueOptions?: TasksQueueOptions
 292   ): void {
 293     if (this.opts.enableTasksQueue === true && !enable) {
 294       this.flushTasksQueues()
 295     }
 296     this.opts.enableTasksQueue = enable
 297     this.setTasksQueueOptions(tasksQueueOptions as TasksQueueOptions)
 298   }
 299
 300   /** @inheritDoc */
 301   public setTasksQueueOptions (tasksQueueOptions: TasksQueueOptions): void {
 302     if (this.opts.enableTasksQueue === true) {
 303       this.checkValidTasksQueueOptions(tasksQueueOptions)
 304       this.opts.tasksQueueOptions =
 305         this.buildTasksQueueOptions(tasksQueueOptions)
 306     } else {
 307       delete this.opts.tasksQueueOptions
 308     }
 309   }
 310
 311   private buildTasksQueueOptions (
 312     tasksQueueOptions: TasksQueueOptions
 313   ): TasksQueueOptions {
 314     return {
 315       concurrency: tasksQueueOptions?.concurrency ?? 1
 316     }
 317   }
 318
 319   /**
 320    * Whether the pool is full or not.
 321    *
 322    * The pool filling boolean status.
 323    */
 324   protected abstract get full (): boolean
 325
 326   /**
 327    * Whether the pool is busy or not.
 328    *
 329    * The pool busyness boolean status.
 330    */
 331   protected abstract get busy (): boolean
 332
 333   protected internalBusy (): boolean {
 334     return (
 335       this.workerNodes.findIndex(workerNode => {
 336         return workerNode.tasksUsage.running === 0
 337       }) === -1
 338     )
 339   }
 340
 341   /** @inheritDoc */
 342   public async execute (data?: Data, name?: string): Promise<Response> {
 343     const [workerNodeKey, workerNode] = this.chooseWorkerNode()
 344     const submittedTask: Task<Data> = {
 345       name,
 346       // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
 347       data: data ?? ({} as Data),
 348       id: crypto.randomUUID()
 349     }
 350     const res = new Promise<Response>((resolve, reject) => {
 351       this.promiseResponseMap.set(submittedTask.id as string, {
 352         resolve,
 353         reject,
 354         worker: workerNode.worker
 355       })
 356     })
 357     if (
 358       this.opts.enableTasksQueue === true &&
 359       (this.busy ||
 360         this.workerNodes[workerNodeKey].tasksUsage.running >=
 361           ((this.opts.tasksQueueOptions as TasksQueueOptions)
 362             .concurrency as number))
 363     ) {
 364       this.enqueueTask(workerNodeKey, submittedTask)
 365     } else {
 366       this.executeTask(workerNodeKey, submittedTask)
 367     }
 368     this.workerChoiceStrategyContext.update(workerNodeKey)
 369     this.checkAndEmitEvents()
 370     // eslint-disable-next-line @typescript-eslint/return-await
 371     return res
 372   }
 373
 374   /** @inheritDoc */
 375   public async destroy (): Promise<void> {
 376     await Promise.all(
 377       this.workerNodes.map(async (workerNode, workerNodeKey) => {
 378         this.flushTasksQueue(workerNodeKey)
 379         await this.destroyWorker(workerNode.worker)
 380       })
 381     )
 382   }
 383
 384   /**
 385    * Shutdowns the given worker.
 386    *
 387    * @param worker - A worker within `workerNodes`.
 388    */
 389   protected abstract destroyWorker (worker: Worker): void | Promise<void>
 390
 391   /**
 392    * Setup hook to execute code before worker node are created in the abstract constructor.
 393    * Can be overridden
 394    *
 395    * @virtual
 396    */
 397   protected setupHook (): void {
 398     // Intentionally empty
 399   }
 400
 401   /**
 402    * Should return whether the worker is the main worker or not.
 403    */
 404   protected abstract isMain (): boolean
 405
 406   /**
 407    * Hook executed before the worker task execution.
 408    * Can be overridden.
 409    *
 410    * @param workerNodeKey - The worker node key.
 411    */
 412   protected beforeTaskExecutionHook (workerNodeKey: number): void {
 413     ++this.workerNodes[workerNodeKey].tasksUsage.running
 414   }
 415
 416   /**
 417    * Hook executed after the worker task execution.
 418    * Can be overridden.
 419    *
 420    * @param worker - The worker.
 421    * @param message - The received message.
 422    */
 423   protected afterTaskExecutionHook (
 424     worker: Worker,
 425     message: MessageValue<Response>
 426   ): void {
 427     const workerNodeKey = this.getWorkerNodeKey(worker)
 428     const workerTasksUsage = this.workerNodes[workerNodeKey].tasksUsage
 429     --workerTasksUsage.running
 430     ++workerTasksUsage.run
 431     if (message.error != null) {
 432       ++workerTasksUsage.error
 433     }
 434     if (this.workerChoiceStrategyContext.getRequiredStatistics().runTime) {
 435       workerTasksUsage.runTime += message.runTime ?? 0
 436       if (
 437         this.workerChoiceStrategyContext.getRequiredStatistics().avgRunTime &&
 438         workerTasksUsage.run !== 0
 439       ) {
 440         workerTasksUsage.avgRunTime =
 441           workerTasksUsage.runTime / workerTasksUsage.run
 442       }
 443       if (
 444         this.workerChoiceStrategyContext.getRequiredStatistics().medRunTime &&
 445         message.runTime != null
 446       ) {
 447         workerTasksUsage.runTimeHistory.push(message.runTime)
 448         workerTasksUsage.medRunTime = median(workerTasksUsage.runTimeHistory)
 449       }
 450     }
 451   }
 452
 453   /**
 454    * Chooses a worker node for the next task.
 455    *
 456    * The default uses a round robin algorithm to distribute the load.
 457    *
 458    * @returns [worker node key, worker node].
 459    */
 460   protected chooseWorkerNode (): [number, WorkerNode<Worker, Data>] {
 461     let workerNodeKey: number
 462     if (this.type === PoolType.DYNAMIC && !this.full && this.internalBusy()) {
 463       const workerCreated = this.createAndSetupWorker()
 464       this.registerWorkerMessageListener(workerCreated, message => {
 465         const currentWorkerNodeKey = this.getWorkerNodeKey(workerCreated)
 466         if (
 467           isKillBehavior(KillBehaviors.HARD, message.kill) ||
 468           (message.kill != null &&
 469             this.workerNodes[currentWorkerNodeKey].tasksUsage.running === 0)
 470         ) {
 471           // Kill message received from the worker: no new tasks are submitted to that worker for a while ( > maxInactiveTime)
 472           this.flushTasksQueue(currentWorkerNodeKey)
 473           void (this.destroyWorker(workerCreated) as Promise<void>)
 474         }
 475       })
 476       workerNodeKey = this.getWorkerNodeKey(workerCreated)
 477     } else {
 478       workerNodeKey = this.workerChoiceStrategyContext.execute()
 479     }
 480     return [workerNodeKey, this.workerNodes[workerNodeKey]]
 481   }
 482
 483   /**
 484    * Sends a message to the given worker.
 485    *
 486    * @param worker - The worker which should receive the message.
 487    * @param message - The message.
 488    */
 489   protected abstract sendToWorker (
 490     worker: Worker,
 491     message: MessageValue<Data>
 492   ): void
 493
 494   /**
 495    * Registers a listener callback on the given worker.
 496    *
 497    * @param worker - The worker which should register a listener.
 498    * @param listener - The message listener callback.
 499    */
 500   protected abstract registerWorkerMessageListener<
 501     Message extends Data | Response
 502   >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
 503
 504   /**
 505    * Returns a newly created worker.
 506    */
 507   protected abstract createWorker (): Worker
 508
 509   /**
 510    * Function that can be hooked up when a worker has been newly created and moved to the pool worker nodes.
 511    *
 512    * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
 513    *
 514    * @param worker - The newly created worker.
 515    */
 516   protected abstract afterWorkerSetup (worker: Worker): void
 517
 518   /**
 519    * Creates a new worker and sets it up completely in the pool worker nodes.
 520    *
 521    * @returns New, completely set up worker.
 522    */
 523   protected createAndSetupWorker (): Worker {
 524     const worker = this.createWorker()
 525
 526     worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
 527     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 528     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 529     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 530     worker.once('exit', () => {
 531       this.removeWorkerNode(worker)
 532     })
 533
 534     this.pushWorkerNode(worker)
 535
 536     this.afterWorkerSetup(worker)
 537
 538     return worker
 539   }
 540
 541   /**
 542    * This function is the listener registered for each worker message.
 543    *
 544    * @returns The listener function to execute when a message is received from a worker.
 545    */
 546   protected workerListener (): (message: MessageValue<Response>) => void {
 547     return message => {
 548       if (message.id != null) {
 549         // Task execution response received
 550         const promiseResponse = this.promiseResponseMap.get(message.id)
 551         if (promiseResponse != null) {
 552           if (message.error != null) {
 553             promiseResponse.reject(message.error)
 554           } else {
 555             promiseResponse.resolve(message.data as Response)
 556           }
 557           this.afterTaskExecutionHook(promiseResponse.worker, message)
 558           this.promiseResponseMap.delete(message.id)
 559           const workerNodeKey = this.getWorkerNodeKey(promiseResponse.worker)
 560           if (
 561             this.opts.enableTasksQueue === true &&
 562             this.tasksQueueSize(workerNodeKey) > 0
 563           ) {
 564             this.executeTask(
 565               workerNodeKey,
 566               this.dequeueTask(workerNodeKey) as Task<Data>
 567             )
 568           }
 569         }
 570       }
 571     }
 572   }
 573
 574   private checkAndEmitEvents (): void {
 575     if (this.opts.enableEvents === true) {
 576       if (this.busy) {
 577         this.emitter?.emit(PoolEvents.busy)
 578       }
 579       if (this.type === PoolType.DYNAMIC && this.full) {
 580         this.emitter?.emit(PoolEvents.full)
 581       }
 582     }
 583   }
 584
 585   /**
 586    * Sets the given worker node its tasks usage in the pool.
 587    *
 588    * @param workerNode - The worker node.
 589    * @param tasksUsage - The worker node tasks usage.
 590    */
 591   private setWorkerNodeTasksUsage (
 592     workerNode: WorkerNode<Worker, Data>,
 593     tasksUsage: TasksUsage
 594   ): void {
 595     workerNode.tasksUsage = tasksUsage
 596   }
 597
 598   /**
 599    * Pushes the given worker in the pool worker nodes.
 600    *
 601    * @param worker - The worker.
 602    * @returns The worker nodes length.
 603    */
 604   private pushWorkerNode (worker: Worker): number {
 605     return this.workerNodes.push({
 606       worker,
 607       tasksUsage: {
 608         run: 0,
 609         running: 0,
 610         runTime: 0,
 611         runTimeHistory: new CircularArray(),
 612         avgRunTime: 0,
 613         medRunTime: 0,
 614         error: 0
 615       },
 616       tasksQueue: new Queue<Task<Data>>()
 617     })
 618   }
 619
 620   /**
 621    * Sets the given worker in the pool worker nodes.
 622    *
 623    * @param workerNodeKey - The worker node key.
 624    * @param worker - The worker.
 625    * @param tasksUsage - The worker tasks usage.
 626    * @param tasksQueue - The worker task queue.
 627    */
 628   private setWorkerNode (
 629     workerNodeKey: number,
 630     worker: Worker,
 631     tasksUsage: TasksUsage,
 632     tasksQueue: Queue<Task<Data>>
 633   ): void {
 634     this.workerNodes[workerNodeKey] = {
 635       worker,
 636       tasksUsage,
 637       tasksQueue
 638     }
 639   }
 640
 641   /**
 642    * Removes the given worker from the pool worker nodes.
 643    *
 644    * @param worker - The worker.
 645    */
 646   private removeWorkerNode (worker: Worker): void {
 647     const workerNodeKey = this.getWorkerNodeKey(worker)
 648     this.workerNodes.splice(workerNodeKey, 1)
 649     this.workerChoiceStrategyContext.remove(workerNodeKey)
 650   }
 651
 652   private executeTask (workerNodeKey: number, task: Task<Data>): void {
 653     this.beforeTaskExecutionHook(workerNodeKey)
 654     this.sendToWorker(this.workerNodes[workerNodeKey].worker, task)
 655   }
 656
 657   private enqueueTask (workerNodeKey: number, task: Task<Data>): number {
 658     return this.workerNodes[workerNodeKey].tasksQueue.enqueue(task)
 659   }
 660
 661   private dequeueTask (workerNodeKey: number): Task<Data> | undefined {
 662     return this.workerNodes[workerNodeKey].tasksQueue.dequeue()
 663   }
 664
 665   private tasksQueueSize (workerNodeKey: number): number {
 666     return this.workerNodes[workerNodeKey].tasksQueue.size
 667   }
 668
 669   private flushTasksQueue (workerNodeKey: number): void {
 670     if (this.tasksQueueSize(workerNodeKey) > 0) {
 671       for (let i = 0; i < this.tasksQueueSize(workerNodeKey); i++) {
 672         this.executeTask(
 673           workerNodeKey,
 674           this.dequeueTask(workerNodeKey) as Task<Data>
 675         )
 676       }
 677     }
 678   }
 679
 680   private flushTasksQueues (): void {
 681     for (const [workerNodeKey] of this.workerNodes.entries()) {
 682       this.flushTasksQueue(workerNodeKey)
 683     }
 684   }
 685 }