src/worker/abstract-worker.ts

   1 import { AsyncResource } from 'node:async_hooks'
   2 import type { Worker } from 'node:cluster'
   3 import type { MessagePort } from 'node:worker_threads'
   4 import { performance } from 'node:perf_hooks'
   5 import type {
   6   MessageValue,
   7   Task,
   8   TaskPerformance,
   9   WorkerStatistics
  10 } from '../utility-types'
  11 import {
  12   DEFAULT_TASK_NAME,
  13   EMPTY_FUNCTION,
  14   isAsyncFunction,
  15   isPlainObject
  16 } from '../utils'
  17 import { KillBehaviors, type WorkerOptions } from './worker-options'
  18 import type {
  19   TaskAsyncFunction,
  20   TaskFunction,
  21   TaskFunctions,
  22   TaskSyncFunction
  23 } from './task-functions'
  24
  25 const DEFAULT_MAX_INACTIVE_TIME = 60000
  26 const DEFAULT_WORKER_OPTIONS: WorkerOptions = {
  27   /**
  28    * The kill behavior option on this worker or its default value.
  29    */
  30   killBehavior: KillBehaviors.SOFT,
  31   /**
  32    * The maximum time to keep this worker active while idle.
  33    * The pool automatically checks and terminates this worker when the time expires.
  34    */
  35   maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME,
  36   /**
  37    * The function to call when the worker is killed.
  38    */
  39   killHandler: EMPTY_FUNCTION
  40 }
  41
  42 /**
  43  * Base class that implements some shared logic for all poolifier workers.
  44  *
  45  * @typeParam MainWorker - Type of main worker.
  46  * @typeParam Data - Type of data this worker receives from pool's execution. This can only be structured-cloneable data.
  47  * @typeParam Response - Type of response the worker sends back to the main worker. This can only be structured-cloneable data.
  48  */
  49 export abstract class AbstractWorker<
  50   MainWorker extends Worker | MessagePort,
  51   Data = unknown,
  52   Response = unknown
  53 > extends AsyncResource {
  54   /**
  55    * Worker id.
  56    */
  57   protected abstract id: number
  58   /**
  59    * Task function(s) processed by the worker when the pool's `execution` function is invoked.
  60    */
  61   protected taskFunctions!: Map<string, TaskFunction<Data, Response>>
  62   /**
  63    * Timestamp of the last task processed by this worker.
  64    */
  65   protected lastTaskTimestamp!: number
  66   /**
  67    * Performance statistics computation requirements.
  68    */
  69   protected statistics!: WorkerStatistics
  70   /**
  71    * Handler id of the `activeInterval` worker activity check.
  72    */
  73   protected activeInterval?: NodeJS.Timeout
  74   /**
  75    * Constructs a new poolifier worker.
  76    *
  77    * @param type - The type of async event.
  78    * @param isMain - Whether this is the main worker or not.
  79    * @param mainWorker - Reference to main worker.
  80    * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked. The first function is the default function.
  81    * @param opts - Options for the worker.
  82    */
  83   public constructor (
  84     type: string,
  85     protected readonly isMain: boolean,
  86     private readonly mainWorker: MainWorker,
  87     taskFunctions: TaskFunction<Data, Response> | TaskFunctions<Data, Response>,
  88     protected opts: WorkerOptions = DEFAULT_WORKER_OPTIONS
  89   ) {
  90     super(type)
  91     if (this.isMain == null) {
  92       throw new Error('isMain parameter is mandatory')
  93     }
  94     this.checkTaskFunctions(taskFunctions)
  95     this.checkWorkerOptions(this.opts)
  96     if (!this.isMain) {
  97       this.getMainWorker().on('message', this.handleReadyMessage.bind(this))
  98     }
  99   }
 100
 101   private checkWorkerOptions (opts: WorkerOptions): void {
 102     if (opts != null && !isPlainObject(opts)) {
 103       throw new TypeError('opts worker options parameter is not a plain object')
 104     }
 105     if (
 106       opts?.killBehavior != null &&
 107       !Object.values(KillBehaviors).includes(opts.killBehavior)
 108     ) {
 109       throw new TypeError(
 110         `killBehavior option '${opts.killBehavior}' is not valid`
 111       )
 112     }
 113     if (
 114       opts?.maxInactiveTime != null &&
 115       !Number.isSafeInteger(opts.maxInactiveTime)
 116     ) {
 117       throw new TypeError('maxInactiveTime option is not an integer')
 118     }
 119     if (opts?.maxInactiveTime != null && opts.maxInactiveTime < 5) {
 120       throw new TypeError(
 121         'maxInactiveTime option is not a positive integer greater or equal than 5'
 122       )
 123     }
 124     if (opts?.killHandler != null && typeof opts.killHandler !== 'function') {
 125       throw new TypeError('killHandler option is not a function')
 126     }
 127     if (opts?.async != null) {
 128       throw new Error('async option is deprecated')
 129     }
 130     this.opts = { ...DEFAULT_WORKER_OPTIONS, ...opts }
 131   }
 132
 133   private checkValidTaskFunction (
 134     name: string,
 135     fn: TaskFunction<Data, Response>
 136   ): void {
 137     if (typeof name !== 'string') {
 138       throw new TypeError(
 139         'A taskFunctions parameter object key is not a string'
 140       )
 141     }
 142     if (typeof name === 'string' && name.trim().length === 0) {
 143       throw new TypeError(
 144         'A taskFunctions parameter object key is an empty string'
 145       )
 146     }
 147     if (typeof fn !== 'function') {
 148       throw new TypeError(
 149         'A taskFunctions parameter object value is not a function'
 150       )
 151     }
 152   }
 153
 154   /**
 155    * Checks if the `taskFunctions` parameter is passed to the constructor and valid.
 156    *
 157    * @param taskFunctions - The task function(s) parameter that should be checked.
 158    */
 159   private checkTaskFunctions (
 160     taskFunctions: TaskFunction<Data, Response> | TaskFunctions<Data, Response>
 161   ): void {
 162     if (taskFunctions == null) {
 163       throw new Error('taskFunctions parameter is mandatory')
 164     }
 165     this.taskFunctions = new Map<string, TaskFunction<Data, Response>>()
 166     if (typeof taskFunctions === 'function') {
 167       const boundFn = taskFunctions.bind(this)
 168       this.taskFunctions.set(DEFAULT_TASK_NAME, boundFn)
 169       this.taskFunctions.set(
 170         typeof taskFunctions.name === 'string' &&
 171         taskFunctions.name.trim().length > 0
 172           ? taskFunctions.name
 173           : 'fn1',
 174         boundFn
 175       )
 176     } else if (isPlainObject(taskFunctions)) {
 177       let firstEntry = true
 178       for (const [name, fn] of Object.entries(taskFunctions)) {
 179         this.checkValidTaskFunction(name, fn)
 180         const boundFn = fn.bind(this)
 181         if (firstEntry) {
 182           this.taskFunctions.set(DEFAULT_TASK_NAME, boundFn)
 183           firstEntry = false
 184         }
 185         this.taskFunctions.set(name, boundFn)
 186       }
 187       if (firstEntry) {
 188         throw new Error('taskFunctions parameter object is empty')
 189       }
 190     } else {
 191       throw new TypeError(
 192         'taskFunctions parameter is not a function or a plain object'
 193       )
 194     }
 195   }
 196
 197   /**
 198    * Checks if the worker has a task function with the given name.
 199    *
 200    * @param name - The name of the task function to check.
 201    * @returns Whether the worker has a task function with the given name or not.
 202    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string or an empty string.
 203    */
 204   public hasTaskFunction (name: string): boolean {
 205     this.checkTaskFunctionName(name)
 206     return this.taskFunctions.has(name)
 207   }
 208
 209   /**
 210    * Adds a task function to the worker.
 211    * If a task function with the same name already exists, it is replaced.
 212    *
 213    * @param name - The name of the task function to add.
 214    * @param fn - The task function to add.
 215    * @returns Whether the task function was added or not.
 216    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string or an empty string.
 217    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 218    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `fn` parameter is not a function.
 219    */
 220   public addTaskFunction (
 221     name: string,
 222     fn: TaskFunction<Data, Response>
 223   ): boolean {
 224     this.checkTaskFunctionName(name)
 225     if (name === DEFAULT_TASK_NAME) {
 226       throw new Error(
 227         'Cannot add a task function with the default reserved name'
 228       )
 229     }
 230     if (typeof fn !== 'function') {
 231       throw new TypeError('fn parameter is not a function')
 232     }
 233     try {
 234       const boundFn = fn.bind(this)
 235       if (
 236         this.taskFunctions.get(name) ===
 237         this.taskFunctions.get(DEFAULT_TASK_NAME)
 238       ) {
 239         this.taskFunctions.set(DEFAULT_TASK_NAME, boundFn)
 240       }
 241       this.taskFunctions.set(name, boundFn)
 242       this.sendTaskFunctionsListToMainWorker()
 243       return true
 244     } catch {
 245       return false
 246     }
 247   }
 248
 249   /**
 250    * Removes a task function from the worker.
 251    *
 252    * @param name - The name of the task function to remove.
 253    * @returns Whether the task function existed and was removed or not.
 254    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string or an empty string.
 255    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 256    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the task function used as default task function.
 257    */
 258   public removeTaskFunction (name: string): boolean {
 259     this.checkTaskFunctionName(name)
 260     if (name === DEFAULT_TASK_NAME) {
 261       throw new Error(
 262         'Cannot remove the task function with the default reserved name'
 263       )
 264     }
 265     if (
 266       this.taskFunctions.get(name) === this.taskFunctions.get(DEFAULT_TASK_NAME)
 267     ) {
 268       throw new Error(
 269         'Cannot remove the task function used as the default task function'
 270       )
 271     }
 272     const deleteStatus = this.taskFunctions.delete(name)
 273     this.sendTaskFunctionsListToMainWorker()
 274     return deleteStatus
 275   }
 276
 277   /**
 278    * Lists the names of the worker's task functions.
 279    *
 280    * @returns The names of the worker's task functions.
 281    */
 282   public listTaskFunctions (): string[] {
 283     const names: string[] = [...this.taskFunctions.keys()]
 284     let defaultTaskFunctionName: string = DEFAULT_TASK_NAME
 285     for (const [name, fn] of this.taskFunctions) {
 286       if (
 287         name !== DEFAULT_TASK_NAME &&
 288         fn === this.taskFunctions.get(DEFAULT_TASK_NAME)
 289       ) {
 290         defaultTaskFunctionName = name
 291         break
 292       }
 293     }
 294     return [
 295       names[names.indexOf(DEFAULT_TASK_NAME)],
 296       defaultTaskFunctionName,
 297       ...names.filter(
 298         name => name !== DEFAULT_TASK_NAME && name !== defaultTaskFunctionName
 299       )
 300     ]
 301   }
 302
 303   /**
 304    * Sets the default task function to use in the worker.
 305    *
 306    * @param name - The name of the task function to use as default task function.
 307    * @returns Whether the default task function was set or not.
 308    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string or an empty string.
 309    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 310    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is a non-existing task function.
 311    */
 312   public setDefaultTaskFunction (name: string): boolean {
 313     this.checkTaskFunctionName(name)
 314     if (name === DEFAULT_TASK_NAME) {
 315       throw new Error(
 316         'Cannot set the default task function reserved name as the default task function'
 317       )
 318     }
 319     if (!this.taskFunctions.has(name)) {
 320       throw new Error(
 321         'Cannot set the default task function to a non-existing task function'
 322       )
 323     }
 324     try {
 325       this.taskFunctions.set(
 326         DEFAULT_TASK_NAME,
 327         this.taskFunctions.get(name) as TaskFunction<Data, Response>
 328       )
 329       return true
 330     } catch {
 331       return false
 332     }
 333   }
 334
 335   private checkTaskFunctionName (name: string): void {
 336     if (typeof name !== 'string') {
 337       throw new TypeError('name parameter is not a string')
 338     }
 339     if (typeof name === 'string' && name.trim().length === 0) {
 340       throw new TypeError('name parameter is an empty string')
 341     }
 342   }
 343
 344   /**
 345    * Handles the ready message sent by the main worker.
 346    *
 347    * @param message - The ready message.
 348    */
 349   protected abstract handleReadyMessage (message: MessageValue<Data>): void
 350
 351   /**
 352    * Worker message listener.
 353    *
 354    * @param message - The received message.
 355    */
 356   protected messageListener (message: MessageValue<Data>): void {
 357     this.checkMessageWorkerId(message)
 358     if (message.statistics != null) {
 359       // Statistics message received
 360       this.statistics = message.statistics
 361     } else if (message.checkActive != null) {
 362       // Check active message received
 363       message.checkActive ? this.startCheckActive() : this.stopCheckActive()
 364     } else if (message.taskId != null && message.data != null) {
 365       // Task message received
 366       this.run(message)
 367     } else if (message.kill === true) {
 368       // Kill message received
 369       this.handleKillMessage(message)
 370     }
 371   }
 372
 373   /**
 374    * Handles a kill message sent by the main worker.
 375    *
 376    * @param message - The kill message.
 377    */
 378   protected handleKillMessage (message: MessageValue<Data>): void {
 379     this.stopCheckActive()
 380     if (isAsyncFunction(this.opts.killHandler)) {
 381       (this.opts.killHandler?.() as Promise<void>)
 382         .then(() => {
 383           this.sendToMainWorker({ kill: 'success', workerId: this.id })
 384           return null
 385         })
 386         .catch(() => {
 387           this.sendToMainWorker({ kill: 'failure', workerId: this.id })
 388         })
 389         .finally(() => {
 390           this.emitDestroy()
 391         })
 392         .catch(EMPTY_FUNCTION)
 393     } else {
 394       try {
 395         // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
 396         this.opts.killHandler?.() as void
 397         this.sendToMainWorker({ kill: 'success', workerId: this.id })
 398       } catch {
 399         this.sendToMainWorker({ kill: 'failure', workerId: this.id })
 400       } finally {
 401         this.emitDestroy()
 402       }
 403     }
 404   }
 405
 406   /**
 407    * Check if the message worker id is set and matches the worker id.
 408    *
 409    * @param message - The message to check.
 410    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the message worker id is not set or does not match the worker id.
 411    */
 412   private checkMessageWorkerId (message: MessageValue<Data>): void {
 413     if (message.workerId == null) {
 414       throw new Error('Message worker id is not set')
 415     } else if (message.workerId != null && message.workerId !== this.id) {
 416       throw new Error(
 417         `Message worker id ${message.workerId} does not match the worker id ${this.id}`
 418       )
 419     }
 420   }
 421
 422   /**
 423    * Starts the worker check active interval.
 424    */
 425   private startCheckActive (): void {
 426     this.lastTaskTimestamp = performance.now()
 427     this.activeInterval = setInterval(
 428       this.checkActive.bind(this),
 429       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) / 2
 430     )
 431   }
 432
 433   /**
 434    * Stops the worker check active interval.
 435    */
 436   private stopCheckActive (): void {
 437     if (this.activeInterval != null) {
 438       clearInterval(this.activeInterval)
 439       delete this.activeInterval
 440     }
 441   }
 442
 443   /**
 444    * Checks if the worker should be terminated, because its living too long.
 445    */
 446   private checkActive (): void {
 447     if (
 448       performance.now() - this.lastTaskTimestamp >
 449       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME)
 450     ) {
 451       this.sendToMainWorker({ kill: this.opts.killBehavior, workerId: this.id })
 452     }
 453   }
 454
 455   /**
 456    * Returns the main worker.
 457    *
 458    * @returns Reference to the main worker.
 459    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the main worker is not set.
 460    */
 461   protected getMainWorker (): MainWorker {
 462     if (this.mainWorker == null) {
 463       throw new Error('Main worker not set')
 464     }
 465     return this.mainWorker
 466   }
 467
 468   /**
 469    * Sends a message to main worker.
 470    *
 471    * @param message - The response message.
 472    */
 473   protected abstract sendToMainWorker (
 474     message: MessageValue<Response, Data>
 475   ): void
 476
 477   /**
 478    * Sends the list of task function names to the main worker.
 479    */
 480   protected sendTaskFunctionsListToMainWorker (): void {
 481     this.sendToMainWorker({
 482       taskFunctions: this.listTaskFunctions(),
 483       workerId: this.id
 484     })
 485   }
 486
 487   /**
 488    * Handles an error and convert it to a string so it can be sent back to the main worker.
 489    *
 490    * @param error - The error raised by the worker.
 491    * @returns The error message.
 492    */
 493   protected handleError (error: Error | string): string {
 494     return error instanceof Error ? error.message : error
 495   }
 496
 497   /**
 498    * Runs the given task.
 499    *
 500    * @param task - The task to execute.
 501    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the task function is not found.
 502    */
 503   protected run (task: Task<Data>): void {
 504     const { name, taskId, data } = task
 505     const fn = this.taskFunctions.get(name ?? DEFAULT_TASK_NAME)
 506     if (fn == null) {
 507       this.sendToMainWorker({
 508         taskError: {
 509           name: name as string,
 510           message: `Task function '${name as string}' not found`,
 511           data
 512         },
 513         workerId: this.id,
 514         taskId
 515       })
 516       return
 517     }
 518     if (isAsyncFunction(fn)) {
 519       this.runInAsyncScope(this.runAsync.bind(this), this, fn, task)
 520     } else {
 521       this.runInAsyncScope(this.runSync.bind(this), this, fn, task)
 522     }
 523   }
 524
 525   /**
 526    * Runs the given task function synchronously.
 527    *
 528    * @param fn - Task function that will be executed.
 529    * @param task - Input data for the task function.
 530    */
 531   protected runSync (
 532     fn: TaskSyncFunction<Data, Response>,
 533     task: Task<Data>
 534   ): void {
 535     const { name, taskId, data } = task
 536     try {
 537       let taskPerformance = this.beginTaskPerformance(name)
 538       const res = fn(data)
 539       taskPerformance = this.endTaskPerformance(taskPerformance)
 540       this.sendToMainWorker({
 541         data: res,
 542         taskPerformance,
 543         workerId: this.id,
 544         taskId
 545       })
 546     } catch (error) {
 547       this.sendToMainWorker({
 548         taskError: {
 549           name: name as string,
 550           message: this.handleError(error as Error | string),
 551           data
 552         },
 553         workerId: this.id,
 554         taskId
 555       })
 556     } finally {
 557       this.updateLastTaskTimestamp()
 558     }
 559   }
 560
 561   /**
 562    * Runs the given task function asynchronously.
 563    *
 564    * @param fn - Task function that will be executed.
 565    * @param task - Input data for the task function.
 566    */
 567   protected runAsync (
 568     fn: TaskAsyncFunction<Data, Response>,
 569     task: Task<Data>
 570   ): void {
 571     const { name, taskId, data } = task
 572     let taskPerformance = this.beginTaskPerformance(name)
 573     fn(data)
 574       .then(res => {
 575         taskPerformance = this.endTaskPerformance(taskPerformance)
 576         this.sendToMainWorker({
 577           data: res,
 578           taskPerformance,
 579           workerId: this.id,
 580           taskId
 581         })
 582         return null
 583       })
 584       .catch(error => {
 585         this.sendToMainWorker({
 586           taskError: {
 587             name: name as string,
 588             message: this.handleError(error as Error | string),
 589             data
 590           },
 591           workerId: this.id,
 592           taskId
 593         })
 594       })
 595       .finally(() => {
 596         this.updateLastTaskTimestamp()
 597       })
 598       .catch(EMPTY_FUNCTION)
 599   }
 600
 601   private beginTaskPerformance (name?: string): TaskPerformance {
 602     this.checkStatistics()
 603     return {
 604       name: name ?? DEFAULT_TASK_NAME,
 605       timestamp: performance.now(),
 606       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 607     }
 608   }
 609
 610   private endTaskPerformance (
 611     taskPerformance: TaskPerformance
 612   ): TaskPerformance {
 613     this.checkStatistics()
 614     return {
 615       ...taskPerformance,
 616       ...(this.statistics.runTime && {
 617         runTime: performance.now() - taskPerformance.timestamp
 618       }),
 619       ...(this.statistics.elu && {
 620         elu: performance.eventLoopUtilization(taskPerformance.elu)
 621       })
 622     }
 623   }
 624
 625   private checkStatistics (): void {
 626     if (this.statistics == null) {
 627       throw new Error('Performance statistics computation requirements not set')
 628     }
 629   }
 630
 631   private updateLastTaskTimestamp (): void {
 632     if (this.activeInterval != null) {
 633       this.lastTaskTimestamp = performance.now()
 634     }
 635   }
 636 }