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, 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<
  43   string,
  44   PromiseResponseWrapper<Worker, Response>
  45   > = new Map<string, PromiseResponseWrapper<Worker, Response>>()
  46
  47   /**
  48    * Worker choice strategy instance implementing the worker choice algorithm.
  49    *
  50    * Default to a strategy implementing a round robin algorithm.
  51    */
  52   protected workerChoiceStrategyContext: WorkerChoiceStrategyContext<
  53   Worker,
  54   Data,
  55   Response
  56   >
  57
  58   /**
  59    * Constructs a new poolifier pool.
  60    *
  61    * @param numberOfWorkers - Number of workers that this pool should manage.
  62    * @param filePath - Path to the worker-file.
  63    * @param opts - Options for the pool.
  64    */
  65   public constructor (
  66     public readonly numberOfWorkers: number,
  67     public readonly filePath: string,
  68     public readonly opts: PoolOptions<Worker>
  69   ) {
  70     if (!this.isMain()) {
  71       throw new Error('Cannot start a pool from a worker!')
  72     }
  73     this.checkNumberOfWorkers(this.numberOfWorkers)
  74     this.checkFilePath(this.filePath)
  75     this.checkPoolOptions(this.opts)
  76
  77     this.chooseWorker.bind(this)
  78     this.internalExecute.bind(this)
  79     this.checkAndEmitBusy.bind(this)
  80     this.sendToWorker.bind(this)
  81
  82     this.setupHook()
  83
  84     for (let i = 1; i <= this.numberOfWorkers; i++) {
  85       this.createAndSetupWorker()
  86     }
  87
  88     if (this.opts.enableEvents === true) {
  89       this.emitter = new PoolEmitter()
  90     }
  91     this.workerChoiceStrategyContext = new WorkerChoiceStrategyContext<
  92     Worker,
  93     Data,
  94     Response
  95     >(
  96       this,
  97       () => {
  98         const createdWorker = this.createAndSetupWorker()
  99         this.registerWorkerMessageListener(createdWorker, message => {
 100           if (
 101             isKillBehavior(KillBehaviors.HARD, message.kill) ||
 102             this.getWorkerTasksUsage(createdWorker)?.running === 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(createdWorker)
 106           }
 107         })
 108         return this.getWorkerKey(createdWorker)
 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   /**
 151    * Number of tasks concurrently running.
 152    */
 153   private get numberOfRunningTasks (): number {
 154     return this.promiseResponseMap.size
 155   }
 156
 157   /**
 158    * Gets the given worker key.
 159    *
 160    * @param worker - The worker.
 161    * @returns The worker key if the worker is found in the pool, `-1` otherwise.
 162    */
 163   private getWorkerKey (worker: Worker): number {
 164     return this.workers.findIndex(workerItem => workerItem.worker === worker)
 165   }
 166
 167   /** {@inheritDoc} */
 168   public setWorkerChoiceStrategy (
 169     workerChoiceStrategy: WorkerChoiceStrategy
 170   ): void {
 171     this.opts.workerChoiceStrategy = workerChoiceStrategy
 172     for (const [index, workerItem] of this.workers.entries()) {
 173       this.setWorker(index, workerItem.worker, {
 174         run: 0,
 175         running: 0,
 176         runTime: 0,
 177         avgRunTime: 0,
 178         error: 0
 179       })
 180     }
 181     this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
 182       workerChoiceStrategy
 183     )
 184   }
 185
 186   /** {@inheritDoc} */
 187   public abstract get full (): boolean
 188
 189   /** {@inheritDoc} */
 190   public abstract get busy (): boolean
 191
 192   protected internalBusy (): boolean {
 193     return (
 194       this.numberOfRunningTasks >= this.numberOfWorkers &&
 195       this.findFreeWorkerKey() === -1
 196     )
 197   }
 198
 199   /** {@inheritDoc} */
 200   public findFreeWorkerKey (): number {
 201     return this.workers.findIndex(workerItem => {
 202       return workerItem.tasksUsage.running === 0
 203     })
 204   }
 205
 206   /** {@inheritDoc} */
 207   public async execute (data: Data): Promise<Response> {
 208     const [workerKey, worker] = this.chooseWorker()
 209     const messageId = crypto.randomUUID()
 210     const res = this.internalExecute(workerKey, worker, messageId)
 211     this.checkAndEmitBusy()
 212     this.sendToWorker(worker, {
 213       // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
 214       data: data ?? ({} as Data),
 215       id: messageId
 216     })
 217     // eslint-disable-next-line @typescript-eslint/return-await
 218     return res
 219   }
 220
 221   /** {@inheritDoc} */
 222   public async destroy (): Promise<void> {
 223     await Promise.all(
 224       this.workers.map(async workerItem => {
 225         await this.destroyWorker(workerItem.worker)
 226       })
 227     )
 228   }
 229
 230   /**
 231    * Shutdowns given worker.
 232    *
 233    * @param worker - A worker within `workers`.
 234    */
 235   protected abstract destroyWorker (worker: Worker): void | Promise<void>
 236
 237   /**
 238    * Setup hook that can be overridden by a Poolifier pool implementation
 239    * to run code before workers are created in the abstract constructor.
 240    */
 241   protected setupHook (): void {
 242     // Can be overridden
 243   }
 244
 245   /**
 246    * Should return whether the worker is the main worker or not.
 247    */
 248   protected abstract isMain (): boolean
 249
 250   /**
 251    * Hook executed before the worker task promise resolution.
 252    * Can be overridden.
 253    *
 254    * @param workerKey - The worker key.
 255    */
 256   protected beforePromiseResponseHook (workerKey: number): void {
 257     ++this.workers[workerKey].tasksUsage.running
 258   }
 259
 260   /**
 261    * Hook executed after the worker task promise resolution.
 262    * Can be overridden.
 263    *
 264    * @param worker - The worker.
 265    * @param message - The received message.
 266    */
 267   protected afterPromiseResponseHook (
 268     worker: Worker,
 269     message: MessageValue<Response>
 270   ): void {
 271     const workerTasksUsage = this.getWorkerTasksUsage(worker) as TasksUsage
 272     --workerTasksUsage.running
 273     ++workerTasksUsage.run
 274     if (message.error != null) {
 275       ++workerTasksUsage.error
 276     }
 277     if (this.workerChoiceStrategyContext.getRequiredStatistics().runTime) {
 278       workerTasksUsage.runTime += message.taskRunTime ?? 0
 279       if (
 280         this.workerChoiceStrategyContext.getRequiredStatistics().avgRunTime &&
 281         workerTasksUsage.run !== 0
 282       ) {
 283         workerTasksUsage.avgRunTime =
 284           workerTasksUsage.runTime / workerTasksUsage.run
 285       }
 286     }
 287   }
 288
 289   /**
 290    * Removes the given worker from the pool.
 291    *
 292    * @param worker - The worker that will be removed.
 293    */
 294   protected removeWorker (worker: Worker): void {
 295     const workerKey = this.getWorkerKey(worker)
 296     this.workers.splice(workerKey, 1)
 297     this.workerChoiceStrategyContext.remove(workerKey)
 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 key, worker].
 306    */
 307   protected chooseWorker (): [number, Worker] {
 308     const workerKey = this.workerChoiceStrategyContext.execute()
 309     return [workerKey, this.workers[workerKey].worker]
 310   }
 311
 312   /**
 313    * Sends a message to the given worker.
 314    *
 315    * @param worker - The worker which should receive the message.
 316    * @param message - The message.
 317    */
 318   protected abstract sendToWorker (
 319     worker: Worker,
 320     message: MessageValue<Data>
 321   ): void
 322
 323   /**
 324    * Registers a listener callback on a given worker.
 325    *
 326    * @param worker - The worker which should register a listener.
 327    * @param listener - The message listener callback.
 328    */
 329   protected abstract registerWorkerMessageListener<
 330     Message extends Data | Response
 331   >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
 332
 333   /**
 334    * Returns a newly created worker.
 335    */
 336   protected abstract createWorker (): Worker
 337
 338   /**
 339    * Function that can be hooked up when a worker has been newly created and moved to the workers registry.
 340    *
 341    * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
 342    *
 343    * @param worker - The newly created worker.
 344    */
 345   protected abstract afterWorkerSetup (worker: Worker): void
 346
 347   /**
 348    * Creates a new worker for this pool and sets it up completely.
 349    *
 350    * @returns New, completely set up worker.
 351    */
 352   protected createAndSetupWorker (): Worker {
 353     const worker = this.createWorker()
 354
 355     worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
 356     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 357     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 358     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 359     worker.once('exit', () => {
 360       this.removeWorker(worker)
 361     })
 362
 363     this.pushWorker(worker, {
 364       run: 0,
 365       running: 0,
 366       runTime: 0,
 367       avgRunTime: 0,
 368       error: 0
 369     })
 370
 371     this.afterWorkerSetup(worker)
 372
 373     return worker
 374   }
 375
 376   /**
 377    * This function is the listener registered for each worker.
 378    *
 379    * @returns The listener function to execute when a message is received from a worker.
 380    */
 381   protected workerListener (): (message: MessageValue<Response>) => void {
 382     return message => {
 383       if (message.id !== undefined) {
 384         const promiseResponse = this.promiseResponseMap.get(message.id)
 385         if (promiseResponse !== undefined) {
 386           if (message.error != null) {
 387             promiseResponse.reject(message.error)
 388           } else {
 389             promiseResponse.resolve(message.data as Response)
 390           }
 391           this.afterPromiseResponseHook(promiseResponse.worker, message)
 392           this.promiseResponseMap.delete(message.id)
 393         }
 394       }
 395     }
 396   }
 397
 398   private async internalExecute (
 399     workerKey: number,
 400     worker: Worker,
 401     messageId: string
 402   ): Promise<Response> {
 403     this.beforePromiseResponseHook(workerKey)
 404     return await new Promise<Response>((resolve, reject) => {
 405       this.promiseResponseMap.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    * Gets worker tasks usage.
 417    *
 418    * @param worker - The worker.
 419    * @returns The worker tasks usage.
 420    */
 421   private getWorkerTasksUsage (worker: Worker): TasksUsage | undefined {
 422     const workerKey = this.getWorkerKey(worker)
 423     if (workerKey !== -1) {
 424       return this.workers[workerKey].tasksUsage
 425     }
 426     throw new Error('Worker could not be found in the pool')
 427   }
 428
 429   /**
 430    * Pushes the given worker.
 431    *
 432    * @param worker - The worker.
 433    * @param tasksUsage - The worker tasks usage.
 434    */
 435   private pushWorker (worker: Worker, tasksUsage: TasksUsage): void {
 436     this.workers.push({
 437       worker,
 438       tasksUsage
 439     })
 440   }
 441
 442   /**
 443    * Sets the given worker.
 444    *
 445    * @param workerKey - The worker key.
 446    * @param worker - The worker.
 447    * @param tasksUsage - The worker tasks usage.
 448    */
 449   private setWorker (
 450     workerKey: number,
 451     worker: Worker,
 452     tasksUsage: TasksUsage
 453   ): void {
 454     this.workers[workerKey] = {
 455       worker,
 456       tasksUsage
 457     }
 458   }
 459 }