src/pools/abstract-pool.ts

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