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 no-new-func
 359           taskFunction: new Function(
 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     this.sendToMainWorker({
 381       taskFunctionOperation,
 382       taskFunctionOperationStatus: response.status,
 383       taskFunctionProperties,
 384       ...(!response.status &&
 385         response.error != null && {
 386         workerError: {
 387           name: taskFunctionProperties.name,
 388           message: this.handleError(response.error as Error | string),
 389         },
 390       }),
 391     })
 392   }
 393
 394   /**
 395    * Handles a kill message sent by the main worker.
 396    * @param message - The kill message.
 397    */
 398   protected handleKillMessage (message: MessageValue<Data>): void {
 399     this.stopCheckActive()
 400     if (isAsyncFunction(this.opts.killHandler)) {
 401       ;(this.opts.killHandler() as Promise<void>)
 402         .then(() => {
 403           this.sendToMainWorker({ kill: 'success' })
 404           return undefined
 405         })
 406         .catch(() => {
 407           this.sendToMainWorker({ kill: 'failure' })
 408         })
 409     } else {
 410       try {
 411         this.opts.killHandler?.()
 412         this.sendToMainWorker({ kill: 'success' })
 413       } catch {
 414         this.sendToMainWorker({ kill: 'failure' })
 415       }
 416     }
 417   }
 418
 419   /**
 420    * Check if the message worker id is set and matches the worker id.
 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 !== 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 })
 464     }
 465   }
 466
 467   /**
 468    * Returns the main worker.
 469    * @returns Reference to the main worker.
 470    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the main worker is not set.
 471    */
 472   protected getMainWorker (): MainWorker {
 473     if (this.mainWorker == null) {
 474       throw new Error('Main worker not set')
 475     }
 476     return this.mainWorker
 477   }
 478
 479   /**
 480    * Sends a message to main worker.
 481    * @param message - The response message.
 482    */
 483   protected abstract sendToMainWorker (
 484     message: MessageValue<Response, Data>
 485   ): void
 486
 487   /**
 488    * Sends task functions properties to the main worker.
 489    */
 490   protected sendTaskFunctionsPropertiesToMainWorker (): void {
 491     this.sendToMainWorker({
 492       taskFunctionsProperties: this.listTaskFunctionsProperties(),
 493     })
 494   }
 495
 496   /**
 497    * Handles an error and convert it to a string so it can be sent back to the main worker.
 498    * @param error - The error raised by the worker.
 499    * @returns The error message.
 500    */
 501   protected handleError (error: Error | string): string {
 502     return error instanceof Error ? error.message : error
 503   }
 504
 505   /**
 506    * Runs the given task.
 507    * @param task - The task to execute.
 508    */
 509   protected readonly run = (task: Task<Data>): void => {
 510     const { name, taskId, data } = task
 511     const taskFunctionName = name ?? DEFAULT_TASK_NAME
 512     if (!this.taskFunctions.has(taskFunctionName)) {
 513       this.sendToMainWorker({
 514         workerError: {
 515           // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 516           name: name!,
 517           message: `Task function '${name}' not found`,
 518           data,
 519         },
 520         taskId,
 521       })
 522       return
 523     }
 524     const fn = this.taskFunctions.get(taskFunctionName)?.taskFunction
 525     if (isAsyncFunction(fn)) {
 526       this.runAsync(fn as TaskAsyncFunction<Data, Response>, task)
 527     } else {
 528       this.runSync(fn as TaskSyncFunction<Data, Response>, task)
 529     }
 530   }
 531
 532   /**
 533    * Runs the given task function synchronously.
 534    * @param fn - Task function that will be executed.
 535    * @param task - Input data for the task function.
 536    */
 537   protected readonly runSync = (
 538     fn: TaskSyncFunction<Data, Response>,
 539     task: Task<Data>
 540   ): void => {
 541     const { name, taskId, data } = task
 542     try {
 543       let taskPerformance = this.beginTaskPerformance(name)
 544       const res = fn(data)
 545       taskPerformance = this.endTaskPerformance(taskPerformance)
 546       this.sendToMainWorker({
 547         data: res,
 548         taskPerformance,
 549         taskId,
 550       })
 551     } catch (error) {
 552       this.sendToMainWorker({
 553         workerError: {
 554           // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 555           name: name!,
 556           message: this.handleError(error as Error | string),
 557           data,
 558         },
 559         taskId,
 560       })
 561     } finally {
 562       this.updateLastTaskTimestamp()
 563     }
 564   }
 565
 566   /**
 567    * Runs the given task function asynchronously.
 568    * @param fn - Task function that will be executed.
 569    * @param task - Input data for the task function.
 570    */
 571   protected readonly runAsync = (
 572     fn: TaskAsyncFunction<Data, Response>,
 573     task: Task<Data>
 574   ): void => {
 575     const { name, taskId, data } = task
 576     let taskPerformance = this.beginTaskPerformance(name)
 577     fn(data)
 578       .then(res => {
 579         taskPerformance = this.endTaskPerformance(taskPerformance)
 580         this.sendToMainWorker({
 581           data: res,
 582           taskPerformance,
 583           taskId,
 584         })
 585         return undefined
 586       })
 587       .catch((error: unknown) => {
 588         this.sendToMainWorker({
 589           workerError: {
 590             // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 591             name: name!,
 592             message: this.handleError(error as Error | string),
 593             data,
 594           },
 595           taskId,
 596         })
 597       })
 598       .finally(() => {
 599         this.updateLastTaskTimestamp()
 600       })
 601       .catch(EMPTY_FUNCTION)
 602   }
 603
 604   private beginTaskPerformance (name?: string): TaskPerformance {
 605     if (this.statistics == null) {
 606       throw new Error('Performance statistics computation requirements not set')
 607     }
 608     return {
 609       name: name ?? DEFAULT_TASK_NAME,
 610       timestamp: performance.now(),
 611       ...(this.statistics.elu && {
 612         elu: performance.eventLoopUtilization(),
 613       }),
 614     }
 615   }
 616
 617   private endTaskPerformance (
 618     taskPerformance: TaskPerformance
 619   ): TaskPerformance {
 620     if (this.statistics == null) {
 621       throw new Error('Performance statistics computation requirements not set')
 622     }
 623     return {
 624       ...taskPerformance,
 625       ...(this.statistics.runTime && {
 626         runTime: performance.now() - taskPerformance.timestamp,
 627       }),
 628       ...(this.statistics.elu && {
 629         elu: performance.eventLoopUtilization(taskPerformance.elu),
 630       }),
 631     }
 632   }
 633
 634   private updateLastTaskTimestamp (): void {
 635     if (this.activeInterval != null) {
 636       this.lastTaskTimestamp = performance.now()
 637     }
 638   }
 639 }