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