src/worker/WorkerSet.ts

   1 // Partial Copyright Jerome Benoit. 2021-2024. All Rights Reserved.
   2
   3 import { EventEmitterAsyncResource } from 'node:events'
   4 import { SHARE_ENV, Worker } from 'node:worker_threads'
   5
   6 import { WorkerAbstract } from './WorkerAbstract.js'
   7 import { EMPTY_FUNCTION, workerSetVersion } from './WorkerConstants.js'
   8 import {
   9   type SetInfo,
  10   type WorkerData,
  11   type WorkerMessage,
  12   WorkerMessageEvents,
  13   type WorkerOptions,
  14   type WorkerSetElement,
  15   WorkerSetEvents
  16 } from './WorkerTypes.js'
  17 import { randomizeDelay, sleep } from './WorkerUtils.js'
  18
  19 export class WorkerSet extends WorkerAbstract<WorkerData> {
  20   public readonly emitter: EventEmitterAsyncResource | undefined
  21   private readonly workerSet: Set<WorkerSetElement>
  22   private started: boolean
  23   private workerStartup: boolean
  24
  25   /**
  26    * Creates a new `WorkerSet`.
  27    *
  28    * @param workerScript -
  29    * @param workerOptions -
  30    */
  31   constructor (workerScript: string, workerOptions: WorkerOptions) {
  32     super(workerScript, workerOptions)
  33     if (this.workerOptions.elementsPerWorker == null) {
  34       throw new TypeError('Elements per worker is not defined')
  35     }
  36     if (!Number.isSafeInteger(this.workerOptions.elementsPerWorker)) {
  37       throw new TypeError('Elements per worker must be an integer')
  38     }
  39     if (this.workerOptions.elementsPerWorker <= 0) {
  40       throw new RangeError('Elements per worker must be greater than zero')
  41     }
  42     this.workerSet = new Set<WorkerSetElement>()
  43     if (this.workerOptions.poolOptions?.enableEvents === true) {
  44       this.emitter = new EventEmitterAsyncResource({ name: 'workerset' })
  45     }
  46     this.started = false
  47     this.workerStartup = false
  48   }
  49
  50   get info (): SetInfo {
  51     return {
  52       version: workerSetVersion,
  53       type: 'set',
  54       worker: 'thread',
  55       started: this.started,
  56       size: this.size,
  57       elementsExecuting: [...this.workerSet].reduce(
  58         (accumulator, workerSetElement) => accumulator + workerSetElement.numberOfWorkerElements,
  59         0
  60       ),
  61       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  62       elementsPerWorker: this.maxElementsPerWorker!
  63     }
  64   }
  65
  66   get size (): number {
  67     return this.workerSet.size
  68   }
  69
  70   get maxElementsPerWorker (): number | undefined {
  71     return this.workerOptions.elementsPerWorker
  72   }
  73
  74   /** @inheritDoc */
  75   public async start (): Promise<void> {
  76     this.addWorkerSetElement()
  77     // Add worker set element sequentially to optimize memory at startup
  78     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  79     this.workerOptions.workerStartDelay! > 0 &&
  80       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
  81       (await sleep(randomizeDelay(this.workerOptions.workerStartDelay!)))
  82     this.emitter?.emit(WorkerSetEvents.started, this.info)
  83     this.started = true
  84   }
  85
  86   /** @inheritDoc */
  87   public async stop (): Promise<void> {
  88     for (const workerSetElement of this.workerSet) {
  89       const worker = workerSetElement.worker
  90       const waitWorkerExit = new Promise<void>(resolve => {
  91         worker.once('exit', () => {
  92           resolve()
  93         })
  94       })
  95       worker.unref()
  96       await worker.terminate()
  97       await waitWorkerExit
  98     }
  99     this.emitter?.emit(WorkerSetEvents.stopped, this.info)
 100     this.started = false
 101     this.emitter?.emitDestroy()
 102     this.emitter?.removeAllListeners()
 103   }
 104
 105   /** @inheritDoc */
 106   public async addElement (elementData: WorkerData): Promise<void> {
 107     if (!this.started) {
 108       throw new Error('Cannot add a WorkerSet element: not started')
 109     }
 110     const workerSetElement = await this.getWorkerSetElement()
 111     workerSetElement.worker.postMessage({
 112       event: WorkerMessageEvents.addWorkerElement,
 113       data: elementData
 114     })
 115     ++workerSetElement.numberOfWorkerElements
 116     // Add element sequentially to optimize memory at startup
 117     // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 118     if (this.workerOptions.elementStartDelay! > 0) {
 119       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 120       await sleep(randomizeDelay(this.workerOptions.elementStartDelay!))
 121     }
 122   }
 123
 124   /**
 125    * Adds a new `WorkerSetElement`.
 126    */
 127   private addWorkerSetElement (): WorkerSetElement {
 128     this.workerStartup = true
 129     const worker = new Worker(this.workerScript, {
 130       env: SHARE_ENV,
 131       ...this.workerOptions.poolOptions?.workerOptions
 132     })
 133     worker.on('message', this.workerOptions.poolOptions?.messageHandler ?? EMPTY_FUNCTION)
 134     worker.on('message', (message: WorkerMessage<WorkerData>) => {
 135       if (message.event === WorkerMessageEvents.addedWorkerElement) {
 136         this.emitter?.emit(WorkerSetEvents.elementAdded, this.info)
 137       } else if (message.event === WorkerMessageEvents.workerElementError) {
 138         this.emitter?.emit(WorkerSetEvents.elementError, message.data)
 139       }
 140     })
 141     worker.on('error', this.workerOptions.poolOptions?.errorHandler ?? EMPTY_FUNCTION)
 142     worker.on('error', error => {
 143       this.emitter?.emit(WorkerSetEvents.error, error)
 144       if (
 145         this.workerOptions.poolOptions?.restartWorkerOnError === true &&
 146         this.started &&
 147         !this.workerStartup
 148       ) {
 149         this.addWorkerSetElement()
 150       }
 151       worker.unref()
 152       worker.terminate().catch((error: unknown) => this.emitter?.emit(WorkerSetEvents.error, error))
 153     })
 154     worker.on('online', this.workerOptions.poolOptions?.onlineHandler ?? EMPTY_FUNCTION)
 155     worker.on('exit', this.workerOptions.poolOptions?.exitHandler ?? EMPTY_FUNCTION)
 156     worker.once('exit', () => {
 157       this.removeWorkerSetElement(this.getWorkerSetElementByWorker(worker))
 158     })
 159     const workerSetElement: WorkerSetElement = { worker, numberOfWorkerElements: 0 }
 160     this.workerSet.add(workerSetElement)
 161     this.workerStartup = false
 162     return workerSetElement
 163   }
 164
 165   private removeWorkerSetElement (workerSetElement: WorkerSetElement | undefined): void {
 166     if (workerSetElement == null) {
 167       return
 168     }
 169     this.workerSet.delete(workerSetElement)
 170   }
 171
 172   private async getWorkerSetElement (): Promise<WorkerSetElement> {
 173     let chosenWorkerSetElement: WorkerSetElement | undefined
 174     for (const workerSetElement of this.workerSet) {
 175       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 176       if (workerSetElement.numberOfWorkerElements < this.workerOptions.elementsPerWorker!) {
 177         chosenWorkerSetElement = workerSetElement
 178         break
 179       }
 180     }
 181     if (chosenWorkerSetElement == null) {
 182       chosenWorkerSetElement = this.addWorkerSetElement()
 183       // Add worker set element sequentially to optimize memory at startup
 184       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 185       this.workerOptions.workerStartDelay! > 0 &&
 186         // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 187         (await sleep(randomizeDelay(this.workerOptions.workerStartDelay!)))
 188     }
 189     return chosenWorkerSetElement
 190   }
 191
 192   private getWorkerSetElementByWorker (worker: Worker): WorkerSetElement | undefined {
 193     let workerSetElt: WorkerSetElement | undefined
 194     for (const workerSetElement of this.workerSet) {
 195       if (workerSetElement.worker.threadId === worker.threadId) {
 196         workerSetElt = workerSetElement
 197         break
 198       }
 199     }
 200     return workerSetElt
 201   }
 202 }