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 {
  18   type KillBehavior,
  19   KillBehaviors,
  20   type WorkerOptions
  21 } from './worker-options'
  22 import type {
  23   TaskFunctions,
  24   WorkerAsyncFunction,
  25   WorkerFunction,
  26   WorkerSyncFunction
  27 } from './worker-functions'
  28
  29 const DEFAULT_MAX_INACTIVE_TIME = 60000
  30 const DEFAULT_KILL_BEHAVIOR: KillBehavior = KillBehaviors.SOFT
  31
  32 /**
  33  * Base class that implements some shared logic for all poolifier workers.
  34  *
  35  * @typeParam MainWorker - Type of main worker.
  36  * @typeParam Data - Type of data this worker receives from pool's execution. This can only be structured-cloneable data.
  37  * @typeParam Response - Type of response the worker sends back to the main worker. This can only be structured-cloneable data.
  38  */
  39 export abstract class AbstractWorker<
  40   MainWorker extends Worker | MessagePort,
  41   Data = unknown,
  42   Response = unknown
  43 > extends AsyncResource {
  44   /**
  45    * Worker id.
  46    */
  47   protected abstract id: number
  48   /**
  49    * Task function(s) processed by the worker when the pool's `execution` function is invoked.
  50    */
  51   protected taskFunctions!: Map<string, WorkerFunction<Data, Response>>
  52   /**
  53    * Timestamp of the last task processed by this worker.
  54    */
  55   protected lastTaskTimestamp!: number
  56   /**
  57    * Performance statistics computation requirements.
  58    */
  59   protected statistics!: WorkerStatistics
  60   /**
  61    * Handler id of the `aliveInterval` worker alive check.
  62    */
  63   protected aliveInterval?: NodeJS.Timeout
  64   /**
  65    * Constructs a new poolifier worker.
  66    *
  67    * @param type - The type of async event.
  68    * @param isMain - Whether this is the main worker or not.
  69    * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked. The first function is the default function.
  70    * @param mainWorker - Reference to main worker.
  71    * @param opts - Options for the worker.
  72    */
  73   public constructor (
  74     type: string,
  75     protected readonly isMain: boolean,
  76     taskFunctions:
  77     | WorkerFunction<Data, Response>
  78     | TaskFunctions<Data, Response>,
  79     protected readonly mainWorker: MainWorker,
  80     protected readonly opts: WorkerOptions = {
  81       /**
  82        * The kill behavior option on this worker or its default value.
  83        */
  84       killBehavior: DEFAULT_KILL_BEHAVIOR,
  85       /**
  86        * The maximum time to keep this worker alive while idle.
  87        * The pool automatically checks and terminates this worker when the time expires.
  88        */
  89       maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME
  90     }
  91   ) {
  92     super(type)
  93     this.checkWorkerOptions(this.opts)
  94     this.checkTaskFunctions(taskFunctions)
  95     if (!this.isMain) {
  96       this.mainWorker?.on('message', this.messageListener.bind(this))
  97     }
  98   }
  99
 100   private checkWorkerOptions (opts: WorkerOptions): void {
 101     this.opts.killBehavior = opts.killBehavior ?? DEFAULT_KILL_BEHAVIOR
 102     this.opts.maxInactiveTime =
 103       opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME
 104     delete this.opts.async
 105   }
 106
 107   /**
 108    * Checks if the `taskFunctions` parameter is passed to the constructor.
 109    *
 110    * @param taskFunctions - The task function(s) parameter that should be checked.
 111    */
 112   private checkTaskFunctions (
 113     taskFunctions:
 114     | WorkerFunction<Data, Response>
 115     | TaskFunctions<Data, Response>
 116   ): void {
 117     if (taskFunctions == null) {
 118       throw new Error('taskFunctions parameter is mandatory')
 119     }
 120     this.taskFunctions = new Map<string, WorkerFunction<Data, Response>>()
 121     if (typeof taskFunctions === 'function') {
 122       this.taskFunctions.set(DEFAULT_TASK_NAME, taskFunctions.bind(this))
 123     } else if (isPlainObject(taskFunctions)) {
 124       let firstEntry = true
 125       for (const [name, fn] of Object.entries(taskFunctions)) {
 126         if (typeof fn !== 'function') {
 127           throw new TypeError(
 128             'A taskFunctions parameter object value is not a function'
 129           )
 130         }
 131         this.taskFunctions.set(name, fn.bind(this))
 132         if (firstEntry) {
 133           this.taskFunctions.set(DEFAULT_TASK_NAME, fn.bind(this))
 134           firstEntry = false
 135         }
 136       }
 137       if (firstEntry) {
 138         throw new Error('taskFunctions parameter object is empty')
 139       }
 140     } else {
 141       throw new TypeError(
 142         'taskFunctions parameter is not a function or a plain object'
 143       )
 144     }
 145   }
 146
 147   /**
 148    * Checks if the worker has a task function with the given name.
 149    *
 150    * @param name - The name of the task function to check.
 151    * @returns Whether the worker has a task function with the given name or not.
 152    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string.
 153    */
 154   public hasTaskFunction (name: string): boolean {
 155     if (typeof name !== 'string') {
 156       throw new TypeError('name parameter is not a string')
 157     }
 158     return this.taskFunctions.has(name)
 159   }
 160
 161   /**
 162    * Adds a task function to the worker.
 163    * If a task function with the same name already exists, it is replaced.
 164    *
 165    * @param name - The name of the task function to add.
 166    * @param fn - The task function to add.
 167    * @returns Whether the task function was added or not.
 168    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string.
 169    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 170    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `fn` parameter is not a function.
 171    */
 172   public addTaskFunction (
 173     name: string,
 174     fn: WorkerFunction<Data, Response>
 175   ): boolean {
 176     if (typeof name !== 'string') {
 177       throw new TypeError('name parameter is not a string')
 178     }
 179     if (name === DEFAULT_TASK_NAME) {
 180       throw new Error(
 181         'Cannot add a task function with the default reserved name'
 182       )
 183     }
 184     if (typeof fn !== 'function') {
 185       throw new TypeError('fn parameter is not a function')
 186     }
 187     try {
 188       if (
 189         this.taskFunctions.get(name) ===
 190         this.taskFunctions.get(DEFAULT_TASK_NAME)
 191       ) {
 192         this.taskFunctions.set(DEFAULT_TASK_NAME, fn.bind(this))
 193       }
 194       this.taskFunctions.set(name, fn.bind(this))
 195       return true
 196     } catch {
 197       return false
 198     }
 199   }
 200
 201   /**
 202    * Removes a task function from the worker.
 203    *
 204    * @param name - The name of the task function to remove.
 205    * @returns Whether the task function existed and was removed or not.
 206    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string.
 207    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 208    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the task function used as default task function.
 209    */
 210   public removeTaskFunction (name: string): boolean {
 211     if (typeof name !== 'string') {
 212       throw new TypeError('name parameter is not a string')
 213     }
 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) === this.taskFunctions.get(DEFAULT_TASK_NAME)
 221     ) {
 222       throw new Error(
 223         'Cannot remove the task function used as the default task function'
 224       )
 225     }
 226     return this.taskFunctions.delete(name)
 227   }
 228
 229   /**
 230    * Sets the default task function to use when no task function name is specified.
 231    *
 232    * @param name - The name of the task function to use as default task function.
 233    * @returns Whether the default task function was set or not.
 234    * @throws {@link https://nodejs.org/api/errors.html#class-typeerror} If the `name` parameter is not a string.
 235    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is the default task function reserved name.
 236    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the `name` parameter is a non-existing task function.
 237    */
 238   public setDefaultTaskFunction (name: string): boolean {
 239     if (typeof name !== 'string') {
 240       throw new TypeError('name parameter is not a string')
 241     }
 242     if (name === DEFAULT_TASK_NAME) {
 243       throw new Error(
 244         'Cannot set the default task function reserved name as the default task function'
 245       )
 246     }
 247     if (!this.taskFunctions.has(name)) {
 248       throw new Error(
 249         'Cannot set the default task function to a non-existing task function'
 250       )
 251     }
 252     try {
 253       this.taskFunctions.set(
 254         DEFAULT_TASK_NAME,
 255         this.taskFunctions.get(name)?.bind(this) as WorkerFunction<
 256         Data,
 257         Response
 258         >
 259       )
 260       return true
 261     } catch {
 262       return false
 263     }
 264   }
 265
 266   /**
 267    * Worker message listener.
 268    *
 269    * @param message - Message received.
 270    */
 271   protected messageListener (message: MessageValue<Data, Data>): void {
 272     if (message.workerId === this.id) {
 273       if (message.ready != null) {
 274         // Startup message received
 275         this.workerReady()
 276       } else if (message.statistics != null) {
 277         // Statistics message received
 278         this.statistics = message.statistics
 279       } else if (message.checkAlive != null) {
 280         // Check alive message received
 281         message.checkAlive ? this.startCheckAlive() : this.stopCheckAlive()
 282       } else if (message.id != null && message.data != null) {
 283         // Task message received
 284         this.run(message)
 285       } else if (message.kill === true) {
 286         // Kill message received
 287         this.stopCheckAlive()
 288         this.emitDestroy()
 289       }
 290     }
 291   }
 292
 293   /**
 294    * Notifies the main worker that this worker is ready to process tasks.
 295    */
 296   protected workerReady (): void {
 297     !this.isMain && this.sendToMainWorker({ ready: true, workerId: this.id })
 298   }
 299
 300   /**
 301    * Starts the worker alive check interval.
 302    */
 303   private startCheckAlive (): void {
 304     this.lastTaskTimestamp = performance.now()
 305     this.aliveInterval = setInterval(
 306       this.checkAlive.bind(this),
 307       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) / 2
 308     )
 309     this.checkAlive.bind(this)()
 310   }
 311
 312   /**
 313    * Stops the worker alive check interval.
 314    */
 315   private stopCheckAlive (): void {
 316     this.aliveInterval != null && clearInterval(this.aliveInterval)
 317   }
 318
 319   /**
 320    * Checks if the worker should be terminated, because its living too long.
 321    */
 322   private checkAlive (): void {
 323     if (
 324       performance.now() - this.lastTaskTimestamp >
 325       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME)
 326     ) {
 327       this.sendToMainWorker({ kill: this.opts.killBehavior, workerId: this.id })
 328     }
 329   }
 330
 331   /**
 332    * Returns the main worker.
 333    *
 334    * @returns Reference to the main worker.
 335    */
 336   protected getMainWorker (): MainWorker {
 337     if (this.mainWorker == null) {
 338       throw new Error('Main worker not set')
 339     }
 340     return this.mainWorker
 341   }
 342
 343   /**
 344    * Sends a message to the main worker.
 345    *
 346    * @param message - The response message.
 347    */
 348   protected abstract sendToMainWorker (
 349     message: MessageValue<Response, Data>
 350   ): void
 351
 352   /**
 353    * Handles an error and convert it to a string so it can be sent back to the main worker.
 354    *
 355    * @param e - The error raised by the worker.
 356    * @returns The error message.
 357    */
 358   protected handleError (e: Error | string): string {
 359     return e instanceof Error ? e.message : e
 360   }
 361
 362   /**
 363    * Runs the given task.
 364    *
 365    * @param task - The task to execute.
 366    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the task function is not found.
 367    */
 368   protected run (task: Task<Data>): void {
 369     const fn = this.getTaskFunction(task.name)
 370     if (isAsyncFunction(fn)) {
 371       this.runInAsyncScope(this.runAsync.bind(this), this, fn, task)
 372     } else {
 373       this.runInAsyncScope(this.runSync.bind(this), this, fn, task)
 374     }
 375   }
 376
 377   /**
 378    * Runs the given function synchronously.
 379    *
 380    * @param fn - Task function that will be executed.
 381    * @param task - Input data for the task function.
 382    */
 383   protected runSync (
 384     fn: WorkerSyncFunction<Data, Response>,
 385     task: Task<Data>
 386   ): void {
 387     try {
 388       let taskPerformance = this.beginTaskPerformance(task.name)
 389       const res = fn(task.data)
 390       taskPerformance = this.endTaskPerformance(taskPerformance)
 391       this.sendToMainWorker({
 392         data: res,
 393         taskPerformance,
 394         workerId: this.id,
 395         id: task.id
 396       })
 397     } catch (e) {
 398       const errorMessage = this.handleError(e as Error | string)
 399       this.sendToMainWorker({
 400         taskError: {
 401           name: task.name ?? DEFAULT_TASK_NAME,
 402           message: errorMessage,
 403           data: task.data
 404         },
 405         workerId: this.id,
 406         id: task.id
 407       })
 408     } finally {
 409       if (!this.isMain && this.aliveInterval != null) {
 410         this.lastTaskTimestamp = performance.now()
 411       }
 412     }
 413   }
 414
 415   /**
 416    * Runs the given function asynchronously.
 417    *
 418    * @param fn - Task function that will be executed.
 419    * @param task - Input data for the task function.
 420    */
 421   protected runAsync (
 422     fn: WorkerAsyncFunction<Data, Response>,
 423     task: Task<Data>
 424   ): void {
 425     let taskPerformance = this.beginTaskPerformance(task.name)
 426     fn(task.data)
 427       .then(res => {
 428         taskPerformance = this.endTaskPerformance(taskPerformance)
 429         this.sendToMainWorker({
 430           data: res,
 431           taskPerformance,
 432           workerId: this.id,
 433           id: task.id
 434         })
 435         return null
 436       })
 437       .catch(e => {
 438         const errorMessage = this.handleError(e as Error | string)
 439         this.sendToMainWorker({
 440           taskError: {
 441             name: task.name ?? DEFAULT_TASK_NAME,
 442             message: errorMessage,
 443             data: task.data
 444           },
 445           workerId: this.id,
 446           id: task.id
 447         })
 448       })
 449       .finally(() => {
 450         if (!this.isMain && this.aliveInterval != null) {
 451           this.lastTaskTimestamp = performance.now()
 452         }
 453       })
 454       .catch(EMPTY_FUNCTION)
 455   }
 456
 457   /**
 458    * Gets the task function with the given name.
 459    *
 460    * @param name - Name of the task function that will be returned.
 461    * @returns The task function.
 462    * @throws {@link https://nodejs.org/api/errors.html#class-error} If the task function is not found.
 463    */
 464   private getTaskFunction (name?: string): WorkerFunction<Data, Response> {
 465     name = name ?? DEFAULT_TASK_NAME
 466     const fn = this.taskFunctions.get(name)
 467     if (fn == null) {
 468       throw new Error(`Task function '${name}' not found`)
 469     }
 470     return fn
 471   }
 472
 473   private beginTaskPerformance (name?: string): TaskPerformance {
 474     this.checkStatistics()
 475     return {
 476       name: name ?? DEFAULT_TASK_NAME,
 477       timestamp: performance.now(),
 478       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 479     }
 480   }
 481
 482   private endTaskPerformance (
 483     taskPerformance: TaskPerformance
 484   ): TaskPerformance {
 485     this.checkStatistics()
 486     return {
 487       ...taskPerformance,
 488       ...(this.statistics.runTime && {
 489         runTime: performance.now() - taskPerformance.timestamp
 490       }),
 491       ...(this.statistics.elu && {
 492         elu: performance.eventLoopUtilization(taskPerformance.elu)
 493       })
 494     }
 495   }
 496
 497   private checkStatistics (): void {
 498     if (this.statistics == null) {
 499       throw new Error('Performance statistics computation requirements not set')
 500     }
 501   }
 502 }