src/worker/abstract-worker.ts

   1 import { AsyncResource } from 'node:async_hooks'
   2 import type { Worker } from 'node:cluster'
   3 import type { MessagePort } from 'node:worker_threads'
   4 import { performance } from 'node:perf_hooks'
   5 import type {
   6   MessageValue,
   7   Task,
   8   TaskPerformance,
   9   WorkerStatistics
  10 } from '../utility-types'
  11 import {
  12   DEFAULT_TASK_NAME,
  13   EMPTY_FUNCTION,
  14   isAsyncFunction,
  15   isPlainObject
  16 } from '../utils'
  17 import { KillBehaviors, type WorkerOptions } from './worker-options'
  18 import type {
  19   TaskAsyncFunction,
  20   TaskFunction,
  21   TaskFunctions,
  22   TaskSyncFunction
  23 } from './task-functions'
  24
  25 const DEFAULT_MAX_INACTIVE_TIME = 60000
  26 const DEFAULT_WORKER_OPTIONS: WorkerOptions = {
  27   /**
  28    * The kill behavior option on this worker or its default value.
  29    */
  30   killBehavior: KillBehaviors.SOFT,
  31   /**
  32    * The maximum time to keep this worker active while idle.
  33    * The pool automatically checks and terminates this worker when the time expires.
  34    */
  35   maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME,
  36   /**
  37    * The function to call when the worker is killed.
  38    */
  39   killHandler: EMPTY_FUNCTION
  40 }
  41
  42 /**
  43  * Base class that implements some shared logic for all poolifier workers.
  44  *
  45  * @typeParam MainWorker - Type of main worker.
  46  * @typeParam Data - Type of data this worker receives from pool's execution. This can only be structured-cloneable data.
  47  * @typeParam Response - Type of response the worker sends back to the main worker. This can only be structured-cloneable data.
  48  */
  49 export abstract class AbstractWorker<
  50   MainWorker extends Worker | MessagePort,
  51   Data = unknown,
  52   Response = unknown
  53 > extends AsyncResource {
  54   /**
  55    * Worker id.
  56    */
  57   protected abstract id: number
  58   /**
  59    * Task function(s) processed by the worker when the pool's `execution` function is invoked.
  60    */
  61   protected taskFunctions!: Map<string, TaskFunction<Data, Response>>
  62   /**
  63    * Timestamp of the last task processed by this worker.
  64    */
  65   protected lastTaskTimestamp!: number
  66   /**
  67    * Performance statistics computation requirements.
  68    */
  69   protected statistics!: WorkerStatistics
  70   /**
  71    * Handler id of the `activeInterval` worker activity check.
  72    */
  73   protected activeInterval?: NodeJS.Timeout
  74   /**
  75    * Constructs a new poolifier worker.
  76    *
  77    * @param type - The type of async event.
  78    * @param isMain - Whether this is the main worker or not.
  79    * @param mainWorker - Reference to main worker.
  80    * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked. The first function is the default function.
  81    * @param opts - Options for the worker.
  82    */
  83   public constructor (
  84     type: string,
  85     protected readonly isMain: boolean,
  86     private readonly mainWorker: MainWorker,
  87     taskFunctions: TaskFunction<Data, Response> | TaskFunctions<Data, Response>,
  88     protected opts: WorkerOptions = DEFAULT_WORKER_OPTIONS
  89   ) {
  90     super(type)
  91     if (this.isMain == null) {
  92       throw new Error('isMain parameter is mandatory')
  93     }
  94     this.checkTaskFunctions(taskFunctions)
  95     this.checkWorkerOptions(this.opts)
  96     if (!this.isMain) {
  97       this.getMainWorker().on('message', this.handleReadyMessage.bind(this))
  98     }
  99   }
 100
 101   private checkWorkerOptions (opts: WorkerOptions): void {
 102     this.opts = { ...DEFAULT_WORKER_OPTIONS, ...opts }
 103     delete this.opts.async
 104   }
 105
 106   /**
 107    * Checks if the `taskFunctions` parameter is passed to the constructor.
 108    *
 109    * @param taskFunctions - The task function(s) parameter that should be checked.
 110    */
 111   private checkTaskFunctions (
 112     taskFunctions: TaskFunction<Data, Response> | TaskFunctions<Data, Response>
 113   ): void {
 114     if (taskFunctions == null) {
 115       throw new Error('taskFunctions parameter is mandatory')
 116     }
 117     this.taskFunctions = new Map<string, TaskFunction<Data, Response>>()
 118     if (typeof taskFunctions === 'function') {
 119       const boundFn = taskFunctions.bind(this)
 120       this.taskFunctions.set(DEFAULT_TASK_NAME, boundFn)
 121       this.taskFunctions.set(
 122         typeof taskFunctions.name === 'string' &&
 123         taskFunctions.name.trim().length > 0
 124           ? taskFunctions.name
 125           : 'fn1',
 126         boundFn
 127       )
 128     } else if (isPlainObject(taskFunctions)) {
 129       let firstEntry = true
 130       for (const [name, fn] of Object.entries(taskFunctions)) {
 131         if (typeof name !== 'string') {
 132           throw new TypeError(
 133             'A taskFunctions parameter object key is not a string'
 134           )
 135         }
 136         if (typeof name === 'string' && name.trim().length === 0) {
 137           throw new TypeError(
 138             'A taskFunctions parameter object key is an empty string'
 139           )
 140         }
 141         if (typeof fn !== 'function') {
 142           throw new TypeError(
 143             'A taskFunctions parameter object value is not a function'
 144           )
 145         }
 146         const boundFn = fn.bind(this)
 147         if (firstEntry) {
 148           this.taskFunctions.set(DEFAULT_TASK_NAME, boundFn)
 149           firstEntry = false
 150         }
 151         this.taskFunctions.set(name, boundFn)
 152       }
 153       if (firstEntry) {
 154         throw new Error('taskFunctions parameter object is empty')
 155       }
 156     } else {
 157       throw new TypeError(
 158         'taskFunctions parameter is not a function or a plain object'
 159       )
 160     }
 161   }
 162
 163   /**
 164    * Checks if the worker has a task function with the given name.
 165    *
 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    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string or an empty string.
 169    */
 170   public hasTaskFunction (name: string): boolean {
 171     this.checkTaskFunctionName(name)
 172     return this.taskFunctions.has(name)
 173   }
 174
 175   /**
 176    * Adds a task function to the worker.
 177    * If a task function with the same name already exists, it is replaced.
 178    *
 179    * @param name - The name of the task function to add.
 180    * @param fn - The task function to add.
 181    * @returns Whether the task function was added or not.
 182    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string or an empty string.
 183    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 184    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `fn` parameter is not a function.
 185    */
 186   public addTaskFunction (
 187     name: string,
 188     fn: TaskFunction<Data, Response>
 189   ): boolean {
 190     this.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       throw new TypeError('fn parameter is not a function')
 198     }
 199     try {
 200       const boundFn = fn.bind(this)
 201       if (
 202         this.taskFunctions.get(name) ===
 203         this.taskFunctions.get(DEFAULT_TASK_NAME)
 204       ) {
 205         this.taskFunctions.set(DEFAULT_TASK_NAME, boundFn)
 206       }
 207       this.taskFunctions.set(name, boundFn)
 208       this.sendTaskFunctionsListToMainWorker()
 209       return true
 210     } catch {
 211       return false
 212     }
 213   }
 214
 215   /**
 216    * Removes a task function from the worker.
 217    *
 218    * @param name - The name of the task function to remove.
 219    * @returns Whether the task function existed and was removed or not.
 220    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string or an empty string.
 221    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 222    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the task function used as default task function.
 223    */
 224   public removeTaskFunction (name: string): boolean {
 225     this.checkTaskFunctionName(name)
 226     if (name === DEFAULT_TASK_NAME) {
 227       throw new Error(
 228         'Cannot remove the task function with the default reserved name'
 229       )
 230     }
 231     if (
 232       this.taskFunctions.get(name) === this.taskFunctions.get(DEFAULT_TASK_NAME)
 233     ) {
 234       throw new Error(
 235         'Cannot remove the task function used as the default task function'
 236       )
 237     }
 238     const deleteStatus = this.taskFunctions.delete(name)
 239     this.sendTaskFunctionsListToMainWorker()
 240     return deleteStatus
 241   }
 242
 243   /**
 244    * Lists the names of the worker's task functions.
 245    *
 246    * @returns The names of the worker's task functions.
 247    */
 248   public listTaskFunctions (): string[] {
 249     const names: string[] = [...this.taskFunctions.keys()]
 250     let defaultTaskFunctionName: string = DEFAULT_TASK_NAME
 251     for (const [name, fn] of this.taskFunctions) {
 252       if (
 253         name !== DEFAULT_TASK_NAME &&
 254         fn === this.taskFunctions.get(DEFAULT_TASK_NAME)
 255       ) {
 256         defaultTaskFunctionName = name
 257         break
 258       }
 259     }
 260     return [
 261       names[names.indexOf(DEFAULT_TASK_NAME)],
 262       defaultTaskFunctionName,
 263       ...names.filter(
 264         (name) => name !== DEFAULT_TASK_NAME && name !== defaultTaskFunctionName
 265       )
 266     ]
 267   }
 268
 269   /**
 270    * Sets the default task function to use in the worker.
 271    *
 272    * @param name - The name of the task function to use as default task function.
 273    * @returns Whether the default task function was set or not.
 274    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string or an empty string.
 275    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 276    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is a non-existing task function.
 277    */
 278   public setDefaultTaskFunction (name: string): boolean {
 279     this.checkTaskFunctionName(name)
 280     if (name === DEFAULT_TASK_NAME) {
 281       throw new Error(
 282         'Cannot set the default task function reserved name as the default task function'
 283       )
 284     }
 285     if (!this.taskFunctions.has(name)) {
 286       throw new Error(
 287         'Cannot set the default task function to a non-existing task function'
 288       )
 289     }
 290     try {
 291       this.taskFunctions.set(
 292         DEFAULT_TASK_NAME,
 293         this.taskFunctions.get(name) as TaskFunction<Data, Response>
 294       )
 295       return true
 296     } catch {
 297       return false
 298     }
 299   }
 300
 301   private checkTaskFunctionName (name: string): void {
 302     if (typeof name !== 'string') {
 303       throw new TypeError('name parameter is not a string')
 304     }
 305     if (typeof name === 'string' && name.trim().length === 0) {
 306       throw new TypeError('name parameter is an empty string')
 307     }
 308   }
 309
 310   /**
 311    * Handles the ready message sent by the main worker.
 312    *
 313    * @param message - The ready message.
 314    */
 315   protected abstract handleReadyMessage (message: MessageValue<Data>): void
 316
 317   /**
 318    * Worker message listener.
 319    *
 320    * @param message - The received message.
 321    */
 322   protected messageListener (message: MessageValue<Data>): void {
 323     this.checkMessageWorkerId(message)
 324     if (message.statistics != null) {
 325       // Statistics message received
 326       this.statistics = message.statistics
 327     } else if (message.checkActive != null) {
 328       // Check active message received
 329       message.checkActive ? this.startCheckActive() : this.stopCheckActive()
 330     } else if (message.taskId != null && message.data != null) {
 331       // Task message received
 332       this.run(message)
 333     } else if (message.kill === true) {
 334       // Kill message received
 335       this.handleKillMessage(message)
 336     }
 337   }
 338
 339   /**
 340    * Handles a kill message sent by the main worker.
 341    *
 342    * @param message - The kill message.
 343    */
 344   protected handleKillMessage (message: MessageValue<Data>): void {
 345     this.stopCheckActive()
 346     if (isAsyncFunction(this.opts.killHandler)) {
 347       (this.opts.killHandler?.() as Promise<void>)
 348         .then(() => {
 349           this.sendToMainWorker({ kill: 'success', workerId: this.id })
 350           return null
 351         })
 352         .catch(() => {
 353           this.sendToMainWorker({ kill: 'failure', workerId: this.id })
 354         })
 355         .finally(() => {
 356           this.emitDestroy()
 357         })
 358         .catch(EMPTY_FUNCTION)
 359     } else {
 360       try {
 361         // eslint-disable-next-line @typescript-eslint/no-invalid-void-type
 362         this.opts.killHandler?.() as void
 363         this.sendToMainWorker({ kill: 'success', workerId: this.id })
 364       } catch {
 365         this.sendToMainWorker({ kill: 'failure', workerId: this.id })
 366       } finally {
 367         this.emitDestroy()
 368       }
 369     }
 370   }
 371
 372   /**
 373    * Check if the message worker id is set and matches the worker id.
 374    *
 375    * @param message - The message to check.
 376    * @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.
 377    */
 378   private checkMessageWorkerId (message: MessageValue<Data>): void {
 379     if (message.workerId == null) {
 380       throw new Error('Message worker id is not set')
 381     } else if (message.workerId != null && message.workerId !== this.id) {
 382       throw new Error(
 383         `Message worker id ${message.workerId} does not match the worker id ${this.id}`
 384       )
 385     }
 386   }
 387
 388   /**
 389    * Starts the worker check active interval.
 390    */
 391   private startCheckActive (): void {
 392     this.lastTaskTimestamp = performance.now()
 393     this.activeInterval = setInterval(
 394       this.checkActive.bind(this),
 395       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) / 2
 396     )
 397   }
 398
 399   /**
 400    * Stops the worker check active interval.
 401    */
 402   private stopCheckActive (): void {
 403     if (this.activeInterval != null) {
 404       clearInterval(this.activeInterval)
 405       delete this.activeInterval
 406     }
 407   }
 408
 409   /**
 410    * Checks if the worker should be terminated, because its living too long.
 411    */
 412   private checkActive (): void {
 413     if (
 414       performance.now() - this.lastTaskTimestamp >
 415       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME)
 416     ) {
 417       this.sendToMainWorker({ kill: this.opts.killBehavior, workerId: this.id })
 418     }
 419   }
 420
 421   /**
 422    * Returns the main worker.
 423    *
 424    * @returns Reference to the main worker.
 425    */
 426   protected getMainWorker (): MainWorker {
 427     if (this.mainWorker == null) {
 428       throw new Error('Main worker not set')
 429     }
 430     return this.mainWorker
 431   }
 432
 433   /**
 434    * Sends a message to main worker.
 435    *
 436    * @param message - The response message.
 437    */
 438   protected abstract sendToMainWorker (
 439     message: MessageValue<Response, Data>
 440   ): void
 441
 442   /**
 443    * Sends the list of task function names to the main worker.
 444    */
 445   protected sendTaskFunctionsListToMainWorker (): void {
 446     this.sendToMainWorker({
 447       taskFunctions: this.listTaskFunctions(),
 448       workerId: this.id
 449     })
 450   }
 451
 452   /**
 453    * Handles an error and convert it to a string so it can be sent back to the main worker.
 454    *
 455    * @param e - The error raised by the worker.
 456    * @returns The error message.
 457    */
 458   protected handleError (e: Error | string): string {
 459     return e instanceof Error ? e.message : e
 460   }
 461
 462   /**
 463    * Runs the given task.
 464    *
 465    * @param task - The task to execute.
 466    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the task function is not found.
 467    */
 468   protected run (task: Task<Data>): void {
 469     const { name, taskId, data } = task
 470     const fn = this.taskFunctions.get(name ?? DEFAULT_TASK_NAME)
 471     if (fn == null) {
 472       this.sendToMainWorker({
 473         taskError: {
 474           name: name as string,
 475           message: `Task function '${name as string}' not found`,
 476           data
 477         },
 478         workerId: this.id,
 479         taskId
 480       })
 481       return
 482     }
 483     if (isAsyncFunction(fn)) {
 484       this.runInAsyncScope(this.runAsync.bind(this), this, fn, task)
 485     } else {
 486       this.runInAsyncScope(this.runSync.bind(this), this, fn, task)
 487     }
 488   }
 489
 490   /**
 491    * Runs the given task function synchronously.
 492    *
 493    * @param fn - Task function that will be executed.
 494    * @param task - Input data for the task function.
 495    */
 496   protected runSync (
 497     fn: TaskSyncFunction<Data, Response>,
 498     task: Task<Data>
 499   ): void {
 500     const { name, taskId, data } = task
 501     try {
 502       let taskPerformance = this.beginTaskPerformance(name)
 503       const res = fn(data)
 504       taskPerformance = this.endTaskPerformance(taskPerformance)
 505       this.sendToMainWorker({
 506         data: res,
 507         taskPerformance,
 508         workerId: this.id,
 509         taskId
 510       })
 511     } catch (e) {
 512       const errorMessage = this.handleError(e as Error | string)
 513       this.sendToMainWorker({
 514         taskError: {
 515           name: name as string,
 516           message: errorMessage,
 517           data
 518         },
 519         workerId: this.id,
 520         taskId
 521       })
 522     } finally {
 523       this.updateLastTaskTimestamp()
 524     }
 525   }
 526
 527   /**
 528    * Runs the given task function asynchronously.
 529    *
 530    * @param fn - Task function that will be executed.
 531    * @param task - Input data for the task function.
 532    */
 533   protected runAsync (
 534     fn: TaskAsyncFunction<Data, Response>,
 535     task: Task<Data>
 536   ): void {
 537     const { name, taskId, data } = task
 538     let taskPerformance = this.beginTaskPerformance(name)
 539     fn(data)
 540       .then((res) => {
 541         taskPerformance = this.endTaskPerformance(taskPerformance)
 542         this.sendToMainWorker({
 543           data: res,
 544           taskPerformance,
 545           workerId: this.id,
 546           taskId
 547         })
 548         return null
 549       })
 550       .catch((e) => {
 551         const errorMessage = this.handleError(e as Error | string)
 552         this.sendToMainWorker({
 553           taskError: {
 554             name: name as string,
 555             message: errorMessage,
 556             data
 557           },
 558           workerId: this.id,
 559           taskId
 560         })
 561       })
 562       .finally(() => {
 563         this.updateLastTaskTimestamp()
 564       })
 565       .catch(EMPTY_FUNCTION)
 566   }
 567
 568   private beginTaskPerformance (name?: string): TaskPerformance {
 569     this.checkStatistics()
 570     return {
 571       name: name ?? DEFAULT_TASK_NAME,
 572       timestamp: performance.now(),
 573       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 574     }
 575   }
 576
 577   private endTaskPerformance (
 578     taskPerformance: TaskPerformance
 579   ): TaskPerformance {
 580     this.checkStatistics()
 581     return {
 582       ...taskPerformance,
 583       ...(this.statistics.runTime && {
 584         runTime: performance.now() - taskPerformance.timestamp
 585       }),
 586       ...(this.statistics.elu && {
 587         elu: performance.eventLoopUtilization(taskPerformance.elu)
 588       })
 589     }
 590   }
 591
 592   private checkStatistics (): void {
 593     if (this.statistics == null) {
 594       throw new Error('Performance statistics computation requirements not set')
 595     }
 596   }
 597
 598   private updateLastTaskTimestamp (): void {
 599     if (this.activeInterval != null) {
 600       this.lastTaskTimestamp = performance.now()
 601     }
 602   }
 603 }