test: ensure the worker choice context class is properly initialized
[poolifier.git] / src / pools / abstract-pool.ts
CommitLineData
fc3e6586 1import crypto from 'node:crypto'
2740a743 2import type { MessageValue, PromiseResponseWrapper } from '../utility-types'
ed6dd37f 3import { EMPTY_FUNCTION } from '../utils'
34a0cfab 4import { KillBehaviors, isKillBehavior } from '../worker/worker-options'
bdaf31cd 5import type { PoolOptions } from './pool'
b4904890 6import { PoolEmitter } from './pool'
ffcbbad8 7import type { IPoolInternal, TasksUsage, WorkerType } from './pool-internal'
b4904890 8import { PoolType } from './pool-internal'
ea7a90d3 9import type { IPoolWorker } from './pool-worker'
a35560ba
S
10import {
11 WorkerChoiceStrategies,
63220255 12 type WorkerChoiceStrategy
bdaf31cd
JB
13} from './selection-strategies/selection-strategies-types'
14import { WorkerChoiceStrategyContext } from './selection-strategies/worker-choice-strategy-context'
c97c7edb 15
729c563d 16/**
ea7a90d3 17 * Base class that implements some shared logic for all poolifier pools.
729c563d 18 *
38e795c1
JB
19 * @typeParam Worker - Type of worker which manages this pool.
20 * @typeParam Data - Type of data sent to the worker. This can only be serializable data.
21 * @typeParam Response - Type of response of execution. This can only be serializable data.
729c563d 22 */
c97c7edb 23export abstract class AbstractPool<
ea7a90d3 24 Worker extends IPoolWorker,
d3c8a1a8
S
25 Data = unknown,
26 Response = unknown
9b2fdd9f 27> implements IPoolInternal<Worker, Data, Response> {
38e795c1 28 /** {@inheritDoc} */
e65c6cd9 29 public readonly workers: Array<WorkerType<Worker>> = []
4a6952ff 30
38e795c1 31 /** {@inheritDoc} */
7c0ba920
JB
32 public readonly emitter?: PoolEmitter
33
be0676b3 34 /**
2740a743 35 * The promise response map.
be0676b3 36 *
2740a743 37 * - `key`: The message id of each submitted task.
c923ce56 38 * - `value`: An object that contains the worker, the promise resolve and reject callbacks.
be0676b3 39 *
2740a743 40 * When we receive a message from the worker we get a map entry with the promise resolve/reject bound to the message.
be0676b3 41 */
c923ce56
JB
42 protected promiseResponseMap: Map<
43 string,
44 PromiseResponseWrapper<Worker, Response>
45 > = new Map<string, PromiseResponseWrapper<Worker, Response>>()
c97c7edb 46
a35560ba 47 /**
51fe3d3c 48 * Worker choice strategy context referencing a worker choice algorithm implementation.
a35560ba 49 *
51fe3d3c 50 * Default to a round robin algorithm.
a35560ba
S
51 */
52 protected workerChoiceStrategyContext: WorkerChoiceStrategyContext<
78cea37e
JB
53 Worker,
54 Data,
55 Response
a35560ba
S
56 >
57
729c563d
S
58 /**
59 * Constructs a new poolifier pool.
60 *
38e795c1
JB
61 * @param numberOfWorkers - Number of workers that this pool should manage.
62 * @param filePath - Path to the worker-file.
63 * @param opts - Options for the pool.
729c563d 64 */
c97c7edb 65 public constructor (
5c5a1fb7 66 public readonly numberOfWorkers: number,
c97c7edb 67 public readonly filePath: string,
1927ee67 68 public readonly opts: PoolOptions<Worker>
c97c7edb 69 ) {
78cea37e 70 if (!this.isMain()) {
c97c7edb
S
71 throw new Error('Cannot start a pool from a worker!')
72 }
8d3782fa 73 this.checkNumberOfWorkers(this.numberOfWorkers)
c510fea7 74 this.checkFilePath(this.filePath)
7c0ba920 75 this.checkPoolOptions(this.opts)
1086026a
JB
76
77 this.chooseWorker.bind(this)
78 this.internalExecute.bind(this)
164d950a 79 this.checkAndEmitFull.bind(this)
1086026a
JB
80 this.checkAndEmitBusy.bind(this)
81 this.sendToWorker.bind(this)
82
c97c7edb
S
83 this.setupHook()
84
5c5a1fb7 85 for (let i = 1; i <= this.numberOfWorkers; i++) {
280c2a77 86 this.createAndSetupWorker()
c97c7edb
S
87 }
88
6bd72cd0 89 if (this.opts.enableEvents === true) {
7c0ba920
JB
90 this.emitter = new PoolEmitter()
91 }
d59df138
JB
92 this.workerChoiceStrategyContext = new WorkerChoiceStrategyContext<
93 Worker,
94 Data,
95 Response
96 >(
a35560ba 97 this,
4a6952ff 98 () => {
c923ce56
JB
99 const createdWorker = this.createAndSetupWorker()
100 this.registerWorkerMessageListener(createdWorker, message => {
4a6952ff
JB
101 if (
102 isKillBehavior(KillBehaviors.HARD, message.kill) ||
c923ce56 103 this.getWorkerTasksUsage(createdWorker)?.running === 0
4a6952ff
JB
104 ) {
105 // Kill received from the worker, means that no new tasks are submitted to that worker for a while ( > maxInactiveTime)
c923ce56 106 void this.destroyWorker(createdWorker)
4a6952ff
JB
107 }
108 })
c923ce56 109 return this.getWorkerKey(createdWorker)
4a6952ff 110 },
e843b904 111 this.opts.workerChoiceStrategy
a35560ba 112 )
c97c7edb
S
113 }
114
a35560ba 115 private checkFilePath (filePath: string): void {
ffcbbad8
JB
116 if (
117 filePath == null ||
118 (typeof filePath === 'string' && filePath.trim().length === 0)
119 ) {
c510fea7
APA
120 throw new Error('Please specify a file with a worker implementation')
121 }
122 }
123
8d3782fa
JB
124 private checkNumberOfWorkers (numberOfWorkers: number): void {
125 if (numberOfWorkers == null) {
126 throw new Error(
127 'Cannot instantiate a pool without specifying the number of workers'
128 )
78cea37e 129 } else if (!Number.isSafeInteger(numberOfWorkers)) {
473c717a 130 throw new TypeError(
8d3782fa
JB
131 'Cannot instantiate a pool with a non integer number of workers'
132 )
133 } else if (numberOfWorkers < 0) {
473c717a 134 throw new RangeError(
8d3782fa
JB
135 'Cannot instantiate a pool with a negative number of workers'
136 )
7c0ba920 137 } else if (this.type === PoolType.FIXED && numberOfWorkers === 0) {
8d3782fa
JB
138 throw new Error('Cannot instantiate a fixed pool with no worker')
139 }
140 }
141
7c0ba920 142 private checkPoolOptions (opts: PoolOptions<Worker>): void {
e843b904
JB
143 this.opts.workerChoiceStrategy =
144 opts.workerChoiceStrategy ?? WorkerChoiceStrategies.ROUND_ROBIN
b529c323
JB
145 if (
146 !Object.values(WorkerChoiceStrategies).includes(
147 this.opts.workerChoiceStrategy
148 )
149 ) {
150 throw new Error(
151 `Invalid worker choice strategy '${this.opts.workerChoiceStrategy}'`
152 )
153 }
7c0ba920
JB
154 this.opts.enableEvents = opts.enableEvents ?? true
155 }
156
38e795c1 157 /** {@inheritDoc} */
7c0ba920
JB
158 public abstract get type (): PoolType
159
c2ade475 160 /**
51fe3d3c 161 * Number of tasks concurrently running in the pool.
c2ade475
JB
162 */
163 private get numberOfRunningTasks (): number {
2740a743 164 return this.promiseResponseMap.size
a35560ba
S
165 }
166
ffcbbad8 167 /**
b4e75778 168 * Gets the given worker key.
ffcbbad8
JB
169 *
170 * @param worker - The worker.
7cf00f70 171 * @returns The worker key if the worker is found in the pool, `-1` otherwise.
ffcbbad8 172 */
e65c6cd9
JB
173 private getWorkerKey (worker: Worker): number {
174 return this.workers.findIndex(workerItem => workerItem.worker === worker)
bf9549ae
JB
175 }
176
38e795c1 177 /** {@inheritDoc} */
a35560ba
S
178 public setWorkerChoiceStrategy (
179 workerChoiceStrategy: WorkerChoiceStrategy
180 ): void {
b98ec2e6 181 this.opts.workerChoiceStrategy = workerChoiceStrategy
c923ce56
JB
182 for (const [index, workerItem] of this.workers.entries()) {
183 this.setWorker(index, workerItem.worker, {
ffcbbad8
JB
184 run: 0,
185 running: 0,
186 runTime: 0,
2740a743
JB
187 avgRunTime: 0,
188 error: 0
ffcbbad8 189 })
ea7a90d3 190 }
a35560ba
S
191 this.workerChoiceStrategyContext.setWorkerChoiceStrategy(
192 workerChoiceStrategy
193 )
194 }
195
c2ade475
JB
196 /** {@inheritDoc} */
197 public abstract get full (): boolean
198
38e795c1 199 /** {@inheritDoc} */
7c0ba920
JB
200 public abstract get busy (): boolean
201
c2ade475 202 protected internalBusy (): boolean {
7c0ba920
JB
203 return (
204 this.numberOfRunningTasks >= this.numberOfWorkers &&
bf90656c 205 this.findFreeWorkerKey() === -1
7c0ba920
JB
206 )
207 }
208
38e795c1 209 /** {@inheritDoc} */
bf90656c
JB
210 public findFreeWorkerKey (): number {
211 return this.workers.findIndex(workerItem => {
c923ce56
JB
212 return workerItem.tasksUsage.running === 0
213 })
7c0ba920
JB
214 }
215
38e795c1 216 /** {@inheritDoc} */
78cea37e 217 public async execute (data: Data): Promise<Response> {
c923ce56 218 const [workerKey, worker] = this.chooseWorker()
b4e75778 219 const messageId = crypto.randomUUID()
c923ce56 220 const res = this.internalExecute(workerKey, worker, messageId)
164d950a 221 this.checkAndEmitFull()
14916bf9 222 this.checkAndEmitBusy()
a05c10de 223 this.sendToWorker(worker, {
e5a5c0fc
JB
224 // eslint-disable-next-line @typescript-eslint/consistent-type-assertions
225 data: data ?? ({} as Data),
b4e75778 226 id: messageId
a05c10de 227 })
78cea37e 228 // eslint-disable-next-line @typescript-eslint/return-await
280c2a77
S
229 return res
230 }
c97c7edb 231
38e795c1 232 /** {@inheritDoc} */
c97c7edb 233 public async destroy (): Promise<void> {
1fbcaa7c 234 await Promise.all(
e65c6cd9
JB
235 this.workers.map(async workerItem => {
236 await this.destroyWorker(workerItem.worker)
1fbcaa7c
JB
237 })
238 )
c97c7edb
S
239 }
240
4a6952ff 241 /**
675bb809 242 * Shutdowns given worker.
4a6952ff 243 *
38e795c1 244 * @param worker - A worker within `workers`.
4a6952ff
JB
245 */
246 protected abstract destroyWorker (worker: Worker): void | Promise<void>
c97c7edb 247
729c563d 248 /**
280c2a77
S
249 * Setup hook that can be overridden by a Poolifier pool implementation
250 * to run code before workers are created in the abstract constructor.
729c563d 251 */
280c2a77
S
252 protected setupHook (): void {
253 // Can be overridden
254 }
c97c7edb 255
729c563d 256 /**
280c2a77
S
257 * Should return whether the worker is the main worker or not.
258 */
259 protected abstract isMain (): boolean
260
261 /**
bf9549ae
JB
262 * Hook executed before the worker task promise resolution.
263 * Can be overridden.
729c563d 264 *
2740a743 265 * @param workerKey - The worker key.
729c563d 266 */
2740a743
JB
267 protected beforePromiseResponseHook (workerKey: number): void {
268 ++this.workers[workerKey].tasksUsage.running
c97c7edb
S
269 }
270
c01733f1 271 /**
bf9549ae
JB
272 * Hook executed after the worker task promise resolution.
273 * Can be overridden.
c01733f1 274 *
c923ce56 275 * @param worker - The worker.
38e795c1 276 * @param message - The received message.
c01733f1 277 */
2740a743 278 protected afterPromiseResponseHook (
c923ce56 279 worker: Worker,
2740a743 280 message: MessageValue<Response>
bf9549ae 281 ): void {
c923ce56 282 const workerTasksUsage = this.getWorkerTasksUsage(worker) as TasksUsage
3032893a
JB
283 --workerTasksUsage.running
284 ++workerTasksUsage.run
2740a743
JB
285 if (message.error != null) {
286 ++workerTasksUsage.error
287 }
97a2abc3 288 if (this.workerChoiceStrategyContext.getRequiredStatistics().runTime) {
3032893a 289 workerTasksUsage.runTime += message.taskRunTime ?? 0
c6bd2650
JB
290 if (
291 this.workerChoiceStrategyContext.getRequiredStatistics().avgRunTime &&
292 workerTasksUsage.run !== 0
293 ) {
3032893a
JB
294 workerTasksUsage.avgRunTime =
295 workerTasksUsage.runTime / workerTasksUsage.run
296 }
297 }
c01733f1 298 }
299
280c2a77 300 /**
675bb809 301 * Chooses a worker for the next task.
280c2a77 302 *
51fe3d3c 303 * The default uses a round robin algorithm to distribute the load.
280c2a77 304 *
c923ce56 305 * @returns [worker key, worker].
280c2a77 306 */
c923ce56
JB
307 protected chooseWorker (): [number, Worker] {
308 const workerKey = this.workerChoiceStrategyContext.execute()
309 return [workerKey, this.workers[workerKey].worker]
c97c7edb
S
310 }
311
280c2a77 312 /**
675bb809 313 * Sends a message to the given worker.
280c2a77 314 *
38e795c1
JB
315 * @param worker - The worker which should receive the message.
316 * @param message - The message.
280c2a77
S
317 */
318 protected abstract sendToWorker (
319 worker: Worker,
320 message: MessageValue<Data>
321 ): void
322
4a6952ff 323 /**
bdede008 324 * Registers a listener callback on a given worker.
4a6952ff 325 *
38e795c1
JB
326 * @param worker - The worker which should register a listener.
327 * @param listener - The message listener callback.
4a6952ff
JB
328 */
329 protected abstract registerWorkerMessageListener<
4f7fa42a 330 Message extends Data | Response
78cea37e 331 >(worker: Worker, listener: (message: MessageValue<Message>) => void): void
c97c7edb 332
729c563d
S
333 /**
334 * Returns a newly created worker.
335 */
280c2a77 336 protected abstract createWorker (): Worker
c97c7edb 337
729c563d
S
338 /**
339 * Function that can be hooked up when a worker has been newly created and moved to the workers registry.
340 *
38e795c1 341 * Can be used to update the `maxListeners` or binding the `main-worker`\<-\>`worker` connection if not bind by default.
729c563d 342 *
38e795c1 343 * @param worker - The newly created worker.
729c563d 344 */
280c2a77 345 protected abstract afterWorkerSetup (worker: Worker): void
c97c7edb 346
4a6952ff
JB
347 /**
348 * Creates a new worker for this pool and sets it up completely.
349 *
350 * @returns New, completely set up worker.
351 */
352 protected createAndSetupWorker (): Worker {
bdacc2d2 353 const worker = this.createWorker()
280c2a77 354
35cf1c03 355 worker.on('message', this.opts.messageHandler ?? EMPTY_FUNCTION)
a35560ba
S
356 worker.on('error', this.opts.errorHandler ?? EMPTY_FUNCTION)
357 worker.on('online', this.opts.onlineHandler ?? EMPTY_FUNCTION)
358 worker.on('exit', this.opts.exitHandler ?? EMPTY_FUNCTION)
a974afa6
JB
359 worker.once('exit', () => {
360 this.removeWorker(worker)
361 })
280c2a77 362
c923ce56 363 this.pushWorker(worker, {
ffcbbad8
JB
364 run: 0,
365 running: 0,
366 runTime: 0,
2740a743
JB
367 avgRunTime: 0,
368 error: 0
ffcbbad8 369 })
280c2a77
S
370
371 this.afterWorkerSetup(worker)
372
c97c7edb
S
373 return worker
374 }
be0676b3
APA
375
376 /**
377 * This function is the listener registered for each worker.
378 *
bdacc2d2 379 * @returns The listener function to execute when a message is received from a worker.
be0676b3
APA
380 */
381 protected workerListener (): (message: MessageValue<Response>) => void {
4a6952ff 382 return message => {
b1989cfd 383 if (message.id != null) {
2740a743 384 const promiseResponse = this.promiseResponseMap.get(message.id)
b1989cfd 385 if (promiseResponse != null) {
78cea37e 386 if (message.error != null) {
2740a743 387 promiseResponse.reject(message.error)
a05c10de 388 } else {
2740a743 389 promiseResponse.resolve(message.data as Response)
a05c10de 390 }
c923ce56 391 this.afterPromiseResponseHook(promiseResponse.worker, message)
2740a743 392 this.promiseResponseMap.delete(message.id)
be0676b3
APA
393 }
394 }
395 }
be0676b3 396 }
7c0ba920 397
78cea37e 398 private async internalExecute (
2740a743 399 workerKey: number,
c923ce56 400 worker: Worker,
b4e75778 401 messageId: string
78cea37e 402 ): Promise<Response> {
2740a743 403 this.beforePromiseResponseHook(workerKey)
78cea37e 404 return await new Promise<Response>((resolve, reject) => {
c923ce56 405 this.promiseResponseMap.set(messageId, { resolve, reject, worker })
78cea37e
JB
406 })
407 }
408
7c0ba920 409 private checkAndEmitBusy (): void {
78cea37e 410 if (this.opts.enableEvents === true && this.busy) {
7c0ba920
JB
411 this.emitter?.emit('busy')
412 }
413 }
bf9549ae 414
164d950a
JB
415 private checkAndEmitFull (): void {
416 if (
417 this.type === PoolType.DYNAMIC &&
418 this.opts.enableEvents === true &&
419 this.full
420 ) {
421 this.emitter?.emit('full')
422 }
423 }
424
c923ce56
JB
425 /**
426 * Gets worker tasks usage.
427 *
428 * @param worker - The worker.
429 * @returns The worker tasks usage.
430 */
431 private getWorkerTasksUsage (worker: Worker): TasksUsage | undefined {
3032893a 432 const workerKey = this.getWorkerKey(worker)
e65c6cd9
JB
433 if (workerKey !== -1) {
434 return this.workers[workerKey].tasksUsage
ffcbbad8 435 }
3032893a 436 throw new Error('Worker could not be found in the pool')
a05c10de
JB
437 }
438
439 /**
51fe3d3c 440 * Pushes the given worker in the pool.
ea7a90d3 441 *
38e795c1 442 * @param worker - The worker.
ffcbbad8 443 * @param tasksUsage - The worker tasks usage.
ea7a90d3 444 */
c923ce56 445 private pushWorker (worker: Worker, tasksUsage: TasksUsage): void {
e65c6cd9 446 this.workers.push({
ffcbbad8
JB
447 worker,
448 tasksUsage
ea7a90d3
JB
449 })
450 }
c923ce56
JB
451
452 /**
51fe3d3c 453 * Sets the given worker in the pool.
c923ce56
JB
454 *
455 * @param workerKey - The worker key.
456 * @param worker - The worker.
457 * @param tasksUsage - The worker tasks usage.
458 */
459 private setWorker (
460 workerKey: number,
461 worker: Worker,
462 tasksUsage: TasksUsage
463 ): void {
464 this.workers[workerKey] = {
465 worker,
466 tasksUsage
467 }
468 }
51fe3d3c
JB
469
470 /**
471 * Removes the given worker from the pool.
472 *
473 * @param worker - The worker that will be removed.
474 */
475 protected removeWorker (worker: Worker): void {
476 const workerKey = this.getWorkerKey(worker)
477 this.workers.splice(workerKey, 1)
478 this.workerChoiceStrategyContext.remove(workerKey)
479 }
c97c7edb 480}