src/pools/abstract-pool.ts

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