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 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.beforeTaskRunHook(message)
 230       const res = fn(message.data)
 231       const { runTime, waitTime, elu } = this.afterTaskRunHook(taskPerformance)
 232       this.sendToMainWorker({
 233         data: res,
 234         runTime,
 235         waitTime,
 236         elu,
 237         id: message.id
 238       })
 239     } catch (e) {
 240       const err = this.handleError(e as Error)
 241       this.sendToMainWorker({
 242         error: err,
 243         errorData: message.data,
 244         id: message.id
 245       })
 246     } finally {
 247       !this.isMain && (this.lastTaskTimestamp = performance.now())
 248     }
 249   }
 250
 251   /**
 252    * Runs the given function asynchronously.
 253    *
 254    * @param fn - Function that will be executed.
 255    * @param message - Input data for the given function.
 256    */
 257   protected runAsync (
 258     fn: WorkerAsyncFunction<Data, Response>,
 259     message: MessageValue<Data>
 260   ): void {
 261     const taskPerformance = this.beforeTaskRunHook(message)
 262     fn(message.data)
 263       .then(res => {
 264         const { runTime, waitTime, elu } =
 265           this.afterTaskRunHook(taskPerformance)
 266         this.sendToMainWorker({
 267           data: res,
 268           runTime,
 269           waitTime,
 270           elu,
 271           id: message.id
 272         })
 273         return null
 274       })
 275       .catch(e => {
 276         const err = this.handleError(e as Error)
 277         this.sendToMainWorker({
 278           error: err,
 279           errorData: message.data,
 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 beforeTaskRunHook (message: MessageValue<Data>): TaskPerformance {
 304     const timestamp = performance.now()
 305     return {
 306       timestamp,
 307       ...(this.statistics.waitTime && {
 308         waitTime: timestamp - (message.timestamp ?? timestamp)
 309       }),
 310       ...(this.statistics.elu && { elu: performance.eventLoopUtilization() })
 311     }
 312   }
 313
 314   private afterTaskRunHook (taskPerformance: TaskPerformance): TaskPerformance {
 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 }