src/worker/abstract-worker.ts

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