src/pools/abstract-pool.ts

   1 import type {
   2   MessageValue,
   3   PromiseWorkerResponseWrapper
   4 } from '../utility-types'
   5 import { isKillBehavior, KillBehaviors } from '../worker/worker-options'
   6 import type { IPoolInternal } from './pool-internal'
   7 import { PoolEmitter } from './pool-internal'
   8 import type { WorkerChoiceStrategy } from './selection-strategies'
   9 import {
  10   WorkerChoiceStrategies,
  11   WorkerChoiceStrategyContext
  12 } from './selection-strategies'
  13
  14 /**
  15  * An intentional empty function.
  16  */
  17 const EMPTY_FUNCTION: () => void = () => {
  18   /* Intentionally empty */
  19 }
  20
  21 /**
  22  * Callback invoked if the worker raised an error.
  23  */
  24 export type ErrorHandler<Worker> = (this: Worker, e: Error) => void
  25
  26 /**
  27  * Callback invoked when the worker has started successfully.
  28  */
  29 export type OnlineHandler<Worker> = (this: Worker) => void
  30
  31 /**
  32  * Callback invoked when the worker exits successfully.
  33  */
  34 export type ExitHandler<Worker> = (this: Worker, code: number) => void
  35
  36 /**
  37  * Basic interface that describes the minimum required implementation of listener events for a pool-worker.
  38  */
  39 export interface IWorker {
  40   /**
  41    * Register a listener to the error event.
  42    *
  43    * @param event `'error'`.
  44    * @param handler The error handler.
  45    */
  46   on(event: 'error', handler: ErrorHandler<this>): void
  47   /**
  48    * Register a listener to the online event.
  49    *
  50    * @param event `'online'`.
  51    * @param handler The online handler.
  52    */
  53   on(event: 'online', handler: OnlineHandler<this>): void
  54   /**
  55    * Register a listener to the exit event.
  56    *
  57    * @param event `'exit'`.
  58    * @param handler The exit handler.
  59    */
  60   on(event: 'exit', handler: ExitHandler<this>): void
  61   /**
  62    * Register a listener to the exit event that will only performed once.
  63    *
  64    * @param event `'exit'`.
  65    * @param handler The exit handler.
  66    */
  67   once(event: 'exit', handler: ExitHandler<this>): void
  68 }
  69
  70 /**
  71  * Options for a poolifier pool.
  72  */
  73 export interface PoolOptions<Worker> {
  74   /**
  75    * A function that will listen for error event on each worker.
  76    */
  77   errorHandler?: ErrorHandler<Worker>
  78   /**
  79    * A function that will listen for online event on each worker.
  80    */
  81   onlineHandler?: OnlineHandler<Worker>
  82   /**
  83    * A function that will listen for exit event on each worker.
  84    */
  85   exitHandler?: ExitHandler<Worker>
  86   /**
  87    * The work choice strategy to use in this pool.
  88    */
  89   workerChoiceStrategy?: WorkerChoiceStrategy
  90 }
  91
  92 /**
  93  * Base class containing some shared logic for all poolifier pools.
  94  *
  95  * @template Worker Type of worker which manages this pool.
  96  * @template Data Type of data sent to the worker. This can only be serializable data.
  97  * @template Response Type of response of execution. This can only be serializable data.
  98  */
  99 export abstract class AbstractPool<
 100   Worker extends IWorker,
 101   Data = unknown,
 102   Response = unknown
 103 > implements IPoolInternal<Worker, Data, Response> {
 104   /** @inheritdoc */
 105   public readonly workers: Worker[] = []
 106
 107   /** @inheritdoc */
 108   public readonly tasks: Map<Worker, number> = new Map<Worker, number>()
 109
 110   /** @inheritdoc */
 111   public readonly emitter: PoolEmitter
 112
 113   /**
 114    * The promise map.
 115    *
 116    * - `key`: This is the message ID of each submitted task.
 117    * - `value`: An object that contains the worker, the resolve function and the reject function.
 118    *
 119    * When we receive a message from the worker we get a map entry and resolve/reject the promise based on the message.
 120    */
 121   protected promiseMap: Map<
 122     number,
 123     PromiseWorkerResponseWrapper<Worker, Response>
 124   > = new Map<number, PromiseWorkerResponseWrapper<Worker, Response>>()
 125
 126   /**
 127    * ID of the next message.
 128    */
 129   protected nextMessageId: number = 0
 130
 131   /**
 132    * Worker choice strategy instance implementing the worker choice algorithm.
 133    *
 134    * Default to a strategy implementing a round robin algorithm.
 135    */
 136   protected workerChoiceStrategyContext: WorkerChoiceStrategyContext<
 137     Worker,
 138     Data,
 139     Response
 140   >
 141
 142   /**
 143    * Constructs a new poolifier pool.
 144    *
 145    * @param numberOfWorkers Number of workers that this pool should manage.
 146    * @param filePath Path to the worker-file.
 147    * @param opts Options for the pool.
 148    */
 149   public constructor (
 150     public readonly numberOfWorkers: number,
 151     public readonly filePath: string,
 152     public readonly opts: PoolOptions<Worker>
 153   ) {
 154     if (!this.isMain()) {
 155       throw new Error('Cannot start a pool from a worker!')
 156     }
 157     this.checkNumberOfWorkers(this.numberOfWorkers)
 158     this.checkFilePath(this.filePath)
 159     this.setupHook()
 160
 161     for (let i = 1; i <= this.numberOfWorkers; i++) {
 162       this.createAndSetupWorker()
 163     }
 164
 165     this.emitter = new PoolEmitter()
 166     this.workerChoiceStrategyContext = new WorkerChoiceStrategyContext(
 167       this,
 168       () => {
 169         const workerCreated = this.createAndSetupWorker()
 170         this.registerWorkerMessageListener(workerCreated, message => {
 171           const tasksInProgress = this.tasks.get(workerCreated)
 172           if (
 173             isKillBehavior(KillBehaviors.HARD, message.kill) ||
 174             tasksInProgress === 0
 175           ) {
 176             // Kill received from the worker, means that no new tasks are submitted to that worker for a while ( > maxInactiveTime)
 177             void this.destroyWorker(workerCreated)
 178           }
 179         })
 180         return workerCreated
 181       },
 182       opts.workerChoiceStrategy ?? WorkerChoiceStrategies.ROUND_ROBIN
 183     )
 184   }
 185
 186   private checkFilePath (filePath: string): void {
 187     if (!filePath) {
 188       throw new Error('Please specify a file with a worker implementation')
 189     }
 190   }
 191
 192   private checkNumberOfWorkers (numberOfWorkers: number): void {
 193     if (numberOfWorkers == null) {
 194       throw new Error(
 195         'Cannot instantiate a pool without specifying the number of workers'
 196       )
 197     } else if (!Number.isSafeInteger(numberOfWorkers)) {
 198       throw new Error(
 199         'Cannot instantiate a pool with a non integer number of workers'
 200       )
 201     } else if (numberOfWorkers < 0) {
 202       throw new Error(
 203         'Cannot instantiate a pool with a negative number of workers'
 204       )
 205     } else if (!this.dynamic && numberOfWorkers === 0) {
 206       throw new Error('Cannot instantiate a fixed pool with no worker')
 207     }
 208   }
 209
 210   /** @inheritdoc */
 211   public get dynamic (): boolean {
 212     return false
 213   }
 214
 215   /** @inheritdoc */
 216   public setWorkerChoiceStrategy (
 217     workerChoiceStrategy: WorkerChoiceStrategy
 218   ): void {
 219     this.opts.workerChoiceStrategy = workerChoiceStrategy
 220     this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
 221       workerChoiceStrategy
 222     )
 223   }
 224
 225   /** @inheritdoc */
 226   public execute (data: Data): Promise<Response> {
 227     // Configure worker to handle message with the specified task
 228     const worker = this.chooseWorker()
 229     this.increaseWorkersTask(worker)
 230     const messageId = ++this.nextMessageId
 231     const res = this.internalExecute(worker, messageId)
 232     this.sendToWorker(worker, { data: data || ({} as Data), id: messageId })
 233     return res
 234   }
 235
 236   /** @inheritdoc */
 237   public async destroy (): Promise<void> {
 238     await Promise.all(this.workers.map(worker => this.destroyWorker(worker)))
 239   }
 240
 241   /**
 242    * Shut down given worker.
 243    *
 244    * @param worker A worker within `workers`.
 245    */
 246   protected abstract destroyWorker (worker: Worker): void | Promise<void>
 247
 248   /**
 249    * Setup hook that can be overridden by a Poolifier pool implementation
 250    * to run code before workers are created in the abstract constructor.
 251    */
 252   protected setupHook (): void {
 253     // Can be overridden
 254   }
 255
 256   /**
 257    * Should return whether the worker is the main worker or not.
 258    */
 259   protected abstract isMain (): boolean
 260
 261   /**
 262    * Increase the number of tasks that the given workers has done.
 263    *
 264    * @param worker Worker whose tasks are increased.
 265    */
 266   protected increaseWorkersTask (worker: Worker): void {
 267     this.stepWorkerNumberOfTasks(worker, 1)
 268   }
 269
 270   /**
 271    * Decrease the number of tasks that the given workers has done.
 272    *
 273    * @param worker Worker whose tasks are decreased.
 274    */
 275   protected decreaseWorkersTasks (worker: Worker): void {
 276     this.stepWorkerNumberOfTasks(worker, -1)
 277   }
 278
 279   /**
 280    * Step the number of tasks that the given workers has done.
 281    *
 282    * @param worker Worker whose tasks are set.
 283    * @param step Worker number of tasks step.
 284    */
 285   private stepWorkerNumberOfTasks (worker: Worker, step: number): void {
 286     const numberOfTasksInProgress = this.tasks.get(worker)
 287     if (numberOfTasksInProgress !== undefined) {
 288       this.tasks.set(worker, numberOfTasksInProgress + step)
 289     } else {
 290       throw Error('Worker could not be found in tasks map')
 291     }
 292   }
 293
 294   /**
 295    * Removes the given worker from the pool.
 296    *
 297    * @param worker Worker that will be removed.
 298    */
 299   protected removeWorker (worker: Worker): void {
 300     // Clean worker from data structure
 301     const workerIndex = this.workers.indexOf(worker)
 302     this.workers.splice(workerIndex, 1)
 303     this.tasks.delete(worker)
 304   }
 305
 306   /**
 307    * Choose a worker for the next task.
 308    *
 309    * The default implementation uses a round robin algorithm to distribute the load.
 310    *
 311    * @returns Worker.
 312    */
 313   protected chooseWorker (): Worker {
 314     return this.workerChoiceStrategyContext.execute()
 315   }
 316
 317   /**
 318    * Send a message to the given worker.
 319    *
 320    * @param worker The worker which should receive the message.
 321    * @param message The message.
 322    */
 323   protected abstract sendToWorker (
 324     worker: Worker,
 325     message: MessageValue<Data>
 326   ): void
 327
 328   /**
 329    * Register a listener callback on a given worker.
 330    *
 331    * @param worker A worker.
 332    * @param listener A message listener callback.
 333    */
 334   protected abstract registerWorkerMessageListener<
 335     Message extends Data | Response
 336   > (worker: Worker, listener: (message: MessageValue<Message>) => void): void
 337
 338   protected internalExecute (
 339     worker: Worker,
 340     messageId: number
 341   ): Promise<Response> {
 342     return new Promise<Response>((resolve, reject) => {
 343       this.promiseMap.set(messageId, { resolve, reject, worker })
 344     })
 345   }
 346
 347   /**
 348    * Returns a newly created worker.
 349    */
 350   protected abstract createWorker (): Worker
 351
 352   /**
 353    * Function that can be hooked up when a worker has been newly created and moved to the workers registry.
 354    *
 355    * Can be used to update the `maxListeners` or binding the `main-worker`<->`worker` connection if not bind by default.
 356    *
 357    * @param worker The newly created worker.
 358    */
 359   protected abstract afterWorkerSetup (worker: Worker): void
 360
 361   /**
 362    * Creates a new worker for this pool and sets it up completely.
 363    *
 364    * @returns New, completely set up worker.
 365    */
 366   protected createAndSetupWorker (): Worker {
 367     const worker: Worker = this.createWorker()
 368
 369     worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
 370     worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
 371     worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
 372     worker.once('exit', () => this.removeWorker(worker))
 373
 374     this.workers.push(worker)
 375
 376     // Init tasks map
 377     this.tasks.set(worker, 0)
 378
 379     this.afterWorkerSetup(worker)
 380
 381     return worker
 382   }
 383
 384   /**
 385    * This function is the listener registered for each worker.
 386    *
 387    * @returns The listener function to execute when a message is sent from a worker.
 388    */
 389   protected workerListener (): (message: MessageValue<Response>) => void {
 390     return message => {
 391       if (message.id) {
 392         const value = this.promiseMap.get(message.id)
 393         if (value) {
 394           this.decreaseWorkersTasks(value.worker)
 395           if (message.error) value.reject(message.error)
 396           else value.resolve(message.data as Response)
 397           this.promiseMap.delete(message.id)
 398         }
 399       }
 400     }
 401   }
 402 }