src/pools/cluster/fixed.ts

   1 import cluster, { type ClusterSettings, type Worker } from 'node:cluster'
   2 import type { MessageValue } from '../../utility-types'
   3 import { AbstractPool } from '../abstract-pool'
   4 import { type PoolOptions, type PoolType, PoolTypes } from '../pool'
   5 import { type WorkerType, WorkerTypes } from '../worker'
   6
   7 /**
   8  * Options for a poolifier cluster pool.
   9  */
  10 export interface ClusterPoolOptions extends PoolOptions<Worker> {
  11   /**
  12    * Key/value pairs to add to worker process environment.
  13    *
  14    * @see https://nodejs.org/api/cluster.html#cluster_cluster_fork_env
  15    */
  16   env?: Record<string, unknown>
  17   /**
  18    * Cluster settings.
  19    *
  20    * @see https://nodejs.org/api/cluster.html#cluster_cluster_settings
  21    */
  22   settings?: ClusterSettings
  23 }
  24
  25 /**
  26  * A cluster pool with a fixed number of workers.
  27  *
  28  * @typeParam Data - Type of data sent to the worker. This can only be structured-cloneable data.
  29  * @typeParam Response - Type of execution response. This can only be structured-cloneable data.
  30  * @author [Christopher Quadflieg](https://github.com/Shinigami92)
  31  * @since 2.0.0
  32  */
  33 export class FixedClusterPool<
  34   Data = unknown,
  35   Response = unknown
  36 > extends AbstractPool<Worker, Data, Response> {
  37   /**
  38    * Constructs a new poolifier fixed cluster pool.
  39    *
  40    * @param numberOfWorkers - Number of workers for this pool.
  41    * @param filePath - Path to an implementation of a `ClusterWorker` file, which can be relative or absolute.
  42    * @param opts - Options for this fixed cluster pool.
  43    */
  44   public constructor (
  45     numberOfWorkers: number,
  46     filePath: string,
  47     protected readonly opts: ClusterPoolOptions = {}
  48   ) {
  49     super(numberOfWorkers, filePath, opts)
  50   }
  51
  52   /** @inheritDoc */
  53   protected setupHook (): void {
  54     cluster.setupPrimary({ ...this.opts.settings, exec: this.filePath })
  55   }
  56
  57   /** @inheritDoc */
  58   protected isMain (): boolean {
  59     return cluster.isPrimary
  60   }
  61
  62   /** @inheritDoc */
  63   protected async destroyWorkerNode (workerNodeKey: number): Promise<void> {
  64     this.flagWorkerNodeAsNotReady(workerNodeKey)
  65     this.flushTasksQueue(workerNodeKey)
  66     // FIXME: wait for tasks to be finished
  67     const workerNode = this.workerNodes[workerNodeKey]
  68     const worker = workerNode.worker
  69     const waitWorkerExit = new Promise<void>(resolve => {
  70       worker.once('exit', () => {
  71         resolve()
  72       })
  73     })
  74     worker.once('disconnect', () => {
  75       worker.kill()
  76     })
  77     await this.sendKillMessageToWorker(workerNodeKey)
  78     workerNode.removeAllListeners()
  79     worker.disconnect()
  80     await waitWorkerExit
  81   }
  82
  83   /** @inheritDoc */
  84   protected sendToWorker (
  85     workerNodeKey: number,
  86     message: MessageValue<Data>
  87   ): void {
  88     this.workerNodes[workerNodeKey].worker.send({
  89       ...message,
  90       workerId: this.getWorkerInfo(workerNodeKey).id as number
  91     })
  92   }
  93
  94   /** @inheritDoc */
  95   protected sendStartupMessageToWorker (workerNodeKey: number): void {
  96     this.sendToWorker(workerNodeKey, {
  97       ready: false
  98     })
  99   }
 100
 101   /** @inheritDoc */
 102   protected registerWorkerMessageListener<Message extends Data | Response>(
 103     workerNodeKey: number,
 104     listener: (message: MessageValue<Message>) => void
 105   ): void {
 106     this.workerNodes[workerNodeKey].worker.on('message', listener)
 107   }
 108
 109   /** @inheritDoc */
 110   protected registerOnceWorkerMessageListener<Message extends Data | Response>(
 111     workerNodeKey: number,
 112     listener: (message: MessageValue<Message>) => void
 113   ): void {
 114     this.workerNodes[workerNodeKey].worker.once('message', listener)
 115   }
 116
 117   /** @inheritDoc */
 118   protected deregisterWorkerMessageListener<Message extends Data | Response>(
 119     workerNodeKey: number,
 120     listener: (message: MessageValue<Message>) => void
 121   ): void {
 122     this.workerNodes[workerNodeKey].worker.off('message', listener)
 123   }
 124
 125   /** @inheritDoc */
 126   protected createWorker (): Worker {
 127     return cluster.fork(this.opts.env)
 128   }
 129
 130   /** @inheritDoc */
 131   protected get type (): PoolType {
 132     return PoolTypes.fixed
 133   }
 134
 135   /** @inheritDoc */
 136   protected get worker (): WorkerType {
 137     return WorkerTypes.cluster
 138   }
 139
 140   /** @inheritDoc */
 141   protected get busy (): boolean {
 142     return this.internalBusy()
 143   }
 144 }