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