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