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 (this.workerChoiceStrategyContext.getRequiredStatistics().medRunTime) {
 444         workerTasksUsage.runTimeHistory.push(message.runTime ?? 0)
 445         workerTasksUsage.medRunTime = median(workerTasksUsage.runTimeHistory)
 446       }
 447     }
 448   }
 449
 450   /**
 451    * Chooses a worker node for the next task.
 452    *
 453    * The default uses a round robin algorithm to distribute the load.
 454    *
 455    * @returns [worker node key, worker node].
 456    */
 457   protected chooseWorkerNode (): [number, WorkerNode<Worker, Data>] {
 458     let workerNodeKey: number
 459     if (this.type === PoolType.DYNAMIC && !this.full && this.internalBusy()) {
 460       const workerCreated = this.createAndSetupWorker()
 461       this.registerWorkerMessageListener(workerCreated, message => {
 462         const currentWorkerNodeKey = this.getWorkerNodeKey(workerCreated)
 463         if (
 464           isKillBehavior(KillBehaviors.HARD, message.kill) ||
 465           (message.kill != null &&
 466             this.workerNodes[currentWorkerNodeKey].tasksUsage.running === 0)
 467         ) {
 468           // Kill message received from the worker: no new tasks are submitted to that worker for a while ( > maxInactiveTime)
 469           this.flushTasksQueue(currentWorkerNodeKey)
 470           void (this.destroyWorker(workerCreated) as Promise<void>)
 471         }
 472       })
 473       workerNodeKey = this.getWorkerNodeKey(workerCreated)
 474     } else {
 475       workerNodeKey = this.workerChoiceStrategyContext.execute()
 476     }
 477     return [workerNodeKey, this.workerNodes[workerNodeKey]]
 478   }
 479
 480   /**
 481    * Sends a message to the given worker.
 482    *
 483    * @param worker - The worker which should receive the message.
 484    * @param message - The message.
 485    */
 486   protected abstract sendToWorker (
 487     worker: Worker,
 488     message: MessageValue<Data>
 489   ): void
 490
 491   /**
 492    * Registers a listener callback on the given worker.
 493    *
 494    * @param worker - The worker which should register a listener.
 495    * @param listener - The message listener callback.
 496    */
 497   protected abstract registerWorkerMessageListener<
 498     Message extends Data | Response
 499   >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
 500
 501   /**
 502    * Returns a newly created worker.
 503    */
 504   protected abstract createWorker (): Worker
 505
 506   /**
 507    * Function that can be hooked up when a worker has been newly created and moved to the pool worker nodes.
 508    *
 509    * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
 510    *
 511    * @param worker - The newly created worker.
 512    */
 513   protected abstract afterWorkerSetup (worker: Worker): void
 514
 515   /**
 516    * Creates a new worker and sets it up completely in the pool worker nodes.
 517    *
 518    * @returns New, completely set up worker.
 519    */
 520   protected createAndSetupWorker (): Worker {
 521     const worker = this.createWorker()
 522
 523     worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
 524     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 525     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 526     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 527     worker.once('exit', () => {
 528       this.removeWorkerNode(worker)
 529     })
 530
 531     this.pushWorkerNode(worker)
 532
 533     this.afterWorkerSetup(worker)
 534
 535     return worker
 536   }
 537
 538   /**
 539    * This function is the listener registered for each worker message.
 540    *
 541    * @returns The listener function to execute when a message is received from a worker.
 542    */
 543   protected workerListener (): (message: MessageValue<Response>) => void {
 544     return message => {
 545       if (message.id != null) {
 546         // Task execution response received
 547         const promiseResponse = this.promiseResponseMap.get(message.id)
 548         if (promiseResponse != null) {
 549           if (message.error != null) {
 550             promiseResponse.reject(message.error)
 551           } else {
 552             promiseResponse.resolve(message.data as Response)
 553           }
 554           this.afterTaskExecutionHook(promiseResponse.worker, message)
 555           this.promiseResponseMap.delete(message.id)
 556           const workerNodeKey = this.getWorkerNodeKey(promiseResponse.worker)
 557           if (
 558             this.opts.enableTasksQueue === true &&
 559             this.tasksQueueSize(workerNodeKey) > 0
 560           ) {
 561             this.executeTask(
 562               workerNodeKey,
 563               this.dequeueTask(workerNodeKey) as Task<Data>
 564             )
 565           }
 566         }
 567       }
 568     }
 569   }
 570
 571   private checkAndEmitEvents (): void {
 572     if (this.opts.enableEvents === true) {
 573       if (this.busy) {
 574         this.emitter?.emit(PoolEvents.busy)
 575       }
 576       if (this.type === PoolType.DYNAMIC && this.full) {
 577         this.emitter?.emit(PoolEvents.full)
 578       }
 579     }
 580   }
 581
 582   /**
 583    * Sets the given worker node its tasks usage in the pool.
 584    *
 585    * @param workerNode - The worker node.
 586    * @param tasksUsage - The worker node tasks usage.
 587    */
 588   private setWorkerNodeTasksUsage (
 589     workerNode: WorkerNode<Worker, Data>,
 590     tasksUsage: TasksUsage
 591   ): void {
 592     workerNode.tasksUsage = tasksUsage
 593   }
 594
 595   /**
 596    * Pushes the given worker in the pool worker nodes.
 597    *
 598    * @param worker - The worker.
 599    * @returns The worker nodes length.
 600    */
 601   private pushWorkerNode (worker: Worker): number {
 602     return this.workerNodes.push({
 603       worker,
 604       tasksUsage: {
 605         run: 0,
 606         running: 0,
 607         runTime: 0,
 608         runTimeHistory: new CircularArray(),
 609         avgRunTime: 0,
 610         medRunTime: 0,
 611         error: 0
 612       },
 613       tasksQueue: new Queue<Task<Data>>()
 614     })
 615   }
 616
 617   /**
 618    * Sets the given worker in the pool worker nodes.
 619    *
 620    * @param workerNodeKey - The worker node key.
 621    * @param worker - The worker.
 622    * @param tasksUsage - The worker tasks usage.
 623    * @param tasksQueue - The worker task queue.
 624    */
 625   private setWorkerNode (
 626     workerNodeKey: number,
 627     worker: Worker,
 628     tasksUsage: TasksUsage,
 629     tasksQueue: Queue<Task<Data>>
 630   ): void {
 631     this.workerNodes[workerNodeKey] = {
 632       worker,
 633       tasksUsage,
 634       tasksQueue
 635     }
 636   }
 637
 638   /**
 639    * Removes the given worker from the pool worker nodes.
 640    *
 641    * @param worker - The worker.
 642    */
 643   private removeWorkerNode (worker: Worker): void {
 644     const workerNodeKey = this.getWorkerNodeKey(worker)
 645     this.workerNodes.splice(workerNodeKey, 1)
 646     this.workerChoiceStrategyContext.remove(workerNodeKey)
 647   }
 648
 649   private executeTask (workerNodeKey: number, task: Task<Data>): void {
 650     this.beforeTaskExecutionHook(workerNodeKey)
 651     this.sendToWorker(this.workerNodes[workerNodeKey].worker, task)
 652   }
 653
 654   private enqueueTask (workerNodeKey: number, task: Task<Data>): number {
 655     return this.workerNodes[workerNodeKey].tasksQueue.enqueue(task)
 656   }
 657
 658   private dequeueTask (workerNodeKey: number): Task<Data> | undefined {
 659     return this.workerNodes[workerNodeKey].tasksQueue.dequeue()
 660   }
 661
 662   private tasksQueueSize (workerNodeKey: number): number {
 663     return this.workerNodes[workerNodeKey].tasksQueue.size
 664   }
 665
 666   private flushTasksQueue (workerNodeKey: number): void {
 667     if (this.tasksQueueSize(workerNodeKey) > 0) {
 668       for (let i = 0; i < this.tasksQueueSize(workerNodeKey); i++) {
 669         this.executeTask(
 670           workerNodeKey,
 671           this.dequeueTask(workerNodeKey) as Task<Data>
 672         )
 673       }
 674     }
 675   }
 676
 677   private flushTasksQueues (): void {
 678     for (const [workerNodeKey] of this.workerNodes.entries()) {
 679       this.flushTasksQueue(workerNodeKey)
 680     }
 681   }
 682 }