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