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