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 type { KillBehavior, WorkerOptions } from './worker-options'
  13 import { KillBehaviors } from './worker-options'
  14
  15 const DEFAULT_FUNCTION_NAME = 'default'
  16 const DEFAULT_MAX_INACTIVE_TIME = 60000
  17 const DEFAULT_KILL_BEHAVIOR: KillBehavior = KillBehaviors.SOFT
  18
  19 /**
  20  * Base class that implements some shared logic for all poolifier workers.
  21  *
  22  * @typeParam MainWorker - Type of main worker.
  23  * @typeParam Data - Type of data this worker receives from pool's execution. This can only be serializable data.
  24  * @typeParam Response - Type of response the worker sends back to the main worker. This can only be serializable data.
  25  */
  26 export abstract class AbstractWorker<
  27   MainWorker extends Worker | MessagePort,
  28   Data = unknown,
  29   Response = unknown
  30 > extends AsyncResource {
  31   /**
  32    * Task function(s) processed by the worker when the pool's `execution` function is invoked.
  33    */
  34   protected taskFunctions!: Map<string, WorkerFunction<Data, Response>>
  35   /**
  36    * Timestamp of the last task processed by this worker.
  37    */
  38   protected lastTaskTimestamp!: number
  39   /**
  40    * Handler id of the `aliveInterval` worker alive check.
  41    */
  42   protected readonly aliveInterval?: NodeJS.Timeout
  43   /**
  44    * Constructs a new poolifier worker.
  45    *
  46    * @param type - The type of async event.
  47    * @param isMain - Whether this is the main worker or not.
  48    * @param taskFunctions - Task function(s) processed by the worker when the pool's `execution` function is invoked. The first function is the default function.
  49    * @param mainWorker - Reference to main worker.
  50    * @param opts - Options for the worker.
  51    */
  52   public constructor (
  53     type: string,
  54     protected readonly isMain: boolean,
  55     taskFunctions:
  56     | WorkerFunction<Data, Response>
  57     | TaskFunctions<Data, Response>,
  58     protected mainWorker: MainWorker | undefined | null,
  59     protected readonly opts: WorkerOptions = {
  60       /**
  61        * The kill behavior option on this worker or its default value.
  62        */
  63       killBehavior: DEFAULT_KILL_BEHAVIOR,
  64       /**
  65        * The maximum time to keep this worker alive while idle.
  66        * The pool automatically checks and terminates this worker when the time expires.
  67        */
  68       maxInactiveTime: DEFAULT_MAX_INACTIVE_TIME
  69     }
  70   ) {
  71     super(type)
  72     this.checkWorkerOptions(this.opts)
  73     this.checkTaskFunctions(taskFunctions)
  74     if (
  75       !this.isMain &&
  76       (this.opts.maxInactiveTime ?? DEFAULT_MAX_INACTIVE_TIME) > 0
  77     ) {
  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         id: message.id,
 219         runTime,
 220         waitTime
 221       })
 222     } catch (e) {
 223       const err = this.handleError(e as Error)
 224       this.sendToMainWorker({ error: err, id: message.id })
 225     } finally {
 226       !this.isMain && (this.lastTaskTimestamp = performance.now())
 227     }
 228   }
 229
 230   /**
 231    * Runs the given function asynchronously.
 232    *
 233    * @param fn - Function that will be executed.
 234    * @param message - Input data for the given function.
 235    */
 236   protected runAsync (
 237     fn: WorkerAsyncFunction<Data, Response>,
 238     message: MessageValue<Data>
 239   ): void {
 240     const startTimestamp = performance.now()
 241     const waitTime = startTimestamp - (message.submissionTimestamp ?? 0)
 242     fn(message.data)
 243       .then(res => {
 244         const runTime = performance.now() - startTimestamp
 245         this.sendToMainWorker({
 246           data: res,
 247           id: message.id,
 248           runTime,
 249           waitTime
 250         })
 251         return null
 252       })
 253       .catch(e => {
 254         const err = this.handleError(e as Error)
 255         this.sendToMainWorker({ error: err, id: message.id })
 256       })
 257       .finally(() => {
 258         !this.isMain && (this.lastTaskTimestamp = performance.now())
 259       })
 260       .catch(EMPTY_FUNCTION)
 261   }
 262
 263   /**
 264    * Gets the task function in the given scope.
 265    *
 266    * @param name - Name of the function that will be returned.
 267    */
 268   private getTaskFunction (name?: string): WorkerFunction<Data, Response> {
 269     name = name ?? DEFAULT_FUNCTION_NAME
 270     const fn = this.taskFunctions.get(name)
 271     if (fn == null) {
 272       throw new Error(`Task function '${name}' not found`)
 273     }
 274     return fn
 275   }
 276 }