src/workers.ts

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