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