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