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