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