src/workers.ts

   1 import { isMainThread, parentPort } from 'worker_threads'
   2
   3 import { AsyncResource } from 'async_hooks'
   4
   5 export interface ThreadWorkerOptions {
   6   /**
   7    * Max time to wait tasks to work on (in ms), after this period the new worker threads will die.
   8    *
   9    * @default 60.000 ms
  10    */
  11   maxInactiveTime?: number
  12   /**
  13    * `true` if your function contains async pieces, else `false`.
  14    *
  15    * @default false
  16    */
  17   async?: boolean
  18 }
  19
  20 /**
  21  * An example worker that will be always alive, you just need to **extend** this class if you want a static pool.
  22  *
  23  * When this worker is inactive for more than 1 minute, it will send this info to the main thread,
  24  * if you are using DynamicThreadPool, the workers created after will be killed, the min num of thread will be guaranteed.
  25  *
  26  * @author [Alessandro Pio Ardizio](https://github.com/pioardi)
  27  * @since 0.0.1
  28  */
  29 // eslint-disable-next-line @typescript-eslint/no-explicit-any
  30 export class ThreadWorker<Data = any, Response = any> extends AsyncResource {
  31   protected readonly maxInactiveTime: number
  32   protected readonly async: boolean
  33   protected lastTask: number
  34   protected readonly interval?: NodeJS.Timeout
  35   protected parent?: MessagePort
  36
  37   public constructor (
  38     fn: (data: Data) => Response,
  39     public readonly opts: ThreadWorkerOptions = {}
  40   ) {
  41     super('worker-thread-pool:pioardi')
  42
  43     this.maxInactiveTime = this.opts.maxInactiveTime ?? 1000 * 60
  44     this.async = !!this.opts.async
  45     this.lastTask = Date.now()
  46     if (!fn) throw new Error('Fn parameter is mandatory')
  47     // keep the worker active
  48     if (!isMainThread) {
  49       this.interval = setInterval(
  50         this.checkAlive.bind(this),
  51         this.maxInactiveTime / 2
  52       )
  53       this.checkAlive.bind(this)()
  54     }
  55     parentPort?.on(
  56       'message',
  57       (value: {
  58         data?: Response
  59         id?: number
  60         parent?: MessagePort
  61         kill?: number
  62       }) => {
  63         if (value?.data && value.id) {
  64           // here you will receive messages
  65           // console.log('This is the main thread ' + isMainThread)
  66           if (this.async) {
  67             this.runInAsyncScope(this.runAsync.bind(this), this, fn, value)
  68           } else {
  69             this.runInAsyncScope(this.run.bind(this), this, fn, value)
  70           }
  71         } else if (value.parent) {
  72           // save the port to communicate with the main thread
  73           // this will be received once
  74           this.parent = value.parent
  75         } else if (value.kill) {
  76           // here is time to kill this thread, just clearing the interval
  77           if (this.interval) clearInterval(this.interval)
  78           this.emitDestroy()
  79         }
  80       }
  81     )
  82   }
  83
  84   protected checkAlive (): void {
  85     if (Date.now() - this.lastTask > this.maxInactiveTime) {
  86       this.parent?.postMessage({ kill: 1 })
  87     }
  88   }
  89
  90   protected run (
  91     fn: (data: Data) => Response,
  92     value: { readonly data: Data; readonly id: number }
  93   ): void {
  94     try {
  95       const res = fn(value.data)
  96       this.parent?.postMessage({ data: res, id: value.id })
  97       this.lastTask = Date.now()
  98     } catch (e) {
  99       this.parent?.postMessage({ error: e, id: value.id })
 100       this.lastTask = Date.now()
 101     }
 102   }
 103
 104   protected runAsync (
 105     fn: (data: Data) => Promise<Response>,
 106     value: { readonly data: Data; readonly id: number }
 107   ): void {
 108     fn(value.data)
 109       .then(res => {
 110         this.parent?.postMessage({ data: res, id: value.id })
 111         this.lastTask = Date.now()
 112       })
 113       .catch(e => {
 114         this.parent?.postMessage({ error: e, id: value.id })
 115         this.lastTask = Date.now()
 116       })
 117   }
 118 }