src/worker/WorkerSet.ts

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