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