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