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