src/pools/abstract-pool.ts

   1 import type {
   2   MessageValue,
   3   PromiseWorkerResponseWrapper
   4 } from '../utility-types'
   5 import { EMPTY_FUNCTION } from '../utils'
   6 import { isKillBehavior, KillBehaviors } from '../worker/worker-options'
   7 import type { AbstractPoolWorker } from './abstract-pool-worker'
   8 import type { PoolOptions } from './pool'
   9 import type { IPoolInternal, TasksUsage } from './pool-internal'
  10 import { PoolEmitter, PoolType } from './pool-internal'
  11 import {
  12   WorkerChoiceStrategies,
  13   WorkerChoiceStrategy
  14 } from './selection-strategies/selection-strategies-types'
  15 import { WorkerChoiceStrategyContext } from './selection-strategies/worker-choice-strategy-context'
  16
  17 const WORKER_NOT_FOUND_TASKS_USAGE_MAP =
  18   'Worker could not be found in worker tasks usage map'
  19
  20 /**
  21  * Base class containing some shared logic for all poolifier pools.
  22  *
  23  * @template Worker Type of worker which manages this pool.
  24  * @template Data Type of data sent to the worker. This can only be serializable data.
  25  * @template Response Type of response of execution. This can only be serializable data.
  26  */
  27 export abstract class AbstractPool<
  28   Worker extends AbstractPoolWorker,
  29   Data = unknown,
  30   Response = unknown
  31 > implements IPoolInternal<Worker, Data, Response> {
  32   /** @inheritDoc */
  33   public readonly workers: Worker[] = []
  34
  35   /**
  36    * The workers tasks usage map.
  37    *
  38    *  `key`: The `Worker`
  39    *  `value`: Worker tasks usage statistics.
  40    */
  41   protected workersTasksUsage: Map<Worker, TasksUsage> = new Map<
  42     Worker,
  43     TasksUsage
  44   >()
  45
  46   /** @inheritDoc */
  47   public readonly emitter?: PoolEmitter
  48
  49   /** @inheritDoc */
  50   public readonly max?: number
  51
  52   /**
  53    * The promise map.
  54    *
  55    * - `key`: This is the message Id of each submitted task.
  56    * - `value`: An object that contains the worker, the resolve function and the reject function.
  57    *
  58    * When we receive a message from the worker we get a map entry and resolve/reject the promise based on the message.
  59    */
  60   protected promiseMap: Map<
  61     number,
  62     PromiseWorkerResponseWrapper<Worker, Response>
  63   > = new Map<number, PromiseWorkerResponseWrapper<Worker, Response>>()
  64
  65   /**
  66    * Id of the next message.
  67    */
  68   protected nextMessageId: number = 0
  69
  70   /**
  71    * Worker choice strategy instance implementing the worker choice algorithm.
  72    *
  73    * Default to a strategy implementing a round robin algorithm.
  74    */
  75   protected workerChoiceStrategyContext: WorkerChoiceStrategyContext<
  76     Worker,
  77     Data,
  78     Response
  79   >
  80
  81   /**
  82    * Constructs a new poolifier pool.
  83    *
  84    * @param numberOfWorkers Number of workers that this pool should manage.
  85    * @param filePath Path to the worker-file.
  86    * @param opts Options for the pool.
  87    */
  88   public constructor (
  89     public readonly numberOfWorkers: number,
  90     public readonly filePath: string,
  91     public readonly opts: PoolOptions<Worker>
  92   ) {
  93     if (!this.isMain()) {
  94       throw new Error('Cannot start a pool from a worker!')
  95     }
  96     this.checkNumberOfWorkers(this.numberOfWorkers)
  97     this.checkFilePath(this.filePath)
  98     this.checkPoolOptions(this.opts)
  99     this.setupHook()
 100
 101     for (let i = 1; i <= this.numberOfWorkers; i++) {
 102       this.createAndSetupWorker()
 103     }
 104
 105     if (this.opts.enableEvents) {
 106       this.emitter = new PoolEmitter()
 107     }
 108     this.workerChoiceStrategyContext = new WorkerChoiceStrategyContext(
 109       this,
 110       () => {
 111         const workerCreated = this.createAndSetupWorker()
 112         this.registerWorkerMessageListener(workerCreated, message => {
 113           if (
 114             isKillBehavior(KillBehaviors.HARD, message.kill) ||
 115             this.getWorkerRunningTasks(workerCreated) === 0
 116           ) {
 117             // Kill received from the worker, means that no new tasks are submitted to that worker for a while ( > maxInactiveTime)
 118             this.destroyWorker(workerCreated) as void
 119           }
 120         })
 121         return workerCreated
 122       },
 123       this.opts.workerChoiceStrategy
 124     )
 125   }
 126
 127   private checkFilePath (filePath: string): void {
 128     if (!filePath) {
 129       throw new Error('Please specify a file with a worker implementation')
 130     }
 131   }
 132
 133   private checkNumberOfWorkers (numberOfWorkers: number): void {
 134     if (numberOfWorkers == null) {
 135       throw new Error(
 136         'Cannot instantiate a pool without specifying the number of workers'
 137       )
 138     } else if (!Number.isSafeInteger(numberOfWorkers)) {
 139       throw new Error(
 140         'Cannot instantiate a pool with a non integer number of workers'
 141       )
 142     } else if (numberOfWorkers < 0) {
 143       throw new Error(
 144         'Cannot instantiate a pool with a negative number of workers'
 145       )
 146     } else if (this.type === PoolType.FIXED && numberOfWorkers === 0) {
 147       throw new Error('Cannot instantiate a fixed pool with no worker')
 148     }
 149   }
 150
 151   private checkPoolOptions (opts: PoolOptions<Worker>): void {
 152     this.opts.workerChoiceStrategy =
 153       opts.workerChoiceStrategy ?? WorkerChoiceStrategies.ROUND_ROBIN
 154     this.opts.enableEvents = opts.enableEvents ?? true
 155   }
 156
 157   /** @inheritDoc */
 158   public abstract get type (): PoolType
 159
 160   /** @inheritDoc */
 161   public get numberOfRunningTasks (): number {
 162     return this.promiseMap.size
 163   }
 164
 165   /** @inheritDoc */
 166   public getWorkerIndex (worker: Worker): number {
 167     return this.workers.indexOf(worker)
 168   }
 169
 170   /** @inheritDoc */
 171   public getWorkerRunningTasks (worker: Worker): number | undefined {
 172     return this.workersTasksUsage.get(worker)?.running
 173   }
 174
 175   /** @inheritDoc */
 176   public getWorkerAverageTasksRunTime (worker: Worker): number | undefined {
 177     return this.workersTasksUsage.get(worker)?.avgRunTime
 178   }
 179
 180   /** @inheritDoc */
 181   public setWorkerChoiceStrategy (
 182     workerChoiceStrategy: WorkerChoiceStrategy
 183   ): void {
 184     this.opts.workerChoiceStrategy = workerChoiceStrategy
 185     this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
 186       workerChoiceStrategy
 187     )
 188   }
 189
 190   /** @inheritDoc */
 191   public abstract get busy (): boolean
 192
 193   protected internalGetBusyStatus (): boolean {
 194     return (
 195       this.numberOfRunningTasks >= this.numberOfWorkers &&
 196       this.findFreeWorker() === false
 197     )
 198   }
 199
 200   /** @inheritDoc */
 201   public findFreeWorker (): Worker | false {
 202     for (const worker of this.workers) {
 203       if (this.getWorkerRunningTasks(worker) === 0) {
 204         // A worker is free, return the matching worker
 205         return worker
 206       }
 207     }
 208     return false
 209   }
 210
 211   /** @inheritDoc */
 212   public execute (data: Data): Promise<Response> {
 213     // Configure worker to handle message with the specified task
 214     const worker = this.chooseWorker()
 215     const messageId = ++this.nextMessageId
 216     const res = this.internalExecute(worker, messageId)
 217     this.checkAndEmitBusy()
 218     data = data ?? ({} as Data)
 219     this.sendToWorker(worker, { data, id: messageId })
 220     return res
 221   }
 222
 223   /** @inheritDoc */
 224   public async destroy (): Promise<void> {
 225     await Promise.all(this.workers.map(worker => this.destroyWorker(worker)))
 226   }
 227
 228   /**
 229    * Shut down given worker.
 230    *
 231    * @param worker A worker within `workers`.
 232    */
 233   protected abstract destroyWorker (worker: Worker): void | Promise<void>
 234
 235   /**
 236    * Setup hook that can be overridden by a Poolifier pool implementation
 237    * to run code before workers are created in the abstract constructor.
 238    */
 239   protected setupHook (): void {
 240     // Can be overridden
 241   }
 242
 243   /**
 244    * Should return whether the worker is the main worker or not.
 245    */
 246   protected abstract isMain (): boolean
 247
 248   /**
 249    * Hook executed before the worker task promise resolution.
 250    * Can be overridden.
 251    *
 252    * @param worker The worker.
 253    */
 254   protected beforePromiseWorkerResponseHook (worker: Worker): void {
 255     this.increaseWorkerRunningTasks(worker)
 256   }
 257
 258   /**
 259    * Hook executed after the worker task promise resolution.
 260    * Can be overridden.
 261    *
 262    * @param message The received message.
 263    * @param promise The Promise response.
 264    */
 265   protected afterPromiseWorkerResponseHook (
 266     message: MessageValue<Response>,
 267     promise: PromiseWorkerResponseWrapper<Worker, Response>
 268   ): void {
 269     this.decreaseWorkerRunningTasks(promise.worker)
 270     this.stepWorkerRunTasks(promise.worker, 1)
 271     this.updateWorkerTasksRunTime(promise.worker, message.taskRunTime)
 272   }
 273
 274   /**
 275    * Removes the given worker from the pool.
 276    *
 277    * @param worker Worker that will be removed.
 278    */
 279   protected removeWorker (worker: Worker): void {
 280     // Clean worker from data structure
 281     this.workers.splice(this.getWorkerIndex(worker), 1)
 282     this.resetWorkerTasksUsage(worker)
 283   }
 284
 285   /**
 286    * Choose 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    * Send 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    * Register a listener callback on a given worker.
 309    *
 310    * @param worker A worker.
 311    * @param listener A message listener callback.
 312    */
 313   protected abstract registerWorkerMessageListener<
 314     Message extends Data | Response
 315   > (worker: Worker, listener: (message: MessageValue<Message>) => void): void
 316
 317   protected internalExecute (
 318     worker: Worker,
 319     messageId: number
 320   ): Promise<Response> {
 321     this.beforePromiseWorkerResponseHook(worker)
 322     return new Promise<Response>((resolve, reject) => {
 323       this.promiseMap.set(messageId, { resolve, reject, worker })
 324     })
 325   }
 326
 327   /**
 328    * Returns a newly created worker.
 329    */
 330   protected abstract createWorker (): Worker
 331
 332   /**
 333    * Function that can be hooked up when a worker has been newly created and moved to the workers registry.
 334    *
 335    * Can be used to update the `maxListeners` or binding the `main-worker`<->`worker` connection if not bind by default.
 336    *
 337    * @param worker The newly created worker.
 338    */
 339   protected abstract afterWorkerSetup (worker: Worker): void
 340
 341   /**
 342    * Creates a new worker for this pool and sets it up completely.
 343    *
 344    * @returns New, completely set up worker.
 345    */
 346   protected createAndSetupWorker (): Worker {
 347     const worker = this.createWorker()
 348
 349     worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
 350     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 351     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 352     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 353     worker.once('exit', () => this.removeWorker(worker))
 354
 355     this.workers.push(worker)
 356
 357     // Init worker tasks usage map
 358     this.workersTasksUsage.set(worker, {
 359       run: 0,
 360       running: 0,
 361       runTime: 0,
 362       avgRunTime: 0
 363     })
 364
 365     this.afterWorkerSetup(worker)
 366
 367     return worker
 368   }
 369
 370   /**
 371    * This function is the listener registered for each worker.
 372    *
 373    * @returns The listener function to execute when a message is received from a worker.
 374    */
 375   protected workerListener (): (message: MessageValue<Response>) => void {
 376     return message => {
 377       if (message.id !== undefined) {
 378         const promise = this.promiseMap.get(message.id)
 379         if (promise !== undefined) {
 380           this.afterPromiseWorkerResponseHook(message, promise)
 381           if (message.error) promise.reject(message.error)
 382           else promise.resolve(message.data as Response)
 383           this.promiseMap.delete(message.id)
 384         }
 385       }
 386     }
 387   }
 388
 389   private checkAndEmitBusy (): void {
 390     if (this.opts.enableEvents && this.busy) {
 391       this.emitter?.emit('busy')
 392     }
 393   }
 394
 395   /**
 396    * Increase the number of tasks that the given worker has applied.
 397    *
 398    * @param worker Worker which running tasks is increased.
 399    */
 400   private increaseWorkerRunningTasks (worker: Worker): void {
 401     this.stepWorkerRunningTasks(worker, 1)
 402   }
 403
 404   /**
 405    * Decrease the number of tasks that the given worker has applied.
 406    *
 407    * @param worker Worker which running tasks is decreased.
 408    */
 409   private decreaseWorkerRunningTasks (worker: Worker): void {
 410     this.stepWorkerRunningTasks(worker, -1)
 411   }
 412
 413   /**
 414    * Step the number of tasks that the given worker has applied.
 415    *
 416    * @param worker Worker which running tasks are stepped.
 417    * @param step Number of running tasks step.
 418    */
 419   private stepWorkerRunningTasks (worker: Worker, step: number): void {
 420     const tasksUsage = this.workersTasksUsage.get(worker)
 421     if (tasksUsage !== undefined) {
 422       tasksUsage.running = tasksUsage.running + step
 423       this.workersTasksUsage.set(worker, tasksUsage)
 424     } else {
 425       throw new Error(WORKER_NOT_FOUND_TASKS_USAGE_MAP)
 426     }
 427   }
 428
 429   /**
 430    * Step the number of tasks that the given worker has run.
 431    *
 432    * @param worker Worker which has run tasks.
 433    * @param step Number of run tasks step.
 434    */
 435   private stepWorkerRunTasks (worker: Worker, step: number) {
 436     const tasksUsage = this.workersTasksUsage.get(worker)
 437     if (tasksUsage !== undefined) {
 438       tasksUsage.run = tasksUsage.run + step
 439       this.workersTasksUsage.set(worker, tasksUsage)
 440     } else {
 441       throw new Error(WORKER_NOT_FOUND_TASKS_USAGE_MAP)
 442     }
 443   }
 444
 445   /**
 446    * Update tasks run time for the given worker.
 447    *
 448    * @param worker Worker which run the task.
 449    * @param taskRunTime Worker task run time.
 450    */
 451   private updateWorkerTasksRunTime (
 452     worker: Worker,
 453     taskRunTime: number | undefined
 454   ) {
 455     const tasksUsage = this.workersTasksUsage.get(worker)
 456     if (tasksUsage !== undefined && tasksUsage.run !== 0) {
 457       tasksUsage.runTime += taskRunTime ?? 0
 458       tasksUsage.avgRunTime = tasksUsage.runTime / tasksUsage.run
 459       this.workersTasksUsage.set(worker, tasksUsage)
 460     } else {
 461       throw new Error(WORKER_NOT_FOUND_TASKS_USAGE_MAP)
 462     }
 463   }
 464
 465   /**
 466    * Reset worker tasks usage statistics.
 467    *
 468    * @param worker The worker.
 469    */
 470   private resetWorkerTasksUsage (worker: Worker): void {
 471     this.workersTasksUsage.delete(worker)
 472   }
 473 }