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