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, 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 readonly 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.lastTaskTimestamp = performance.now()
  92       this.aliveInterval = setInterval(
  93         this.checkAlive.bind(this),
  94         (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) / 2
  95       )
  96       this.checkAlive.bind(this)()
  97       this.mainWorker?.on('message', this.messageListener.bind(this))
  98     }
  99   }
 100
 101   private checkWorkerOptions (opts: WorkerOptions): void {
 102     this.opts.killBehavior = opts.killBehavior ?? DEFAULT_KILL_BEHAVIOR
 103     this.opts.maxInactiveTime =
 104       opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME
 105     delete this.opts.async
 106   }
 107
 108   /**
 109    * Checks if the `taskFunctions` parameter is passed to the constructor.
 110    *
 111    * @param taskFunctions - The task function(s) parameter that should be checked.
 112    */
 113   private checkTaskFunctions (
 114     taskFunctions:
 115     | WorkerFunction<Data, Response>
 116     | TaskFunctions<Data, Response>
 117   ): void {
 118     if (taskFunctions == null) {
 119       throw new Error('taskFunctions parameter is mandatory')
 120     }
 121     this.taskFunctions = new Map<string, WorkerFunction<Data, Response>>()
 122     if (typeof taskFunctions === 'function') {
 123       this.taskFunctions.set(DEFAULT_FUNCTION_NAME, taskFunctions.bind(this))
 124     } else if (isPlainObject(taskFunctions)) {
 125       let firstEntry = true
 126       for (const [name, fn] of Object.entries(taskFunctions)) {
 127         if (typeof fn !== 'function') {
 128           throw new TypeError(
 129             'A taskFunctions parameter object value is not a function'
 130           )
 131         }
 132         this.taskFunctions.set(name, fn.bind(this))
 133         if (firstEntry) {
 134           this.taskFunctions.set(DEFAULT_FUNCTION_NAME, fn.bind(this))
 135           firstEntry = false
 136         }
 137       }
 138       if (firstEntry) {
 139         throw new Error('taskFunctions parameter object is empty')
 140       }
 141     } else {
 142       throw new TypeError(
 143         'taskFunctions parameter is not a function or a plain object'
 144       )
 145     }
 146   }
 147
 148   /**
 149    * Worker message listener.
 150    *
 151    * @param message - Message received.
 152    */
 153   protected messageListener (message: MessageValue<Data, Data>): void {
 154     if (message.id != null && message.data != null) {
 155       // Task message received
 156       const fn = this.getTaskFunction(message.name)
 157       if (fn?.constructor.name === 'AsyncFunction') {
 158         this.runInAsyncScope(this.runAsync.bind(this), this, fn, message)
 159       } else {
 160         this.runInAsyncScope(this.runSync.bind(this), this, fn, message)
 161       }
 162     } else if (message.statistics != null) {
 163       // Statistics message received
 164       this.statistics = message.statistics
 165     } else if (message.kill != null) {
 166       // Kill message received
 167       this.aliveInterval != null && clearInterval(this.aliveInterval)
 168       this.emitDestroy()
 169     }
 170   }
 171
 172   /**
 173    * Returns the main worker.
 174    *
 175    * @returns Reference to the main worker.
 176    */
 177   protected getMainWorker (): MainWorker {
 178     if (this.mainWorker == null) {
 179       throw new Error('Main worker not set')
 180     }
 181     return this.mainWorker
 182   }
 183
 184   /**
 185    * Sends a message to the main worker.
 186    *
 187    * @param message - The response message.
 188    */
 189   protected abstract sendToMainWorker (
 190     message: MessageValue<Response, Data>
 191   ): void
 192
 193   /**
 194    * Checks if the worker should be terminated, because its living too long.
 195    */
 196   protected checkAlive (): void {
 197     if (
 198       performance.now() - this.lastTaskTimestamp >
 199       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME)
 200     ) {
 201       this.sendToMainWorker({ kill: this.opts.killBehavior })
 202     }
 203   }
 204
 205   /**
 206    * Handles an error and convert it to a string so it can be sent back to the main worker.
 207    *
 208    * @param e - The error raised by the worker.
 209    * @returns The error message.
 210    */
 211   protected handleError (e: Error | string): string {
 212     return e instanceof Error ? e.message : e
 213   }
 214
 215   /**
 216    * Runs the given function synchronously.
 217    *
 218    * @param fn - Function that will be executed.
 219    * @param message - Input data for the given function.
 220    */
 221   protected runSync (
 222     fn: WorkerSyncFunction<Data, Response>,
 223     message: MessageValue<Data>
 224   ): void {
 225     try {
 226       let taskPerformance = this.beginTaskPerformance()
 227       const res = fn(message.data)
 228       taskPerformance = this.endTaskPerformance(taskPerformance)
 229       this.sendToMainWorker({
 230         data: res,
 231         taskPerformance,
 232         workerId: this.id,
 233         id: message.id
 234       })
 235     } catch (e) {
 236       const errorMessage = this.handleError(e as Error | string)
 237       this.sendToMainWorker({
 238         taskError: {
 239           workerId: this.id,
 240           message: errorMessage,
 241           data: message.data
 242         },
 243         id: message.id
 244       })
 245     } finally {
 246       !this.isMain && (this.lastTaskTimestamp = performance.now())
 247     }
 248   }
 249
 250   /**
 251    * Runs the given function asynchronously.
 252    *
 253    * @param fn - Function that will be executed.
 254    * @param message - Input data for the given function.
 255    */
 256   protected runAsync (
 257     fn: WorkerAsyncFunction<Data, Response>,
 258     message: MessageValue<Data>
 259   ): void {
 260     let taskPerformance = this.beginTaskPerformance()
 261     fn(message.data)
 262       .then(res => {
 263         taskPerformance = this.endTaskPerformance(taskPerformance)
 264         this.sendToMainWorker({
 265           data: res,
 266           taskPerformance,
 267           workerId: this.id,
 268           id: message.id
 269         })
 270         return null
 271       })
 272       .catch(e => {
 273         const errorMessage = this.handleError(e as Error | string)
 274         this.sendToMainWorker({
 275           taskError: {
 276             workerId: this.id,
 277             message: errorMessage,
 278             data: message.data
 279           },
 280           id: message.id
 281         })
 282       })
 283       .finally(() => {
 284         !this.isMain && (this.lastTaskTimestamp = performance.now())
 285       })
 286       .catch(EMPTY_FUNCTION)
 287   }
 288
 289   /**
 290    * Gets the task function in the given scope.
 291    *
 292    * @param name - Name of the function that will be returned.
 293    */
 294   private getTaskFunction (name?: string): WorkerFunction<Data, Response> {
 295     name = name ?? DEFAULT_FUNCTION_NAME
 296     const fn = this.taskFunctions.get(name)
 297     if (fn == null) {
 298       throw new Error(`Task function '${name}' not found`)
 299     }
 300     return fn
 301   }
 302
 303   private beginTaskPerformance (): TaskPerformance {
 304     this.checkStatistics()
 305     return {
 306       timestamp: performance.now(),
 307       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 308     }
 309   }
 310
 311   private endTaskPerformance (
 312     taskPerformance: TaskPerformance
 313   ): TaskPerformance {
 314     this.checkStatistics()
 315     return {
 316       ...taskPerformance,
 317       ...(this.statistics.runTime && {
 318         runTime: performance.now() - taskPerformance.timestamp
 319       }),
 320       ...(this.statistics.elu && {
 321         elu: performance.eventLoopUtilization(taskPerformance.elu)
 322       })
 323     }
 324   }
 325
 326   private checkStatistics (): void {
 327     if (this.statistics == null) {
 328       throw new Error('Performance statistics computation requirements not set')
 329     }
 330   }
 331 }