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 === 'has') {
 359       response = this.hasTaskFunction(taskFunctionName as string)
 360     } else if (taskFunctionOperation === 'add') {
 361       response = this.addTaskFunction(
 362         taskFunctionName as string,
 363         // eslint-disable-next-line @typescript-eslint/no-implied-eval, no-new-func
 364         new Function(`return ${taskFunction as string}`)() as TaskFunction<
 365         Data,
 366         Response
 367         >
 368       )
 369     } else if (taskFunctionOperation === 'remove') {
 370       response = this.removeTaskFunction(taskFunctionName as string)
 371     } else if (taskFunctionOperation === 'default') {
 372       response = this.setDefaultTaskFunction(taskFunctionName as string)
 373     }
 374     this.sendToMainWorker({
 375       taskFunctionOperation,
 376       taskFunctionOperationStatus: response.status,
 377       workerError: {
 378         name: taskFunctionName as string,
 379         message: this.handleError(response.error as Error | string)
 380       },
 381       workerId: this.id
 382     })
 383   }
 384
 385   /**
 386    * Handles a kill message sent by the main worker.
 387    *
 388    * @param message - The kill message.
 389    */
 390   protected handleKillMessage (message: MessageValue<Data>): void {
 391     this.stopCheckActive()
 392     if (isAsyncFunction(this.opts.killHandler)) {
 393       (this.opts.killHandler?.() as Promise<void>)
 394         .then(() => {
 395           this.sendToMainWorker({ kill: 'success', workerId: this.id })
 396           return null
 397         })
 398         .catch(() => {
 399           this.sendToMainWorker({ kill: 'failure', workerId: this.id })
 400         })
 401         .finally(() => {
 402           this.emitDestroy()
 403         })
 404         .catch(EMPTY_FUNCTION)
 405     } else {
 406       try {
 407         // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
 408         this.opts.killHandler?.() as void
 409         this.sendToMainWorker({ kill: 'success', workerId: this.id })
 410       } catch {
 411         this.sendToMainWorker({ kill: 'failure', workerId: this.id })
 412       } finally {
 413         this.emitDestroy()
 414       }
 415     }
 416   }
 417
 418   /**
 419    * Check if the message worker id is set and matches the worker id.
 420    *
 421    * @param message - The message to check.
 422    * @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.
 423    */
 424   private checkMessageWorkerId (message: MessageValue<Data>): void {
 425     if (message.workerId == null) {
 426       throw new Error('Message worker id is not set')
 427     } else if (message.workerId != null && message.workerId !== this.id) {
 428       throw new Error(
 429         `Message worker id ${message.workerId} does not match the worker id ${this.id}`
 430       )
 431     }
 432   }
 433
 434   /**
 435    * Starts the worker check active interval.
 436    */
 437   private startCheckActive (): void {
 438     this.lastTaskTimestamp = performance.now()
 439     this.activeInterval = setInterval(
 440       this.checkActive.bind(this),
 441       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) / 2
 442     )
 443   }
 444
 445   /**
 446    * Stops the worker check active interval.
 447    */
 448   private stopCheckActive (): void {
 449     if (this.activeInterval != null) {
 450       clearInterval(this.activeInterval)
 451       delete this.activeInterval
 452     }
 453   }
 454
 455   /**
 456    * Checks if the worker should be terminated, because its living too long.
 457    */
 458   private checkActive (): void {
 459     if (
 460       performance.now() - this.lastTaskTimestamp >
 461       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME)
 462     ) {
 463       this.sendToMainWorker({ kill: this.opts.killBehavior, workerId: this.id })
 464     }
 465   }
 466
 467   /**
 468    * Returns the main worker.
 469    *
 470    * @returns Reference to the main worker.
 471    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the main worker is not set.
 472    */
 473   protected getMainWorker (): MainWorker {
 474     if (this.mainWorker == null) {
 475       throw new Error('Main worker not set')
 476     }
 477     return this.mainWorker
 478   }
 479
 480   /**
 481    * Sends a message to main worker.
 482    *
 483    * @param message - The response message.
 484    */
 485   protected abstract sendToMainWorker (
 486     message: MessageValue<Response, Data>
 487   ): void
 488
 489   /**
 490    * Sends the list of task function names to the main worker.
 491    */
 492   protected sendTaskFunctionsListToMainWorker (): void {
 493     this.sendToMainWorker({
 494       taskFunctionNames: this.listTaskFunctionNames(),
 495       workerId: this.id
 496     })
 497   }
 498
 499   /**
 500    * Handles an error and convert it to a string so it can be sent back to the main worker.
 501    *
 502    * @param error - The error raised by the worker.
 503    * @returns The error message.
 504    */
 505   protected handleError (error: Error | string): string {
 506     return error instanceof Error ? error.message : error
 507   }
 508
 509   /**
 510    * Runs the given task.
 511    *
 512    * @param task - The task to execute.
 513    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the task function is not found.
 514    */
 515   protected run (task: Task<Data>): void {
 516     const { name, taskId, data } = task
 517     const fn = this.taskFunctions.get(name ?? DEFAULT_TASK_NAME)
 518     if (fn == null) {
 519       this.sendToMainWorker({
 520         workerError: {
 521           name: name as string,
 522           message: `Task function '${name as string}' not found`,
 523           data
 524         },
 525         workerId: this.id,
 526         taskId
 527       })
 528       return
 529     }
 530     if (isAsyncFunction(fn)) {
 531       this.runInAsyncScope(this.runAsync.bind(this), this, fn, task)
 532     } else {
 533       this.runInAsyncScope(this.runSync.bind(this), this, fn, task)
 534     }
 535   }
 536
 537   /**
 538    * Runs the given task function synchronously.
 539    *
 540    * @param fn - Task function that will be executed.
 541    * @param task - Input data for the task function.
 542    */
 543   protected runSync (
 544     fn: TaskSyncFunction<Data, Response>,
 545     task: Task<Data>
 546   ): void {
 547     const { name, taskId, data } = task
 548     try {
 549       let taskPerformance = this.beginTaskPerformance(name)
 550       const res = fn(data)
 551       taskPerformance = this.endTaskPerformance(taskPerformance)
 552       this.sendToMainWorker({
 553         data: res,
 554         taskPerformance,
 555         workerId: this.id,
 556         taskId
 557       })
 558     } catch (error) {
 559       this.sendToMainWorker({
 560         workerError: {
 561           name: name as string,
 562           message: this.handleError(error as Error | string),
 563           data
 564         },
 565         workerId: this.id,
 566         taskId
 567       })
 568     } finally {
 569       this.updateLastTaskTimestamp()
 570     }
 571   }
 572
 573   /**
 574    * Runs the given task function asynchronously.
 575    *
 576    * @param fn - Task function that will be executed.
 577    * @param task - Input data for the task function.
 578    */
 579   protected runAsync (
 580     fn: TaskAsyncFunction<Data, Response>,
 581     task: Task<Data>
 582   ): void {
 583     const { name, taskId, data } = task
 584     let taskPerformance = this.beginTaskPerformance(name)
 585     fn(data)
 586       .then(res => {
 587         taskPerformance = this.endTaskPerformance(taskPerformance)
 588         this.sendToMainWorker({
 589           data: res,
 590           taskPerformance,
 591           workerId: this.id,
 592           taskId
 593         })
 594         return null
 595       })
 596       .catch(error => {
 597         this.sendToMainWorker({
 598           workerError: {
 599             name: name as string,
 600             message: this.handleError(error as Error | string),
 601             data
 602           },
 603           workerId: this.id,
 604           taskId
 605         })
 606       })
 607       .finally(() => {
 608         this.updateLastTaskTimestamp()
 609       })
 610       .catch(EMPTY_FUNCTION)
 611   }
 612
 613   private beginTaskPerformance (name?: string): TaskPerformance {
 614     this.checkStatistics()
 615     return {
 616       name: name ?? DEFAULT_TASK_NAME,
 617       timestamp: performance.now(),
 618       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 619     }
 620   }
 621
 622   private endTaskPerformance (
 623     taskPerformance: TaskPerformance
 624   ): TaskPerformance {
 625     this.checkStatistics()
 626     return {
 627       ...taskPerformance,
 628       ...(this.statistics.runTime && {
 629         runTime: performance.now() - taskPerformance.timestamp
 630       }),
 631       ...(this.statistics.elu && {
 632         elu: performance.eventLoopUtilization(taskPerformance.elu)
 633       })
 634     }
 635   }
 636
 637   private checkStatistics (): void {
 638     if (this.statistics == null) {
 639       throw new Error('Performance statistics computation requirements not set')
 640     }
 641   }
 642
 643   private updateLastTaskTimestamp (): void {
 644     if (this.activeInterval != null) {
 645       this.lastTaskTimestamp = performance.now()
 646     }
 647   }
 648 }