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 { EMPTY_FUNCTION, isAsyncFunction, isPlainObject } from '../utils'
  11 import {
  12   type KillBehavior,
  13   KillBehaviors,
  14   type WorkerOptions
  15 } from './worker-options'
  16 import type {
  17   TaskFunctions,
  18   WorkerAsyncFunction,
  19   WorkerFunction,
  20   WorkerSyncFunction
  21 } from './worker-functions'
  22
  23 const DEFAULT_FUNCTION_NAME = 'default'
  24 const DEFAULT_MAX_INACTIVE_TIME = 60000
  25 const DEFAULT_KILL_BEHAVIOR: KillBehavior = KillBehaviors.SOFT
  26
  27 /**
  28  * Base class that implements some shared logic for all poolifier workers.
  29  *
  30  * @typeParam MainWorker - Type of main worker.
  31  * @typeParam Data - Type of data this worker receives from pool's execution. This can only be structured-cloneable data.
  32  * @typeParam Response - Type of response the worker sends back to the main worker. This can only be structured-cloneable data.
  33  */
  34 export abstract class AbstractWorker<
  35   MainWorker extends Worker | MessagePort,
  36   Data = unknown,
  37   Response = unknown
  38 > extends AsyncResource {
  39   /**
  40    * Worker id.
  41    */
  42   protected abstract id: number
  43   /**
  44    * Task function(s) processed by the worker when the pool's `execution` function is invoked.
  45    */
  46   protected taskFunctions!: Map<string, WorkerFunction<Data, Response>>
  47   /**
  48    * Timestamp of the last task processed by this worker.
  49    */
  50   protected lastTaskTimestamp!: number
  51   /**
  52    * Performance statistics computation requirements.
  53    */
  54   protected statistics!: WorkerStatistics
  55   /**
  56    * Handler id of the `aliveInterval` worker alive check.
  57    */
  58   protected aliveInterval?: NodeJS.Timeout
  59   /**
  60    * Constructs a new poolifier worker.
  61    *
  62    * @param type - The type of async event.
  63    * @param isMain - Whether this is the main worker or not.
  64    * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked. The first function is the default function.
  65    * @param mainWorker - Reference to main worker.
  66    * @param opts - Options for the worker.
  67    */
  68   public constructor (
  69     type: string,
  70     protected readonly isMain: boolean,
  71     taskFunctions:
  72     | WorkerFunction<Data, Response>
  73     | TaskFunctions<Data, Response>,
  74     protected readonly mainWorker: MainWorker,
  75     protected readonly opts: WorkerOptions = {
  76       /**
  77        * The kill behavior option on this worker or its default value.
  78        */
  79       killBehavior: DEFAULT_KILL_BEHAVIOR,
  80       /**
  81        * The maximum time to keep this worker alive while idle.
  82        * The pool automatically checks and terminates this worker when the time expires.
  83        */
  84       maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME
  85     }
  86   ) {
  87     super(type)
  88     this.checkWorkerOptions(this.opts)
  89     this.checkTaskFunctions(taskFunctions)
  90     if (!this.isMain) {
  91       this.mainWorker?.on('message', this.messageListener.bind(this))
  92     }
  93   }
  94
  95   private checkWorkerOptions (opts: WorkerOptions): void {
  96     this.opts.killBehavior = opts.killBehavior ?? DEFAULT_KILL_BEHAVIOR
  97     this.opts.maxInactiveTime =
  98       opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME
  99     delete this.opts.async
 100   }
 101
 102   /**
 103    * Checks if the `taskFunctions` parameter is passed to the constructor.
 104    *
 105    * @param taskFunctions - The task function(s) parameter that should be checked.
 106    */
 107   private checkTaskFunctions (
 108     taskFunctions:
 109     | WorkerFunction<Data, Response>
 110     | TaskFunctions<Data, Response>
 111   ): void {
 112     if (taskFunctions == null) {
 113       throw new Error('taskFunctions parameter is mandatory')
 114     }
 115     this.taskFunctions = new Map<string, WorkerFunction<Data, Response>>()
 116     if (typeof taskFunctions === 'function') {
 117       this.taskFunctions.set(DEFAULT_FUNCTION_NAME, taskFunctions.bind(this))
 118     } else if (isPlainObject(taskFunctions)) {
 119       let firstEntry = true
 120       for (const [name, fn] of Object.entries(taskFunctions)) {
 121         if (typeof fn !== 'function') {
 122           throw new TypeError(
 123             'A taskFunctions parameter object value is not a function'
 124           )
 125         }
 126         this.taskFunctions.set(name, fn.bind(this))
 127         if (firstEntry) {
 128           this.taskFunctions.set(DEFAULT_FUNCTION_NAME, fn.bind(this))
 129           firstEntry = false
 130         }
 131       }
 132       if (firstEntry) {
 133         throw new Error('taskFunctions parameter object is empty')
 134       }
 135     } else {
 136       throw new TypeError(
 137         'taskFunctions parameter is not a function or a plain object'
 138       )
 139     }
 140   }
 141
 142   /**
 143    * Worker message listener.
 144    *
 145    * @param message - Message received.
 146    */
 147   protected messageListener (message: MessageValue<Data, Data>): void {
 148     if (message.ready != null && message.workerId === this.id) {
 149       // Startup message received
 150       this.workerReady()
 151     } else if (message.statistics != null && message.workerId === this.id) {
 152       // Statistics message received
 153       this.statistics = message.statistics
 154     } else if (message.checkAlive != null && message.workerId === this.id) {
 155       // Check alive message received
 156       message.checkAlive ? this.startCheckAlive() : this.stopCheckAlive()
 157     } else if (
 158       message.id != null &&
 159       message.data != null &&
 160       message.workerId === this.id
 161     ) {
 162       // Task message received
 163       const fn = this.getTaskFunction(message.name)
 164       if (isAsyncFunction(fn)) {
 165         this.runInAsyncScope(this.runAsync.bind(this), this, fn, message)
 166       } else {
 167         this.runInAsyncScope(this.runSync.bind(this), this, fn, message)
 168       }
 169     } else if (message.kill === true && message.workerId === this.id) {
 170       // Kill message received
 171       this.stopCheckAlive()
 172       this.emitDestroy()
 173     }
 174   }
 175
 176   /**
 177    * Notifies the main worker that this worker is ready to process tasks.
 178    */
 179   protected workerReady (): void {
 180     !this.isMain && this.sendToMainWorker({ ready: true, workerId: this.id })
 181   }
 182
 183   /**
 184    * Starts the worker alive check interval.
 185    */
 186   private startCheckAlive (): void {
 187     this.lastTaskTimestamp = performance.now()
 188     this.aliveInterval = setInterval(
 189       this.checkAlive.bind(this),
 190       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) / 2
 191     )
 192     this.checkAlive.bind(this)()
 193   }
 194
 195   /**
 196    * Stops the worker alive check interval.
 197    */
 198   private stopCheckAlive (): void {
 199     this.aliveInterval != null && clearInterval(this.aliveInterval)
 200   }
 201
 202   /**
 203    * Checks if the worker should be terminated, because its living too long.
 204    */
 205   private checkAlive (): void {
 206     if (
 207       performance.now() - this.lastTaskTimestamp >
 208       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME)
 209     ) {
 210       this.sendToMainWorker({ kill: this.opts.killBehavior, workerId: this.id })
 211     }
 212   }
 213
 214   /**
 215    * Returns the main worker.
 216    *
 217    * @returns Reference to the main worker.
 218    */
 219   protected getMainWorker (): MainWorker {
 220     if (this.mainWorker == null) {
 221       throw new Error('Main worker not set')
 222     }
 223     return this.mainWorker
 224   }
 225
 226   /**
 227    * Sends a message to the main worker.
 228    *
 229    * @param message - The response message.
 230    */
 231   protected abstract sendToMainWorker (
 232     message: MessageValue<Response, Data>
 233   ): void
 234
 235   /**
 236    * Handles an error and convert it to a string so it can be sent back to the main worker.
 237    *
 238    * @param e - The error raised by the worker.
 239    * @returns The error message.
 240    */
 241   protected handleError (e: Error | string): string {
 242     return e instanceof Error ? e.message : e
 243   }
 244
 245   /**
 246    * Runs the given function synchronously.
 247    *
 248    * @param fn - Function that will be executed.
 249    * @param message - Input data for the given function.
 250    */
 251   protected runSync (
 252     fn: WorkerSyncFunction<Data, Response>,
 253     message: MessageValue<Data>
 254   ): void {
 255     try {
 256       let taskPerformance = this.beginTaskPerformance()
 257       const res = fn(message.data)
 258       taskPerformance = this.endTaskPerformance(taskPerformance)
 259       this.sendToMainWorker({
 260         data: res,
 261         taskPerformance,
 262         workerId: this.id,
 263         id: message.id
 264       })
 265     } catch (e) {
 266       const errorMessage = this.handleError(e as Error | string)
 267       this.sendToMainWorker({
 268         taskError: {
 269           message: errorMessage,
 270           data: message.data
 271         },
 272         workerId: this.id,
 273         id: message.id
 274       })
 275     } finally {
 276       if (!this.isMain && this.aliveInterval != null) {
 277         this.lastTaskTimestamp = performance.now()
 278       }
 279     }
 280   }
 281
 282   /**
 283    * Runs the given function asynchronously.
 284    *
 285    * @param fn - Function that will be executed.
 286    * @param message - Input data for the given function.
 287    */
 288   protected runAsync (
 289     fn: WorkerAsyncFunction<Data, Response>,
 290     message: MessageValue<Data>
 291   ): void {
 292     let taskPerformance = this.beginTaskPerformance()
 293     fn(message.data)
 294       .then(res => {
 295         taskPerformance = this.endTaskPerformance(taskPerformance)
 296         this.sendToMainWorker({
 297           data: res,
 298           taskPerformance,
 299           workerId: this.id,
 300           id: message.id
 301         })
 302         return null
 303       })
 304       .catch(e => {
 305         const errorMessage = this.handleError(e as Error | string)
 306         this.sendToMainWorker({
 307           taskError: {
 308             message: errorMessage,
 309             data: message.data
 310           },
 311           workerId: this.id,
 312           id: message.id
 313         })
 314       })
 315       .finally(() => {
 316         if (!this.isMain && this.aliveInterval != null) {
 317           this.lastTaskTimestamp = performance.now()
 318         }
 319       })
 320       .catch(EMPTY_FUNCTION)
 321   }
 322
 323   /**
 324    * Gets the task function in the given scope.
 325    *
 326    * @param name - Name of the function that will be returned.
 327    */
 328   private getTaskFunction (name?: string): WorkerFunction<Data, Response> {
 329     name = name ?? DEFAULT_FUNCTION_NAME
 330     const fn = this.taskFunctions.get(name)
 331     if (fn == null) {
 332       throw new Error(`Task function '${name}' not found`)
 333     }
 334     return fn
 335   }
 336
 337   private beginTaskPerformance (): TaskPerformance {
 338     this.checkStatistics()
 339     return {
 340       timestamp: performance.now(),
 341       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 342     }
 343   }
 344
 345   private endTaskPerformance (
 346     taskPerformance: TaskPerformance
 347   ): TaskPerformance {
 348     this.checkStatistics()
 349     return {
 350       ...taskPerformance,
 351       ...(this.statistics.runTime && {
 352         runTime: performance.now() - taskPerformance.timestamp
 353       }),
 354       ...(this.statistics.elu && {
 355         elu: performance.eventLoopUtilization(taskPerformance.elu)
 356       })
 357     }
 358   }
 359
 360   private checkStatistics (): void {
 361     if (this.statistics == null) {
 362       throw new Error('Performance statistics computation requirements not set')
 363     }
 364   }
 365 }