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       const { uuid, event, data } = message
 157       if (this.promiseResponseMap.has(uuid)) {
 158         // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 159         const { resolve, reject, workerSetElement } = this.promiseResponseMap.get(uuid)!
 160         if (event === WorkerMessageEvents.addedWorkerElement) {
 161           this.emitter?.emit(WorkerSetEvents.elementAdded, this.info)
 162           ++workerSetElement.numberOfWorkerElements
 163           resolve(data)
 164         } else if (event === WorkerMessageEvents.workerElementError) {
 165           this.emitter?.emit(WorkerSetEvents.elementError, data)
 166           reject(data)
 167         }
 168         this.promiseResponseMap.delete(uuid)
 169       }
 170     })
 171     worker.on('error', this.workerOptions.poolOptions?.errorHandler ?? EMPTY_FUNCTION)
 172     worker.on('error', error => {
 173       this.emitter?.emit(WorkerSetEvents.error, error)
 174       if (
 175         this.workerOptions.poolOptions?.restartWorkerOnError === true &&
 176         this.started &&
 177         !this.workerStartup
 178       ) {
 179         this.addWorkerSetElement()
 180       }
 181       worker.unref()
 182       worker.terminate().catch((error: unknown) => this.emitter?.emit(WorkerSetEvents.error, error))
 183     })
 184     worker.on('online', this.workerOptions.poolOptions?.onlineHandler ?? EMPTY_FUNCTION)
 185     worker.on('exit', this.workerOptions.poolOptions?.exitHandler ?? EMPTY_FUNCTION)
 186     worker.once('exit', () => {
 187       this.removeWorkerSetElement(this.getWorkerSetElementByWorker(worker))
 188     })
 189     const workerSetElement: WorkerSetElement = { worker, numberOfWorkerElements: 0 }
 190     this.workerSet.add(workerSetElement)
 191     this.workerStartup = false
 192     return workerSetElement
 193   }
 194
 195   private removeWorkerSetElement (workerSetElement: WorkerSetElement | undefined): void {
 196     if (workerSetElement == null) {
 197       return
 198     }
 199     this.workerSet.delete(workerSetElement)
 200   }
 201
 202   private async getWorkerSetElement (): Promise<WorkerSetElement> {
 203     let chosenWorkerSetElement: WorkerSetElement | undefined
 204     for (const workerSetElement of this.workerSet) {
 205       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 206       if (workerSetElement.numberOfWorkerElements < this.workerOptions.elementsPerWorker!) {
 207         chosenWorkerSetElement = workerSetElement
 208         break
 209       }
 210     }
 211     if (chosenWorkerSetElement == null) {
 212       chosenWorkerSetElement = this.addWorkerSetElement()
 213       // Add worker set element sequentially to optimize memory at startup
 214       // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 215       this.workerOptions.workerStartDelay! > 0 &&
 216         // eslint-disable-next-line @typescript-eslint/no-non-null-assertion
 217         (await sleep(randomizeDelay(this.workerOptions.workerStartDelay!)))
 218     }
 219     return chosenWorkerSetElement
 220   }
 221
 222   private getWorkerSetElementByWorker (worker: Worker): WorkerSetElement | undefined {
 223     let workerSetElt: WorkerSetElement | undefined
 224     for (const workerSetElement of this.workerSet) {
 225       if (workerSetElement.worker.threadId === worker.threadId) {
 226         workerSetElt = workerSetElement
 227         break
 228       }
 229     }
 230     return workerSetElt
 231   }
 232 }