src/worker/abstract-worker.ts

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