src/pools/abstract-pool.ts

   1 import crypto from 'node:crypto'
   2 import type {
   3   MessageValue,
   4   PromiseWorkerResponseWrapper
   5 } from '../utility-types'
   6 import { EMPTY_FUNCTION } from '../utils'
   7 import { KillBehaviors, isKillBehavior } from '../worker/worker-options'
   8 import type { PoolOptions } from './pool'
   9 import { PoolEmitter } from './pool'
  10 import type { IPoolInternal, TasksUsage, WorkerType } from './pool-internal'
  11 import { PoolType } from './pool-internal'
  12 import type { IPoolWorker } from './pool-worker'
  13 import {
  14   WorkerChoiceStrategies,
  15   type WorkerChoiceStrategy
  16 } from './selection-strategies/selection-strategies-types'
  17 import { WorkerChoiceStrategyContext } from './selection-strategies/worker-choice-strategy-context'
  18
  19 /**
  20  * Base class that implements some shared logic for all poolifier pools.
  21  *
  22  * @typeParam Worker - Type of worker which manages this pool.
  23  * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
  24  * @typeParam Response - Type of response of execution. This can only be serializable data.
  25  */
  26 export abstract class AbstractPool<
  27   Worker extends IPoolWorker,
  28   Data = unknown,
  29   Response = unknown
  30 > implements IPoolInternal<Worker, Data, Response> {
  31   /** {@inheritDoc} */
  32   public readonly workers: Map<number, WorkerType<Worker>> = new Map<
  33   number,
  34   WorkerType<Worker>
  35   >()
  36
  37   /** {@inheritDoc} */
  38   public readonly emitter?: PoolEmitter
  39
  40   /**
  41    * Id of the next worker.
  42    */
  43   protected nextWorkerId: number = 0
  44
  45   /**
  46    * The promise map.
  47    *
  48    * - `key`: This is the message id of each submitted task.
  49    * - `value`: An object that contains the worker, the resolve function and the reject function.
  50    *
  51    * When we receive a message from the worker we get a map entry and resolve/reject the promise based on the message.
  52    */
  53   protected promiseMap: Map<
  54   string,
  55   PromiseWorkerResponseWrapper<Worker, Response>
  56   > = new Map<string, PromiseWorkerResponseWrapper<Worker, Response>>()
  57
  58   /**
  59    * Worker choice strategy instance implementing the worker choice algorithm.
  60    *
  61    * Default to a strategy implementing a round robin algorithm.
  62    */
  63   protected workerChoiceStrategyContext: WorkerChoiceStrategyContext<
  64   Worker,
  65   Data,
  66   Response
  67   >
  68
  69   /**
  70    * Constructs a new poolifier pool.
  71    *
  72    * @param numberOfWorkers - Number of workers that this pool should manage.
  73    * @param filePath - Path to the worker-file.
  74    * @param opts - Options for the pool.
  75    */
  76   public constructor (
  77     public readonly numberOfWorkers: number,
  78     public readonly filePath: string,
  79     public readonly opts: PoolOptions<Worker>
  80   ) {
  81     if (!this.isMain()) {
  82       throw new Error('Cannot start a pool from a worker!')
  83     }
  84     this.checkNumberOfWorkers(this.numberOfWorkers)
  85     this.checkFilePath(this.filePath)
  86     this.checkPoolOptions(this.opts)
  87     this.setupHook()
  88
  89     for (let i = 1; i <= this.numberOfWorkers; i++) {
  90       this.createAndSetupWorker()
  91     }
  92
  93     if (this.opts.enableEvents === true) {
  94       this.emitter = new PoolEmitter()
  95     }
  96     this.workerChoiceStrategyContext = new WorkerChoiceStrategyContext(
  97       this,
  98       () => {
  99         const workerCreated = this.createAndSetupWorker()
 100         this.registerWorkerMessageListener(workerCreated, message => {
 101           if (
 102             isKillBehavior(KillBehaviors.HARD, message.kill) ||
 103             this.getWorkerRunningTasks(workerCreated) === 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(workerCreated)
 107           }
 108         })
 109         return workerCreated
 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     this.opts.enableEvents = opts.enableEvents ?? true
 146   }
 147
 148   /** {@inheritDoc} */
 149   public abstract get type (): PoolType
 150
 151   /** {@inheritDoc} */
 152   public get numberOfRunningTasks (): number {
 153     return this.promiseMap.size
 154   }
 155
 156   /**
 157    * Gets the given worker key.
 158    *
 159    * @param worker - The worker.
 160    * @returns The worker key.
 161    */
 162   private getWorkerKey (worker: Worker): number | undefined {
 163     return [...this.workers].find(([, value]) => value.worker === worker)?.[0]
 164   }
 165
 166   /** {@inheritDoc} */
 167   public getWorkerRunningTasks (worker: Worker): number | undefined {
 168     return this.getWorkerTasksUsage(worker)?.running
 169   }
 170
 171   /** {@inheritDoc} */
 172   public getWorkerRunTasks (worker: Worker): number | undefined {
 173     return this.getWorkerTasksUsage(worker)?.run
 174   }
 175
 176   /** {@inheritDoc} */
 177   public getWorkerAverageTasksRunTime (worker: Worker): number | undefined {
 178     return this.getWorkerTasksUsage(worker)?.avgRunTime
 179   }
 180
 181   /** {@inheritDoc} */
 182   public setWorkerChoiceStrategy (
 183     workerChoiceStrategy: WorkerChoiceStrategy
 184   ): void {
 185     this.opts.workerChoiceStrategy = workerChoiceStrategy
 186     for (const [key, value] of this.workers) {
 187       this.setWorker(key, value.worker, {
 188         run: 0,
 189         running: 0,
 190         runTime: 0,
 191         avgRunTime: 0
 192       })
 193     }
 194     this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
 195       workerChoiceStrategy
 196     )
 197   }
 198
 199   /** {@inheritDoc} */
 200   public abstract get busy (): boolean
 201
 202   protected internalGetBusyStatus (): boolean {
 203     return (
 204       this.numberOfRunningTasks >= this.numberOfWorkers &&
 205       this.findFreeWorker() === false
 206     )
 207   }
 208
 209   /** {@inheritDoc} */
 210   public findFreeWorker (): Worker | false {
 211     for (const value of this.workers.values()) {
 212       if (value.tasksUsage.running === 0) {
 213         // A worker is free, return the matching worker
 214         return value.worker
 215       }
 216     }
 217     return false
 218   }
 219
 220   /** {@inheritDoc} */
 221   public async execute (data: Data): Promise<Response> {
 222     const worker = this.chooseWorker()
 223     const messageId = crypto.randomUUID()
 224     const res = this.internalExecute(worker, messageId)
 225     this.checkAndEmitBusy()
 226     this.sendToWorker(worker, {
 227       // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
 228       data: data ?? ({} as Data),
 229       id: messageId
 230     })
 231     // eslint-disable-next-line @typescript-eslint/return-await
 232     return res
 233   }
 234
 235   /** {@inheritDoc} */
 236   public async destroy (): Promise<void> {
 237     await Promise.all(
 238       [...this.workers].map(async ([, value]) => {
 239         await this.destroyWorker(value.worker)
 240       })
 241     )
 242   }
 243
 244   /**
 245    * Shutdowns given worker.
 246    *
 247    * @param worker - A worker within `workers`.
 248    */
 249   protected abstract destroyWorker (worker: Worker): void | Promise<void>
 250
 251   /**
 252    * Setup hook that can be overridden by a Poolifier pool implementation
 253    * to run code before workers are created in the abstract constructor.
 254    */
 255   protected setupHook (): void {
 256     // Can be overridden
 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 worker - The worker.
 269    */
 270   protected beforePromiseWorkerResponseHook (worker: Worker): void {
 271     this.increaseWorkerRunningTasks(worker)
 272   }
 273
 274   /**
 275    * Hook executed after the worker task promise resolution.
 276    * Can be overridden.
 277    *
 278    * @param message - The received message.
 279    * @param promise - The Promise response.
 280    */
 281   protected afterPromiseWorkerResponseHook (
 282     message: MessageValue<Response>,
 283     promise: PromiseWorkerResponseWrapper<Worker, Response>
 284   ): void {
 285     this.decreaseWorkerRunningTasks(promise.worker)
 286     this.stepWorkerRunTasks(promise.worker, 1)
 287     this.updateWorkerTasksRunTime(promise.worker, message.taskRunTime)
 288   }
 289
 290   /**
 291    * Removes the given worker from the pool.
 292    *
 293    * @param worker - The worker that will be removed.
 294    */
 295   protected removeWorker (worker: Worker): void {
 296     this.workers.delete(this.getWorkerKey(worker) as number)
 297     --this.nextWorkerId
 298   }
 299
 300   /**
 301    * Chooses a worker for the next task.
 302    *
 303    * The default implementation uses a round robin algorithm to distribute the load.
 304    *
 305    * @returns Worker.
 306    */
 307   protected chooseWorker (): Worker {
 308     return this.workerChoiceStrategyContext.execute()
 309   }
 310
 311   /**
 312    * Sends a message to the given worker.
 313    *
 314    * @param worker - The worker which should receive the message.
 315    * @param message - The message.
 316    */
 317   protected abstract sendToWorker (
 318     worker: Worker,
 319     message: MessageValue<Data>
 320   ): void
 321
 322   /**
 323    * Registers a listener callback on a given worker.
 324    *
 325    * @param worker - The worker which should register a listener.
 326    * @param listener - The message listener callback.
 327    */
 328   protected abstract registerWorkerMessageListener<
 329     Message extends Data | Response
 330   >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
 331
 332   /**
 333    * Returns a newly created worker.
 334    */
 335   protected abstract createWorker (): Worker
 336
 337   /**
 338    * Function that can be hooked up when a worker has been newly created and moved to the workers registry.
 339    *
 340    * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
 341    *
 342    * @param worker - The newly created worker.
 343    */
 344   protected abstract afterWorkerSetup (worker: Worker): void
 345
 346   /**
 347    * Creates a new worker for this pool and sets it up completely.
 348    *
 349    * @returns New, completely set up worker.
 350    */
 351   protected createAndSetupWorker (): Worker {
 352     const worker = this.createWorker()
 353
 354     worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
 355     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 356     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 357     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 358     worker.once('exit', () => {
 359       this.removeWorker(worker)
 360     })
 361
 362     this.setWorker(this.nextWorkerId, worker, {
 363       run: 0,
 364       running: 0,
 365       runTime: 0,
 366       avgRunTime: 0
 367     })
 368     ++this.nextWorkerId
 369
 370     this.afterWorkerSetup(worker)
 371
 372     return worker
 373   }
 374
 375   /**
 376    * This function is the listener registered for each worker.
 377    *
 378    * @returns The listener function to execute when a message is received from a worker.
 379    */
 380   protected workerListener (): (message: MessageValue<Response>) => void {
 381     return message => {
 382       if (message.id !== undefined) {
 383         const promise = this.promiseMap.get(message.id)
 384         if (promise !== undefined) {
 385           if (message.error != null) {
 386             promise.reject(message.error)
 387           } else {
 388             promise.resolve(message.data as Response)
 389           }
 390           this.afterPromiseWorkerResponseHook(message, promise)
 391           this.promiseMap.delete(message.id)
 392         }
 393       }
 394     }
 395   }
 396
 397   private async internalExecute (
 398     worker: Worker,
 399     messageId: string
 400   ): Promise<Response> {
 401     this.beforePromiseWorkerResponseHook(worker)
 402     return await new Promise<Response>((resolve, reject) => {
 403       this.promiseMap.set(messageId, { resolve, reject, worker })
 404     })
 405   }
 406
 407   private checkAndEmitBusy (): void {
 408     if (this.opts.enableEvents === true && this.busy) {
 409       this.emitter?.emit('busy')
 410     }
 411   }
 412
 413   /**
 414    * Increases the number of tasks that the given worker has applied.
 415    *
 416    * @param worker - Worker which running tasks is increased.
 417    */
 418   private increaseWorkerRunningTasks (worker: Worker): void {
 419     this.stepWorkerRunningTasks(worker, 1)
 420   }
 421
 422   /**
 423    * Decreases the number of tasks that the given worker has applied.
 424    *
 425    * @param worker - Worker which running tasks is decreased.
 426    */
 427   private decreaseWorkerRunningTasks (worker: Worker): void {
 428     this.stepWorkerRunningTasks(worker, -1)
 429   }
 430
 431   /**
 432    * Gets tasks usage of the given worker.
 433    *
 434    * @param worker - Worker which tasks usage is returned.
 435    */
 436   private getWorkerTasksUsage (worker: Worker): TasksUsage | undefined {
 437     if (this.checkWorker(worker)) {
 438       const workerKey = this.getWorkerKey(worker) as number
 439       const workerEntry = this.workers.get(workerKey) as WorkerType<Worker>
 440       return workerEntry.tasksUsage
 441     }
 442   }
 443
 444   /**
 445    * Steps the number of tasks that the given worker has applied.
 446    *
 447    * @param worker - Worker which running tasks are stepped.
 448    * @param step - Number of running tasks step.
 449    */
 450   private stepWorkerRunningTasks (worker: Worker, step: number): void {
 451     // prettier-ignore
 452     (this.getWorkerTasksUsage(worker) as TasksUsage).running += step
 453   }
 454
 455   /**
 456    * Steps the number of tasks that the given worker has run.
 457    *
 458    * @param worker - Worker which has run tasks.
 459    * @param step - Number of run tasks step.
 460    */
 461   private stepWorkerRunTasks (worker: Worker, step: number): void {
 462     // prettier-ignore
 463     (this.getWorkerTasksUsage(worker) as TasksUsage).run += step
 464   }
 465
 466   /**
 467    * Updates tasks runtime for the given worker.
 468    *
 469    * @param worker - Worker which run the task.
 470    * @param taskRunTime - Worker task runtime.
 471    */
 472   private updateWorkerTasksRunTime (
 473     worker: Worker,
 474     taskRunTime: number | undefined
 475   ): void {
 476     if (
 477       this.workerChoiceStrategyContext.getWorkerChoiceStrategy()
 478         .requiredStatistics.runTime
 479     ) {
 480       const workerTasksUsage = this.getWorkerTasksUsage(worker) as TasksUsage
 481       workerTasksUsage.runTime += taskRunTime ?? 0
 482       if (workerTasksUsage.run !== 0) {
 483         workerTasksUsage.avgRunTime =
 484           workerTasksUsage.runTime / workerTasksUsage.run
 485       }
 486     }
 487   }
 488
 489   /**
 490    * Sets the given worker.
 491    *
 492    * @param workerKey - The worker key.
 493    * @param worker - The worker.
 494    * @param tasksUsage - The worker tasks usage.
 495    */
 496   private setWorker (
 497     workerKey: number,
 498     worker: Worker,
 499     tasksUsage: TasksUsage
 500   ): void {
 501     this.workers.set(workerKey, {
 502       worker,
 503       tasksUsage
 504     })
 505   }
 506
 507   /**
 508    * Checks if the given worker is registered in the pool.
 509    *
 510    * @param worker - Worker to check.
 511    * @returns `true` if the worker is registered in the pool.
 512    */
 513   private checkWorker (worker: Worker): boolean {
 514     if (this.getWorkerKey(worker) == null) {
 515       throw new Error('Worker could not be found in the pool')
 516     }
 517     return true
 518   }
 519 }