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