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    * Worker message listener.
 148    *
 149    * @param message - Message received.
 150    */
 151   protected messageListener (message: MessageValue<Data, Data>): void {
 152     if (message.workerId === this.id) {
 153       if (message.ready != null) {
 154         // Startup message received
 155         this.workerReady()
 156       } else if (message.statistics != null) {
 157         // Statistics message received
 158         this.statistics = message.statistics
 159       } else if (message.checkAlive != null) {
 160         // Check alive message received
 161         message.checkAlive ? this.startCheckAlive() : this.stopCheckAlive()
 162       } else if (message.id != null && message.data != null) {
 163         // Task message received
 164         const fn = this.getTaskFunction(message.name)
 165         if (isAsyncFunction(fn)) {
 166           this.runInAsyncScope(this.runAsync.bind(this), this, fn, message)
 167         } else {
 168           this.runInAsyncScope(this.runSync.bind(this), this, fn, message)
 169         }
 170       } else if (message.kill === true) {
 171         // Kill message received
 172         this.stopCheckAlive()
 173         this.emitDestroy()
 174       }
 175     }
 176   }
 177
 178   /**
 179    * Notifies the main worker that this worker is ready to process tasks.
 180    */
 181   protected workerReady (): void {
 182     !this.isMain && this.sendToMainWorker({ ready: true, workerId: this.id })
 183   }
 184
 185   /**
 186    * Starts the worker alive check interval.
 187    */
 188   private startCheckAlive (): void {
 189     this.lastTaskTimestamp = performance.now()
 190     this.aliveInterval = setInterval(
 191       this.checkAlive.bind(this),
 192       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) / 2
 193     )
 194     this.checkAlive.bind(this)()
 195   }
 196
 197   /**
 198    * Stops the worker alive check interval.
 199    */
 200   private stopCheckAlive (): void {
 201     this.aliveInterval != null && clearInterval(this.aliveInterval)
 202   }
 203
 204   /**
 205    * Checks if the worker should be terminated, because its living too long.
 206    */
 207   private checkAlive (): void {
 208     if (
 209       performance.now() - this.lastTaskTimestamp >
 210       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME)
 211     ) {
 212       this.sendToMainWorker({ kill: this.opts.killBehavior, workerId: this.id })
 213     }
 214   }
 215
 216   /**
 217    * Returns the main worker.
 218    *
 219    * @returns Reference to the main worker.
 220    */
 221   protected getMainWorker (): MainWorker {
 222     if (this.mainWorker == null) {
 223       throw new Error('Main worker not set')
 224     }
 225     return this.mainWorker
 226   }
 227
 228   /**
 229    * Sends a message to the main worker.
 230    *
 231    * @param message - The response message.
 232    */
 233   protected abstract sendToMainWorker (
 234     message: MessageValue<Response, Data>
 235   ): void
 236
 237   /**
 238    * Handles an error and convert it to a string so it can be sent back to the main worker.
 239    *
 240    * @param e - The error raised by the worker.
 241    * @returns The error message.
 242    */
 243   protected handleError (e: Error | string): string {
 244     return e instanceof Error ? e.message : e
 245   }
 246
 247   /**
 248    * Runs the given function synchronously.
 249    *
 250    * @param fn - Function that will be executed.
 251    * @param message - Input data for the given function.
 252    */
 253   protected runSync (
 254     fn: WorkerSyncFunction<Data, Response>,
 255     message: MessageValue<Data>
 256   ): void {
 257     try {
 258       let taskPerformance = this.beginTaskPerformance(message.name)
 259       const res = fn(message.data)
 260       taskPerformance = this.endTaskPerformance(taskPerformance)
 261       this.sendToMainWorker({
 262         data: res,
 263         taskPerformance,
 264         workerId: this.id,
 265         id: message.id
 266       })
 267     } catch (e) {
 268       const errorMessage = this.handleError(e as Error | string)
 269       this.sendToMainWorker({
 270         taskError: {
 271           name: message.name ?? DEFAULT_TASK_NAME,
 272           message: errorMessage,
 273           data: message.data
 274         },
 275         workerId: this.id,
 276         id: message.id
 277       })
 278     } finally {
 279       if (!this.isMain && this.aliveInterval != null) {
 280         this.lastTaskTimestamp = performance.now()
 281       }
 282     }
 283   }
 284
 285   /**
 286    * Runs the given function asynchronously.
 287    *
 288    * @param fn - Function that will be executed.
 289    * @param message - Input data for the given function.
 290    */
 291   protected runAsync (
 292     fn: WorkerAsyncFunction<Data, Response>,
 293     message: MessageValue<Data>
 294   ): void {
 295     let taskPerformance = this.beginTaskPerformance(message.name)
 296     fn(message.data)
 297       .then(res => {
 298         taskPerformance = this.endTaskPerformance(taskPerformance)
 299         this.sendToMainWorker({
 300           data: res,
 301           taskPerformance,
 302           workerId: this.id,
 303           id: message.id
 304         })
 305         return null
 306       })
 307       .catch(e => {
 308         const errorMessage = this.handleError(e as Error | string)
 309         this.sendToMainWorker({
 310           taskError: {
 311             name: message.name ?? DEFAULT_TASK_NAME,
 312             message: errorMessage,
 313             data: message.data
 314           },
 315           workerId: this.id,
 316           id: message.id
 317         })
 318       })
 319       .finally(() => {
 320         if (!this.isMain && this.aliveInterval != null) {
 321           this.lastTaskTimestamp = performance.now()
 322         }
 323       })
 324       .catch(EMPTY_FUNCTION)
 325   }
 326
 327   /**
 328    * Gets the task function in the given scope.
 329    *
 330    * @param name - Name of the task function that will be returned.
 331    */
 332   private getTaskFunction (name?: string): WorkerFunction<Data, Response> {
 333     name = name ?? DEFAULT_TASK_NAME
 334     const fn = this.taskFunctions.get(name)
 335     if (fn == null) {
 336       throw new Error(`Task function '${name}' not found`)
 337     }
 338     return fn
 339   }
 340
 341   private beginTaskPerformance (name?: string): TaskPerformance {
 342     this.checkStatistics()
 343     return {
 344       name: name ?? DEFAULT_TASK_NAME,
 345       timestamp: performance.now(),
 346       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 347     }
 348   }
 349
 350   private endTaskPerformance (
 351     taskPerformance: TaskPerformance
 352   ): TaskPerformance {
 353     this.checkStatistics()
 354     return {
 355       ...taskPerformance,
 356       ...(this.statistics.runTime && {
 357         runTime: performance.now() - taskPerformance.timestamp
 358       }),
 359       ...(this.statistics.elu && {
 360         elu: performance.eventLoopUtilization(taskPerformance.elu)
 361       })
 362     }
 363   }
 364
 365   private checkStatistics (): void {
 366     if (this.statistics == null) {
 367       throw new Error('Performance statistics computation requirements not set')
 368     }
 369   }
 370 }