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       this.taskFunctions.set(
 281         DEFAULT_TASK_NAME,
 282         this.taskFunctions.get(name) as TaskFunction<Data, Response>
 283       )
 284       this.sendTaskFunctionNamesToMainWorker()
 285       return { status: true }
 286     } catch (error) {
 287       return { status: false, error: error as Error }
 288     }
 289   }
 290
 291   /**
 292    * Handles the ready message sent by the main worker.
 293    *
 294    * @param message - The ready message.
 295    */
 296   protected abstract handleReadyMessage (message: MessageValue<Data>): void
 297
 298   /**
 299    * Worker message listener.
 300    *
 301    * @param message - The received message.
 302    */
 303   protected messageListener (message: MessageValue<Data>): void {
 304     this.checkMessageWorkerId(message)
 305     if (message.statistics != null) {
 306       // Statistics message received
 307       this.statistics = message.statistics
 308     } else if (message.checkActive != null) {
 309       // Check active message received
 310       message.checkActive ? this.startCheckActive() : this.stopCheckActive()
 311     } else if (message.taskFunctionOperation != null) {
 312       // Task function operation message received
 313       this.handleTaskFunctionOperationMessage(message)
 314     } else if (message.taskId != null && message.data != null) {
 315       // Task message received
 316       this.run(message)
 317     } else if (message.kill === true) {
 318       // Kill message received
 319       this.handleKillMessage(message)
 320     }
 321   }
 322
 323   protected handleTaskFunctionOperationMessage (
 324     message: MessageValue<Data>
 325   ): void {
 326     const { taskFunctionOperation, taskFunctionName, taskFunction } = message
 327     let response!: TaskFunctionOperationResult
 328     switch (taskFunctionOperation) {
 329       case 'add':
 330         response = this.addTaskFunction(
 331           taskFunctionName as string,
 332           // eslint-disable-next-line @typescript-eslint/no-implied-eval, no-new-func
 333           new Function(`return ${taskFunction as string}`)() as TaskFunction<
 334           Data,
 335           Response
 336           >
 337         )
 338         break
 339       case 'remove':
 340         response = this.removeTaskFunction(taskFunctionName as string)
 341         break
 342       case 'default':
 343         response = this.setDefaultTaskFunction(taskFunctionName as string)
 344         break
 345       default:
 346         response = { status: false, error: new Error('Unknown task operation') }
 347         break
 348     }
 349     this.sendToMainWorker({
 350       taskFunctionOperation,
 351       taskFunctionOperationStatus: response.status,
 352       taskFunctionName,
 353       ...(!response.status &&
 354         response?.error != null && {
 355         workerError: {
 356           name: taskFunctionName as string,
 357           message: this.handleError(response.error as Error | string)
 358         }
 359       })
 360     })
 361   }
 362
 363   /**
 364    * Handles a kill message sent by the main worker.
 365    *
 366    * @param message - The kill message.
 367    */
 368   protected handleKillMessage (_message: MessageValue<Data>): void {
 369     this.stopCheckActive()
 370     if (isAsyncFunction(this.opts.killHandler)) {
 371       (this.opts.killHandler?.() as Promise<void>)
 372         .then(() => {
 373           this.sendToMainWorker({ kill: 'success' })
 374           return undefined
 375         })
 376         .catch(() => {
 377           this.sendToMainWorker({ kill: 'failure' })
 378         })
 379     } else {
 380       try {
 381         // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
 382         this.opts.killHandler?.() as void
 383         this.sendToMainWorker({ kill: 'success' })
 384       } catch {
 385         this.sendToMainWorker({ kill: 'failure' })
 386       }
 387     }
 388   }
 389
 390   /**
 391    * Check if the message worker id is set and matches the worker id.
 392    *
 393    * @param message - The message to check.
 394    * @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.
 395    */
 396   private checkMessageWorkerId (message: MessageValue<Data>): void {
 397     if (message.workerId == null) {
 398       throw new Error('Message worker id is not set')
 399     } else if (message.workerId !== this.id) {
 400       throw new Error(
 401         `Message worker id ${message.workerId} does not match the worker id ${this.id}`
 402       )
 403     }
 404   }
 405
 406   /**
 407    * Starts the worker check active interval.
 408    */
 409   private startCheckActive (): void {
 410     this.lastTaskTimestamp = performance.now()
 411     this.activeInterval = setInterval(
 412       this.checkActive.bind(this),
 413       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) / 2
 414     )
 415   }
 416
 417   /**
 418    * Stops the worker check active interval.
 419    */
 420   private stopCheckActive (): void {
 421     if (this.activeInterval != null) {
 422       clearInterval(this.activeInterval)
 423       delete this.activeInterval
 424     }
 425   }
 426
 427   /**
 428    * Checks if the worker should be terminated, because its living too long.
 429    */
 430   private checkActive (): void {
 431     if (
 432       performance.now() - this.lastTaskTimestamp >
 433       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME)
 434     ) {
 435       this.sendToMainWorker({ kill: this.opts.killBehavior })
 436     }
 437   }
 438
 439   /**
 440    * Returns the main worker.
 441    *
 442    * @returns Reference to the main worker.
 443    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the main worker is not set.
 444    */
 445   protected getMainWorker (): MainWorker {
 446     if (this.mainWorker == null) {
 447       throw new Error('Main worker not set')
 448     }
 449     return this.mainWorker
 450   }
 451
 452   /**
 453    * Sends a message to main worker.
 454    *
 455    * @param message - The response message.
 456    */
 457   protected abstract sendToMainWorker (
 458     message: MessageValue<Response, Data>
 459   ): void
 460
 461   /**
 462    * Sends task function names to the main worker.
 463    */
 464   protected sendTaskFunctionNamesToMainWorker (): void {
 465     this.sendToMainWorker({
 466       taskFunctionNames: this.listTaskFunctionNames()
 467     })
 468   }
 469
 470   /**
 471    * Handles an error and convert it to a string so it can be sent back to the main worker.
 472    *
 473    * @param error - The error raised by the worker.
 474    * @returns The error message.
 475    */
 476   protected handleError (error: Error | string): string {
 477     return error instanceof Error ? error.message : error
 478   }
 479
 480   /**
 481    * Runs the given task.
 482    *
 483    * @param task - The task to execute.
 484    */
 485   protected readonly run = (task: Task<Data>): void => {
 486     const { name, taskId, data } = task
 487     const taskFunctionName = name ?? DEFAULT_TASK_NAME
 488     if (!this.taskFunctions.has(taskFunctionName)) {
 489       this.sendToMainWorker({
 490         workerError: {
 491           name: name as string,
 492           message: `Task function '${name as string}' not found`,
 493           data
 494         },
 495         taskId
 496       })
 497       return
 498     }
 499     const fn = this.taskFunctions.get(taskFunctionName)
 500     if (isAsyncFunction(fn)) {
 501       this.runAsync(fn as TaskAsyncFunction<Data, Response>, task)
 502     } else {
 503       this.runSync(fn as TaskSyncFunction<Data, Response>, task)
 504     }
 505   }
 506
 507   /**
 508    * Runs the given task function synchronously.
 509    *
 510    * @param fn - Task function that will be executed.
 511    * @param task - Input data for the task function.
 512    */
 513   protected readonly runSync = (
 514     fn: TaskSyncFunction<Data, Response>,
 515     task: Task<Data>
 516   ): void => {
 517     const { name, taskId, data } = task
 518     try {
 519       let taskPerformance = this.beginTaskPerformance(name)
 520       const res = fn(data)
 521       taskPerformance = this.endTaskPerformance(taskPerformance)
 522       this.sendToMainWorker({
 523         data: res,
 524         taskPerformance,
 525         taskId
 526       })
 527     } catch (error) {
 528       this.sendToMainWorker({
 529         workerError: {
 530           name: name as string,
 531           message: this.handleError(error as Error | string),
 532           data
 533         },
 534         taskId
 535       })
 536     } finally {
 537       this.updateLastTaskTimestamp()
 538     }
 539   }
 540
 541   /**
 542    * Runs the given task function asynchronously.
 543    *
 544    * @param fn - Task function that will be executed.
 545    * @param task - Input data for the task function.
 546    */
 547   protected readonly runAsync = (
 548     fn: TaskAsyncFunction<Data, Response>,
 549     task: Task<Data>
 550   ): void => {
 551     const { name, taskId, data } = task
 552     let taskPerformance = this.beginTaskPerformance(name)
 553     fn(data)
 554       .then(res => {
 555         taskPerformance = this.endTaskPerformance(taskPerformance)
 556         this.sendToMainWorker({
 557           data: res,
 558           taskPerformance,
 559           taskId
 560         })
 561         return undefined
 562       })
 563       .catch(error => {
 564         this.sendToMainWorker({
 565           workerError: {
 566             name: name as string,
 567             message: this.handleError(error as Error | string),
 568             data
 569           },
 570           taskId
 571         })
 572       })
 573       .finally(() => {
 574         this.updateLastTaskTimestamp()
 575       })
 576       .catch(EMPTY_FUNCTION)
 577   }
 578
 579   private beginTaskPerformance (name?: string): TaskPerformance {
 580     this.checkStatistics()
 581     return {
 582       name: name ?? DEFAULT_TASK_NAME,
 583       timestamp: performance.now(),
 584       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 585     }
 586   }
 587
 588   private endTaskPerformance (
 589     taskPerformance: TaskPerformance
 590   ): TaskPerformance {
 591     this.checkStatistics()
 592     return {
 593       ...taskPerformance,
 594       ...(this.statistics.runTime && {
 595         runTime: performance.now() - taskPerformance.timestamp
 596       }),
 597       ...(this.statistics.elu && {
 598         elu: performance.eventLoopUtilization(taskPerformance.elu)
 599       })
 600     }
 601   }
 602
 603   private checkStatistics (): void {
 604     if (this.statistics == null) {
 605       throw new Error('Performance statistics computation requirements not set')
 606     }
 607   }
 608
 609   private updateLastTaskTimestamp (): void {
 610     if (this.activeInterval != null) {
 611       this.lastTaskTimestamp = performance.now()
 612     }
 613   }
 614 }