src/worker/abstract-worker.ts

   1 import { AsyncResource } from 'async_hooks'
   2 import type { Worker } from 'cluster'
   3 import type { MessagePort } from 'worker_threads'
   4 import type { MessageValue } from '../utility-types'
   5 import type { WorkerOptions } from './worker-options'
   6
   7 /**
   8  * Base class containing some shared logic for all poolifier workers.
   9  *
  10  * @template MainWorker Type of main worker.
  11  * @template Data Type of data this worker receives from pool's execution.
  12  * @template Response Type of response the worker sends back to the main worker.
  13  */
  14 export abstract class AbstractWorker<
  15   MainWorker extends Worker | MessagePort,
  16   Data = unknown,
  17   Response = unknown
  18 > extends AsyncResource {
  19   /**
  20    * The maximum time to keep this worker alive while idle. The pool automatically checks and terminates this worker when the time expires.
  21    */
  22   protected readonly maxInactiveTime: number
  23   /**
  24    * Whether the worker is working asynchronously or not.
  25    */
  26   protected readonly async: boolean
  27   /**
  28    * Timestamp of the last task processed by this worker.
  29    */
  30   protected lastTask: number
  31   /**
  32    * Handler ID of the `interval` alive check.
  33    */
  34   protected readonly interval?: NodeJS.Timeout
  35
  36   /**
  37    * Constructs a new poolifier worker.
  38    *
  39    * @param type The type of async event.
  40    * @param isMain Whether this is the main worker or not.
  41    * @param fn Function processed by the worker when the pool's `execution` function is invoked.
  42    * @param mainWorker Reference to main worker.
  43    * @param opts Options for the worker.
  44    */
  45   public constructor (
  46     type: string,
  47     isMain: boolean,
  48     fn: (data: Data) => Response,
  49     protected mainWorker?: MainWorker | null,
  50     public readonly opts: WorkerOptions = {}
  51   ) {
  52     super(type)
  53
  54     this.maxInactiveTime = this.opts.maxInactiveTime ?? 1000 * 60
  55     this.async = !!this.opts.async
  56     this.lastTask = Date.now()
  57     if (!fn) throw new Error('fn parameter is mandatory')
  58     // Keep the worker active
  59     if (!isMain) {
  60       this.interval = setInterval(
  61         this.checkAlive.bind(this),
  62         this.maxInactiveTime / 2
  63       )
  64       this.checkAlive.bind(this)()
  65     }
  66
  67     this.mainWorker?.on('message', (value: MessageValue<Data, MainWorker>) => {
  68       if (value?.data && value.id) {
  69         // Here you will receive messages
  70         if (this.async) {
  71           this.runInAsyncScope(this.runAsync.bind(this), this, fn, value)
  72         } else {
  73           this.runInAsyncScope(this.run.bind(this), this, fn, value)
  74         }
  75       } else if (value.parent) {
  76         // Save a reference of the main worker to communicate with it
  77         // This will be received once
  78         this.mainWorker = value.parent
  79       } else if (value.kill) {
  80         // Here is time to kill this worker, just clearing the interval
  81         if (this.interval) clearInterval(this.interval)
  82         this.emitDestroy()
  83       }
  84     })
  85   }
  86
  87   /**
  88    * Returns the main worker.
  89    *
  90    * @returns Reference to the main worker.
  91    */
  92   protected getMainWorker (): MainWorker {
  93     if (!this.mainWorker) {
  94       throw new Error('Main worker was not set')
  95     }
  96     return this.mainWorker
  97   }
  98
  99   /**
 100    * Send a message to the main worker.
 101    *
 102    * @param message The response message.
 103    */
 104   protected abstract sendToMainWorker (message: MessageValue<Response>): void
 105
 106   /**
 107    * Check to see if the worker should be terminated, because its living too long.
 108    */
 109   protected checkAlive (): void {
 110     if (Date.now() - this.lastTask > this.maxInactiveTime) {
 111       this.sendToMainWorker({ kill: 1 })
 112     }
 113   }
 114
 115   /**
 116    * Handle an error and convert it to a string so it can be sent back to the main worker.
 117    *
 118    * @param e The error raised by the worker.
 119    * @returns Message of the error.
 120    */
 121   protected handleError (e: Error | string): string {
 122     return (e as unknown) as string
 123   }
 124
 125   /**
 126    * Run the given function synchronously.
 127    *
 128    * @param fn Function that will be executed.
 129    * @param value Input data for the given function.
 130    */
 131   protected run (
 132     fn: (data?: Data) => Response,
 133     value: MessageValue<Data>
 134   ): void {
 135     try {
 136       const res = fn(value.data)
 137       this.sendToMainWorker({ data: res, id: value.id })
 138       this.lastTask = Date.now()
 139     } catch (e) {
 140       const err = this.handleError(e)
 141       this.sendToMainWorker({ error: err, id: value.id })
 142       this.lastTask = Date.now()
 143     }
 144   }
 145
 146   /**
 147    * Run the given function asynchronously.
 148    *
 149    * @param fn Function that will be executed.
 150    * @param value Input data for the given function.
 151    */
 152   protected runAsync (
 153     fn: (data?: Data) => Promise<Response>,
 154     value: MessageValue<Data>
 155   ): void {
 156     fn(value.data)
 157       .then(res => {
 158         this.sendToMainWorker({ data: res, id: value.id })
 159         this.lastTask = Date.now()
 160         return null
 161       })
 162       .catch(e => {
 163         const err = this.handleError(e)
 164         this.sendToMainWorker({ error: err, id: value.id })
 165         this.lastTask = Date.now()
 166       })
 167   }
 168 }