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