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