src/pools/thread/fixed.ts

   1 import { isMainThread, MessageChannel, SHARE_ENV, Worker } from 'worker_threads'
   2 import type { Draft, JSONValue, MessageValue } from '../../utility-types'
   3 import type { PoolOptions } from '../abstract-pool'
   4 import { AbstractPool } from '../abstract-pool'
   5
   6 /**
   7  * A thread worker with message channels for communication between main thread and thread worker.
   8  */
   9 export type ThreadWorkerWithMessageChannel = Worker & Draft<MessageChannel>
  10
  11 /**
  12  * A thread pool with a fixed number of threads.
  13  *
  14  * It is possible to perform tasks in sync or asynchronous mode as you prefer.
  15  *
  16  * This pool selects the threads in a round robin fashion.
  17  *
  18  * @template Data Type of data sent to the worker.
  19  * @template Response Type of response of execution.
  20  *
  21  * @author [Alessandro Pio Ardizio](https://github.com/pioardi)
  22  * @since 0.0.1
  23  */
  24 export class FixedThreadPool<
  25   Data extends JSONValue = JSONValue,
  26   Response extends JSONValue = JSONValue
  27 > extends AbstractPool<ThreadWorkerWithMessageChannel, Data, Response> {
  28   /**
  29    * Constructs a new poolifier fixed thread pool.
  30    *
  31    * @param numberOfThreads Number of threads for this pool.
  32    * @param filePath Path to an implementation of a `ThreadWorker` file, which can be relative or absolute.
  33    * @param opts Options for this fixed thread pool. Default: `{ maxTasks: 1000 }`
  34    */
  35   public constructor (
  36     numberOfThreads: number,
  37     filePath: string,
  38     opts: PoolOptions<ThreadWorkerWithMessageChannel> = { maxTasks: 1000 }
  39   ) {
  40     super(numberOfThreads, filePath, opts)
  41   }
  42
  43   protected isMain (): boolean {
  44     return isMainThread
  45   }
  46
  47   protected async destroyWorker (
  48     worker: ThreadWorkerWithMessageChannel
  49   ): Promise<void> {
  50     await worker.terminate()
  51   }
  52
  53   protected sendToWorker (
  54     worker: ThreadWorkerWithMessageChannel,
  55     message: MessageValue<Data>
  56   ): void {
  57     worker.postMessage(message)
  58   }
  59
  60   protected registerWorkerMessageListener<Message extends Data | Response> (
  61     messageChannel: ThreadWorkerWithMessageChannel,
  62     listener: (message: MessageValue<Message>) => void
  63   ): void {
  64     messageChannel.port2?.on('message', listener)
  65   }
  66
  67   protected unregisterWorkerMessageListener<Message extends Data | Response> (
  68     messageChannel: ThreadWorkerWithMessageChannel,
  69     listener: (message: MessageValue<Message>) => void
  70   ): void {
  71     messageChannel.port2?.removeListener('message', listener)
  72   }
  73
  74   protected createWorker (): ThreadWorkerWithMessageChannel {
  75     return new Worker(this.filePath, {
  76       env: SHARE_ENV
  77     })
  78   }
  79
  80   protected afterWorkerSetup (worker: ThreadWorkerWithMessageChannel): void {
  81     const { port1, port2 } = new MessageChannel()
  82     worker.postMessage({ parent: port1 }, [port1])
  83     worker.port1 = port1
  84     worker.port2 = port2
  85     // we will attach a listener for every task,
  86     // when task is completed the listener will be removed but to avoid warnings we are increasing the max listeners size
  87     worker.port2.setMaxListeners(this.opts.maxTasks ?? 1000)
  88   }
  89 }