src/worker/abstract-worker.ts

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