src/worker/abstract-worker.ts

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