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