src/pools/abstract-pool.ts

   1 import crypto from 'node:crypto'
   2 import type { MessageValue, PromiseResponseWrapper } from '../utility-types'
   3 import { EMPTY_FUNCTION } from '../utils'
   4 import { KillBehaviors, isKillBehavior } from '../worker/worker-options'
   5 import type { PoolOptions } from './pool'
   6 import { PoolEmitter } from './pool'
   7 import type { IPoolInternal, TasksUsage, WorkerType } from './pool-internal'
   8 import { PoolType } from './pool-internal'
   9 import type { IPoolWorker } from './pool-worker'
  10 import {
  11   WorkerChoiceStrategies,
  12   type WorkerChoiceStrategy
  13 } from './selection-strategies/selection-strategies-types'
  14 import { WorkerChoiceStrategyContext } from './selection-strategies/worker-choice-strategy-context'
  15
  16 /**
  17  * Base class that implements some shared logic for all poolifier pools.
  18  *
  19  * @typeParam Worker - Type of worker which manages this pool.
  20  * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
  21  * @typeParam Response - Type of response of execution. This can only be serializable data.
  22  */
  23 export abstract class AbstractPool<
  24   Worker extends IPoolWorker,
  25   Data = unknown,
  26   Response = unknown
  27 > implements IPoolInternal<Worker, Data, Response> {
  28   /** @inheritDoc */
  29   public readonly workers: Array<WorkerType<Worker>> = []
  30
  31   /** @inheritDoc */
  32   public readonly emitter?: PoolEmitter
  33
  34   /**
  35    * The promise response map.
  36    *
  37    * - `key`: The message id of each submitted task.
  38    * - `value`: An object that contains the worker, the promise resolve and reject callbacks.
  39    *
  40    * When we receive a message from the worker we get a map entry with the promise resolve/reject bound to the message.
  41    */
  42   protected promiseResponseMap: Map<
  43   string,
  44   PromiseResponseWrapper<Worker, Response>
  45   > = new Map<string, PromiseResponseWrapper<Worker, Response>>()
  46
  47   /**
  48    * Worker choice strategy context referencing a worker choice algorithm implementation.
  49    *
  50    * Default to a round robin algorithm.
  51    */
  52   protected workerChoiceStrategyContext: WorkerChoiceStrategyContext<
  53   Worker,
  54   Data,
  55   Response
  56   >
  57
  58   /**
  59    * Constructs a new poolifier pool.
  60    *
  61    * @param numberOfWorkers - Number of workers that this pool should manage.
  62    * @param filePath - Path to the worker-file.
  63    * @param opts - Options for the pool.
  64    */
  65   public constructor (
  66     public readonly numberOfWorkers: number,
  67     public readonly filePath: string,
  68     public readonly opts: PoolOptions<Worker>
  69   ) {
  70     if (!this.isMain()) {
  71       throw new Error('Cannot start a pool from a worker!')
  72     }
  73     this.checkNumberOfWorkers(this.numberOfWorkers)
  74     this.checkFilePath(this.filePath)
  75     this.checkPoolOptions(this.opts)
  76
  77     this.chooseWorker.bind(this)
  78     this.internalExecute.bind(this)
  79     this.checkAndEmitFull.bind(this)
  80     this.checkAndEmitBusy.bind(this)
  81     this.sendToWorker.bind(this)
  82
  83     this.setupHook()
  84
  85     for (let i = 1; i <= this.numberOfWorkers; i++) {
  86       this.createAndSetupWorker()
  87     }
  88
  89     if (this.opts.enableEvents === true) {
  90       this.emitter = new PoolEmitter()
  91     }
  92     this.workerChoiceStrategyContext = new WorkerChoiceStrategyContext<
  93     Worker,
  94     Data,
  95     Response
  96     >(
  97       this,
  98       () => {
  99         const createdWorker = this.createAndSetupWorker()
 100         this.registerWorkerMessageListener(createdWorker, message => {
 101           if (
 102             isKillBehavior(KillBehaviors.HARD, message.kill) ||
 103             this.getWorkerTasksUsage(createdWorker)?.running === 0
 104           ) {
 105             // Kill received from the worker, means that no new tasks are submitted to that worker for a while ( > maxInactiveTime)
 106             void this.destroyWorker(createdWorker)
 107           }
 108         })
 109         return this.getWorkerKey(createdWorker)
 110       },
 111       this.opts.workerChoiceStrategy
 112     )
 113   }
 114
 115   private checkFilePath (filePath: string): void {
 116     if (
 117       filePath == null ||
 118       (typeof filePath === 'string' && filePath.trim().length === 0)
 119     ) {
 120       throw new Error('Please specify a file with a worker implementation')
 121     }
 122   }
 123
 124   private checkNumberOfWorkers (numberOfWorkers: number): void {
 125     if (numberOfWorkers == null) {
 126       throw new Error(
 127         'Cannot instantiate a pool without specifying the number of workers'
 128       )
 129     } else if (!Number.isSafeInteger(numberOfWorkers)) {
 130       throw new TypeError(
 131         'Cannot instantiate a pool with a non integer number of workers'
 132       )
 133     } else if (numberOfWorkers < 0) {
 134       throw new RangeError(
 135         'Cannot instantiate a pool with a negative number of workers'
 136       )
 137     } else if (this.type === PoolType.FIXED && numberOfWorkers === 0) {
 138       throw new Error('Cannot instantiate a fixed pool with no worker')
 139     }
 140   }
 141
 142   private checkPoolOptions (opts: PoolOptions<Worker>): void {
 143     this.opts.workerChoiceStrategy =
 144       opts.workerChoiceStrategy ?? WorkerChoiceStrategies.ROUND_ROBIN
 145     if (
 146       !Object.values(WorkerChoiceStrategies).includes(
 147         this.opts.workerChoiceStrategy
 148       )
 149     ) {
 150       throw new Error(
 151         `Invalid worker choice strategy '${this.opts.workerChoiceStrategy}'`
 152       )
 153     }
 154     this.opts.enableEvents = opts.enableEvents ?? true
 155   }
 156
 157   /** @inheritDoc */
 158   public abstract get type (): PoolType
 159
 160   /**
 161    * Number of tasks concurrently running in the pool.
 162    */
 163   private get numberOfRunningTasks (): number {
 164     return this.promiseResponseMap.size
 165   }
 166
 167   /**
 168    * Gets the given worker key.
 169    *
 170    * @param worker - The worker.
 171    * @returns The worker key if the worker is found in the pool, `-1` otherwise.
 172    */
 173   private getWorkerKey (worker: Worker): number {
 174     return this.workers.findIndex(workerItem => workerItem.worker === worker)
 175   }
 176
 177   /** @inheritDoc */
 178   public setWorkerChoiceStrategy (
 179     workerChoiceStrategy: WorkerChoiceStrategy
 180   ): void {
 181     this.opts.workerChoiceStrategy = workerChoiceStrategy
 182     for (const [index, workerItem] of this.workers.entries()) {
 183       this.setWorker(index, workerItem.worker, {
 184         run: 0,
 185         running: 0,
 186         runTime: 0,
 187         avgRunTime: 0,
 188         error: 0
 189       })
 190     }
 191     this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
 192       workerChoiceStrategy
 193     )
 194   }
 195
 196   /** @inheritDoc */
 197   public abstract get full (): boolean
 198
 199   /** @inheritDoc */
 200   public abstract get busy (): boolean
 201
 202   protected internalBusy (): boolean {
 203     return (
 204       this.numberOfRunningTasks >= this.numberOfWorkers &&
 205       this.findFreeWorkerKey() === -1
 206     )
 207   }
 208
 209   /** @inheritDoc */
 210   public findFreeWorkerKey (): number {
 211     return this.workers.findIndex(workerItem => {
 212       return workerItem.tasksUsage.running === 0
 213     })
 214   }
 215
 216   /** @inheritDoc */
 217   public async execute (data: Data): Promise<Response> {
 218     const [workerKey, worker] = this.chooseWorker()
 219     const messageId = crypto.randomUUID()
 220     const res = this.internalExecute(workerKey, worker, messageId)
 221     this.checkAndEmitFull()
 222     this.checkAndEmitBusy()
 223     this.sendToWorker(worker, {
 224       // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
 225       data: data ?? ({} as Data),
 226       id: messageId
 227     })
 228     // eslint-disable-next-line @typescript-eslint/return-await
 229     return res
 230   }
 231
 232   /** @inheritDoc */
 233   public async destroy (): Promise<void> {
 234     await Promise.all(
 235       this.workers.map(async workerItem => {
 236         await this.destroyWorker(workerItem.worker)
 237       })
 238     )
 239   }
 240
 241   /**
 242    * Shutdowns given worker in the pool.
 243    *
 244    * @param worker - A worker within `workers`.
 245    */
 246   protected abstract destroyWorker (worker: Worker): void | Promise<void>
 247
 248   /**
 249    * Setup hook that can be overridden by a Poolifier pool implementation
 250    * to run code before workers are created in the abstract constructor.
 251    * Can be overridden
 252    *
 253    * @virtual
 254    */
 255   protected setupHook (): void {
 256     // Intentionally empty
 257   }
 258
 259   /**
 260    * Should return whether the worker is the main worker or not.
 261    */
 262   protected abstract isMain (): boolean
 263
 264   /**
 265    * Hook executed before the worker task promise resolution.
 266    * Can be overridden.
 267    *
 268    * @param workerKey - The worker key.
 269    */
 270   protected beforePromiseResponseHook (workerKey: number): void {
 271     ++this.workers[workerKey].tasksUsage.running
 272   }
 273
 274   /**
 275    * Hook executed after the worker task promise resolution.
 276    * Can be overridden.
 277    *
 278    * @param worker - The worker.
 279    * @param message - The received message.
 280    */
 281   protected afterPromiseResponseHook (
 282     worker: Worker,
 283     message: MessageValue<Response>
 284   ): void {
 285     const workerTasksUsage = this.getWorkerTasksUsage(worker) as TasksUsage
 286     --workerTasksUsage.running
 287     ++workerTasksUsage.run
 288     if (message.error != null) {
 289       ++workerTasksUsage.error
 290     }
 291     if (this.workerChoiceStrategyContext.getRequiredStatistics().runTime) {
 292       workerTasksUsage.runTime += message.taskRunTime ?? 0
 293       if (
 294         this.workerChoiceStrategyContext.getRequiredStatistics().avgRunTime &&
 295         workerTasksUsage.run !== 0
 296       ) {
 297         workerTasksUsage.avgRunTime =
 298           workerTasksUsage.runTime / workerTasksUsage.run
 299       }
 300     }
 301   }
 302
 303   /**
 304    * Chooses a worker for the next task.
 305    *
 306    * The default uses a round robin algorithm to distribute the load.
 307    *
 308    * @returns [worker key, worker].
 309    */
 310   protected chooseWorker (): [number, Worker] {
 311     const workerKey = this.workerChoiceStrategyContext.execute()
 312     return [workerKey, this.workers[workerKey].worker]
 313   }
 314
 315   /**
 316    * Sends a message to the given worker.
 317    *
 318    * @param worker - The worker which should receive the message.
 319    * @param message - The message.
 320    */
 321   protected abstract sendToWorker (
 322     worker: Worker,
 323     message: MessageValue<Data>
 324   ): void
 325
 326   /**
 327    * Registers a listener callback on a given worker.
 328    *
 329    * @param worker - The worker which should register a listener.
 330    * @param listener - The message listener callback.
 331    */
 332   protected abstract registerWorkerMessageListener<
 333     Message extends Data | Response
 334   >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
 335
 336   /**
 337    * Returns a newly created worker.
 338    */
 339   protected abstract createWorker (): Worker
 340
 341   /**
 342    * Function that can be hooked up when a worker has been newly created and moved to the workers registry.
 343    *
 344    * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
 345    *
 346    * @param worker - The newly created worker.
 347    * @virtual
 348    */
 349   protected abstract afterWorkerSetup (worker: Worker): void
 350
 351   /**
 352    * Creates a new worker for this pool and sets it up completely.
 353    *
 354    * @returns New, completely set up worker.
 355    */
 356   protected createAndSetupWorker (): Worker {
 357     const worker = this.createWorker()
 358
 359     worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
 360     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 361     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 362     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 363     worker.once('exit', () => {
 364       this.removeWorker(worker)
 365     })
 366
 367     this.pushWorker(worker, {
 368       run: 0,
 369       running: 0,
 370       runTime: 0,
 371       avgRunTime: 0,
 372       error: 0
 373     })
 374
 375     this.afterWorkerSetup(worker)
 376
 377     return worker
 378   }
 379
 380   /**
 381    * This function is the listener registered for each worker.
 382    *
 383    * @returns The listener function to execute when a message is received from a worker.
 384    */
 385   protected workerListener (): (message: MessageValue<Response>) => void {
 386     return message => {
 387       if (message.id != null) {
 388         const promiseResponse = this.promiseResponseMap.get(message.id)
 389         if (promiseResponse != null) {
 390           if (message.error != null) {
 391             promiseResponse.reject(message.error)
 392           } else {
 393             promiseResponse.resolve(message.data as Response)
 394           }
 395           this.afterPromiseResponseHook(promiseResponse.worker, message)
 396           this.promiseResponseMap.delete(message.id)
 397         }
 398       }
 399     }
 400   }
 401
 402   private async internalExecute (
 403     workerKey: number,
 404     worker: Worker,
 405     messageId: string
 406   ): Promise<Response> {
 407     this.beforePromiseResponseHook(workerKey)
 408     return await new Promise<Response>((resolve, reject) => {
 409       this.promiseResponseMap.set(messageId, { resolve, reject, worker })
 410     })
 411   }
 412
 413   private checkAndEmitBusy (): void {
 414     if (this.opts.enableEvents === true && this.busy) {
 415       this.emitter?.emit('busy')
 416     }
 417   }
 418
 419   private checkAndEmitFull (): void {
 420     if (
 421       this.type === PoolType.DYNAMIC &&
 422       this.opts.enableEvents === true &&
 423       this.full
 424     ) {
 425       this.emitter?.emit('full')
 426     }
 427   }
 428
 429   /**
 430    * Gets the given worker tasks usage in the pool.
 431    *
 432    * @param worker - The worker.
 433    * @returns The worker tasks usage.
 434    */
 435   private getWorkerTasksUsage (worker: Worker): TasksUsage | undefined {
 436     const workerKey = this.getWorkerKey(worker)
 437     if (workerKey !== -1) {
 438       return this.workers[workerKey].tasksUsage
 439     }
 440     throw new Error('Worker could not be found in the pool')
 441   }
 442
 443   /**
 444    * Pushes the given worker in the pool.
 445    *
 446    * @param worker - The worker.
 447    * @param tasksUsage - The worker tasks usage.
 448    */
 449   private pushWorker (worker: Worker, tasksUsage: TasksUsage): void {
 450     this.workers.push({
 451       worker,
 452       tasksUsage
 453     })
 454   }
 455
 456   /**
 457    * Sets the given worker in the pool.
 458    *
 459    * @param workerKey - The worker key.
 460    * @param worker - The worker.
 461    * @param tasksUsage - The worker tasks usage.
 462    */
 463   private setWorker (
 464     workerKey: number,
 465     worker: Worker,
 466     tasksUsage: TasksUsage
 467   ): void {
 468     this.workers[workerKey] = {
 469       worker,
 470       tasksUsage
 471     }
 472   }
 473
 474   /**
 475    * Removes the given worker from the pool.
 476    *
 477    * @param worker - The worker that will be removed.
 478    */
 479   protected removeWorker (worker: Worker): void {
 480     const workerKey = this.getWorkerKey(worker)
 481     this.workers.splice(workerKey, 1)
 482     this.workerChoiceStrategyContext.remove(workerKey)
 483   }
 484 }