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