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   // eslint-disable-next-line no-undef
  82   protected activeInterval?: NodeJS.Timeout
  83
  84   /**
  85    * Constructs a new poolifier worker.
  86    * @param isMain - Whether this is the main worker or not.
  87    * @param mainWorker - Reference to main worker.
  88    * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked. The first function is the default function.
  89    * @param opts - Options for the worker.
  90    */
  91   public constructor (
  92     protected readonly isMain: boolean | undefined,
  93     private readonly mainWorker: MainWorker | undefined | null,
  94     taskFunctions: TaskFunction<Data, Response> | TaskFunctions<Data, Response>,
  95     protected opts: WorkerOptions = DEFAULT_WORKER_OPTIONS
  96   ) {
  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       // Should be once() but Node.js on windows has a bug that prevents it from working
 104       this.getMainWorker().on('message', this.handleReadyMessage.bind(this))
 105     }
 106   }
 107
 108   private checkWorkerOptions (opts: WorkerOptions): void {
 109     checkValidWorkerOptions(opts)
 110     this.opts = { ...DEFAULT_WORKER_OPTIONS, ...opts }
 111   }
 112
 113   /**
 114    * Checks if the `taskFunctions` parameter is passed to the constructor and valid.
 115    * @param taskFunctions - The task function(s) parameter that should be checked.
 116    */
 117   private checkTaskFunctions (
 118     taskFunctions:
 119       | TaskFunction<Data, Response>
 120       | TaskFunctions<Data, Response>
 121       | undefined
 122   ): void {
 123     if (taskFunctions == null) {
 124       throw new Error('taskFunctions parameter is mandatory')
 125     }
 126     this.taskFunctions = new Map<string, TaskFunctionObject<Data, Response>>()
 127     if (typeof taskFunctions === 'function') {
 128       const fnObj = { taskFunction: taskFunctions.bind(this) }
 129       this.taskFunctions.set(DEFAULT_TASK_NAME, fnObj)
 130       this.taskFunctions.set(
 131         typeof taskFunctions.name === 'string' &&
 132           taskFunctions.name.trim().length > 0
 133           ? taskFunctions.name
 134           : 'fn1',
 135         fnObj
 136       )
 137     } else if (isPlainObject(taskFunctions)) {
 138       let firstEntry = true
 139       for (let [name, fnObj] of Object.entries(taskFunctions)) {
 140         if (typeof fnObj === 'function') {
 141           fnObj = { taskFunction: fnObj } satisfies TaskFunctionObject<
 142             Data,
 143             Response
 144           >
 145         }
 146         checkValidTaskFunctionObjectEntry<Data, Response>(name, fnObj)
 147         fnObj.taskFunction = fnObj.taskFunction.bind(this)
 148         if (firstEntry) {
 149           this.taskFunctions.set(DEFAULT_TASK_NAME, fnObj)
 150           firstEntry = false
 151         }
 152         this.taskFunctions.set(name, fnObj)
 153       }
 154       if (firstEntry) {
 155         throw new Error('taskFunctions parameter object is empty')
 156       }
 157     } else {
 158       throw new TypeError(
 159         'taskFunctions parameter is not a function or a plain object'
 160       )
 161     }
 162   }
 163
 164   /**
 165    * Checks if the worker has a task function with the given name.
 166    * @param name - The name of the task function to check.
 167    * @returns Whether the worker has a task function with the given name or not.
 168    */
 169   public hasTaskFunction (name: string): TaskFunctionOperationResult {
 170     try {
 171       checkTaskFunctionName(name)
 172     } catch (error) {
 173       return { status: false, error: error as Error }
 174     }
 175     return { status: this.taskFunctions.has(name) }
 176   }
 177
 178   /**
 179    * Adds a task function to the worker.
 180    * If a task function with the same name already exists, it is replaced.
 181    * @param name - The name of the task function to add.
 182    * @param fn - The task function to add.
 183    * @returns Whether the task function was added or not.
 184    */
 185   public addTaskFunction (
 186     name: string,
 187     fn: TaskFunction<Data, Response> | TaskFunctionObject<Data, Response>
 188   ): TaskFunctionOperationResult {
 189     try {
 190       checkTaskFunctionName(name)
 191       if (name === DEFAULT_TASK_NAME) {
 192         throw new Error(
 193           'Cannot add a task function with the default reserved name'
 194         )
 195       }
 196       if (typeof fn === 'function') {
 197         fn = { taskFunction: fn } satisfies TaskFunctionObject<Data, Response>
 198       }
 199       checkValidTaskFunctionObjectEntry<Data, Response>(name, fn)
 200       fn.taskFunction = fn.taskFunction.bind(this)
 201       if (
 202         this.taskFunctions.get(name) ===
 203         this.taskFunctions.get(DEFAULT_TASK_NAME)
 204       ) {
 205         this.taskFunctions.set(DEFAULT_TASK_NAME, fn)
 206       }
 207       this.taskFunctions.set(name, fn)
 208       this.sendTaskFunctionsPropertiesToMainWorker()
 209       return { status: true }
 210     } catch (error) {
 211       return { status: false, error: error as Error }
 212     }
 213   }
 214
 215   /**
 216    * Removes a task function from the worker.
 217    * @param name - The name of the task function to remove.
 218    * @returns Whether the task function existed and was removed or not.
 219    */
 220   public removeTaskFunction (name: string): TaskFunctionOperationResult {
 221     try {
 222       checkTaskFunctionName(name)
 223       if (name === DEFAULT_TASK_NAME) {
 224         throw new Error(
 225           'Cannot remove the task function with the default reserved name'
 226         )
 227       }
 228       if (
 229         this.taskFunctions.get(name) ===
 230         this.taskFunctions.get(DEFAULT_TASK_NAME)
 231       ) {
 232         throw new Error(
 233           'Cannot remove the task function used as the default task function'
 234         )
 235       }
 236       const deleteStatus = this.taskFunctions.delete(name)
 237       this.sendTaskFunctionsPropertiesToMainWorker()
 238       return { status: deleteStatus }
 239     } catch (error) {
 240       return { status: false, error: error as Error }
 241     }
 242   }
 243
 244   /**
 245    * Lists the properties of the worker's task functions.
 246    * @returns The properties of the worker's task functions.
 247    */
 248   public listTaskFunctionsProperties (): TaskFunctionProperties[] {
 249     let defaultTaskFunctionName = DEFAULT_TASK_NAME
 250     for (const [name, fnObj] of this.taskFunctions) {
 251       if (
 252         name !== DEFAULT_TASK_NAME &&
 253         fnObj === this.taskFunctions.get(DEFAULT_TASK_NAME)
 254       ) {
 255         defaultTaskFunctionName = name
 256         break
 257       }
 258     }
 259     const taskFunctionsProperties: TaskFunctionProperties[] = []
 260     for (const [name, fnObj] of this.taskFunctions) {
 261       if (name === DEFAULT_TASK_NAME || name === defaultTaskFunctionName) {
 262         continue
 263       }
 264       taskFunctionsProperties.push(buildTaskFunctionProperties(name, fnObj))
 265     }
 266     return [
 267       buildTaskFunctionProperties(
 268         DEFAULT_TASK_NAME,
 269         this.taskFunctions.get(DEFAULT_TASK_NAME)
 270       ),
 271       buildTaskFunctionProperties(
 272         defaultTaskFunctionName,
 273         this.taskFunctions.get(defaultTaskFunctionName)
 274       ),
 275       ...taskFunctionsProperties,
 276     ]
 277   }
 278
 279   /**
 280    * Sets the default task function to use in the worker.
 281    * @param name - The name of the task function to use as default task function.
 282    * @returns Whether the default task function was set or not.
 283    */
 284   public setDefaultTaskFunction (name: string): TaskFunctionOperationResult {
 285     try {
 286       checkTaskFunctionName(name)
 287       if (name === DEFAULT_TASK_NAME) {
 288         throw new Error(
 289           'Cannot set the default task function reserved name as the default task function'
 290         )
 291       }
 292       if (!this.taskFunctions.has(name)) {
 293         throw new Error(
 294           'Cannot set the default task function to a non-existing task function'
 295         )
 296       }
 297       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 298       this.taskFunctions.set(DEFAULT_TASK_NAME, this.taskFunctions.get(name)!)
 299       this.sendTaskFunctionsPropertiesToMainWorker()
 300       return { status: true }
 301     } catch (error) {
 302       return { status: false, error: error as Error }
 303     }
 304   }
 305
 306   /**
 307    * Handles the ready message sent by the main worker.
 308    * @param message - The ready message.
 309    */
 310   protected abstract handleReadyMessage (message: MessageValue<Data>): void
 311
 312   /**
 313    * Worker message listener.
 314    * @param message - The received message.
 315    */
 316   protected messageListener (message: MessageValue<Data>): void {
 317     this.checkMessageWorkerId(message)
 318     const {
 319       statistics,
 320       checkActive,
 321       taskFunctionOperation,
 322       taskId,
 323       data,
 324       kill,
 325     } = message
 326     if (statistics != null) {
 327       // Statistics message received
 328       this.statistics = statistics
 329     } else if (checkActive != null) {
 330       // Check active message received
 331       checkActive ? this.startCheckActive() : this.stopCheckActive()
 332     } else if (taskFunctionOperation != null) {
 333       // Task function operation message received
 334       this.handleTaskFunctionOperationMessage(message)
 335     } else if (taskId != null && data != null) {
 336       // Task message received
 337       this.run(message)
 338     } else if (kill === true) {
 339       // Kill message received
 340       this.handleKillMessage(message)
 341     }
 342   }
 343
 344   protected handleTaskFunctionOperationMessage (
 345     message: MessageValue<Data>
 346   ): void {
 347     const { taskFunctionOperation, taskFunctionProperties, taskFunction } =
 348       message
 349     if (taskFunctionProperties == null) {
 350       throw new Error(
 351         'Cannot handle task function operation message without task function properties'
 352       )
 353     }
 354     let response: TaskFunctionOperationResult
 355     switch (taskFunctionOperation) {
 356       case 'add':
 357         response = this.addTaskFunction(taskFunctionProperties.name, {
 358           // eslint-disable-next-line @typescript-eslint/no-implied-eval, no-new-func
 359           taskFunction: new Function(
 360             // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 361             `return ${taskFunction!}`
 362           )() as TaskFunction<Data, Response>,
 363           ...(taskFunctionProperties.priority != null && {
 364             priority: taskFunctionProperties.priority,
 365           }),
 366           ...(taskFunctionProperties.strategy != null && {
 367             strategy: taskFunctionProperties.strategy,
 368           }),
 369         })
 370         break
 371       case 'remove':
 372         response = this.removeTaskFunction(taskFunctionProperties.name)
 373         break
 374       case 'default':
 375         response = this.setDefaultTaskFunction(taskFunctionProperties.name)
 376         break
 377       default:
 378         response = { status: false, error: new Error('Unknown task operation') }
 379         break
 380     }
 381     this.sendToMainWorker({
 382       taskFunctionOperation,
 383       taskFunctionOperationStatus: response.status,
 384       taskFunctionProperties,
 385       ...(!response.status &&
 386         response.error != null && {
 387         workerError: {
 388           name: taskFunctionProperties.name,
 389           message: this.handleError(response.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 }