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