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