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