src/pools/abstract-pool.ts

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